diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 1dda1f29fb9..1426d100713 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -196,6 +196,11 @@ LAZY_FEATURES: Tuple[LazyFeature, ...] = ( module_path="litellm.proxy.anthropic_endpoints.skills_endpoints", path_prefixes=("/v1/skills", "/skills"), ), + LazyFeature( + name="claude_code_gateway", + module_path="litellm.proxy.anthropic_endpoints.gateway_endpoints", + path_prefixes=("/claude_code_gateway",), + ), LazyFeature( name="langfuse_passthrough", module_path="litellm.proxy.vertex_ai_endpoints.langfuse_endpoints", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7df725bf965..5577c5caff4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2223,6 +2223,14 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="opt-in to RFC 8628 verification_uri_complete for the CLI SSO device flow, pre-filling the user_code in the browser. Off by default; intended for same-host clients where the device that starts the flow and the browser run on the same machine", ) + enable_claude_code_gateway: bool | None = Field( + None, + description="serve the Claude Code gateway protocol (https://code.claude.com/docs/en/claude-apps-gateway) under /claude_code_gateway: OAuth device-flow sign-in reusing proxy SSO, plus managed settings and OTLP telemetry ingestion. Off by default", + ) + claude_code_gateway_managed_settings: Dict[str, Any] | None = Field( + None, + description="Claude Code managed-settings.json served verbatim at the gateway's /claude_code_gateway/managed/settings endpoint. When unset the endpoint returns 404 (no managed policy)", + ) database_url: Optional[str] = Field( None, description="connect to a postgres db - needed for generating temporary keys + tracking spend / key", diff --git a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py new file mode 100644 index 00000000000..1ec4cd488c5 --- /dev/null +++ b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py @@ -0,0 +1,289 @@ +""" +Claude Code gateway protocol. + +Implements the wire contract the Claude Code CLI uses to talk to a gateway: +OAuth 2.0 device-authorization sign-in (RFC 8414 / RFC 8628), inference via the +Anthropic Messages API, managed settings, and OTLP telemetry ingestion. See +https://code.claude.com/docs/en/claude-apps-gateway. + +Everything lives under the ``/claude_code_gateway`` base so operators point +Claude Code at ``https:///claude_code_gateway`` via ``/login``. The +device flow reuses the proxy's existing SSO login machinery: the browser leg is +served by ``/sso/key/generate`` and the shared ``cli_sso_session_cache`` flow, +so the bearer token minted here is the same session JWT the LiteLLM CLI uses and +is accepted by every bearer-authenticated proxy route. +""" + +import hashlib +import json +import secrets +from typing import Any + +from fastapi import APIRouter, Depends, Request, Response +from fastapi.responses import JSONResponse + +from litellm.constants import ( + CLI_JWT_EXPIRATION_HOURS, + CLI_SSO_SESSION_TTL_SECONDS, + LITELLM_CLI_SOURCE_IDENTIFIER, +) +from litellm.proxy.anthropic_endpoints.endpoints import anthropic_response, count_tokens +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + +GATEWAY_PREFIX = "/claude_code_gateway" +_DEVICE_CODE_GRANT = "urn:ietf:params:oauth:grant-type:device_code" +_REFRESH_TOKEN_GRANT = "refresh_token" +_DEVICE_POLL_INTERVAL_SECONDS = 5 + + +def _is_gateway_enabled() -> bool: + from litellm.proxy.proxy_server import general_settings + + return bool((general_settings or {}).get("enable_claude_code_gateway", False)) + + +def ensure_gateway_enabled() -> None: + from fastapi import HTTPException + + if not _is_gateway_enabled(): + raise HTTPException(status_code=404, detail="Claude Code gateway is not enabled") + + +def _managed_settings() -> dict[str, Any] | None: + from litellm.proxy.proxy_server import general_settings + + settings = (general_settings or {}).get("claude_code_gateway_managed_settings") + return settings if isinstance(settings, dict) else None + + +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 + + +def _oauth_error_response(err: _OAuthError) -> JSONResponse: + body: dict[str, str] = {"error": err.error} + if err.description is not None: + body["error_description"] = err.description + return JSONResponse(status_code=err.status_code, content=body) + + +router = APIRouter(prefix=GATEWAY_PREFIX, tags=["Claude Code gateway"]) + +router.add_api_route( + "/v1/messages", + anthropic_response, + methods=["POST"], + dependencies=[Depends(ensure_gateway_enabled)], + include_in_schema=False, +) +router.add_api_route( + "/v1/messages/count_tokens", + count_tokens, + methods=["POST"], + dependencies=[Depends(ensure_gateway_enabled)], + include_in_schema=False, +) + + +@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")) + + from litellm.proxy.utils import get_custom_url + + request_base_url = str(request.base_url) + issuer = get_custom_url(request_base_url=request_base_url, route="claude_code_gateway") + return JSONResponse( + content={ + "issuer": issuer, + "device_authorization_endpoint": get_custom_url( + request_base_url=request_base_url, route="claude_code_gateway/oauth/device_authorization" + ), + "token_endpoint": get_custom_url( + request_base_url=request_base_url, route="claude_code_gateway/oauth/token" + ), + "grant_types_supported": [_DEVICE_CODE_GRANT, _REFRESH_TOKEN_GRANT], + } + ) + + +@router.post("/oauth/device_authorization", include_in_schema=False) +async def device_authorization(request: Request) -> JSONResponse: + from urllib.parse import urlencode + + from litellm.proxy.management_endpoints.ui_sso import ( + _check_cli_sso_start_rate_limit, + _generate_cli_sso_user_code, + _hash_cli_sso_secret, + _normalize_cli_sso_user_code, + _set_cli_sso_flow, + ) + from litellm.proxy.proxy_server import cli_sso_session_cache, general_settings + 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")) + + _check_cli_sso_start_rate_limit( + request=request, + cache=cli_sso_session_cache, + use_x_forwarded_for=bool((general_settings or {}).get("use_x_forwarded_for", False)), + ) + + device_code = f"cli-{secrets.token_urlsafe(24)}" + user_code = _generate_cli_sso_user_code() + flow = { + "poll_secret_hash": _hash_cli_sso_secret(device_code), + "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) + + request_base_url = str(request.base_url) + verification_uri = get_custom_url(request_base_url=request_base_url, route="sso/key/generate") + verification_uri_complete = ( + verification_uri + + "?" + + urlencode({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": device_code, "user_code": user_code}) + ) + verification_uri_no_code = ( + verification_uri + "?" + urlencode({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": device_code}) + ) + return JSONResponse( + content={ + "device_code": device_code, + "user_code": user_code, + "verification_uri": verification_uri_no_code, + "verification_uri_complete": verification_uri_complete, + "expires_in": CLI_SSO_SESSION_TTL_SECONDS, + "interval": _DEVICE_POLL_INTERVAL_SECONDS, + } + ) + + +def _mint_access_token_from_flow(flow: dict[str, Any]) -> str: + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken + + session_data = flow.get("session_data") + if not isinstance(session_data, dict): + raise _oauth_error(status_code=400, error="authorization_pending") + + teams = session_data.get("teams") or [] + team_id = teams[0] if isinstance(teams, list) and teams else None + user_info = LiteLLM_UserTable( + user_id=session_data["user_id"], + user_role=session_data["user_role"], + models=session_data.get("models", []), + ) + return ExperimentalUIJWTToken.get_cli_jwt_auth_token(user_info=user_info, team_id=team_id) + + +async def _handle_device_code_grant(device_code: str | None) -> JSONResponse: + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.ui_sso import ( + _get_cli_sso_flow_cache_key, + _get_cli_sso_flow_or_raise, + ) + 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") + ) + + try: + flow = _get_cli_sso_flow_or_raise(login_id=device_code, cache=cli_sso_session_cache) + except HTTPException: + return _oauth_error_response(_oauth_error(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")) + + try: + access_token = _mint_access_token_from_flow(flow) + except _OAuthError as err: + return _oauth_error_response(err) + + cli_sso_session_cache.delete_cache(key=_get_cli_sso_flow_cache_key(device_code)) + return JSONResponse( + content={ + "access_token": access_token, + "token_type": "Bearer", + "expires_in": CLI_JWT_EXPIRATION_HOURS * 3600, + } + ) + + +@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")) + + form = await request.form() + grant_type = form.get("grant_type") + + if grant_type == _DEVICE_CODE_GRANT: + device_code = form.get("device_code") + return await _handle_device_code_grant(device_code if isinstance(device_code, str) else None) + + if grant_type == _REFRESH_TOKEN_GRANT: + return _oauth_error_response( + _oauth_error( + status_code=401, + error="invalid_grant", + description="This gateway does not issue refresh tokens; sign in again", + ) + ) + + return _oauth_error_response( + _oauth_error(status_code=400, error="unsupported_grant_type", description=f"Unsupported grant_type: {grant_type}") + ) + + +@router.get("/managed/settings", include_in_schema=False, dependencies=[Depends(user_api_key_auth)]) +async def managed_settings(request: Request) -> Response: + ensure_gateway_enabled() + + settings = _managed_settings() + if settings is None: + return Response(status_code=404) + + body = json.dumps(settings, sort_keys=True, separators=(",", ":")) + etag = '"' + hashlib.sha256(body.encode("utf-8")).hexdigest() + '"' + if_none_match = request.headers.get("If-None-Match") + if if_none_match is not None and if_none_match == etag: + return Response(status_code=304, headers={"ETag": etag}) + return Response(content=body, media_type="application/json", headers={"ETag": etag}) + + +async def _accept_otlp(request: Request) -> Response: + ensure_gateway_enabled() + await request.body() + return Response(status_code=200) + + +@router.post("/v1/metrics", include_in_schema=False, dependencies=[Depends(user_api_key_auth)]) +async def otlp_metrics(request: Request) -> Response: + return await _accept_otlp(request) + + +@router.post("/v1/logs", include_in_schema=False, dependencies=[Depends(user_api_key_auth)]) +async def otlp_logs(request: Request) -> Response: + return await _accept_otlp(request) + + +@router.post("/v1/traces", include_in_schema=False, dependencies=[Depends(user_api_key_auth)]) +async def otlp_traces(request: Request) -> Response: + return await _accept_otlp(request) diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py new file mode 100644 index 00000000000..8645f4a8680 --- /dev/null +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py @@ -0,0 +1,219 @@ +""" +Tests for the Claude Code gateway protocol (anthropic_endpoints/gateway_endpoints.py). + +Covers the OAuth device-flow surface (RFC 8414 discovery, RFC 8628 device +authorization + token), managed settings, OTLP ingestion, and the enable flag. +""" + +from contextlib import contextmanager +from typing import Any, Iterator, Optional +from unittest.mock import patch + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from litellm.caching.dual_cache import DualCache +from litellm.proxy.anthropic_endpoints import gateway_endpoints +from litellm.proxy.management_endpoints.ui_sso import _get_cli_sso_flow_cache_key + + +@contextmanager +def _gateway_env( + *, + enabled: bool = True, + managed_settings: Optional[dict[str, Any]] = None, +) -> Iterator[tuple[TestClient, DualCache]]: + general_settings: dict[str, Any] = {"enable_claude_code_gateway": enabled} + if managed_settings is not None: + general_settings["claude_code_gateway_managed_settings"] = managed_settings + cache = DualCache(default_in_memory_ttl=600) + + app = FastAPI() + app.include_router(gateway_endpoints.router) + + async def _fake_auth() -> Any: + return object() + + app.dependency_overrides[gateway_endpoints.user_api_key_auth] = _fake_auth + + with patch("litellm.proxy.proxy_server.general_settings", general_settings), patch( + "litellm.proxy.proxy_server.cli_sso_session_cache", cache + ): + with TestClient(app) as client: + yield client, cache + + +def _complete_flow(cache: DualCache, device_code: str) -> None: + key = _get_cli_sso_flow_cache_key(device_code) + flow = cache.get_cache(key=key) + assert isinstance(flow, dict) + flow["sso_complete"] = True + flow["user_code_verified"] = True + flow["session_data"] = { + "user_id": "user-123", + "user_role": "internal_user", + "models": ["claude-sonnet-4-5"], + "teams": ["team-a"], + } + cache.set_cache(key=key, value=flow, ttl=600) + + +def test_discovery_shape(): + with _gateway_env() as (client, _): + resp = client.get("/claude_code_gateway/.well-known/oauth-authorization-server") + assert resp.status_code == 200 + body = resp.json() + assert body["device_authorization_endpoint"].endswith("/claude_code_gateway/oauth/device_authorization") + assert body["token_endpoint"].endswith("/claude_code_gateway/oauth/token") + assert body["grant_types_supported"] == [ + "urn:ietf:params:oauth:grant-type:device_code", + "refresh_token", + ] + # authorization_endpoint is intentionally absent (device flow only). + assert "authorization_endpoint" not in body + # Both endpoints must be same-origin with the issuer. + assert body["device_authorization_endpoint"].startswith(body["issuer"]) + assert body["token_endpoint"].startswith(body["issuer"]) + + +def test_discovery_404_when_disabled(): + with _gateway_env(enabled=False) as (client, _): + resp = client.get("/claude_code_gateway/.well-known/oauth-authorization-server") + assert resp.status_code == 404 + + +def test_device_authorization_returns_rfc8628_shape_and_persists_flow(): + with _gateway_env() as (client, cache): + resp = client.post("/claude_code_gateway/oauth/device_authorization") + assert resp.status_code == 200 + body = resp.json() + device_code = body["device_code"] + assert device_code.startswith("cli-") + 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 isinstance(stored, dict) + assert stored["sso_complete"] is False + + +def test_token_authorization_pending_before_browser_completes(): + with _gateway_env() as (client, _): + device_code = client.post("/claude_code_gateway/oauth/device_authorization").json()["device_code"] + resp = client.post( + "/claude_code_gateway/oauth/token", + data={"grant_type": "urn:ietf:params:oauth:grant-type:device_code", "device_code": device_code}, + ) + assert resp.status_code == 400 + assert resp.json()["error"] == "authorization_pending" + + +def test_token_success_mints_bearer_and_is_single_use(): + with _gateway_env() as (client, cache): + device_code = client.post("/claude_code_gateway/oauth/device_authorization").json()["device_code"] + _complete_flow(cache, device_code) + + with patch( + "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", + return_value="sk-litellm-session-token", + ) as mint: + resp = client.post( + "/claude_code_gateway/oauth/token", + data={"grant_type": "urn:ietf:params:oauth:grant-type:device_code", "device_code": device_code}, + ) + assert resp.status_code == 200 + body = resp.json() + assert body["access_token"] == "sk-litellm-session-token" + assert body["token_type"] == "Bearer" + assert body["expires_in"] > 0 + + called_user = mint.call_args.kwargs["user_info"] + assert called_user.user_id == "user-123" + assert mint.call_args.kwargs["team_id"] == "team-a" + + # Single-use: the flow is deleted, so a replay returns expired_token. + replay = client.post( + "/claude_code_gateway/oauth/token", + data={"grant_type": "urn:ietf:params:oauth:grant-type:device_code", "device_code": device_code}, + ) + assert replay.status_code == 400 + assert replay.json()["error"] == "expired_token" + + +def test_token_unknown_device_code_is_expired_token(): + with _gateway_env() as (client, _): + resp = client.post( + "/claude_code_gateway/oauth/token", + data={"grant_type": "urn:ietf:params:oauth:grant-type:device_code", "device_code": "cli-does-not-exist"}, + ) + assert resp.status_code == 400 + assert resp.json()["error"] == "expired_token" + + +def test_refresh_grant_forces_relogin(): + with _gateway_env() as (client, _): + resp = client.post( + "/claude_code_gateway/oauth/token", + data={"grant_type": "refresh_token", "refresh_token": "whatever"}, + ) + assert resp.status_code == 401 + assert resp.json()["error"] == "invalid_grant" + + +def test_unsupported_grant_type(): + with _gateway_env() as (client, _): + resp = client.post("/claude_code_gateway/oauth/token", data={"grant_type": "password"}) + assert resp.status_code == 400 + assert resp.json()["error"] == "unsupported_grant_type" + + +def test_managed_settings_404_when_unset(): + with _gateway_env() as (client, _): + resp = client.get("/claude_code_gateway/managed/settings") + assert resp.status_code == 404 + + +def test_managed_settings_returns_json_with_etag_and_304(): + settings = {"permissions": {"defaultMode": "acceptEdits"}, "env": {"FOO": "bar"}} + with _gateway_env(managed_settings=settings) as (client, _): + resp = client.get("/claude_code_gateway/managed/settings") + assert resp.status_code == 200 + assert resp.json() == settings + etag = resp.headers["ETag"] + assert etag + + not_modified = client.get("/claude_code_gateway/managed/settings", headers={"If-None-Match": etag}) + assert not_modified.status_code == 304 + assert not_modified.headers["ETag"] == etag + + +def test_managed_settings_404_when_gateway_disabled(): + with _gateway_env(enabled=False, managed_settings={"env": {}}) as (client, _): + resp = client.get("/claude_code_gateway/managed/settings") + assert resp.status_code == 404 + + +@pytest.mark.parametrize("signal", ["metrics", "logs", "traces"]) +def test_otlp_endpoints_accept_and_return_200(signal: str): + with _gateway_env() as (client, _): + resp = client.post(f"/claude_code_gateway/v1/{signal}", content=b"\x00\x01binary-otlp") + assert resp.status_code == 200 + + +@pytest.mark.parametrize("signal", ["metrics", "logs", "traces"]) +def test_otlp_endpoints_404_when_disabled(signal: str): + with _gateway_env(enabled=False) as (client, _): + resp = client.post(f"/claude_code_gateway/v1/{signal}", content=b"payload") + assert resp.status_code == 404 + + +def test_messages_gated_by_enable_flag(): + with _gateway_env(enabled=False) as (client, _): + resp = client.post("/claude_code_gateway/v1/messages", json={"model": "claude-sonnet-4-5", "messages": []}) + assert resp.status_code == 404