Fix Azure AI Anthropic CountTokens 401 auth error (#20069)

Add x-api-key header to CountTokens handler to match chat completion
authentication. Azure AI Anthropic requires this header per Microsoft's
native API format.

litellm/llms/azure_ai/anthropic/count_tokens/transformation.py

@@ -30,30 +30,32 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig):
         """
         Get the required headers for the Azure AI Anthropic CountTokens API.
 
-        Uses Azure authentication (api-key header) instead of Anthropic's x-api-key.
+        Azure AI Anthropic uses Anthropic's native API format, which requires the
+        x-api-key header for authentication (in addition to Azure's api-key header).
 
         Args:
             api_key: The Azure AI API key
             litellm_params: Optional LiteLLM parameters for additional auth config
 
         Returns:
-            Dictionary of required headers with Azure authentication
+            Dictionary of required headers with both x-api-key and Azure authentication
         """
-        # Start with base headers
+        # Start with base headers including x-api-key for Anthropic API compatibility
         headers = {
             "Content-Type": "application/json",
             "anthropic-version": "2023-06-01",
             "anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION,
+            "x-api-key": api_key,  # Azure AI Anthropic requires this header
         }
 
-        # Use Azure authentication
+        # Also set up Azure auth headers for flexibility
         litellm_params = litellm_params or {}
         if "api_key" not in litellm_params:
             litellm_params["api_key"] = api_key
 
         litellm_params_obj = GenericLiteLLMParams(**litellm_params)
 
-        # Get Azure auth headers
+        # Get Azure auth headers (api-key or Authorization)
         azure_headers = BaseAzureLLM._base_validate_azure_environment(
             headers={}, litellm_params=litellm_params_obj
         )
@@ -68,7 +70,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig):
         Get the Azure AI Anthropic CountTokens API endpoint.
 
         Args:
-            api_base: The Azure AI API base URL 
+            api_base: The Azure AI API base URL
                       (e.g., https://my-resource.services.ai.azure.com or
                        https://my-resource.services.ai.azure.com/anthropic)
 

tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py

@@ -0,0 +1,111 @@
+"""
+Tests for Azure AI Anthropic CountTokens transformation.
+
+Verifies that the CountTokens API uses the correct authentication headers.
+"""
+import os
+import sys
+
+sys.path.insert(
+    0, os.path.abspath("../../../../..")
+)  # Adds the parent directory to the system path
+
+
+from litellm.llms.azure_ai.anthropic.count_tokens.transformation import (
+    AzureAIAnthropicCountTokensConfig,
+)
+
+
+class TestAzureAIAnthropicCountTokensConfig:
+    """Test Azure AI Anthropic CountTokens configuration and headers."""
+
+    def test_get_required_headers_includes_x_api_key(self):
+        """
+        Test that get_required_headers includes x-api-key header.
+
+        Azure AI Anthropic uses Anthropic's native API format which requires
+        the x-api-key header for authentication (not just Azure's api-key).
+        """
+        config = AzureAIAnthropicCountTokensConfig()
+        api_key = "test-api-key-12345"
+
+        headers = config.get_required_headers(api_key=api_key)
+
+        # Verify x-api-key header is set
+        assert "x-api-key" in headers
+        assert headers["x-api-key"] == api_key
+
+        # Verify base headers are present
+        assert headers["Content-Type"] == "application/json"
+        assert headers["anthropic-version"] == "2023-06-01"
+        assert "anthropic-beta" in headers
+
+    def test_get_required_headers_includes_azure_api_key(self):
+        """
+        Test that get_required_headers includes Azure api-key header.
+
+        Both x-api-key and api-key headers should be present.
+        """
+        config = AzureAIAnthropicCountTokensConfig()
+        api_key = "test-azure-key-67890"
+
+        headers = config.get_required_headers(api_key=api_key)
+
+        # Verify both authentication headers are set
+        assert "x-api-key" in headers
+        assert "api-key" in headers
+        assert headers["x-api-key"] == api_key
+        assert headers["api-key"] == api_key
+
+    def test_get_required_headers_with_litellm_params(self):
+        """
+        Test that get_required_headers works with litellm_params.
+        """
+        config = AzureAIAnthropicCountTokensConfig()
+        api_key = "test-key"
+        litellm_params = {"api_key": "param-key", "custom_field": "value"}
+
+        headers = config.get_required_headers(
+            api_key=api_key, litellm_params=litellm_params
+        )
+
+        # x-api-key should use the direct api_key parameter
+        assert headers["x-api-key"] == api_key
+        # Azure api-key should come from litellm_params
+        assert headers["api-key"] == "param-key"
+
+    def test_get_count_tokens_endpoint_with_base_url(self):
+        """Test endpoint generation from base URL."""
+        config = AzureAIAnthropicCountTokensConfig()
+
+        api_base = "https://my-resource.services.ai.azure.com"
+        endpoint = config.get_count_tokens_endpoint(api_base)
+
+        assert (
+            endpoint
+            == "https://my-resource.services.ai.azure.com/anthropic/v1/messages/count_tokens"
+        )
+
+    def test_get_count_tokens_endpoint_with_anthropic_path(self):
+        """Test endpoint generation when base URL already includes /anthropic."""
+        config = AzureAIAnthropicCountTokensConfig()
+
+        api_base = "https://my-resource.services.ai.azure.com/anthropic"
+        endpoint = config.get_count_tokens_endpoint(api_base)
+
+        assert (
+            endpoint
+            == "https://my-resource.services.ai.azure.com/anthropic/v1/messages/count_tokens"
+        )
+
+    def test_get_count_tokens_endpoint_with_trailing_slash(self):
+        """Test endpoint generation with trailing slash in base URL."""
+        config = AzureAIAnthropicCountTokensConfig()
+
+        api_base = "https://my-resource.services.ai.azure.com/"
+        endpoint = config.get_count_tokens_endpoint(api_base)
+
+        assert (
+            endpoint
+            == "https://my-resource.services.ai.azure.com/anthropic/v1/messages/count_tokens"
+        )