diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 8e8d10c961b..a9a4bf0358e 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -54,13 +54,19 @@ class AnthropicTokenCounter(BaseTokenCounter): deployment = deployment or {} litellm_params: Final = deployment.get("litellm_params", {}) - # Get Anthropic API key from deployment config or environment - api_key = litellm_params.get("api_key") - if not api_key: - api_key = os.getenv("ANTHROPIC_API_KEY") + from litellm.llms.anthropic.wif import aget_anthropic_wif_token + + static_key: Final = litellm_params.get("api_key") or os.getenv("ANTHROPIC_API_KEY") + # A federated deployment holds no static key by design. Without a minted one this returns + # None and the caller silently falls back to the local tokenizer, so a workload identity + # deployment would never reach Anthropic's authoritative count. The minted token is an + # sk-ant-oat, which get_required_headers already sends as a Bearer rather than x-api-key. + api_key: Final = static_key or await aget_anthropic_wif_token( + litellm_params, litellm_params.get("api_base"), model_to_use + ) if not api_key: - verbose_logger.warning("No Anthropic API key found for token counting") + verbose_logger.warning("No Anthropic credential found for token counting") return None try: diff --git a/tests/test_litellm/llms/anthropic/test_count_tokens_oauth.py b/tests/test_litellm/llms/anthropic/test_count_tokens_oauth.py index 1551292b246..240038b1db2 100644 --- a/tests/test_litellm/llms/anthropic/test_count_tokens_oauth.py +++ b/tests/test_litellm/llms/anthropic/test_count_tokens_oauth.py @@ -10,6 +10,8 @@ Regression test for https://github.com/BerriAI/litellm/issues/22040 import os import sys +import pytest + sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))) from litellm.llms.anthropic.count_tokens.transformation import ( @@ -78,3 +80,53 @@ class TestCountTokensOAuthHeaders: beta_value = headers.get("anthropic-beta", "") assert "token-counting" in beta_value, f"token-counting beta missing from OAuth headers: {beta_value}" assert "oauth-2025-04-20" in beta_value, f"oauth beta missing from OAuth headers: {beta_value}" + + +class TestCountTokensUsesWorkloadIdentity: + """A federated deployment holds no static key. Without minting one, count_tokens returns None + and the caller silently falls back to the local tokenizer, so the number a federated + deployment reports would never come from Anthropic.""" + + @pytest.mark.asyncio + async def test_a_federated_deployment_mints_and_counts(self, monkeypatch): + from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module + + monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False) + minted = "sk-ant-oat01-minted-for-count" + + async def fake_mint(_params, _api_base, _model): + return minted + + monkeypatch.setattr(token_counter_module, "aget_anthropic_wif_token", fake_mint, raising=False) + monkeypatch.setattr("litellm.llms.anthropic.wif.aget_anthropic_wif_token", fake_mint, raising=False) + + seen: dict[str, object] = {} + + async def fake_request(**kwargs): + seen.update(kwargs) + return {"input_tokens": 42} + + monkeypatch.setattr( + token_counter_module.anthropic_count_tokens_handler, + "handle_count_tokens_request", + fake_request, + raising=False, + ) + + result = await token_counter_module.AnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={ + "litellm_params": { + "model": "anthropic/claude-sonnet-4-5", + "anthropic_federation_rule_id": "fdrl_x", + "anthropic_organization_id": "org-x", + } + }, + request_model="claude-sonnet-4-5", + ) + + assert result is not None + assert result.total_tokens == 42 + assert seen["api_key"] == minted