diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 0cceddd9ac..ebabf1b4d7 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -31,6 +31,15 @@ def is_anthropic_oauth_key(value: Optional[str]) -> bool: value = value[7:] return value.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX) +def _merge_beta_headers(existing: Optional[str], new_beta: str) -> str: + """Merge a new beta value into an existing comma-separated anthropic-beta header.""" + if not existing: + return new_beta + betas = {b.strip() for b in existing.split(",") if b.strip()} + betas.add(new_beta) + return ",".join(sorted(betas)) + + def optionally_handle_anthropic_oauth( headers: dict, api_key: Optional[str] ) -> tuple[dict, Optional[str]]: @@ -52,14 +61,18 @@ def optionally_handle_anthropic_oauth( if auth_header and auth_header.startswith(f"Bearer {ANTHROPIC_OAUTH_TOKEN_PREFIX}"): api_key = auth_header.replace("Bearer ", "") headers.pop("x-api-key", None) - headers["anthropic-beta"] = ANTHROPIC_OAUTH_BETA_HEADER + headers["anthropic-beta"] = _merge_beta_headers( + headers.get("anthropic-beta"), ANTHROPIC_OAUTH_BETA_HEADER + ) headers["anthropic-dangerous-direct-browser-access"] = "true" return headers, api_key # Check api_key directly (standard chat/completion flow) if api_key and api_key.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX): headers.pop("x-api-key", None) headers["authorization"] = f"Bearer {api_key}" - headers["anthropic-beta"] = ANTHROPIC_OAUTH_BETA_HEADER + headers["anthropic-beta"] = _merge_beta_headers( + headers.get("anthropic-beta"), ANTHROPIC_OAUTH_BETA_HEADER + ) headers["anthropic-dangerous-direct-browser-access"] = "true" return headers, api_key diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index c3ad72436b..6ecbd54699 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -63,12 +63,20 @@ class AnthropicCountTokensConfig: Returns: Dictionary of required headers """ - return { + from litellm.llms.anthropic.common_utils import ( + optionally_handle_anthropic_oauth, + ) + + headers: Dict[str, str] = { "Content-Type": "application/json", "x-api-key": api_key, "anthropic-version": "2023-06-01", "anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, } + headers, _ = optionally_handle_anthropic_oauth( + headers=headers, api_key=api_key + ) + return headers def validate_request( self, model: str, messages: List[Dict[str, Any]] diff --git a/tests/test_litellm/llms/anthropic/test_count_tokens_oauth.py b/tests/test_litellm/llms/anthropic/test_count_tokens_oauth.py new file mode 100644 index 0000000000..64b9a3c153 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/test_count_tokens_oauth.py @@ -0,0 +1,86 @@ +""" +Tests for Anthropic CountTokens API OAuth token handling. + +Verifies that get_required_headers() correctly handles OAuth tokens +(sk-ant-oat*) by delegating to optionally_handle_anthropic_oauth(). + +Regression test for https://github.com/BerriAI/litellm/issues/22040 +""" + +import os +import sys + +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) +) + +from litellm.llms.anthropic.count_tokens.transformation import ( + AnthropicCountTokensConfig, +) + +# Fake tokens for testing (not real secrets) +FAKE_OAUTH_TOKEN = "sk-ant-oat01-fake-token-for-testing-123456789abcdef" +FAKE_REGULAR_KEY = "sk-ant-api03-regular-key-for-testing-123456789" + + +class TestCountTokensOAuthHeaders: + """Tests that count_tokens headers are correct for both regular and OAuth keys.""" + + def test_regular_api_key_uses_x_api_key(self): + """Regular API keys should be sent via x-api-key header.""" + config = AnthropicCountTokensConfig() + headers = config.get_required_headers(FAKE_REGULAR_KEY) + + assert headers["x-api-key"] == FAKE_REGULAR_KEY + assert "authorization" not in headers + + def test_oauth_key_uses_bearer_authorization(self): + """OAuth tokens (sk-ant-oat*) should be sent via Authorization: Bearer.""" + config = AnthropicCountTokensConfig() + headers = config.get_required_headers(FAKE_OAUTH_TOKEN) + + assert headers.get("authorization") == f"Bearer {FAKE_OAUTH_TOKEN}" + assert "x-api-key" not in headers + + def test_oauth_key_sets_oauth_beta_header(self): + """OAuth tokens should trigger the anthropic-beta oauth header.""" + config = AnthropicCountTokensConfig() + headers = config.get_required_headers(FAKE_OAUTH_TOKEN) + + assert "oauth-2025-04-20" in headers.get("anthropic-beta", "") + + def test_regular_key_preserves_token_counting_beta(self): + """Regular keys should keep the token-counting beta header.""" + config = AnthropicCountTokensConfig() + headers = config.get_required_headers(FAKE_REGULAR_KEY) + + assert "token-counting" in headers.get("anthropic-beta", "") + + def test_headers_always_have_content_type(self): + """Both regular and OAuth paths should have Content-Type.""" + config = AnthropicCountTokensConfig() + + for key in [FAKE_REGULAR_KEY, FAKE_OAUTH_TOKEN]: + headers = config.get_required_headers(key) + assert headers["Content-Type"] == "application/json" + + def test_headers_always_have_anthropic_version(self): + """Both paths should have anthropic-version.""" + config = AnthropicCountTokensConfig() + + for key in [FAKE_REGULAR_KEY, FAKE_OAUTH_TOKEN]: + headers = config.get_required_headers(key) + assert headers["anthropic-version"] == "2023-06-01" + + def test_oauth_key_preserves_token_counting_beta(self): + """OAuth tokens must preserve the token-counting beta alongside the OAuth beta.""" + config = AnthropicCountTokensConfig() + headers = config.get_required_headers(FAKE_OAUTH_TOKEN) + + 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}" + )