fix(proxy): carry used_client_oauth_token onto failure spend rows for litellm_metadata routes

This commit is contained in:
mateo-berri 2026-09-24 16:22:44 -07:00
parent 23a4bc5b6c
commit d688293c81
2 changed files with 48 additions and 7 deletions

View file

@ -83,6 +83,10 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset(
str(CallTypes.aretrieve_batch),
)
)
_FAILURE_ROW_KEYS_LIFTED_FROM_LITELLM_METADATA: Final[tuple[str, ...]] = (
"standard_logging_guardrail_information",
"used_client_oauth_token",
)
def _proxy_spend_writer() -> DBSpendUpdateWriter:
@ -195,13 +199,13 @@ class _ProxyDBLogger(CustomLogger):
existing_metadata.update(_metadata)
litellm_metadata_bucket: Final = request_data.get("litellm_metadata")
if (
isinstance(litellm_metadata_bucket, dict)
and "standard_logging_guardrail_information" not in existing_metadata
):
guardrail_info: Final = litellm_metadata_bucket.get("standard_logging_guardrail_information")
if guardrail_info is not None:
existing_metadata["standard_logging_guardrail_information"] = guardrail_info
existing_metadata.update(
(key, litellm_metadata_bucket[key])
for key in _FAILURE_ROW_KEYS_LIFTED_FROM_LITELLM_METADATA
if isinstance(litellm_metadata_bucket, dict)
and key not in existing_metadata
and litellm_metadata_bucket.get(key) is not None
)
if "litellm_params" not in request_data:
request_data["litellm_params"] = {}

View file

@ -160,6 +160,43 @@ async def test_async_post_call_failure_hook_does_not_clobber_guardrail_info_in_m
assert metadata["standard_logging_guardrail_information"] == metadata_bucket_info
@pytest.mark.asyncio
@pytest.mark.parametrize("used_client_oauth_token", [True, False])
async def test_async_post_call_failure_hook_carries_used_client_oauth_token_from_litellm_metadata(
used_client_oauth_token: bool,
):
"""
/v1/messages and /v1/responses stamp the proxy's own fields into request_data["litellm_metadata"]
and leave request_data["metadata"] to the caller's native metadata, so a failed request on those
routes wrote a spend row whose used_client_oauth_token was null instead of the stamped value
"""
logger = _ProxyDBLogger()
request_data = {
"model": "claude-sonnet-5",
"messages": [{"role": "user", "content": "Hello"}],
"metadata": {"user_id": "anthropic-native-metadata"},
"litellm_metadata": {"used_client_oauth_token": used_client_oauth_token},
"proxy_server_request": {"request_id": "test_request_id"},
}
with patch(
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
new_callable=AsyncMock,
) as mock_update_database:
await logger.async_post_call_failure_hook(
request_data=request_data,
original_exception=Exception("rate limited"),
user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"),
)
call_kwargs = mock_update_database.call_args[1]["kwargs"]
assert call_kwargs["litellm_params"]["metadata"]["user_id"] == "anthropic-native-metadata"
payload = get_logging_payload(
kwargs=call_kwargs, response_obj={}, start_time=datetime.now(), end_time=datetime.now()
)
assert json.loads(payload["metadata"])["used_client_oauth_token"] is used_client_oauth_token
@pytest.mark.asyncio
async def test_async_post_call_failure_hook_bills_guardrail_cost_on_blocked_request():
"""LIT-5651: a request blocked by a guardrail never reaches the LLM, but the