From 0730a74ddabc25f7a94fc986b4e370c1d7b99130 Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Mon, 5 Jan 2026 12:25:13 +0530
Subject: [PATCH] fix: auth header for custom api base in generate Content
 request

---
 .../gemini/google_genai/transformation.py     |  4 +-
 .../test_google_genai_transformation.py       | 81 +++++++++++++++++++
 2 files changed, 84 insertions(+), 1 deletion(-)

diff --git a/litellm/llms/gemini/google_genai/transformation.py b/litellm/llms/gemini/google_genai/transformation.py
index d3b55b13ad..3474c8abe3 100644
--- a/litellm/llms/gemini/google_genai/transformation.py
+++ b/litellm/llms/gemini/google_genai/transformation.py
@@ -153,7 +153,9 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM):
         gemini_api_key = api_key or self._get_google_ai_studio_api_key(
             dict(litellm_params or {})
         )
-        if gemini_api_key is not None:
+        if isinstance(gemini_api_key, dict):
+            default_headers.update(gemini_api_key)
+        elif gemini_api_key is not None:
             default_headers[self.XGOOGLE_API_KEY] = gemini_api_key
         if headers is not None:
             default_headers.update(headers)
diff --git a/tests/test_litellm/google_genai/test_google_genai_transformation.py b/tests/test_litellm/google_genai/test_google_genai_transformation.py
index 0d5468cf27..8943d198dc 100644
--- a/tests/test_litellm/google_genai/test_google_genai_transformation.py
+++ b/tests/test_litellm/google_genai/test_google_genai_transformation.py
@@ -366,3 +366,84 @@ def test_transform_generate_content_request_system_instruction_with_tools():
     assert "tools" in result, "tools should be in request body"
     assert result["tools"] == tools
     assert result["model"] == "gemini-3-flash-preview"
+
+
+def test_validate_environment_with_dict_api_key():
+    """
+    Test that validate_environment correctly handles api_key as a dict.
+    
+    This happens when using custom api_base with Gemini - the auth_header
+    is returned as {"x-goog-api-key": "sk-test"} and should be merged into
+    headers instead of being set as a header value.
+    
+    Regression test for: https://github.com/BerriAI/litellm/issues/xxxxx
+    """
+    config = GoogleGenAIConfig()
+    
+    # Simulate the case where auth_header is a dict (custom api_base scenario)
+    auth_header_dict = {"x-goog-api-key": "sk-test-key-123"}
+    
+    result = config.validate_environment(
+        api_key=auth_header_dict,
+        headers=None,
+        model="gemini-2.5-pro",
+        litellm_params={}
+    )
+    
+    # The dict should be merged into headers, not set as a value
+    assert "x-goog-api-key" in result, "x-goog-api-key should be in headers"
+    assert result["x-goog-api-key"] == "sk-test-key-123", "API key should be the string value, not a dict"
+    assert isinstance(result["x-goog-api-key"], str), "Header value should be a string, not a dict"
+    assert "Content-Type" in result, "Content-Type should be in headers"
+    assert result["Content-Type"] == "application/json"
+
+
+def test_validate_environment_with_string_api_key():
+    """
+    Test that validate_environment correctly handles api_key as a string.
+    
+    This is the normal case when using standard Gemini API.
+    """
+    config = GoogleGenAIConfig()
+    
+    # Normal case: api_key is a string
+    api_key_string = "sk-test-key-456"
+    
+    result = config.validate_environment(
+        api_key=api_key_string,
+        headers=None,
+        model="gemini-2.5-pro",
+        litellm_params={}
+    )
+    
+    # The string should be set as the header value
+    assert "x-goog-api-key" in result, "x-goog-api-key should be in headers"
+    assert result["x-goog-api-key"] == "sk-test-key-456", "API key should match input"
+    assert isinstance(result["x-goog-api-key"], str), "Header value should be a string"
+    assert "Content-Type" in result, "Content-Type should be in headers"
+
+
+def test_validate_environment_with_extra_headers():
+    """
+    Test that validate_environment correctly merges extra headers with dict api_key.
+    """
+    config = GoogleGenAIConfig()
+    
+    # Custom api_base scenario with additional headers
+    auth_header_dict = {"x-goog-api-key": "sk-test-key-789"}
+    extra_headers = {"X-Custom-Header": "custom-value"}
+    
+    result = config.validate_environment(
+        api_key=auth_header_dict,
+        headers=extra_headers,
+        model="gemini-2.5-pro",
+        litellm_params={}
+    )
+    
+    # Both the auth dict and extra headers should be merged
+    assert "x-goog-api-key" in result, "x-goog-api-key should be in headers"
+    assert result["x-goog-api-key"] == "sk-test-key-789", "API key should be correctly set"
+    assert isinstance(result["x-goog-api-key"], str), "Header value should be a string"
+    assert "X-Custom-Header" in result, "Extra headers should be merged"
+    assert result["X-Custom-Header"] == "custom-value"
+    assert "Content-Type" in result