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>
This commit is contained in:
mateo 2026-09-24 22:27:15 +00:00
parent c7938248d2
commit 521b90a5cf
2 changed files with 17 additions and 2 deletions

View file

@ -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,

View file

@ -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")}