fix(anthropic): count tokens through the federated credential

A federated deployment holds no static key, so the token counter found none, returned
None, and the caller silently fell back to the local tokenizer. A workload identity
deployment could therefore never get Anthropic's authoritative count.

It now mints one when no static key is configured. The minted token is an sk-ant-oat,
which the existing header builder already sends as a Bearer rather than an x-api-key,
so nothing else on that path changes.
This commit is contained in:
derhornspieler 2026-08-24 09:17:24 -04:00
parent 0a551902ca
commit 33ed7b14bb
2 changed files with 63 additions and 5 deletions

View file

@ -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:

View file

@ -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