From 4dad6419a9453150287c2f90ef08547ee888511c Mon Sep 17 00:00:00 2001 From: moe-berri Date: Mon, 26 Oct 2026 20:36:56 -0700 Subject: [PATCH] fix(lens): keep route filters and isolate spend exports --- backend/main.py | 1 + gateway/main.py | 1 + .../litellm_proxy_extras/utils.py | 7 +++- .../clickhouse/clickhouse_spend_logger.py | 22 +++++++++- litellm/proxy/_lazy_features.py | 7 +++- litellm/proxy/lens/adapter.py | 3 +- litellm/proxy/tracing_runtime.py | 4 +- .../test_clickhouse_spend_logger.py | 33 ++++++++++++++- .../test_litellm_proxy_extras_utils.py | 31 ++++++++++++-- tests/unit/proxy/lens/test_adapter.py | 33 ++++++++++++++- .../proxy/proxy_server/test_proxy_config.py | 40 +++++++++++++++++++ tests/unit/proxy/test_component_allowlists.py | 10 ++++- 12 files changed, 176 insertions(+), 16 deletions(-) diff --git a/backend/main.py b/backend/main.py index e0cef90c979..b6ec6cb8c44 100644 --- a/backend/main.py +++ b/backend/main.py @@ -55,6 +55,7 @@ async def _backend_lifespan( app_: Starlette, lifespan: Lifespan[Starlette] = _proxy_lifespan ) -> AsyncGenerator[Mapping[str, object], None]: async with lifespan(app_) as state: + app_.state.component_route_filter = _is_backend_route app_.router.routes = [r for r in app_.router.routes if _is_backend_route(r)] yield state if state is not None else {} diff --git a/gateway/main.py b/gateway/main.py index fb4ae830808..3476c08dc7c 100644 --- a/gateway/main.py +++ b/gateway/main.py @@ -66,6 +66,7 @@ async def _gateway_lifespan( app_: Starlette, lifespan: Lifespan[Starlette] = _proxy_lifespan ) -> AsyncGenerator[Mapping[str, object], None]: async with lifespan(app_) as state: + app_.state.component_route_filter = _is_gateway_route app_.router.routes = [r for r in app_.router.routes if _is_gateway_route(r)] yield state if state is not None else {} diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index be785e2fd78..06a9d9f6665 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -501,7 +501,11 @@ class ProxyExtrasDBManager: @staticmethod def _resolve_all_migrations( - migrations_dir: str, schema_path: str, mark_all_applied: bool = True + migrations_dir: str, + schema_path: str, + mark_all_applied: bool = True, + *, + lens_connect: LensCheckConnect | None = None, ): """ 1. Compare the current database state to schema.prisma and generate a migration for the diff. @@ -512,6 +516,7 @@ class ProxyExtrasDBManager: if not database_url: logger.error("DATABASE_URL not set") return + ProxyExtrasDBManager.raise_if_lens_rename_pending(connect=lens_connect) # Prefer DIRECT_URL for schema introspection — pooler URLs (e.g. neon -pooler) # do not support the extended query protocol required by prisma migrate diff. diff_url = os.getenv("DIRECT_URL") or database_url diff --git a/litellm/integrations/clickhouse/clickhouse_spend_logger.py b/litellm/integrations/clickhouse/clickhouse_spend_logger.py index 1d01965cd2c..4db0e2dc6f6 100644 --- a/litellm/integrations/clickhouse/clickhouse_spend_logger.py +++ b/litellm/integrations/clickhouse/clickhouse_spend_logger.py @@ -22,7 +22,7 @@ from litellm.integrations.clickhouse.schema import SPEND_LOGS_TABLE from litellm.litellm_core_utils.llm_response_utils.get_headers import get_provider_request_id from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.tracing.types import SpendLogPayload, SpendLogRecord -from litellm.types.utils import StandardLoggingPayload +from litellm.types.utils import StandardLoggingPayload, StandardLoggingUserAPIKeyMetadata # litellm_logging.py rewrites cache-hit ids as f"{id}_cache_hit{time.time()}" MILLISECONDS_PER_SECOND: Final = 1000 @@ -37,6 +37,24 @@ _METADATA_VALUE: Final = TypeAdapter(JsonValue) _INTERNAL_METADATA_KEYS: Final = frozenset( ("user_api_key", "user_api_key_auth", "user_api_key_budget_reservation", "proxy_server_request") ) +_CUSTOM_METADATA_KEYS: Final = frozenset( + ("spend_logs_metadata", "requester_metadata", "requester_custom_headers", "user_api_key_auth_metadata") +) +_BILLING_METADATA_KEYS: Final = frozenset(StandardLoggingUserAPIKeyMetadata.__annotations__) - _CUSTOM_METADATA_KEYS + + +def _safe_payload_metadata(metadata: Mapping[str, object], private: bool) -> Mapping[str, object]: + retained: Final = { + key: value for key, value in metadata.items() if not (private and key in _CUSTOM_METADATA_KEYS) + } + return { + **redact_credentials_in_payload(retained), + **{ + key: value + for key, value in retained.items() + if key in _BILLING_METADATA_KEYS and (value is None or isinstance(value, (str, int, float, bool))) + }, + } def _request_metadata_fields(value: object) -> Iterator[tuple[str, JsonValue]]: @@ -198,7 +216,7 @@ def spend_log_row_from_payload( MappingProxyType( { **request_metadata, - **metadata, + **_safe_payload_metadata(metadata, redact), "litellm_lens_internal": is_lens_analysis(), **({} if track_end_user else {"user_api_key_end_user_id": None}), } diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 4e54d6ac733..f15ce81bd3e 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -455,7 +455,12 @@ def _route_table( lazy_routes: Mapping[str, tuple[BaseRoute, ...]], features: tuple[LazyFeature, ...], ) -> list[BaseRoute]: # mutable-ok: assigned to Router.routes, a list - ordered: Final = _in_registry_order(routes, lazy_routes, features, _lazy_slots(app)) + allowed: Final[Callable[[BaseRoute], bool] | None] = getattr(app.state, "component_route_filter", None) + ordered: Final = tuple( + route + for route in _in_registry_order(routes, lazy_routes, features, _lazy_slots(app)) + if allowed is None or allowed(route) + ) if builtin_pass_through_routes_first(): return hot_routes_first(ordered) return hot_routes_first( diff --git a/litellm/proxy/lens/adapter.py b/litellm/proxy/lens/adapter.py index 2f34bf661a2..bfd404fd17f 100644 --- a/litellm/proxy/lens/adapter.py +++ b/litellm/proxy/lens/adapter.py @@ -12,7 +12,7 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response from pydantic import BaseModel, ConfigDict from starlette.responses import JSONResponse -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth, hash_token from litellm.proxy.auth.authorization import OwnedRows, resolve_trace_read_scope from litellm.proxy.auth.authorization_dependencies import LogTeamLookupDependency from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -85,6 +85,7 @@ async def delegated_identity( return Identity.model_validate( { **auth.model_dump(include={"user_id", "team_id", "org_id", "token", "models"}), + "token": auth.token if auth.via_virtual_key or auth.token is None else hash_token(auth.token), "user_role": auth.user_role or LitellmUserRoles.INTERNAL_USER, "log_team_ids": scope.team_ids if isinstance(scope, OwnedRows) else (), } diff --git a/litellm/proxy/tracing_runtime.py b/litellm/proxy/tracing_runtime.py index 0ac05216a30..f26d3232da5 100644 --- a/litellm/proxy/tracing_runtime.py +++ b/litellm/proxy/tracing_runtime.py @@ -50,13 +50,13 @@ async def manage_tracing( verbose_proxy_logger.error("Agent tracing store lens failed to start, tracing is disabled: %s", error) yield None return - async with client_factory(connection) as client: + async with client_factory(connection) as client, client_factory(connection) as export_client: tracing: Final = ( receiver_factory() if receiver_factory else TraceReceiver(storage=LensTraceStorage(RemoteTraceStore(client))) ) - async with _export_requests(LensExporter(client)): + async with _export_requests(LensExporter(export_client)): yield tracing diff --git a/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py b/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py index f94d7bbee37..12a93801cb4 100644 --- a/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py +++ b/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py @@ -313,12 +313,43 @@ async def test_turn_off_message_logging_omits_all_custom_request_metadata() -> N metadata: Final = _JSON_OBJECT_ADAPTER.validate_json(metadata_json) standard_metadata: Final = cast(Mapping[str, object], payload["metadata"]) assert metadata == { - **standard_metadata, + **{ + key: value + for key, value in standard_metadata.items() + if key not in ( + "spend_logs_metadata", "requester_metadata", "requester_custom_headers", "user_api_key_auth_metadata" + ) + }, "litellm_lens_internal": False, } assert {"project", "api_key", "auth", "prompt"}.isdisjoint(metadata) +@pytest.mark.parametrize("private", (False, True)) +@pytest.mark.parametrize("carrier", ("spend_logs_metadata", "requester_metadata", "requester_custom_headers")) +def test_standard_metadata_carriers_cannot_restore_credentials( + monkeypatch: pytest.MonkeyPatch, private: bool, carrier: str +) -> None: + monkeypatch.setattr(litellm, "turn_off_message_logging", private) + payload: Final = cast( + StandardLoggingPayload, + _payload(metadata={ + **_payload()["metadata"], + carrier: {"authorization": "dummy-bearer-credential", "nested": [{"api_key": "dummy-provider-key"}]}, + }), + ) + row: Final = spend_log_row_from_payload(payload, {}) + metadata: Final = json.loads(row["metadata"]) + assert "dummy-bearer-credential" not in row["metadata"] + assert "dummy-provider-key" not in row["metadata"] + assert metadata["user_api_key_team_id"] == "team-1" + if private: + assert carrier not in metadata + else: + assert metadata[carrier]["authorization"] == REDACTED + assert metadata[carrier]["nested"][0]["api_key"] == REDACTED + + def test_anthropic_cache_fields_are_used_as_fallback(): usage = {"cache_read_input_tokens": 11, "cache_creation_input_tokens": 3} payload = _payload() diff --git a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py index 02bdf3429ea..65cdfe0e8f6 100644 --- a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py +++ b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py @@ -693,7 +693,10 @@ class TestResolveAllMigrationsLedger: return _FakeCompleted() monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", fake_run) - ProxyExtrasDBManager._resolve_all_migrations(str(tmp_path), "schema.prisma") + ProxyExtrasDBManager._resolve_all_migrations( + str(tmp_path), "schema.prisma", + lens_connect=lambda *args, **kwargs: _FakePsycopgConn([], legacy=False), + ) return calls def _resolved(self, calls): @@ -753,13 +756,17 @@ class TestPartitionedSpendLogsPushGuard: class _FakeCursor: + def __init__(self, legacy=True): + self.legacy = legacy + def fetchone(self): - return (1,) + return (1,) if self.legacy else None class _FakePsycopgConn: - def __init__(self, executed): + def __init__(self, executed, legacy=True): self._executed = executed + self.legacy = legacy def __enter__(self): return self @@ -769,12 +776,28 @@ class _FakePsycopgConn: def execute(self, query, params): self._executed.append((query, params)) - return _FakeCursor() + return _FakeCursor(self.legacy) class TestLensRenamePendingCheck: _DATABASE_URL: Final = "postgresql://litellm:hunter2@localhost:5432/litellm" + def test_legacy_drift_recovery_refuses_to_drop_unrenamed_lens_tables( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + monkeypatch.setenv("DATABASE_URL", "postgresql://example:fake-password@localhost/example") + executed: Final[list[tuple[str, tuple[str]]]] = [] + + def legacy_tables_present(conninfo: str, *, connect_timeout: int, autocommit: bool) -> _FakePsycopgConn: + return _FakePsycopgConn(executed) + + with pytest.raises(RuntimeError, match="Legacy Lens tables exist"): + ProxyExtrasDBManager._resolve_all_migrations( + str(tmp_path), "schema.prisma", lens_connect=legacy_tables_present + ) + assert executed[0][1] == ("public",) + assert not tuple(tmp_path.iterdir()) + def test_database_failure_text_reaches_the_raised_message(self, monkeypatch: pytest.MonkeyPatch) -> None: import psycopg diff --git a/tests/unit/proxy/lens/test_adapter.py b/tests/unit/proxy/lens/test_adapter.py index d6d0ba5d663..55cdf76f0cb 100644 --- a/tests/unit/proxy/lens/test_adapter.py +++ b/tests/unit/proxy/lens/test_adapter.py @@ -7,7 +7,7 @@ import pytest from fastapi import HTTPException from starlette.requests import Request -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth, hash_token from litellm.proxy.lens.adapter import ( Connection, Identity, @@ -116,7 +116,7 @@ async def test_delegated_identity_preserves_authenticated_lens_scope_and_log_per user_id="user", team_id="credential-team", org_id="org", - token="token-hash", + token=hash_token("token-hash"), models=("model",), log_team_ids=() if role in (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) @@ -124,6 +124,35 @@ async def test_delegated_identity_preserves_authenticated_lens_scope_and_log_per ) +@pytest.mark.asyncio +@pytest.mark.parametrize("user_id", (None, "oauth-user")) +@pytest.mark.parametrize("credential", ("opaque-fixture-credential", "a" * 64)) +async def test_opaque_oauth_credentials_never_leave_in_delegated_claims( + connection: Connection, user_id: str | None, credential: str +) -> None: + async def lookup(auth: UserAPIKeyAuth) -> tuple[str, ...]: + return () + + identity: Final = await delegated_identity(UserAPIKeyAuth(user_id=user_id, api_key=credential), lookup) + token: Final = connection.identity_token(identity, 1000) + assert token is not None + claims: Final = jwt.decode(token, options={"verify_signature": False}) + assert claims["identity"]["token"] == hash_token(credential) + assert claims["sub"] == (user_id or hash_token(credential)) + assert credential not in str(claims) + + +@pytest.mark.asyncio +async def test_verified_virtual_key_keeps_its_trace_scope() -> None: + async def lookup(auth: UserAPIKeyAuth) -> tuple[str, ...]: + return () + + auth: Final = UserAPIKeyAuth(token=hash_token("sk-fixture-key")) + auth.via_virtual_key = True + identity: Final = await delegated_identity(auth, lookup) + assert identity.token == auth.token + + @pytest.mark.asyncio async def test_delegation_permission_failure_falls_back_to_own_user_scope() -> None: async def unavailable(auth: UserAPIKeyAuth) -> tuple[str, ...]: diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index e59dd19271c..2bc0aa2b397 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -189,6 +189,46 @@ async def test_tracing_config_automatically_exports_spend_without_a_storage_depe assert logger.task is not None and logger.task.done() and not logger.task.cancelled() +@pytest.mark.asyncio +async def test_trace_read_pool_pressure_cannot_drop_spend_records(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.tracing_runtime import manage_tracing + from litellm.tracing.exporter import LensExporter + from litellm.tracing.generated.types import TraceScope + from litellm.tracing.remote import LensConnection + + monkeypatch.setenv("LITELLM_LENS_URL", "https://lens.example.com") + monkeypatch.setenv("LITELLM_LENS_SERVICE_TOKEN", "dummy-service-token-with-32-characters") + received: Final = asyncio.Queue[httpx.Request]() + + def client(connection: LensConnection) -> httpx.AsyncClient: + congested: Final = asyncio.Event() + + def accept(request: httpx.Request) -> httpx.Response: + if request.url.path == "/internal/read": + congested.set() + if congested.is_set(): + raise httpx.PoolTimeout("fixture read pool occupied") + received.put_nowait(request) + return httpx.Response(204) + + return httpx.AsyncClient(base_url=connection.url, transport=httpx.MockTransport(accept)) + + async with manage_tracing(enabled=True, client_factory=client) as tracing: + assert tracing is not None + with pytest.raises(RuntimeError, match="storage is unavailable"): + await tracing.list_traces(TraceScope(all_teams=0, user_id="fixture-user", team_ids=()), 0, 1) + logger: Final = next( + callback for callback in litellm._async_success_callback if isinstance(callback, LensExporter) + ) + await logger.async_log_success_event( + {"standard_logging_object": {"id": "isolated-spend", "response_cost": 0.25}}, None, None, None + ) + + assert logger.rows_written == 1 + assert logger.rows_dropped == 0 + assert received.get_nowait().url.path == "/internal/spend" + + @pytest.mark.asyncio async def test_tracing_logs_an_error_naming_the_store_when_lens_is_misconfigured( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture diff --git a/tests/unit/proxy/test_component_allowlists.py b/tests/unit/proxy/test_component_allowlists.py index eb058d01770..2e38214d733 100644 --- a/tests/unit/proxy/test_component_allowlists.py +++ b/tests/unit/proxy/test_component_allowlists.py @@ -162,6 +162,7 @@ def test_composed_lifespan_preserves_request_state_and_teardown( def register_trace_route(application: Starlette, module: object) -> None: application.router.routes.append(Route("/v1/traces", trace_state)) + application.router.routes.append(Route("/not-a-component-route/lazy", trace_state)) @asynccontextmanager async def stateful_lifespan(application: Starlette) -> AsyncGenerator[Mapping[str, object], None]: @@ -189,9 +190,14 @@ def test_composed_lifespan_preserves_request_state_and_teardown( assert response.json() == {"keys": [] if state_kind == "stateless" else sorted(state)} filtered: Final = client.get("/not-a-component-route") assert filtered.status_code == (200 if component_lifespan is None else 404), filtered.text - assert events == (["startup", "request", "request"] if component_lifespan is None else ["startup", "request"]) + lazy_filtered: Final = client.get("/not-a-component-route/lazy") + assert lazy_filtered.status_code == (200 if component_lifespan is None else 404), lazy_filtered.text + assert events == ( + ["startup", "request", "request", "request"] if component_lifespan is None else ["startup", "request"] + ) assert events == ( - ["startup", "request", "request", "shutdown"] if component_lifespan is None else ["startup", "request", "shutdown"] + ["startup", "request", "request", "request", "shutdown"] + if component_lifespan is None else ["startup", "request", "shutdown"] )