From fc95d223670541d1cbd43f4ae76fd3d8e87ec51b Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 19:18:38 -0700 Subject: [PATCH] fix(mcp): accept VS Code OAuth registration callbacks --- .../mcp_server/gateway_dcr_flow.py | 8 +- .../mcp_server/test_gateway_dcr_flow.py | 74 +++++++++++++++++-- 2 files changed, 72 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index f4889008e94..f3fdd54b39d 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -122,13 +122,13 @@ _USED_CODE_CACHE_PREFIX: Final = "mcp_gateway_dcr_code_used:" _USED_FLOW_CACHE_PREFIX: Final = "mcp_gateway_dcr_flow_used:" _USED_REFRESH_CACHE_PREFIX: Final = "mcp_gateway_dcr_refresh_used:" -MAX_REDIRECT_URIS: Final = 3 +MAX_REDIRECT_URIS: Final = 4 MAX_REDIRECT_URI_LENGTH: Final = 256 MAX_CLIENT_ID_LENGTH: Final = 2048 """Registration bounds. They exist to bound the sealed client_id, which rides inside -every session-token claim set: 3 URIs of 256 bytes seal to roughly 1.2KB, comfortably -under this cap and under the session token's own 4KB ceiling. Claude Desktop and MCP -Inspector register one or two redirect URIs.""" +every session-token claim set. Four 256-character ASCII URIs seal to roughly 1.5KB; +the encoded client_id is checked against its own cap before registration succeeds. +VS Code registers four callbacks for its web and desktop environments.""" MAX_STATE_LENGTH: Final = 1024 """Bound on the client ``state`` sealed into the flow cookie and echoed on the auth-code diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 8aa4cfc5619..7c80ee77cd7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -6,6 +6,7 @@ import re from base64 import urlsafe_b64encode from datetime import datetime, timedelta, timezone from http.cookies import SimpleCookie +from typing import Final from urllib.parse import parse_qs, urlparse import pytest @@ -18,6 +19,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( GATEWAY_AUTH_CODE_PREFIX, GATEWAY_AUTH_CODE_TTL_SECONDS, MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS, + MAX_CLIENT_ID_LENGTH, ConsentTeam, MintedProxyCredential, _GatewayAuthCode, @@ -53,6 +55,13 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i MASTER_KEY = "sk-gateway-dcr-flow-tests" REDIRECT_URI = "https://claude.ai/api/mcp/auth_callback" +VSCODE_REDIRECT_URIS: Final = ( + "https://insiders.vscode.dev/redirect", + "https://vscode.dev/redirect", + "http://127.0.0.1/", + "http://127.0.0.1:33418/", +) +MAX_LENGTH_REDIRECT_URIS: Final = tuple(f"https://client.example/{index}/".ljust(256, "a") for index in range(4)) CODE_VERIFIER = "verifier-" + "v" * 43 CODE_CHALLENGE = urlsafe_b64encode(hashlib.sha256(CODE_VERIFIER.encode("ascii")).digest()).rstrip(b"=").decode("ascii") @@ -104,6 +113,58 @@ async def test_register_mints_stateless_public_client(): assert record.redirect_uris == (REDIRECT_URI,) +@pytest.mark.asyncio +@pytest.mark.parametrize("redirect_uris", [VSCODE_REDIRECT_URIS, MAX_LENGTH_REDIRECT_URIS]) +async def test_register_four_callbacks_preserves_metadata(redirect_uris: tuple[str, ...]) -> None: + response: Final = await register_aggregate_client( + request=_request(path="/register", method="POST"), + request_body={ + "client_name": "Visual Studio Code", + "client_uri": "https://code.visualstudio.com", + "grant_types": ["authorization_code", "refresh_token"], + "response_types": ["code"], + "redirect_uris": list(redirect_uris), + "token_endpoint_auth_method": "none", + "application_type": "native", + }, + ) + assert response.status_code == 201 + body: Final = json.loads(response.body) + assert body["redirect_uris"] == list(redirect_uris) + assert body["token_endpoint_auth_method"] == "none" + assert "client_secret" not in body + assert len(body["client_id"]) <= MAX_CLIENT_ID_LENGTH + record: Final = open_gateway_dcr_client(body["client_id"]) + assert record is not None + assert record.redirect_uris == redirect_uris + + +@pytest.mark.asyncio +async def test_register_rejects_five_valid_callbacks() -> None: + response: Final = await register_aggregate_client( + request=_request(path="/register", method="POST"), + request_body={"redirect_uris": [*VSCODE_REDIRECT_URIS, "http://127.0.0.1:33419/"]}, + ) + assert response.status_code == 400 + assert json.loads(response.body) == { + "error": "invalid_redirect_uri", + "error_description": "redirect_uris must be a list of 1 to 4 URIs", + } + + +@pytest.mark.asyncio +async def test_register_four_callbacks_preserves_encoded_size_guard() -> None: + response: Final = await register_aggregate_client( + request=_request(path="/register", method="POST"), + request_body={"redirect_uris": [f"https://client.example/{index}/".ljust(256, "é") for index in range(4)]}, + ) + assert response.status_code == 400 + assert json.loads(response.body) == { + "error": "invalid_client_metadata", + "error_description": "registered metadata is too large", + } + + @pytest.mark.asyncio async def test_register_allows_loopback_http_for_dev_clients(): body = await _register(["http://localhost:6274/oauth/callback"]) @@ -162,7 +223,6 @@ async def test_register_rejects_userinfo_spoofed_origin(): ["https://claude.ai/cb#fragment"], ["ftp://claude.ai/cb"], ["https://a.example.com/" + "p" * 300], - ["https://a.example.com/1", "https://a.example.com/2", "https://a.example.com/3", "https://a.example.com/4"], [12345], ], ) @@ -248,12 +308,14 @@ def _flow_cookie_from(response) -> tuple: @pytest.mark.asyncio -async def test_full_walk_register_authorize_complete_token_and_replay(): +@pytest.mark.parametrize("redirect_uris", [(REDIRECT_URI,), VSCODE_REDIRECT_URIS, MAX_LENGTH_REDIRECT_URIS]) +async def test_full_walk_register_authorize_complete_token_and_replay(redirect_uris: tuple[str, ...]): """The whole front door on one deterministic walk: register -> authorize -> complete -> token, then the security edges on the same artifacts (user mismatch, PKCE mismatch, single-use replay, refresh rotation, cross-client refresh).""" - client_id = (await _register([REDIRECT_URI]))["client_id"] - authorize_response = _authorize(client_id, session_user_id="u1") + redirect_uri: Final = redirect_uris[-1] + client_id = (await _register(list(redirect_uris)))["client_id"] + authorize_response = _authorize(client_id, session_user_id="u1", redirect_uri=redirect_uri) handle, cookies = _flow_cookie_from(authorize_response) denied = await complete_connect_flow( @@ -280,7 +342,7 @@ async def test_full_walk_register_authorize_complete_token_and_replay(): ) assert completed.status_code == 303 redirect = urlparse(completed.headers["location"]) - assert f"{redirect.scheme}://{redirect.netloc}{redirect.path}" == REDIRECT_URI + assert f"{redirect.scheme}://{redirect.netloc}{redirect.path}" == redirect_uri params = parse_qs(redirect.query) assert params["state"] == ["client-state-123"] code = params["code"][0] @@ -293,7 +355,7 @@ async def test_full_walk_register_authorize_complete_token_and_replay(): "request": _request("/token", method="POST"), "grant_type": "authorization_code", "code": code, - "redirect_uri": REDIRECT_URI, + "redirect_uri": redirect_uri, "client_id": client_id, "code_verifier": CODE_VERIFIER, "refresh_token": None,