fix(claude_code_gateway): keep the device secret out of the browser URL and validate the login before claiming it

litellm/proxy/anthropic_endpoints/gateway_endpoints.py

@@ -18,6 +18,7 @@ import hashlib
 import json
 import secrets
 from collections.abc import Mapping
+from dataclasses import dataclass
 from types import MappingProxyType
 from typing import Final
 
@@ -32,13 +33,16 @@ from litellm.constants import (
     CLI_SSO_SESSION_TTL_SECONDS,
     LITELLM_CLI_SOURCE_IDENTIFIER,
 )
+from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles
 from litellm.proxy.anthropic_endpoints.endpoints import anthropic_response, count_tokens
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.common_utils.http_parsing_utils import _safe_set_request_parsed_body
+from litellm.proxy.management_endpoints.ui_sso import CliSsoTeamDetail
 
 GATEWAY_PREFIX: Final = "/claude_code_gateway"
 _DEVICE_CODE_GRANT: Final = "urn:ietf:params:oauth:grant-type:device_code"
 _REFRESH_TOKEN_GRANT: Final = "refresh_token"
+_DEVICE_CODE_SEPARATOR: Final = "."
 _DEVICE_POLL_INTERVAL_SECONDS: Final = 5
 _SECONDS_PER_HOUR: Final = 3600
 _MANAGED_SETTINGS_ADAPTER: Final = TypeAdapter(dict[str, object])
@@ -48,12 +52,19 @@ _POST_ONLY: Final = ["POST"]  # mutable-ok: FastAPI's add_api_route only accepts
 
 class _GatewaySessionData(BaseModel):
     user_id: str
-    user_role: str | None
+    user_role: LitellmUserRoles
     models: list[str] = Field(default_factory=list)
     teams: tuple[str, ...] = ()
     team_details: object | None = None
 
 
+@dataclass(frozen=True, slots=True)
+class _GatewayLogin:
+    user_info: LiteLLM_UserTable
+    team_id: str | None
+    team: CliSsoTeamDetail
+
+
 class _OAuthErrorBody(BaseModel):
     error: str
     error_description: str | None = None
@@ -70,7 +81,7 @@ class _DeviceAuthorizationBody(BaseModel):
     device_code: str
     user_code: str
     verification_uri: str
-    verification_uri_complete: str
+    verification_uri_complete: str | None = None
     expires_in: int
     interval: int
 
@@ -111,15 +122,11 @@ def _managed_settings() -> dict[str, object] | None:
     return _MANAGED_SETTINGS_ADAPTER.validate_python(settings)
 
 
-def _oauth_error(*, status_code: int, error: str, description: str | None = None) -> "_OAuthError":
-    return _OAuthError(status_code=status_code, error=error, description=description)
-
-
-class _OAuthError(Exception):
-    def __init__(self, *, status_code: int, error: str, description: str | None) -> None:
-        self.status_code = status_code
-        self.error = error
-        self.description = description
+@dataclass(frozen=True, slots=True)
+class _OAuthError:
+    status_code: int
+    error: str
+    description: str | None = None
 
 
 def _oauth_error_response(err: _OAuthError) -> JSONResponse:
