diff --git a/litellm/proxy/spend_tracking/usage_ingestion_endpoints.py b/litellm/proxy/spend_tracking/usage_ingestion_endpoints.py index 63e0f20a395..9c3e5029173 100644 --- a/litellm/proxy/spend_tracking/usage_ingestion_endpoints.py +++ b/litellm/proxy/spend_tracking/usage_ingestion_endpoints.py @@ -30,7 +30,14 @@ MAX_RECORDS_PER_REQUEST: Final = 1000 class ExternalUsageRecord(BaseModel): - api_key: str = Field(min_length=1, description="Raw virtual key (sk-...) to attribute usage to. Never logged.") + api_key: str | None = Field( + default=None, min_length=1, description="Raw virtual key (sk-...) to attribute usage to. Never logged." + ) + api_key_hash: str | None = Field( + default=None, + min_length=1, + description="SHA-256 hash of the virtual key. Use instead of api_key to avoid submitting raw keys.", + ) model: str = Field(min_length=1) prompt_tokens: int = Field(ge=0) completion_tokens: int = Field(ge=0) @@ -51,6 +58,12 @@ class ExternalUsageRecord(BaseModel): raise ValueError("end_time must not be before start_time") return self + @model_validator(mode="after") + def exactly_one_key_identifier(self) -> "ExternalUsageRecord": + if (self.api_key is None) == (self.api_key_hash is None): + raise ValueError("exactly one of api_key or api_key_hash is required") + return self + class UsageIngestRequest(BaseModel): records: tuple[ExternalUsageRecord, ...] = Field(min_length=1, max_length=MAX_RECORDS_PER_REQUEST) @@ -153,7 +166,7 @@ async def process_external_usage_record( record: ExternalUsageRecord, deps: UsageIngestionDeps ) -> UsageIngestRecordResult: request_id: Final = record.idempotency_key or deps.generate_request_id() - hashed_token: Final = hash_token(record.api_key) + hashed_token: Final = record.api_key_hash if record.api_key_hash is not None else hash_token(record.api_key or "") key: Final = await deps.lookup_key(hashed_token) if key is None: @@ -316,7 +329,8 @@ async def ingest_external_usage( dispatching directly to model gateways), so budgets and spend stay coherent in litellm as the single metering system. - Attribution (user/team/org) is derived from the given virtual key. Records accept an optional + Attribution (user/team/org) is derived from the given virtual key, submitted either raw + (api_key) or pre-hashed (api_key_hash) to keep raw keys out of request bodies. Records accept an optional idempotency_key, stored as the spend-log request_id: the reservation insert, counter updates and dedup are checked atomically at the database primary key, so overlapping retries are safe. The reservation row is always written (even when disable_spend_logs is set), because it is both the diff --git a/tests/test_litellm/proxy/spend_tracking/test_usage_ingestion_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_usage_ingestion_endpoints.py index b5ae287269a..8a477c0ea81 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_usage_ingestion_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_usage_ingestion_endpoints.py @@ -256,3 +256,20 @@ def test_failed_booking_is_retry_safe_error_not_permanent_duplicate(): assert "safe to retry" in (result.error or "") assert result.spend is None assert deps.spend_calls == [] + + +def test_key_hash_resolves_without_raw_key_in_body(): + deps = RecordingDeps() + record = make_record(cost=0.01, idempotency_key="k-11") + record = ExternalUsageRecord(**{**record.model_dump(), "api_key": None, "api_key_hash": hash_token(RAW_KEY)}) + result = run(process_external_usage_record(record, deps.as_deps())) + assert result.status == "recorded" + assert deps.looked_up_hashed == hash_token(RAW_KEY) + assert deps.reserve_calls[0]["hashed_token"] == hash_token(RAW_KEY) + + +def test_exactly_one_key_identifier_required(): + with pytest.raises(ValidationError): + make_record(api_key=None) + with pytest.raises(ValidationError): + make_record(api_key_hash=hash_token(RAW_KEY)) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 3647ef4c396..35bd7151383 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -12876,7 +12876,8 @@ export interface paths { * dispatching directly to model gateways), so budgets and spend stay coherent in litellm as the * single metering system. * - * Attribution (user/team/org) is derived from the given virtual key. Records accept an optional + * Attribution (user/team/org) is derived from the given virtual key, submitted either raw + * (api_key) or pre-hashed (api_key_hash) to keep raw keys out of request bodies. Records accept an optional * idempotency_key, stored as the spend-log request_id: the reservation insert, counter updates and * dedup are checked atomically at the database primary key, so overlapping retries are safe. The * reservation row is always written (even when disable_spend_logs is set), because it is both the @@ -24439,7 +24440,12 @@ export interface components { * Api Key * @description Raw virtual key (sk-...) to attribute usage to. Never logged. */ - api_key: string; + api_key?: string | null; + /** + * Api Key Hash + * @description SHA-256 hash of the virtual key. Use instead of api_key to avoid submitting raw keys. + */ + api_key_hash?: string | null; /** Completion Tokens */ completion_tokens: number; /**