mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-30 01:52:18 +00:00
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:
parent
c7938248d2
commit
521b90a5cf
2 changed files with 17 additions and 2 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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")}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue