mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
Merge pull request #34267 from BerriAI/litellm_claude_code_gateway_protocol
feat(proxy): serve the Claude Code gateway protocol under /claude_code_gateway
This commit is contained in:
Mateo Wang
GitHub
commit
c1de8665ff
10 changed files with 1023 additions and 1 deletions
|
|
@ -233,6 +233,11 @@ LAZY_FEATURES: Final[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",
|
||||
|
|
|
|||
|
|
@ -5235,6 +5235,12 @@
|
|||
}
|
||||
}
|
||||
},
|
||||
"claude_code_gateway": {
|
||||
"components": {
|
||||
"schemas": {}
|
||||
},
|
||||
"paths": {}
|
||||
},
|
||||
"claude_code_marketplace": {
|
||||
"components": {
|
||||
"schemas": {
|
||||
|
|
|
|||
|
|
@ -512,6 +512,8 @@ class LiteLLMRoutes(enum.Enum):
|
|||
anthropic_routes = [
|
||||
"/v1/messages",
|
||||
"/v1/messages/count_tokens",
|
||||
"/claude_code_gateway/v1/messages",
|
||||
"/claude_code_gateway/v1/messages/count_tokens",
|
||||
"/v1/skills",
|
||||
"/v1/skills/{skill_id}",
|
||||
"/claude-code/marketplace.json",
|
||||
|
|
@ -889,6 +891,11 @@ class LiteLLMRoutes(enum.Enum):
|
|||
# of; a caller who administers none gets an empty result set.
|
||||
"/organization/daily/activity",
|
||||
"/user/available_roles", # read-only role metadata; any authenticated user may read
|
||||
# Claude Code gateway: the signed-in CLI fetches its managed settings and posts its own telemetry
|
||||
"/claude_code_gateway/managed/settings",
|
||||
"/claude_code_gateway/v1/metrics",
|
||||
"/claude_code_gateway/v1/logs",
|
||||
"/claude_code_gateway/v1/traces",
|
||||
"/user/list", # org admins checked in endpoint; non-admins get 403
|
||||
"/management/v1/users/bulk_delete", # proxy admins delete anyone, org admins only their orgs' users; others 403
|
||||
"/model/{model_id}/update",
|
||||
|
|
@ -2604,6 +2611,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: str | None = Field(
|
||||
None,
|
||||
description="connect to a postgres db - needed for generating temporary keys + tracking spend / key",
|
||||
|
|
|
|||
397
litellm/proxy/anthropic_endpoints/gateway_endpoints.py
Normal file
397
litellm/proxy/anthropic_endpoints/gateway_endpoints.py
Normal file
|
|
@ -0,0 +1,397 @@
|
|||
"""
|
||||
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://<proxy-host>/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 collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from fastapi import APIRouter, Depends, Request, Response
|
||||
from fastapi.responses import JSONResponse
|
||||
from pydantic import BaseModel, Field, TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.constants import (
|
||||
CLI_JWT_EXPIRATION_HOURS,
|
||||
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])
|
||||
_NO_SETTINGS: Final = MappingProxyType({})
|
||||
_POST_ONLY: Final = ["POST"] # mutable-ok: FastAPI's add_api_route only accepts a list of methods
|
||||
|
||||
|
||||
class _GatewaySessionData(BaseModel):
|
||||
user_id: str
|
||||
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
|
||||
|
||||
|
||||
class _AuthorizationServerMetadata(BaseModel):
|
||||
issuer: str
|
||||
device_authorization_endpoint: str
|
||||
token_endpoint: str
|
||||
grant_types_supported: tuple[str, ...]
|
||||
|
||||
|
||||
class _DeviceAuthorizationBody(BaseModel):
|
||||
device_code: str
|
||||
user_code: str
|
||||
verification_uri: str
|
||||
verification_uri_complete: str | None = None
|
||||
expires_in: int
|
||||
interval: int
|
||||
|
||||
|
||||
class _AccessTokenBody(BaseModel):
|
||||
access_token: str
|
||||
expires_in: int
|
||||
token_type: str = "Bearer"
|
||||
|
||||
|
||||
class _ManagedSettingsBody(BaseModel):
|
||||
uuid: str
|
||||
checksum: str
|
||||
settings: dict[str, object]
|
||||
|
||||
|
||||
def _general_settings() -> Mapping[str, object]:
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
|
||||
return general_settings or _NO_SETTINGS
|
||||
|
||||
|
||||
def _is_gateway_enabled() -> bool:
|
||||
return bool(_general_settings().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, object] | None:
|
||||
settings: Final[object] = _general_settings().get("claude_code_gateway_managed_settings")
|
||||
if not isinstance(settings, dict):
|
||||
return None
|
||||
return _MANAGED_SETTINGS_ADAPTER.validate_python(settings)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _OAuthError:
|
||||
status_code: int
|
||||
error: str
|
||||
description: str | None = None
|
||||
|
||||
|
||||
def _oauth_error_response(err: _OAuthError) -> JSONResponse:
|
||||
body: Final = _OAuthErrorBody(error=err.error, error_description=err.description)
|
||||
return JSONResponse(status_code=err.status_code, content=body.model_dump(exclude_none=True))
|
||||
|
||||
|
||||
router: Final = APIRouter(
|
||||
prefix=GATEWAY_PREFIX,
|
||||
tags=["Claude Code gateway"], # mutable-ok: FastAPI's APIRouter only accepts a list of tags
|
||||
)
|
||||
_GATEWAY_ENABLED: Final = (Depends(ensure_gateway_enabled),)
|
||||
_AUTHENTICATED: Final = (Depends(user_api_key_auth),)
|
||||
|
||||
router.add_api_route(
|
||||
"/v1/messages",
|
||||
anthropic_response,
|
||||
methods=_POST_ONLY,
|
||||
dependencies=_GATEWAY_ENABLED,
|
||||
include_in_schema=False,
|
||||
)
|
||||
router.add_api_route(
|
||||
"/v1/messages/count_tokens",
|
||||
count_tokens,
|
||||
methods=_POST_ONLY,
|
||||
dependencies=_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(_OAuthError(status_code=404, error="not_found"))
|
||||
|
||||
from litellm.proxy.utils import get_custom_url
|
||||
|
||||
request_base_url: Final = str(request.base_url)
|
||||
metadata: Final = _AuthorizationServerMetadata(
|
||||
issuer=get_custom_url(request_base_url=request_base_url, route="claude_code_gateway"),
|
||||
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),
|
||||
)
|
||||
return JSONResponse(content=metadata.model_dump())
|
||||
|
||||
|
||||
@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, # 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
|
||||
_set_cli_sso_flow, # pyright: ignore[reportPrivateUsage] # shared device-flow helper
|
||||
)
|
||||
from litellm.proxy.proxy_server import cli_sso_session_cache
|
||||
from litellm.proxy.utils import get_custom_url
|
||||
|
||||
if not _is_gateway_enabled():
|
||||
return _oauth_error_response(_OAuthError(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().get("use_x_forwarded_for", False)),
|
||||
)
|
||||
|
||||
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(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=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": login_id})
|
||||
body: Final = _DeviceAuthorizationBody(
|
||||
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(exclude_none=True))
|
||||
|
||||
|
||||
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)
|
||||
return _OAuthError(
|
||||
status_code=400, error="invalid_grant", description="The login session is malformed; sign in again"
|
||||
)
|
||||
|
||||
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:
|
||||
return _OAuthError(
|
||||
status_code=400,
|
||||
error="invalid_grant",
|
||||
description=f"Could not resolve the model grants for team {team_id}; sign in again",
|
||||
)
|
||||
|
||||
user_info: Final = LiteLLM_UserTable(
|
||||
user_id=session_data.user_id,
|
||||
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=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(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(login_id)}:claimed",
|
||||
value=1,
|
||||
ttl=CLI_SSO_SESSION_TTL_SECONDS,
|
||||
)
|
||||
return claims == 1
|
||||
|
||||
|
||||
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, # 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(
|
||||
_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=login_id, cache=cli_sso_session_cache)
|
||||
except HTTPException:
|
||||
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(_OAuthError(status_code=400, error="authorization_pending"))
|
||||
|
||||
login: Final = _validate_login(flow)
|
||||
if isinstance(login, _OAuthError):
|
||||
return _oauth_error_response(login)
|
||||
|
||||
access_token: Final = _mint_access_token(login)
|
||||
if not await _claim_device_code(login_id, cli_sso_session_cache):
|
||||
return _oauth_error_response(_OAuthError(status_code=400, error="expired_token"))
|
||||
|
||||
await cli_sso_session_cache.async_delete_cache(key=_get_cli_sso_flow_cache_key(login_id))
|
||||
body: Final = _AccessTokenBody(access_token=access_token, 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(_OAuthError(status_code=404, error="not_found"))
|
||||
|
||||
form: Final = await request.form()
|
||||
grant_type: Final = form.get("grant_type")
|
||||
|
||||
if grant_type == _DEVICE_CODE_GRANT:
|
||||
device_code: Final = 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(
|
||||
_OAuthError(
|
||||
status_code=401,
|
||||
error="invalid_grant",
|
||||
description="This gateway does not issue refresh tokens; sign in again",
|
||||
)
|
||||
)
|
||||
|
||||
return _oauth_error_response(
|
||||
_OAuthError(
|
||||
status_code=400, error="unsupported_grant_type", description=f"Unsupported grant_type: {grant_type}"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@router.get("/managed/settings", include_in_schema=False, dependencies=_AUTHENTICATED)
|
||||
async def managed_settings(request: Request) -> Response:
|
||||
ensure_gateway_enabled()
|
||||
|
||||
settings: Final = _managed_settings()
|
||||
if settings is None:
|
||||
return Response(status_code=404)
|
||||
|
||||
canonical: Final = json.dumps(settings, sort_keys=True, separators=(",", ":"))
|
||||
checksum: Final = "sha256:" + hashlib.sha256(canonical.encode("utf-8")).hexdigest()
|
||||
etag: Final = f'"{checksum}"'
|
||||
headers: Final = MappingProxyType({"ETag": etag})
|
||||
if request.headers.get("If-None-Match") == etag:
|
||||
return Response(status_code=304, headers=headers)
|
||||
body: Final = _ManagedSettingsBody(uuid=checksum, checksum=checksum, settings=settings)
|
||||
return Response(content=body.model_dump_json(), media_type="application/json", headers=headers)
|
||||
|
||||
|
||||
async def _skip_otlp_body_parsing(request: Request) -> None:
|
||||
_safe_set_request_parsed_body(request=request, parsed_body={})
|
||||
|
||||
|
||||
_OTLP_AUTHENTICATED: Final = (Depends(_skip_otlp_body_parsing), *_AUTHENTICATED)
|
||||
|
||||
|
||||
def _accept_otlp() -> Response:
|
||||
ensure_gateway_enabled()
|
||||
return Response(status_code=200)
|
||||
|
||||
|
||||
@router.post("/v1/metrics", include_in_schema=False, dependencies=_OTLP_AUTHENTICATED)
|
||||
async def otlp_metrics() -> Response:
|
||||
return _accept_otlp()
|
||||
|
||||
|
||||
@router.post("/v1/logs", include_in_schema=False, dependencies=_OTLP_AUTHENTICATED)
|
||||
async def otlp_logs() -> Response:
|
||||
return _accept_otlp()
|
||||
|
||||
|
||||
@router.post("/v1/traces", include_in_schema=False, dependencies=_OTLP_AUTHENTICATED)
|
||||
async def otlp_traces() -> Response:
|
||||
return _accept_otlp()
|
||||
|
|
@ -7,6 +7,7 @@ from collections.abc import MutableMapping
|
|||
from typing import Any, Final
|
||||
|
||||
from fastapi import Request
|
||||
from starlette.routing import get_route_path
|
||||
from starlette.types import ASGIApp, Receive, Scope, Send
|
||||
|
||||
import litellm
|
||||
|
|
@ -15,6 +16,12 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
|||
|
||||
# Cache the header name at module level to avoid repeated enum attribute access
|
||||
_AUTHORIZATION_HEADER: Final = SpecialHeaders.openai_authorization.value # "Authorization"
|
||||
_METRICS_MOUNT: Final = "/metrics"
|
||||
|
||||
|
||||
def _is_metrics_route(scope: Scope) -> bool:
|
||||
route_path: Final = get_route_path(scope)
|
||||
return route_path == _METRICS_MOUNT or route_path.startswith(_METRICS_MOUNT + "/")
|
||||
|
||||
|
||||
class PrometheusAuthMiddleware:
|
||||
|
|
@ -36,7 +43,7 @@ class PrometheusAuthMiddleware:
|
|||
|
||||
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
# Fast path: only inspect HTTP requests; pass through websocket/lifespan immediately
|
||||
if scope["type"] != "http" or "/metrics" not in scope.get("path", ""):
|
||||
if scope["type"] != "http" or not _is_metrics_route(scope):
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,475 @@
|
|||
"""
|
||||
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.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Iterator, Mapping
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
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,
|
||||
_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(
|
||||
{
|
||||
"user_id": "user-123",
|
||||
"user_role": "internal_user",
|
||||
"models": ["claude-sonnet-4-5"],
|
||||
"teams": ["team-a"],
|
||||
"team_details": [
|
||||
{
|
||||
"team_id": "team-a",
|
||||
"team_alias": "Team A",
|
||||
"team_models": ["claude-sonnet-4-5"],
|
||||
"team_model_aliases": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class _SharedRedisFake:
|
||||
def __init__(self) -> None:
|
||||
self.values: Mapping[str, object] = MappingProxyType({})
|
||||
self.counters: Mapping[str, float] = MappingProxyType({})
|
||||
|
||||
def set_cache(self, key: str, value: object, **kwargs: object) -> None:
|
||||
self.values = MappingProxyType({**self.values, key: value})
|
||||
|
||||
def get_cache(self, key: str, **kwargs: object) -> object:
|
||||
return self.values.get(key)
|
||||
|
||||
def delete_cache(self, key: str) -> None:
|
||||
self.values = MappingProxyType({name: value for name, value in self.values.items() if name != key})
|
||||
|
||||
async def async_delete_cache(self, key: str) -> None:
|
||||
self.delete_cache(key)
|
||||
|
||||
async def async_increment(self, key: str, value: float, **kwargs: object) -> float:
|
||||
incremented: Final = self.counters.get(key, 0) + value
|
||||
self.counters = MappingProxyType({**self.counters, key: incremented})
|
||||
return incremented
|
||||
|
||||
|
||||
def _replica(redis: _SharedRedisFake) -> DualCache:
|
||||
return DualCache(redis_cache=redis, default_in_memory_ttl=600) # pyright: ignore[reportArgumentType] # duck-typed Redis double
|
||||
|
||||
|
||||
def _real_auth_proxy_attrs() -> Mapping[str, object]:
|
||||
proxy_logging_obj: Final = MagicMock()
|
||||
proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock()
|
||||
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
|
||||
return MappingProxyType(
|
||||
{
|
||||
"master_key": _MASTER_KEY,
|
||||
"prisma_client": None,
|
||||
"user_api_key_cache": DualCache(),
|
||||
"proxy_logging_obj": proxy_logging_obj,
|
||||
"llm_router": None,
|
||||
"llm_model_list": [],
|
||||
"user_custom_auth": None,
|
||||
"litellm_proxy_admin_name": "admin",
|
||||
"jwt_handler": None,
|
||||
"open_telemetry_logger": None,
|
||||
"model_max_budget_limiter": MagicMock(),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _gateway_env(
|
||||
*,
|
||||
enabled: bool = True,
|
||||
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)
|
||||
|
||||
app: Final = FastAPI()
|
||||
app.add_middleware(PrometheusAuthMiddleware)
|
||||
app.include_router(gateway_endpoints.router)
|
||||
|
||||
async def _fake_auth() -> object:
|
||||
return object()
|
||||
|
||||
with ExitStack() as stack:
|
||||
stack.enter_context(
|
||||
patch( # test-quality-ok: the gateway reads this proxy_server module global and has no injection seam
|
||||
"litellm.proxy.proxy_server.general_settings", general_settings
|
||||
)
|
||||
)
|
||||
stack.enter_context(
|
||||
patch( # test-quality-ok: the CLI SSO flow cache is this proxy_server module global shared with ui_sso
|
||||
"litellm.proxy.proxy_server.cli_sso_session_cache", session_cache
|
||||
)
|
||||
)
|
||||
if real_auth:
|
||||
for name, value in _real_auth_proxy_attrs().items():
|
||||
stack.enter_context(patch(f"litellm.proxy.proxy_server.{name}", value))
|
||||
else:
|
||||
app.dependency_overrides[gateway_endpoints.user_api_key_auth] = _fake_auth
|
||||
with TestClient(app) as client:
|
||||
yield client, session_cache
|
||||
|
||||
|
||||
def _start_device_flow(client: TestClient) -> str:
|
||||
return client.post("/claude_code_gateway/oauth/device_authorization").json()["device_code"]
|
||||
|
||||
|
||||
def _request_token(client: TestClient, device_code: str) -> httpx.Response:
|
||||
return client.post(
|
||||
"/claude_code_gateway/oauth/token",
|
||||
data={"grant_type": _DEVICE_CODE_GRANT, "device_code": device_code},
|
||||
)
|
||||
|
||||
|
||||
def _completed_flow(session_data: Mapping[str, object] = _COMPLETED_SESSION) -> dict[str, object]:
|
||||
return {
|
||||
"poll_secret_hash": _hash_cli_sso_secret(_SHARED_POLL_SECRET),
|
||||
"user_code_hash": "unused",
|
||||
"sso_complete": True,
|
||||
"user_code_verified": True,
|
||||
"session_data": dict(session_data),
|
||||
}
|
||||
|
||||
|
||||
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(_login_id(device_code))
|
||||
flow: Final = cache.get_cache(key=key)
|
||||
assert isinstance(flow, dict)
|
||||
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():
|
||||
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"]
|
||||
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
|
||||
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():
|
||||
with _gateway_env() as (client, _):
|
||||
resp = _request_token(client, _start_device_flow(client))
|
||||
assert resp.status_code == 400
|
||||
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)
|
||||
_complete_flow(cache, device_code)
|
||||
|
||||
with patch(_MINT, return_value="sk-litellm-session-token") as mint:
|
||||
resp = _request_token(client, 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"
|
||||
assert mint.call_args.kwargs["team_alias"] == "Team A"
|
||||
assert mint.call_args.kwargs["team_models"] == ("claude-sonnet-4-5",)
|
||||
|
||||
# Single-use: the flow is deleted, so a replay returns expired_token.
|
||||
replay = _request_token(client, device_code)
|
||||
assert replay.status_code == 400
|
||||
assert replay.json()["error"] == "expired_token"
|
||||
|
||||
|
||||
def test_token_teamless_user_mints_without_a_team():
|
||||
with _gateway_env() as (client, cache):
|
||||
device_code = _start_device_flow(client)
|
||||
_complete_flow(cache, device_code, session_data={**_COMPLETED_SESSION, "teams": [], "team_details": []})
|
||||
with patch(_MINT, return_value="sk-litellm-session-token") as mint:
|
||||
resp = _request_token(client, device_code)
|
||||
assert resp.status_code == 200
|
||||
assert mint.call_args.kwargs["team_id"] is None
|
||||
assert mint.call_args.kwargs["team_models"] == ()
|
||||
|
||||
|
||||
@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=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()
|
||||
|
||||
|
||||
def test_token_mint_failure_leaves_the_login_unconsumed():
|
||||
with _gateway_env() as (client, cache):
|
||||
device_code = _start_device_flow(client)
|
||||
_complete_flow(cache, device_code)
|
||||
with patch(_MINT, side_effect=RuntimeError("signing key unavailable")), pytest.raises(RuntimeError):
|
||||
_request_token(client, device_code)
|
||||
with patch(_MINT, return_value="sk-session"):
|
||||
retry = _request_token(client, device_code)
|
||||
assert retry.status_code == 200
|
||||
assert retry.json()["access_token"] == "sk-session"
|
||||
|
||||
|
||||
def test_token_unknown_team_grants_is_invalid_grant():
|
||||
with _gateway_env() as (client, cache):
|
||||
device_code = _start_device_flow(client)
|
||||
_complete_flow(cache, device_code, session_data={**_COMPLETED_SESSION, "team_details": []})
|
||||
with patch(_MINT) as mint:
|
||||
resp = _request_token(client, device_code)
|
||||
assert resp.status_code == 400
|
||||
assert resp.json()["error"] == "invalid_grant"
|
||||
mint.assert_not_called()
|
||||
|
||||
|
||||
def test_token_mints_on_a_replica_that_did_not_start_the_login():
|
||||
redis: Final = _SharedRedisFake()
|
||||
_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, _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)
|
||||
_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, return_value="sk-session"):
|
||||
resp = _request_token(client, _SHARED_DEVICE_CODE)
|
||||
assert resp.status_code == 400
|
||||
assert resp.json() == {"error": "expired_token"}
|
||||
|
||||
|
||||
def test_token_unknown_device_code_is_expired_token():
|
||||
with _gateway_env() as (client, _):
|
||||
resp = _request_token(client, "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_client_envelope_and_304_on_cached_checksum():
|
||||
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
|
||||
body = resp.json()
|
||||
assert body["settings"] == settings
|
||||
checksum = body["checksum"]
|
||||
assert checksum.startswith("sha256:")
|
||||
assert body["uuid"] == checksum
|
||||
assert resp.headers["ETag"] == f'"{checksum}"'
|
||||
|
||||
not_modified = client.get(
|
||||
"/claude_code_gateway/managed/settings", headers={"If-None-Match": f'"{checksum}"'}
|
||||
)
|
||||
assert not_modified.status_code == 304
|
||||
assert not_modified.headers["ETag"] == f'"{checksum}"'
|
||||
|
||||
stale = client.get("/claude_code_gateway/managed/settings", headers={"If-None-Match": '"sha256:stale"'})
|
||||
assert stale.status_code == 200
|
||||
assert stale.json()["checksum"] == checksum
|
||||
|
||||
|
||||
def test_managed_settings_checksum_tracks_policy_content():
|
||||
with _gateway_env(managed_settings={"env": {"FOO": "bar"}}) as (client, _):
|
||||
first = client.get("/claude_code_gateway/managed/settings").json()["checksum"]
|
||||
with _gateway_env(managed_settings={"env": {"FOO": "baz"}}) as (client, _):
|
||||
second = client.get("/claude_code_gateway/managed/settings").json()["checksum"]
|
||||
assert first != second
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
@pytest.mark.parametrize("signal", ["metrics", "logs", "traces"])
|
||||
def test_otlp_protobuf_body_is_accepted_through_real_auth(signal: str):
|
||||
with _gateway_env(real_auth=True) as (client, _):
|
||||
resp = client.post(
|
||||
f"/claude_code_gateway/v1/{signal}",
|
||||
content=_PROTOBUF_BODY,
|
||||
headers={"Authorization": f"Bearer {_MASTER_KEY}", "Content-Type": "application/x-protobuf"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
|
||||
|
||||
def test_otlp_without_a_bearer_is_rejected_by_real_auth():
|
||||
with _gateway_env(real_auth=True) as (client, _), pytest.raises(ProxyException) as exc_info:
|
||||
client.post(
|
||||
"/claude_code_gateway/v1/metrics",
|
||||
content=_PROTOBUF_BODY,
|
||||
headers={"Content-Type": "application/x-protobuf"},
|
||||
)
|
||||
assert exc_info.value.code == "401"
|
||||
|
||||
|
||||
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
|
||||
|
|
@ -910,6 +910,36 @@ def test_anthropic_count_tokens_route_accessible_to_internal_users():
|
|||
assert RouteChecks.is_llm_api_route("/v1/messages") is True
|
||||
|
||||
|
||||
_CLAUDE_CODE_GATEWAY_ROUTES: Final = (
|
||||
"/claude_code_gateway/v1/messages",
|
||||
"/claude_code_gateway/v1/messages/count_tokens",
|
||||
"/claude_code_gateway/managed/settings",
|
||||
"/claude_code_gateway/v1/metrics",
|
||||
"/claude_code_gateway/v1/logs",
|
||||
"/claude_code_gateway/v1/traces",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", _CLAUDE_CODE_GATEWAY_ROUTES)
|
||||
@pytest.mark.parametrize(
|
||||
"role", [LitellmUserRoles.INTERNAL_USER.value, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value]
|
||||
)
|
||||
def test_claude_code_gateway_routes_open_to_signed_in_cli_users(role: str, route: str):
|
||||
user_obj: Final = LiteLLM_UserTable(user_id="test_user", user_email="test@example.com", user_role=role)
|
||||
valid_token: Final = UserAPIKeyAuth(user_id="test_user", user_role=role)
|
||||
request: Final = MagicMock(spec=Request)
|
||||
request.query_params = {}
|
||||
|
||||
RouteChecks.non_proxy_admin_allowed_routes_check(
|
||||
user_obj=user_obj,
|
||||
_user_role=role,
|
||||
route=route,
|
||||
request=request,
|
||||
valid_token=valid_token,
|
||||
request_data={},
|
||||
)
|
||||
|
||||
|
||||
def test_virtual_key_llm_api_routes_allows_registered_pass_through_endpoints():
|
||||
"""
|
||||
Virtual keys with llm_api_routes can access auth=true pass-through endpoints only when
|
||||
|
|
|
|||
|
|
@ -573,6 +573,13 @@ async def test_lone_surrogate_escape_is_rejected_with_400(content: bytes):
|
|||
assert parsed["messages"][0]["content"] == "say ok \U0001F600"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("media_type", ["application/x-protobuf", "application/protobuf", "application/octet-stream"])
|
||||
async def test_json_body_under_a_binary_content_type_is_still_parsed(media_type: str):
|
||||
request = _starlette_request(b'{"model": "claude-sonnet-5"}', media_type)
|
||||
assert await _read_request_body(request) == {"model": "claude-sonnet-5"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_form_data():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -51,6 +51,14 @@ def app_with_middleware():
|
|||
async def embeddings():
|
||||
return {"msg": "embeddings OK"}
|
||||
|
||||
@app.post("/claude_code_gateway/v1/metrics")
|
||||
async def gateway_telemetry():
|
||||
return {"msg": "gateway telemetry OK"}
|
||||
|
||||
@app.get("/metrics/detail")
|
||||
async def metrics_detail():
|
||||
return {"msg": "metrics detail OK"}
|
||||
|
||||
return app
|
||||
|
||||
|
||||
|
|
@ -240,3 +248,63 @@ def test_non_metrics_requests_dont_trigger_auth(app_with_middleware, monkeypatch
|
|||
response = client.get("/embeddings")
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == {"msg": "embeddings OK"}
|
||||
|
||||
|
||||
def test_gateway_telemetry_path_is_not_treated_as_the_metrics_endpoint(app_with_middleware, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "require_auth_for_metrics_endpoint", True)
|
||||
|
||||
def should_not_be_called(*args, **kwargs):
|
||||
raise Exception("Auth should not be called for the gateway telemetry route")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth",
|
||||
should_not_be_called,
|
||||
)
|
||||
|
||||
client = TestClient(app_with_middleware)
|
||||
|
||||
response = client.post("/claude_code_gateway/v1/metrics", content=b"\x0a\x05hello")
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == {"msg": "gateway telemetry OK"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", ["/metrics", "/metrics/", "/metrics/detail"])
|
||||
def test_metrics_paths_still_require_auth(app_with_middleware, monkeypatch, path):
|
||||
monkeypatch.setattr(litellm, "require_auth_for_metrics_endpoint", True)
|
||||
|
||||
async def reject(*args, **kwargs):
|
||||
raise Exception("Invalid API key")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth",
|
||||
reject,
|
||||
)
|
||||
|
||||
client = TestClient(app_with_middleware)
|
||||
|
||||
response = client.get(path)
|
||||
assert response.status_code == 401, response.text
|
||||
|
||||
|
||||
def test_metrics_under_a_root_path_still_requires_auth(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "require_auth_for_metrics_endpoint", True)
|
||||
|
||||
async def reject(*args, **kwargs):
|
||||
raise Exception("Invalid API key")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth",
|
||||
reject,
|
||||
)
|
||||
|
||||
app = FastAPI(root_path="/litellm")
|
||||
app.add_middleware(PrometheusAuthMiddleware)
|
||||
|
||||
@app.get("/metrics")
|
||||
async def metrics():
|
||||
return {"msg": "metrics OK"}
|
||||
|
||||
client = TestClient(app, root_path="/litellm")
|
||||
|
||||
response = client.get("/metrics")
|
||||
assert response.status_code == 401, response.text
|
||||
|
|
|
|||
12
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
12
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -26686,6 +26686,13 @@ export interface components {
|
|||
* @description cancel the in-flight upstream LLM request (non-streaming) when the client disconnects, freeing backend capacity (e.g. a vLLM GPU slot); the request is logged as a 499 failure
|
||||
*/
|
||||
cancel_on_disconnect?: boolean | null;
|
||||
/**
|
||||
* Claude Code Gateway Managed Settings
|
||||
* @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)
|
||||
*/
|
||||
claude_code_gateway_managed_settings?: {
|
||||
[key: string]: unknown;
|
||||
} | null;
|
||||
/**
|
||||
* Completion Model
|
||||
* @description proxy level default model for all chat completion calls
|
||||
|
|
@ -26780,6 +26787,11 @@ export interface components {
|
|||
* @description If True, disables ownership enforcement on Responses API ids. Keys may then retrieve, cancel, delete, and chain from any response id, including ids belonging to another user or team and ids this proxy never issued. WARNING: this removes tenant isolation on /v1/responses
|
||||
*/
|
||||
disable_responses_id_security?: boolean | null;
|
||||
/**
|
||||
* Enable Claude Code Gateway
|
||||
* @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
|
||||
*/
|
||||
enable_claude_code_gateway?: boolean | null;
|
||||
/**
|
||||
* Enable Openai Websocket Passthrough
|
||||
* @description Serve the OpenAI pass-through WebSocket route, which relays frames to OpenAI under the proxy's own provider credential without reading them. Off by default.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue