From 59f7a00cf62fc40aa281eac50a57e02cedd3d0f7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:23:59 -0700 Subject: [PATCH] fix(claude_code_gateway): scope the protobuf body skip to the OTLP routes and match the metrics middleware on the route path --- .../anthropic_endpoints/gateway_endpoints.py | 14 +++- .../proxy/common_utils/http_parsing_utils.py | 10 +-- .../middleware/prometheus_auth_middleware.py | 9 ++- .../test_gateway_endpoints.py | 7 +- .../common_utils/test_http_parsing_utils.py | 8 +-- .../test_prometheus_auth_middleware.py | 68 +++++++++++++++++++ 6 files changed, 97 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py index cc3106fce53..08579186f5e 100644 --- a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py @@ -34,6 +34,7 @@ from litellm.constants import ( ) 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 GATEWAY_PREFIX: Final = "/claude_code_gateway" _DEVICE_CODE_GRANT: Final = "urn:ietf:params:oauth:grant-type:device_code" @@ -349,21 +350,28 @@ async def managed_settings(request: Request) -> Response: 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=_AUTHENTICATED) +@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=_AUTHENTICATED) +@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=_AUTHENTICATED) +@router.post("/v1/traces", include_in_schema=False, dependencies=_OTLP_AUTHENTICATED) async def otlp_traces() -> Response: return _accept_otlp() diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 592060e84ee..f5b6a0a766d 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -18,8 +18,6 @@ from litellm.types.router import Deployment _FORM_CONTENT_TYPES: Final[frozenset[str]] = frozenset({"application/x-www-form-urlencoded", "multipart/form-data"}) -_PROTOBUF_CONTENT_TYPES: Final[frozenset[str]] = frozenset({"application/x-protobuf", "application/protobuf"}) - _ANNOTATION_QUALIFIERS: Final[frozenset[object]] = frozenset({Annotated, NotRequired, ReadOnly, Required}) @@ -46,10 +44,6 @@ def is_json_content_type(content_type: str) -> bool: return _normalize_media_type(content_type) == "application/json" -def _is_protobuf_content_type(content_type: str) -> bool: - return _normalize_media_type(content_type) in _PROTOBUF_CONTENT_TYPES - - def _unqualified(annotation: object) -> object: """Which qualifiers ``get_type_hints`` already stripped varies by interpreter version, so peel them all.""" if get_origin(annotation) not in _ANNOTATION_QUALIFIERS: @@ -139,9 +133,7 @@ async def _read_request_body(request: Request | None) -> dict: _request_headers: Final[dict] = _safe_get_request_headers(request=request) content_type: Final = _request_headers.get("content-type", "") - if _is_protobuf_content_type(content_type): - parsed_body = {} - elif _is_form_content_type(content_type): + if _is_form_content_type(content_type): try: form_data: Final = await request.form() except Exception as e: diff --git a/litellm/proxy/middleware/prometheus_auth_middleware.py b/litellm/proxy/middleware/prometheus_auth_middleware.py index 36818a8cfbd..ebdd3e92bb2 100644 --- a/litellm/proxy/middleware/prometheus_auth_middleware.py +++ b/litellm/proxy/middleware/prometheus_auth_middleware.py @@ -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 diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py index 158fe253796..e49047634bc 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py @@ -21,6 +21,7 @@ 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.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware _DEVICE_CODE_GRANT: Final = "urn:ietf:params:oauth:grant-type:device_code" _MASTER_KEY: Final = "sk-master-key" @@ -107,6 +108,7 @@ def _gateway_env( 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: @@ -378,10 +380,11 @@ def test_otlp_endpoints_404_when_disabled(signal: str): assert resp.status_code == 404 -def test_otlp_protobuf_body_is_accepted_through_real_auth(): +@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( - "/claude_code_gateway/v1/metrics", + f"/claude_code_gateway/v1/{signal}", content=_PROTOBUF_BODY, headers={"Authorization": f"Bearer {_MASTER_KEY}", "Content-Type": "application/x-protobuf"}, ) diff --git a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py index bd9912a96ac..7929a0b21af 100644 --- a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py @@ -574,10 +574,10 @@ async def test_lone_surrogate_escape_is_rejected_with_400(content: bytes): @pytest.mark.asyncio -@pytest.mark.parametrize("media_type", ["application/x-protobuf", "application/protobuf"]) -async def test_protobuf_body_is_left_unparsed(media_type: str): - request = _starlette_request(b"\x0a\x05hello\x12\x03{{{", media_type) - assert await _read_request_body(request) == {} +@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 diff --git a/tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py b/tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py index 1d0c0f90fd1..beb841878d5 100644 --- a/tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py +++ b/tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py @@ -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