From e7b92ce1c2d048122fe5ac1fe26aa066c6ecb667 Mon Sep 17 00:00:00 2001 From: Noobgam Date: Sat, 15 Aug 2026 22:08:48 +0200 Subject: [PATCH] fix(proxy): avoid pinning rotated credentials in token counter --- litellm/proxy/proxy_server.py | 15 ++--- .../proxy/proxy_server/test_routes_utils.py | 61 ++++++++++++++++++- 2 files changed, 66 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 672d94077b2..bc82f3fd349 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -12425,10 +12425,11 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) verbose_proxy_logger.exception( "litellm.proxy.proxy_server.token_counter(): Exception occured while getting deployment" ) - if deployment is not None: - litellm_model_name = deployment.get("litellm_params", {}).get("model") - model_info = deployment.get("model_info", {}) - load_credentials_from_list(deployment.get("litellm_params", {})) + request_deployment: Final = copy.deepcopy(deployment) + if request_deployment is not None: + litellm_model_name = request_deployment.get("litellm_params", {}).get("model") + model_info = request_deployment.get("model_info", {}) + load_credentials_from_list(request_deployment.get("litellm_params", {})) # remove the custom_llm_provider_prefix in the litellm_model_name if "/" in litellm_model_name: litellm_model_name = litellm_model_name.split("/", 1)[1] @@ -12440,9 +12441,9 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) # Try provider-specific token counting first - only for non-direct requests (from provider endpoints) provider_counter: BaseTokenCounter | None = None custom_llm_provider: str | None = None - if call_endpoint is True and deployment is not None: + if call_endpoint is True and request_deployment is not None: # Auto-route to the correct provider based on model - provider_counter, _model, custom_llm_provider = _get_provider_token_counter(deployment, model_to_use) + provider_counter, _model, custom_llm_provider = _get_provider_token_counter(request_deployment, model_to_use) if _model is not None: model_to_use = _model @@ -12453,7 +12454,7 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) model_to_use=model_to_use, messages=messages, contents=contents, - deployment=deployment, + deployment=request_deployment, request_model=request.model, tools=tools, system=system, diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py index 1e1436fcef8..b7d982816cd 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py @@ -10,12 +10,14 @@ from __future__ import annotations import asyncio import json +from unittest.mock import AsyncMock, MagicMock import pytest import litellm from litellm.proxy import proxy_server from litellm.router_utils import pattern_match_deployments +from litellm.types.utils import CredentialItem, TokenCountResponse from .conftest import normalize # type: ignore[import-not-found] @@ -93,6 +95,61 @@ def test_token_counter_missing_input_returns_400( assert "prompt or messages or contents" in response.text +def test_token_counter_named_credentials_do_not_mutate_router_deployment(client, auth_as, monkeypatch): + """Provider token counting must not pin rotating credentials to the router.""" + named_credential = CredentialItem( + credential_name="rotating-bedrock", + credential_values={"aws_session_token": "temporary-session-token"}, + credential_info={"custom_llm_provider": "bedrock"}, + ) + router_deployment = { + "litellm_params": { + "model": "bedrock/anthropic.claude-3-sonnet", + "litellm_credential_name": "rotating-bedrock", + }, + "model_info": {}, + } + + mock_router = MagicMock() + mock_router.async_get_available_deployment = AsyncMock(return_value=router_deployment) + mock_counter = MagicMock() + mock_counter.should_use_token_counting_api.return_value = True + mock_counter.count_tokens = AsyncMock( + return_value=TokenCountResponse( + total_tokens=1, + request_model="claude-bedrock", + model_used="anthropic.claude-3-sonnet", + tokenizer_type="bedrock_api", + ) + ) + + monkeypatch.setattr(litellm, "credential_list", [named_credential]) + monkeypatch.setattr(proxy_server, "llm_router", mock_router) + monkeypatch.setattr( + proxy_server, + "_get_provider_token_counter", + lambda deployment, model: ( + mock_counter, + "anthropic.claude-3-sonnet", + "bedrock", + ), + ) + + with auth_as(): + response = client.post( + "/utils/token_counter?call_endpoint=true", + json={ + "model": "claude-bedrock", + "messages": [{"role": "user", "content": "hello"}], + }, + ) + + assert response.status_code == 200 + hydrated_deployment = mock_counter.count_tokens.await_args.kwargs["deployment"] + assert hydrated_deployment["litellm_params"]["aws_session_token"] == "temporary-session-token" + assert "aws_session_token" not in router_deployment["litellm_params"] + + # --------------------------------------------------------------------------- # GET /utils/supported_openai_params # --------------------------------------------------------------------------- @@ -117,9 +174,7 @@ def patched_supported_params(monkeypatch): def test_supported_openai_params_happy_path(client, auth_as, patched_supported_params): """Pins ``GET /utils/supported_openai_params``.""" with auth_as(): - response = client.get( - "/utils/supported_openai_params", params={"model": "gpt-4"} - ) + response = client.get("/utils/supported_openai_params", params={"model": "gpt-4"}) assert response.status_code == 200 assert normalize(response.json()) == { "supported_openai_params": ["max_tokens", "temperature", "top_p"],