From 521b90a5cfee4e95328f9bc45f676c6e0b22ed7a Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 24 Sep 2026 22:27:15 +0000 Subject: [PATCH] fix(spend): log proxy executed batch rows under the cli-session alias instead of the session token _row_metadata set user_api_key from the raw bearer token while user_api_key_hash carried the alias, so the spend log redaction rejected the alias as untrusted and hashed the random session token instead Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_executed_batches.py | 2 +- .../test_litellm_executed_batches.py | 17 ++++++++++++++++- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/batches_endpoints/litellm_executed_batches.py b/litellm/proxy/batches_endpoints/litellm_executed_batches.py index 67201d99422..caf34404a7d 100644 --- a/litellm/proxy/batches_endpoints/litellm_executed_batches.py +++ b/litellm/proxy/batches_endpoints/litellm_executed_batches.py @@ -656,7 +656,7 @@ class LiteLLMExecutedBatchRunner: def _row_metadata(self, run: _BatchRun) -> dict[str, object]: # mutable-ok: router updates metadata in place return { # mutable-ok: the router updates request metadata in place **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(run.user_api_key_dict), - "user_api_key": run.user_api_key_dict.api_key, + "user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(run.user_api_key_dict), "user_api_end_user_max_budget": run.user_api_key_dict.end_user_max_budget, "tags": list(run.request_tags), # mutable-ok: litellm types request tags as a list "batch_id": run.unified_batch_id, diff --git a/tests/test_litellm/proxy/batches_endpoints/test_litellm_executed_batches.py b/tests/test_litellm/proxy/batches_endpoints/test_litellm_executed_batches.py index 6f2341a578c..cae92ce18bb 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_litellm_executed_batches.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_litellm_executed_batches.py @@ -378,6 +378,7 @@ def make_runner( general_settings: Mapping[str, object] = MappingProxyType({}), heartbeat_seconds: float = 30.0, completion_window_seconds: float = 24 * 60 * 60, + user: UserAPIKeyAuth | None = None, ) -> Harness: store = store_factory({INPUT_FILE_ID: managed_input_file()} if files is None else files) router = FakeRouter() @@ -385,7 +386,7 @@ def make_runner( storage = FakeStorageBackend({STORAGE_URL: content}) storage_factory = FakeStorageBackendFactory(storage, storage_error) prisma = FakePrismaClient(store.objects) - user = UserAPIKeyAuth( + user = user or UserAPIKeyAuth( api_key="sk-batch-key", user_id="user-1", team_id="team-1", key_alias="alias-1", user_email="user@example.com" ) runner = LiteLLMExecutedBatchRunner( @@ -668,6 +669,20 @@ async def test_create_dispatches_each_row_with_the_batch_model_and_the_key_metad assert metadata["user_api_key_user_email"] == "user@example.com" +async def test_create_dispatches_cli_session_rows_under_the_per_user_alias_not_the_login_token() -> None: + session = UserAPIKeyAuth( + api_key="cli-session-Qm7xJ2kP9sLw4vT1nR8yAa", + user_id="alice", + key_alias="cli-session-alice", + is_session_token=True, + ) + harness = make_runner(user=session) + await harness.create_and_finish() + + logged_keys = {call.kwargs["metadata"]["user_api_key"] for call in harness.router.acompletion.await_args_list} + assert logged_keys == {"cli-session-alice"} + + async def test_create_uploads_one_output_line_per_row_with_the_router_response() -> None: harness = make_runner() replies = {"hi 1": chat_response("hi 1"), "hi 2": chat_response("hi 2")}