@@ -153,7 +160,7 @@ router.add_api_route(
 @router.get("/.well-known/oauth-authorization-server", include_in_schema=False)
 async def oauth_authorization_server(request: Request) -> JSONResponse:
     if not _is_gateway_enabled():
-        return _oauth_error_response(_oauth_error(status_code=404, error="not_found"))
+        return _oauth_error_response(_OAuthError(status_code=404, error="not_found"))
 
     from litellm.proxy.utils import get_custom_url
 
@@ -175,6 +182,7 @@ async def device_authorization(request: Request) -> JSONResponse:
 
     from litellm.proxy.management_endpoints.ui_sso import (
         _check_cli_sso_start_rate_limit,  # pyright: ignore[reportPrivateUsage]  # shared device-flow helper
+        _cli_sso_verification_uri_complete_enabled,  # pyright: ignore[reportPrivateUsage]  # shared device-flow helper
         _generate_cli_sso_user_code,  # pyright: ignore[reportPrivateUsage]  # shared device-flow helper
         _hash_cli_sso_secret,  # pyright: ignore[reportPrivateUsage]  # shared device-flow helper
         _normalize_cli_sso_user_code,  # pyright: ignore[reportPrivateUsage]  # shared device-flow helper
@@ -184,7 +192,7 @@ async def device_authorization(request: Request) -> JSONResponse:
     from litellm.proxy.utils import get_custom_url
 
     if not _is_gateway_enabled():
-        return _oauth_error_response(_oauth_error(status_code=404, error="not_found"))
+        return _oauth_error_response(_OAuthError(status_code=404, error="not_found"))
 
     _check_cli_sso_start_rate_limit(
         request=request,
@@ -192,50 +200,51 @@ async def device_authorization(request: Request) -> JSONResponse:
         use_x_forwarded_for=bool(_general_settings().get("use_x_forwarded_for", False)),
     )
 
-    device_code: Final = f"cli-{secrets.token_urlsafe(24)}"
+    login_id: Final = f"cli-{secrets.token_urlsafe(24)}"
+    poll_secret: Final = secrets.token_urlsafe(32)
     user_code: Final = _generate_cli_sso_user_code()
     flow: Final = {  # mutable-ok: the shared CLI SSO cache entry is a dict the browser leg mutates
-        "poll_secret_hash": _hash_cli_sso_secret(device_code),
+        "poll_secret_hash": _hash_cli_sso_secret(poll_secret),
         "user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code(user_code)),
         "sso_complete": False,
         "user_code_verified": False,
         "session_data": None,
     }
-    _set_cli_sso_flow(login_id=device_code, cache=cli_sso_session_cache, flow=flow)
+    _set_cli_sso_flow(login_id=login_id, cache=cli_sso_session_cache, flow=flow)
 
     request_base_url: Final = str(request.base_url)
     verification_uri: Final = get_custom_url(request_base_url=request_base_url, route="sso/key/generate")
-    query: Final = MappingProxyType({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": device_code})
+    query: Final = MappingProxyType({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": login_id})
     body: Final = _DeviceAuthorizationBody(
-        device_code=device_code,
+        device_code=f"{login_id}{_DEVICE_CODE_SEPARATOR}{poll_secret}",
         user_code=user_code,
         verification_uri=f"{verification_uri}?{urlencode(query)}",
         verification_uri_complete=(
             f"{verification_uri}?{urlencode(MappingProxyType({**query, 'user_code': user_code}))}"
+            if _cli_sso_verification_uri_complete_enabled()
+            else None
         ),
         expires_in=CLI_SSO_SESSION_TTL_SECONDS,
         interval=_DEVICE_POLL_INTERVAL_SECONDS,
     )
-    return JSONResponse(content=body.model_dump())
+    return JSONResponse(content=body.model_dump(exclude_none=True))
 
 
-def _mint_access_token_from_flow(flow: Mapping[str, object]) -> str:
-    from litellm.proxy._types import LiteLLM_UserTable
-    from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
+def _validate_login(flow: Mapping[str, object]) -> _GatewayLogin | _OAuthError:
     from litellm.proxy.management_endpoints.ui_sso import selected_cli_sso_team_detail
 
     try:
         session_data: Final = _GatewaySessionData.model_validate(flow.get("session_data"))
     except ValidationError as err:
         verbose_proxy_logger.warning("Claude Code gateway login session is malformed: %s", err)
-        raise _oauth_error(
+        return _OAuthError(
             status_code=400, error="invalid_grant", description="The login session is malformed; sign in again"
-        ) from err
+        )
 
     team_id: Final = session_data.teams[0] if session_data.teams else None
     selected_team: Final = selected_cli_sso_team_detail(team_details=session_data.team_details, team_id=team_id)
     if selected_team is None:
-        raise _oauth_error(
+        return _OAuthError(
             status_code=400,
             error="invalid_grant",
             description=f"Could not resolve the model grants for team {team_id}; sign in again",
@@ -243,26 +252,32 @@ def _mint_access_token_from_flow(flow: Mapping[str, object]) -> str:
 
     user_info: Final = LiteLLM_UserTable(
         user_id=session_data.user_id,
-        user_role=session_data.user_role,
+        user_role=session_data.user_role.value,
         models=session_data.models,
     )
+    return _GatewayLogin(user_info=user_info, team_id=team_id, team=selected_team)
+
+
+def _mint_access_token(login: _GatewayLogin) -> str:
+    from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
+
     return ExperimentalUIJWTToken.get_cli_jwt_auth_token(
-        user_info=user_info,
-        team_id=team_id,
-        team_alias=selected_team.team_alias,
-        team_models=selected_team.team_models,
-        team_model_aliases=selected_team.team_model_aliases,
+        user_info=login.user_info,
+        team_id=login.team_id,
+        team_alias=login.team.team_alias,
+        team_models=login.team.team_models,
+        team_model_aliases=login.team.team_model_aliases,
         max_budget=None,
     )
 
 
-async def _claim_device_code(device_code: str, cache: DualCache) -> bool:
+async def _claim_device_code(login_id: str, cache: DualCache) -> bool:
     from litellm.proxy.management_endpoints.ui_sso import (
         _get_cli_sso_flow_cache_key,  # pyright: ignore[reportPrivateUsage]  # shared device-flow helper
     )
 
     claims: Final = await cache.async_increment_cache(
-        key=f"{_get_cli_sso_flow_cache_key(device_code)}:claimed",
+        key=f"{_get_cli_sso_flow_cache_key(login_id)}:claimed",
         value=1,
         ttl=CLI_SSO_SESSION_TTL_SECONDS,
     )
@@ -275,39 +290,45 @@ async def _handle_device_code_grant(device_code: str | None) -> JSONResponse:
     from litellm.proxy.management_endpoints.ui_sso import (
         _get_cli_sso_flow_cache_key,  # pyright: ignore[reportPrivateUsage]  # shared device-flow helper
         _get_cli_sso_flow_or_raise,  # pyright: ignore[reportPrivateUsage]  # shared device-flow helper
+        _verify_cli_sso_poll_secret,  # pyright: ignore[reportPrivateUsage]  # shared device-flow helper
     )
     from litellm.proxy.proxy_server import cli_sso_session_cache
 
     if not device_code:
         return _oauth_error_response(
-            _oauth_error(status_code=400, error="invalid_request", description="device_code is required")
+            _OAuthError(status_code=400, error="invalid_request", description="device_code is required")
         )
 
+    login_id, _, poll_secret = device_code.partition(_DEVICE_CODE_SEPARATOR)
     try:
-        flow: Final = _get_cli_sso_flow_or_raise(login_id=device_code, cache=cli_sso_session_cache)
+        flow: Final = _get_cli_sso_flow_or_raise(login_id=login_id, cache=cli_sso_session_cache)
     except HTTPException:
-        return _oauth_error_response(_oauth_error(status_code=400, error="expired_token"))
+        return _oauth_error_response(_OAuthError(status_code=400, error="expired_token"))
+
+    if not _verify_cli_sso_poll_secret(flow, poll_secret):
+        return _oauth_error_response(_OAuthError(status_code=400, error="expired_token"))
 
     if not flow.get("sso_complete") or not flow.get("user_code_verified"):
-        return _oauth_error_response(_oauth_error(status_code=400, error="authorization_pending"))
+        return _oauth_error_response(_OAuthError(status_code=400, error="authorization_pending"))
 
-    if not await _claim_device_code(device_code, cli_sso_session_cache):
-        return _oauth_error_response(_oauth_error(status_code=400, error="expired_token"))
+    login: Final = _validate_login(flow)
+    if isinstance(login, _OAuthError):
+        return _oauth_error_response(login)
 
-    await cli_sso_session_cache.async_delete_cache(key=_get_cli_sso_flow_cache_key(device_code))
-    try:
-        access_token: Final = _mint_access_token_from_flow(flow)
-    except _OAuthError as err:
-        return _oauth_error_response(err)
+    if not await _claim_device_code(login_id, cli_sso_session_cache):
+        return _oauth_error_response(_OAuthError(status_code=400, error="expired_token"))
 
-    body: Final = _AccessTokenBody(access_token=access_token, expires_in=CLI_JWT_EXPIRATION_HOURS * _SECONDS_PER_HOUR)
+    await cli_sso_session_cache.async_delete_cache(key=_get_cli_sso_flow_cache_key(login_id))
+    body: Final = _AccessTokenBody(
+        access_token=_mint_access_token(login), expires_in=CLI_JWT_EXPIRATION_HOURS * _SECONDS_PER_HOUR
+    )
     return JSONResponse(content=body.model_dump())
 
 
 @router.post("/oauth/token", include_in_schema=False)
 async def oauth_token(request: Request) -> JSONResponse:
     if not _is_gateway_enabled():
-        return _oauth_error_response(_oauth_error(status_code=404, error="not_found"))
+        return _oauth_error_response(_OAuthError(status_code=404, error="not_found"))
 
     form: Final = await request.form()
     grant_type: Final = form.get("grant_type")
@@ -318,7 +339,7 @@ async def oauth_token(request: Request) -> JSONResponse:
 
     if grant_type == _REFRESH_TOKEN_GRANT:
         return _oauth_error_response(
-            _oauth_error(
+            _OAuthError(
                 status_code=401,
                 error="invalid_grant",
                 description="This gateway does not issue refresh tokens; sign in again",
@@ -326,7 +347,7 @@ async def oauth_token(request: Request) -> JSONResponse:
         )
 
     return _oauth_error_response(
-        _oauth_error(
+        _OAuthError(
             status_code=400, error="unsupported_grant_type", description=f"Unsupported grant_type: {grant_type}"
         )
     )

tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py

@@ -20,11 +20,18 @@ from fastapi.testclient import TestClient
 from litellm.caching.dual_cache import DualCache
 from litellm.proxy._types import ProxyException
 from litellm.proxy.anthropic_endpoints import gateway_endpoints
-from litellm.proxy.management_endpoints.ui_sso import _get_cli_sso_flow_cache_key, _set_cli_sso_flow
+from litellm.proxy.management_endpoints.ui_sso import (
+    _get_cli_sso_flow_cache_key,
+    _hash_cli_sso_secret,
+    _set_cli_sso_flow,
+)
 from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware
 
 _DEVICE_CODE_GRANT: Final = "urn:ietf:params:oauth:grant-type:device_code"
 _MASTER_KEY: Final = "sk-master-key"
+_SHARED_LOGIN_ID: Final = "cli-shared-login-code"
+_SHARED_POLL_SECRET: Final = "shared-poll-secret"
+_SHARED_DEVICE_CODE: Final = f"{_SHARED_LOGIN_ID}.{_SHARED_POLL_SECRET}"
 _MINT: Final = "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token"
 _PROTOBUF_BODY: Final = b"\x0a\x05hello\x12\x03{{{"
 _COMPLETED_SESSION: Final = MappingProxyType(
@@ -100,10 +107,12 @@ def _gateway_env(
     managed_settings: Mapping[str, object] | None = None,
     cache: DualCache | None = None,
     real_auth: bool = False,
+    extra_settings: Mapping[str, object] = MappingProxyType({}),
 ) -> Iterator[tuple[TestClient, DualCache]]:
     general_settings: Final = {
         "enable_claude_code_gateway": enabled,
         **({} if managed_settings is None else {"claude_code_gateway_managed_settings": dict(managed_settings)}),
+        **extra_settings,
     }
     session_cache: Final = cache or DualCache(default_in_memory_ttl=600)
 
@@ -147,7 +156,7 @@ def _request_token(client: TestClient, device_code: str) -> httpx.Response:
 
 def _completed_flow(session_data: Mapping[str, object] = _COMPLETED_SESSION) -> dict[str, object]:
     return {
-        "poll_secret_hash": "unused",
+        "poll_secret_hash": _hash_cli_sso_secret(_SHARED_POLL_SECRET),
         "user_code_hash": "unused",
         "sso_complete": True,
         "user_code_verified": True,
@@ -155,13 +164,18 @@ def _completed_flow(session_data: Mapping[str, object] = _COMPLETED_SESSION) ->
     }
 
 
+def _login_id(device_code: str) -> str:
+    return device_code.partition(".")[0]
+
+
 def _complete_flow(
     cache: DualCache, device_code: str, session_data: Mapping[str, object] = _COMPLETED_SESSION
 ) -> None:
-    key: Final = _get_cli_sso_flow_cache_key(device_code)
+    key: Final = _get_cli_sso_flow_cache_key(_login_id(device_code))
     flow: Final = cache.get_cache(key=key)
     assert isinstance(flow, dict)
-    cache.set_cache(key=key, value={**flow, **_completed_flow(session_data)}, ttl=600)
+    completed: Final = {**flow, **_completed_flow(session_data), "poll_secret_hash": flow["poll_secret_hash"]}
+    cache.set_cache(key=key, value=completed, ttl=600)
 
 
 def test_discovery_shape():
@@ -194,18 +208,35 @@ def test_device_authorization_returns_rfc8628_shape_and_persists_flow():
         assert resp.status_code == 200
         body = resp.json()
         device_code = body["device_code"]
-        assert device_code.startswith("cli-")
+        login_id, separator, poll_secret = device_code.partition(".")
+        assert login_id.startswith("cli-")
+        assert separator == "."
+        assert len(poll_secret) >= 32
         assert body["user_code"]
         assert body["expires_in"] == 600
         assert body["interval"] == 5
-        # verification_uri_complete carries the user_code; the short uri does not.
-        assert f"user_code={body['user_code']}" in body["verification_uri_complete"]
-        assert "user_code=" not in body["verification_uri"]
-        assert f"key={device_code}" in body["verification_uri"]
-        # The device flow is stored under the device_code so the browser SSO leg can complete it.
-        stored = cache.get_cache(key=_get_cli_sso_flow_cache_key(device_code))
+        assert "verification_uri_complete" not in body
+        assert body["verification_uri"].endswith(f"/sso/key/generate?source=litellm-cli&key={login_id}")
+        assert poll_secret not in body["verification_uri"]
+        stored = cache.get_cache(key=_get_cli_sso_flow_cache_key(login_id))
         assert isinstance(stored, dict)
         assert stored["sso_complete"] is False
+        assert stored["poll_secret_hash"] == _hash_cli_sso_secret(poll_secret)
+        assert cache.get_cache(key=_get_cli_sso_flow_cache_key(device_code)) is None
+
+
+@pytest.mark.parametrize("opted_in", [True, False])
+def test_verification_uri_complete_carries_the_user_code_only_when_the_operator_opts_in(opted_in: bool):
+    with _gateway_env(extra_settings={"allow_cli_sso_verification_uri_complete": opted_in}) as (client, _):
+        body = client.post("/claude_code_gateway/oauth/device_authorization").json()
+    login_id = _login_id(body["device_code"])
+    if not opted_in:
+        assert "verification_uri_complete" not in body
+        return
+    assert body["verification_uri_complete"].endswith(
+        f"/sso/key/generate?source=litellm-cli&key={login_id}&user_code={body['user_code']}"
+    )
+    assert "user_code=" not in body["verification_uri"]
 
 
 def test_token_authorization_pending_before_browser_completes():
@@ -215,6 +246,22 @@ def test_token_authorization_pending_before_browser_completes():
     assert resp.json()["error"] == "authorization_pending"
 
 
+@pytest.mark.parametrize("tamper", ["login_id_only", "wrong_secret"])
+def test_token_refuses_the_browser_login_id_without_the_client_secret(tamper: str):
+    with _gateway_env() as (client, cache):
+        device_code = _start_device_flow(client)
+        _complete_flow(cache, device_code)
+        login_id = _login_id(device_code)
+        presented = login_id if tamper == "login_id_only" else f"{login_id}.not-the-secret"
+        with patch(_MINT, return_value="sk-session") as mint:
+            resp = _request_token(client, presented)
+            assert resp.status_code == 400
+            assert resp.json()["error"] == "expired_token"
+            mint.assert_not_called()
+            with_secret = _request_token(client, device_code)
+    assert with_secret.status_code == 200
+
+
 def test_token_success_mints_bearer_and_is_single_use():
     with _gateway_env() as (client, cache):
         device_code = _start_device_flow(client)
@@ -251,14 +298,25 @@ def test_token_teamless_user_mints_without_a_team():
     assert mint.call_args.kwargs["team_models"] == ()
 
 
-def test_token_malformed_session_is_invalid_grant():
+@pytest.mark.parametrize(
+    "session_data",
+    [
+        {"user_role": "internal_user"},
+        {**_COMPLETED_SESSION, "user_role": None},
+        {**_COMPLETED_SESSION, "user_role": "not-a-role"},
+    ],
+    ids=["missing_user_id", "no_role", "unknown_role"],
+)
+def test_token_malformed_session_is_invalid_grant_and_does_not_consume_the_login(session_data: Mapping[str, object]):
     with _gateway_env() as (client, cache):
         device_code = _start_device_flow(client)
-        _complete_flow(cache, device_code, session_data={"user_role": "internal_user"})
+        _complete_flow(cache, device_code, session_data=session_data)
         with patch(_MINT) as mint:
             resp = _request_token(client, device_code)
+            again = _request_token(client, device_code)
     assert resp.status_code == 400
     assert resp.json()["error"] == "invalid_grant"
+    assert again.json()["error"] == "invalid_grant"
     mint.assert_not_called()
 
 
@@ -275,25 +333,24 @@ def test_token_unknown_team_grants_is_invalid_grant():
 
 def test_token_mints_on_a_replica_that_did_not_start_the_login():
     redis: Final = _SharedRedisFake()
-    device_code: Final = "cli-shared-login-code"
-    _set_cli_sso_flow(login_id=device_code, cache=_replica(redis), flow=_completed_flow())
+    _set_cli_sso_flow(login_id=_SHARED_LOGIN_ID, cache=_replica(redis), flow=_completed_flow())
 
     with _gateway_env(cache=_replica(redis)) as (client, _), patch(_MINT, return_value="sk-session") as mint:
-        resp = _request_token(client, device_code)
+        resp = _request_token(client, _SHARED_DEVICE_CODE)
     assert resp.status_code == 200
     assert resp.json()["access_token"] == "sk-session"
     assert mint.call_args.kwargs["team_id"] == "team-a"
+    assert mint.call_args.kwargs["user_info"].user_role == "internal_user"
 
 
 def test_token_refuses_a_device_code_another_replica_already_claimed():
     redis: Final = _SharedRedisFake()
     replica_a: Final = _replica(redis)
-    device_code: Final = "cli-shared-login-code"
-    _set_cli_sso_flow(login_id=device_code, cache=replica_a, flow=_completed_flow())
-    assert asyncio.run(gateway_endpoints._claim_device_code(device_code, replica_a)) is True
+    _set_cli_sso_flow(login_id=_SHARED_LOGIN_ID, cache=replica_a, flow=_completed_flow())
+    assert asyncio.run(gateway_endpoints._claim_device_code(_SHARED_LOGIN_ID, replica_a)) is True
 
     with _gateway_env(cache=_replica(redis)) as (client, _), patch(_MINT) as mint:
-        resp = _request_token(client, device_code)
+        resp = _request_token(client, _SHARED_DEVICE_CODE)
     assert resp.status_code == 400
     assert resp.json()["error"] == "expired_token"
     mint.assert_not_called()