Merge 489af0b8b1 into a2c23d95ae

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

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}),
                 }

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 (),
         }

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
 
 

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()

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
 

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, ...]:

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