From 07b43b583dfbc90b064dcfcdef0c14a235b97bfb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:42:42 -0700 Subject: [PATCH] fix(spend): attribute CLI session spend to the per-user cli-session alias instead of the hashed session token A CLI session token is a fresh random secret on every login, so since v1.99 each login's spend rows carried a different sha256 hash as api_key and the usage APIs could resolve neither key_alias nor user_email for them. Spend rows and logging callbacks now attribute a session request to its stable alias, cli-session-, and the usage endpoints derive that alias and owner from the key itself instead of scanning for a matching digest --- .../proxy/hooks/proxy_track_cost_callback.py | 4 +- litellm/proxy/litellm_pre_call_utils.py | 17 +++++- .../common_daily_activity.py | 6 +- .../spend_tracking/key_metadata_recovery.py | 20 ++++++- .../spend_tracking/spend_tracking_utils.py | 38 +++++++++--- .../test_common_daily_activity.py | 16 +++++ .../test_key_metadata_recovery.py | 24 ++++++++ .../test_spend_tracking_utils.py | 58 +++++++++++++++++++ .../proxy/test_litellm_pre_call_utils.py | 34 +++++++++++ 9 files changed, 201 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index c4fba8ecf9e..3c3a6f502ae 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -109,7 +109,7 @@ class _ProxyDBLogger(CustomLogger): _metadata = dict( LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) ) - _metadata["user_api_key"] = user_api_key_dict.api_key + _metadata["user_api_key"] = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) _metadata["status"] = "failure" _error_information = StandardLoggingPayloadSetup.get_error_information( original_exception=original_exception, @@ -204,7 +204,7 @@ class _ProxyDBLogger(CustomLogger): ) await proxy_logging_obj.db_spend_update_writer.update_database( - token=user_api_key_dict.api_key, + token=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict), response_cost=recovered_response_cost, user_id=user_api_key_dict.user_id, end_user_id=user_api_key_dict.end_user_id, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index da033dc2276..47977624282 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1571,6 +1571,19 @@ class LiteLLMProxyRequestSetup: data[_metadata_variable_name].update(metadata_from_headers) return data + @staticmethod + def get_logged_api_key(user_api_key_dict: UserAPIKeyAuth) -> str | None: + """ + The value spend rows and logging callbacks attribute a request to. + + A CLI session token is a per-login random secret, so the key is the + session's stable alias (cli-session-) instead: every login of + one user rolls up to one key and the live token never lands in a row. + """ + if user_api_key_dict.is_session_token and user_api_key_dict.key_alias: + return user_api_key_dict.key_alias + return user_api_key_dict.api_key + @staticmethod def get_sanitized_user_information_from_key( user_api_key_dict: UserAPIKeyAuth, @@ -1578,7 +1591,7 @@ class LiteLLMProxyRequestSetup: stripped_metadata: Final = strip_callback_config(user_api_key_dict.metadata) auth_metadata: Final = cast("dict[str, str] | None", stripped_metadata) # cast-ok: metadata is free-form JSON user_api_key_logged_metadata: Final = StandardLoggingUserAPIKeyMetadata( - user_api_key_hash=user_api_key_dict.api_key, # just the hashed token + user_api_key_hash=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict), user_api_key_alias=user_api_key_dict.key_alias, user_api_key_spend=user_api_key_dict.spend, user_api_key_max_budget=user_api_key_dict.max_budget, @@ -1616,7 +1629,7 @@ class LiteLLMProxyRequestSetup: user_api_key_dict=user_api_key_dict ) data[_metadata_variable_name].update(user_api_key_logged_metadata) - data[_metadata_variable_name]["user_api_key"] = user_api_key_dict.api_key # this is just the hashed token + data[_metadata_variable_name]["user_api_key"] = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) # Key-owned agent_id for spend attribution; keep existing (e.g. from header) if key has none _key_agent_id: Final = getattr(user_api_key_dict, "agent_id", None) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index a8aef30107c..3eda9669eea 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -13,6 +13,7 @@ from litellm.constants import PTU_SENTINEL_API_KEY from litellm.proxy._types import CommonProxyErrors from litellm.proxy.spend_tracking.key_metadata_recovery import ( attach_user_emails, + recover_cli_session_key_metadata, recover_double_hashed_key_metadata, recover_key_metadata_from_spend_logs, ) @@ -501,11 +502,12 @@ async def get_api_key_metadata( e, ) - still_missing: Final = api_keys - frozenset(result) + from_session_keys: Final = recover_cli_session_key_metadata(api_keys - frozenset(result)) + still_missing: Final = api_keys - frozenset(result) - frozenset(from_session_keys) from_reverse_hash: Final = ( await recover_double_hashed_key_metadata(prisma_client, still_missing) if still_missing else _EMPTY_KEY_METADATA ) - after_token_recovery: Final = MappingProxyType({**result, **from_reverse_hash}) + after_token_recovery: Final = MappingProxyType({**result, **from_session_keys, **from_reverse_hash}) unresolved: Final = api_keys - frozenset(after_token_recovery) from_spend_logs: Final = ( await recover_key_metadata_from_spend_logs(prisma_client, unresolved, spend_logs_window) diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 29688b61b3d..bc28685c0ad 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -11,6 +11,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( + CLI_SESSION_KEY_PREFIX, SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS, SPEND_LOG_KEY_METADATA_CACHE_TTL, SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL, @@ -62,6 +63,7 @@ _SPEND_LOG_STATEMENT_TIMEOUT_SQL: Final = f"SET LOCAL statement_timeout = {SPEND _SPEND_LOG_TRANSACTION_TIMEOUT: Final = timedelta(milliseconds=2 * SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS) _HASHED_JWT_PREFIX: Final = "hashed-jwt-" +_CLI_SESSION_KEY_PREFIX: Final = f"{CLI_SESSION_KEY_PREFIX}-" class KeyMetadataDict(TypedDict, total=False): @@ -198,6 +200,16 @@ async def attach_user_emails( return MappingProxyType({api_key: _meta_with_email(meta, emails) for api_key, meta in recovered.items()}) +def recover_cli_session_key_metadata(missing_keys: AbstractSet[str]) -> Mapping[str, KeyMetadataDict]: + return MappingProxyType( + { + key: KeyMetadataDict(key_alias=key, user_id=key.removeprefix(_CLI_SESSION_KEY_PREFIX)) + for key in missing_keys + if key.startswith(_CLI_SESSION_KEY_PREFIX) and len(key) > len(_CLI_SESSION_KEY_PREFIX) + } + ) + + async def recover_double_hashed_key_metadata( prisma_client: PrismaClient, missing_keys: AbstractSet[str], @@ -383,9 +395,15 @@ async def fill_missing_api_key_aliases( if not missing_keys: return tuple(rows) + from_session_keys: Final = recover_cli_session_key_metadata(missing_keys) recovered: Final = await attach_user_emails( prisma_client, - await recover_double_hashed_key_metadata(prisma_client, missing_keys), + MappingProxyType( + { + **from_session_keys, + **await recover_double_hashed_key_metadata(prisma_client, missing_keys - frozenset(from_session_keys)), + } + ), ) if not recovered: return tuple(rows) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index a21d761996f..fc3ab2cdb8a 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -11,6 +11,7 @@ from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ( + CLI_SESSION_KEY_PREFIX, LITELLM_PROXY_MASTER_KEY_ALIAS, LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, @@ -85,19 +86,28 @@ _NON_SECRET_KEY_ALIASES: Final = frozenset( ) -def _is_non_secret_key_value(value: str) -> bool: +def _is_cli_session_alias(value: str, key_alias: object) -> bool: + return value.startswith(f"{CLI_SESSION_KEY_PREFIX}-") and value == key_alias + + +def _is_non_secret_key_value(value: str, *, key_alias: object = None) -> bool: return ( - value in _NON_SECRET_KEY_ALIASES or is_valid_sha256_hash(value) or _HASHED_JWT_RE.fullmatch(value) is not None + value in _NON_SECRET_KEY_ALIASES + or is_valid_sha256_hash(value) + or _HASHED_JWT_RE.fullmatch(value) is not None + or _is_cli_session_alias(value, key_alias) ) -def _redact_logged_api_key(value: str | None, *, already_redacted: bool = False) -> str | None: +def _redact_logged_api_key( + value: str | None, *, already_redacted: bool = False, key_alias: object = None +) -> str | None: if not isinstance(value, str) or not value: return None stripped: Final = re.sub(r"(?i)^bearer ", "", value) if not stripped: return None - if already_redacted and _is_non_secret_key_value(stripped): + if already_redacted and _is_non_secret_key_value(stripped, key_alias=key_alias): return stripped return hash_token(stripped) @@ -189,10 +199,15 @@ def _get_spend_logs_metadata( ) _raw_key: Final = clean_metadata.get("user_api_key") _trusted_hash: Final = metadata.get("user_api_key_hash") + _key_alias: Final = metadata.get("user_api_key_alias") _already_redacted: Final = ( - isinstance(_trusted_hash, str) and _is_non_secret_key_value(_trusted_hash) and _trusted_hash == _raw_key + isinstance(_trusted_hash, str) + and _is_non_secret_key_value(_trusted_hash, key_alias=_key_alias) + and _trusted_hash == _raw_key + ) + clean_metadata["user_api_key"] = _redact_logged_api_key( + _raw_key, already_redacted=_already_redacted, key_alias=_key_alias ) - clean_metadata["user_api_key"] = _redact_logged_api_key(_raw_key, already_redacted=_already_redacted) clean_metadata["applied_guardrails"] = applied_guardrails clean_metadata["batch_models"] = batch_models clean_metadata["batch_successful_requests"] = batch_successful_requests @@ -388,10 +403,13 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs standard_logging_completion_tokens = standard_logging_payload.get("completion_tokens", 0) standard_logging_total_tokens = standard_logging_payload.get("total_tokens", 0) _trusted_hash = metadata.get("user_api_key_hash") + _key_alias = metadata.get("user_api_key_alias") _key_already_redacted = ( - isinstance(_trusted_hash, str) and _is_non_secret_key_value(_trusted_hash) and _trusted_hash == api_key + isinstance(_trusted_hash, str) + and _is_non_secret_key_value(_trusted_hash, key_alias=_key_alias) + and _trusted_hash == api_key ) - api_key = _redact_logged_api_key(api_key, already_redacted=_key_already_redacted) or "" + api_key = _redact_logged_api_key(api_key, already_redacted=_key_already_redacted, key_alias=_key_alias) or "" if ( standard_logging_payload is not None @@ -399,7 +417,9 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs api_key = ( api_key or _redact_logged_api_key( - standard_logging_payload["metadata"].get("user_api_key_hash"), already_redacted=True + standard_logging_payload["metadata"].get("user_api_key_hash"), + already_redacted=True, + key_alias=standard_logging_payload["metadata"].get("user_api_key_alias"), ) or "" ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 6cd900cb041..32a91bac506 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -2237,3 +2237,19 @@ def test_spend_logs_window_is_none_when_no_date_parses(): from litellm.proxy.management_endpoints.common_daily_activity import _spend_logs_window assert _spend_logs_window({"garbage", ""}) is None + + +@pytest.mark.asyncio +async def test_get_api_key_metadata_resolves_cli_session_keys_from_the_key_itself(): + mock_prisma = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.query_raw = AsyncMock(side_effect=AssertionError("no reverse-hash or spend-log scan expected")) + mock_prisma.db.litellm_usertable.find_many = AsyncMock( + return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com")] + ) + + result = await get_api_key_metadata(prisma_client=mock_prisma, api_keys={"cli-session-alice"}) + + assert result["cli-session-alice"]["key_alias"] == "cli-session-alice" + assert result["cli-session-alice"]["user_email"] == "alice@example.com" diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py index 89be341c87b..ad709a299c7 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -16,6 +16,7 @@ from litellm.constants import ( ) from litellm.proxy.spend_tracking.key_metadata_recovery import ( fill_missing_api_key_aliases, + recover_cli_session_key_metadata, recover_double_hashed_key_metadata, recover_key_metadata_from_spend_logs, ) @@ -588,3 +589,26 @@ async def test_recover_key_metadata_from_spend_logs_bounds_the_scan_with_a_state assert mock_prisma.db.tx.call_args.kwargs["timeout"] == timedelta( milliseconds=2 * SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS ) + + +def test_recover_cli_session_key_metadata_names_the_alias_and_owner_from_the_key(): + result = recover_cli_session_key_metadata({"cli-session-alice", hash_token("sk-other"), "cli-session-", "sk-raw"}) + + assert dict(result) == {"cli-session-alice": {"key_alias": "cli-session-alice", "user_id": "alice"}} + + +@pytest.mark.asyncio +async def test_fill_missing_api_key_aliases_resolves_cli_session_keys_without_a_digest_lookup(): + mock_prisma = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=AssertionError("no reverse-hash lookup expected")) + mock_prisma.db.litellm_usertable.find_many = AsyncMock( + return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com")] + ) + rows = ({"api_key": "cli-session-alice", "api_key_alias": None, "team_id": None, "user_email": None, "spend": 2.0},) + + filled = await fill_missing_api_key_aliases(mock_prisma, rows) + + assert filled[0]["api_key_alias"] == "cli-session-alice" + assert filled[0]["user_email"] == "alice@example.com" + assert filled[0]["spend"] == 2.0 + assert mock_prisma.db.litellm_usertable.find_many.call_args.kwargs["where"] == {"user_id": {"in": ["alice"]}} diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 5a79560b972..a6783fe67bc 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -4432,3 +4432,61 @@ def test_spend_log_request_id_is_the_response_id_a_bridged_messages_caller_recei ) == "resp_01Lit6806Bridged" ) + + +_CLI_SESSION_ALIAS: Final = "cli-session-alice" +_CLI_SESSION_TOKEN: Final = "cli-session-Qm7xJ2kP9sLw4vT1nR8yAa" + + +def _cli_session_request_metadata(logged_key: str) -> dict[str, str]: + return { + "user_api_key": logged_key, + "user_api_key_hash": logged_key, + "user_api_key_alias": _CLI_SESSION_ALIAS, + "user_api_key_user_id": "alice", + } + + +@pytest.mark.parametrize("call_type", ["acompletion", "anthropic_messages", "aresponses"]) +def test_get_logging_payload_attributes_a_cli_session_to_its_alias(call_type: str): + payload = get_logging_payload( + kwargs={ + "call_type": call_type, + "model": "gpt-5.4-nano", + "response_cost": 0.00001, + "litellm_params": {"metadata": _cli_session_request_metadata(_CLI_SESSION_ALIAS)}, + }, + response_obj=litellm.ModelResponse(id=f"{call_type}-1", choices=[], usage=litellm.Usage()), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert payload["api_key"] == _CLI_SESSION_ALIAS + assert json.loads(payload["metadata"])["user_api_key"] == _CLI_SESSION_ALIAS + assert json.loads(payload["metadata"])["user_api_key_alias"] == _CLI_SESSION_ALIAS + + +@pytest.mark.parametrize("call_type", ["acompletion", "anthropic_messages", "aresponses"]) +def test_get_logging_payload_never_lands_a_raw_cli_session_token(call_type: str): + payload = get_logging_payload( + kwargs={ + "call_type": call_type, + "model": "gpt-5.4-nano", + "litellm_params": {"metadata": _cli_session_request_metadata(_CLI_SESSION_TOKEN)}, + }, + response_obj=litellm.ModelResponse(id=f"{call_type}-2", choices=[], usage=litellm.Usage()), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert payload["api_key"] == hash_token(_CLI_SESSION_TOKEN) + assert json.loads(payload["metadata"])["user_api_key"] == hash_token(_CLI_SESSION_TOKEN) + + +def test_redact_logged_api_key_cli_session_alias_needs_alias_provenance(): + assert _redact_logged_api_key(_CLI_SESSION_ALIAS, already_redacted=True, key_alias=_CLI_SESSION_ALIAS) == ( + _CLI_SESSION_ALIAS + ) + assert _redact_logged_api_key(_CLI_SESSION_ALIAS, already_redacted=True) == hash_token(_CLI_SESSION_ALIAS) + assert _redact_logged_api_key(_CLI_SESSION_ALIAS, key_alias=_CLI_SESSION_ALIAS) == hash_token(_CLI_SESSION_ALIAS) + assert _redact_logged_api_key("alice", already_redacted=True, key_alias="alice") == hash_token("alice") diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 94aaa32519e..62fb39cb4d7 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -7943,3 +7943,37 @@ def test_default_team_settings_bool_turn_off_message_logging_redacts(): ) is True ) + + +def test_add_user_api_key_auth_to_request_metadata_attributes_a_cli_session_to_its_alias(): + data = {"model": "gpt-5.4-nano", "messages": [{"role": "user", "content": "hi"}], "litellm_metadata": {}} + session = UserAPIKeyAuth( + api_key="cli-session-Qm7xJ2kP9sLw4vT1nR8yAa", + key_alias="cli-session-alice", + user_id="alice", + is_session_token=True, + ) + + metadata = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( + data=data, user_api_key_dict=session, _metadata_variable_name="litellm_metadata" + )["litellm_metadata"] + + assert metadata["user_api_key"] == "cli-session-alice" + assert metadata["user_api_key_hash"] == "cli-session-alice" + assert metadata["user_api_key_alias"] == "cli-session-alice" + assert "Qm7xJ2kP9sLw4vT1nR8yAa" not in (metadata["user_api_key"], metadata["user_api_key_hash"]) + + +def test_add_user_api_key_auth_to_request_metadata_keeps_the_hashed_token_for_virtual_keys(): + from litellm.proxy._types import hash_token + + data = {"model": "gpt-5.4-nano", "messages": [], "litellm_metadata": {}} + hashed = hash_token("sk-virtual-key") + virtual_key = UserAPIKeyAuth(api_key=hashed, key_alias="cli-session-alice", user_id="alice") + + metadata = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( + data=data, user_api_key_dict=virtual_key, _metadata_variable_name="litellm_metadata" + )["litellm_metadata"] + + assert metadata["user_api_key"] == hashed + assert metadata["user_api_key_hash"] == hashed