From 6b64b98e76869bad94573a9cc4ac2172c7bc1e13 Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 24 Sep 2026 21:53:05 +0000 Subject: [PATCH] fix(spend): attribute CLI session pass-through, queue, and managed batch spend to the cli-session alias Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/common_utils/check_batch_cost.py | 3 ++ .../proxy/hooks/managed_files.py | 3 +- .../pass_through_endpoints.py | 2 +- litellm/proxy/proxy_server.py | 7 +++-- .../proxy_unit_tests/test_check_batch_cost.py | 14 +++++++++ .../test_pass_through_endpoints.py | 31 +++++++++++++++++++ 6 files changed, 55 insertions(+), 5 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 41974c26158..af4d0b0d754 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tupl from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( + CLI_SESSION_KEY_PREFIX, MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, MAX_OBJECTS_PER_POLL_CYCLE, ) @@ -151,6 +152,8 @@ class CheckBatchCost: """Resolve the creating virtual key's alias from its hashed token.""" if not api_key: return None + if api_key.startswith(f"{CLI_SESSION_KEY_PREFIX}-"): + return api_key try: key_row: prisma_models.LiteLLM_VerificationToken | None = await _token_table( self.prisma_client diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 5ac7c1e53c1..21bf7abdc2e 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -50,6 +50,7 @@ from litellm.proxy._types import ( ProxyException, UserAPIKeyAuth, ) +from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.openai_files_endpoints.common_utils import ( BATCH_CREATE_HIDDEN_PARAM, FILE_LIST_CONTINUATION_CHUNK_SIZE, @@ -359,7 +360,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): from prisma import Json - api_key = user_api_key_dict.api_key or None + api_key = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) or None attribution_columns = ( { **({"api_key": api_key} if api_key is not None else {}), diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index a119335ba46..7d6db30e3e3 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -607,7 +607,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): # body that mirrors them cannot clobber the authenticated key, the real # parent span, or the proxy's own session-id decision. _metadata.pop(SESSION_ID_OMITTED_METADATA_KEY, None) - _metadata["user_api_key"] = user_api_key_dict.api_key + _metadata["user_api_key"] = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) _metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span _metadata["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation _metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = user_api_key_dict.matched_model_access_groups diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c0071aa7c81..547c76dc908 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -552,7 +552,7 @@ from litellm.proxy.list_api.common import ( problem_response, request_validation_problem, ) -from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request +from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_litellm_data_to_request from litellm.proxy.logging_endpoints.callback_logs_endpoints import ( rust_control_plane_router, ) @@ -16347,8 +16347,9 @@ async def async_queue_request( # Covers both missing and JSON-string metadata (multipart / # extra_body); see above for the same guard upstream. data["metadata"] = {} - data["metadata"]["user_api_key"] = user_api_key_dict.api_key - data["metadata"]["user_api_key_hash"] = user_api_key_dict.api_key + logged_api_key: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) + data["metadata"]["user_api_key"] = logged_api_key + data["metadata"]["user_api_key_hash"] = logged_api_key data["metadata"]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata) _headers: Final = _safe_get_request_headers(request).copy() _headers.pop("authorization", None) # do not store the original `sk-..` api key in the db diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 6417c7c8aa6..85707da2497 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -2620,6 +2620,20 @@ class TestBatchCostAttribution: assert metadata["user_api_key"] == "hash-alice" assert metadata.get("user_api_key_alias") is None + @pytest.mark.asyncio + async def test_cli_session_batch_keeps_its_alias_without_a_key_row(self): + """A batch created from a CLI session stores the per-user cli-session alias as its api_key. + No verification token row exists for that alias, so the alias is its own key_alias; without + it the spend log redaction would hash the alias back into an unrecoverable value.""" + instance = self._instance(key_row=None) + + metadata = await instance._build_creator_attribution_metadata( + self._job(api_key="cli-session-alice"), "batch-1" + ) + + assert metadata["user_api_key"] == "cli-session-alice" + assert metadata["user_api_key_alias"] == "cli-session-alice" + @pytest.mark.asyncio async def test_unnamed_key_keeps_the_creating_user_alias(self): """Regression: a key generated without key_alias resolves to no alias, and the diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index f70a0a453a4..70f7523ec64 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7417,3 +7417,34 @@ async def test_chat_completion_pass_through_endpoint_failure_carries_the_callers record = next(r for r in caplog.records if "Exception occured" in r.getMessage()) assert record.litellm_call_id == call_id assert call_id in record.getMessage() + + +def test_passthrough_attributes_a_cli_session_to_its_alias_not_the_login_token(): + """A pass-through request signed with a CLI session token must stamp the same stable per-user alias the + unified routes do, so every login rolls up to one key and the token itself never reaches a spend row.""" + from litellm.proxy.spend_tracking.spend_tracking_utils import _get_spend_logs_metadata + + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = "http://0.0.0.0:4000/anthropic/v1/messages" + mock_request.headers = Headers({}) + mock_request.scope = {} + session = UserAPIKeyAuth( + api_key="cli-session-Qm7xJ2kP9sLw4vT1nR8yAa", + key_alias="cli-session-alice", + user_id="alice", + is_session_token=True, + ) + + kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( + request=mock_request, + user_api_key_dict=session, + passthrough_logging_payload=MagicMock(), + logging_obj=MagicMock(), + _parsed_body={}, + litellm_call_id="lit-6852-passthrough-call-id", + ) + + metadata = kwargs["litellm_params"]["metadata"] + assert metadata["user_api_key"] == "cli-session-alice" + assert _get_spend_logs_metadata(metadata)["user_api_key"] == "cli-session-alice"