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 af4d0b0d754..35fc3510414 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -148,11 +148,11 @@ class CheckBatchCost: verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}") return {} - async def _get_key_alias(self, batch_id: str, api_key: str | None) -> str | None: + async def _get_key_alias(self, batch_id: str, api_key: str | None, created_by: str | None) -> str | None: """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}-"): + if created_by and api_key == f"{CLI_SESSION_KEY_PREFIX}-{created_by}": return api_key try: key_row: prisma_models.LiteLLM_VerificationToken | None = await _token_table( @@ -234,7 +234,7 @@ class CheckBatchCost: **(await self._get_user_info(batch_id, job.created_by)), } - key_alias = await self._get_key_alias(batch_id, api_key) + key_alias = await self._get_key_alias(batch_id, api_key, job.created_by) if key_alias is not None: metadata["user_api_key_alias"] = key_alias team_alias = await self._get_team_alias(team_id) diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 85707da2497..95a703e8f11 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -2634,6 +2634,19 @@ class TestBatchCostAttribution: assert metadata["user_api_key"] == "cli-session-alice" assert metadata["user_api_key_alias"] == "cli-session-alice" + @pytest.mark.asyncio + async def test_raw_cli_session_token_on_a_legacy_batch_row_is_not_treated_as_the_alias(self): + """A batch row written by an older build stores the raw per-login session token, which shares + the cli-session- prefix with the alias. Only the exact cli-session- value is the + alias; anything else stays a secret so redaction hashes it instead of persisting the token.""" + instance = self._instance(key_row=None) + + metadata = await instance._build_creator_attribution_metadata( + self._job(api_key="cli-session-Qm7xJ2kP9sLw4vT1nR8yAa"), "batch-1" + ) + + assert metadata.get("user_api_key_alias") is None + @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