From 63a15c5ffdcc17e26e79b51dd8e20d10b8da493a Mon Sep 17 00:00:00 2001
From: Chesars <cesarponce19544@gmail.com>
Date: Fri, 13 Mar 2026 23:54:42 -0300
Subject: [PATCH] fix(vertex_ai): pass extra_headers to vertex embedding
 handler

Forward `extra_headers` to `vertex_embedding.embedding()` in the
embedding call path, consistent with how bedrock, gemini, and vertex
batch embeddings already handle it.

Fixes #21020
---
 litellm/main.py                               |  1 +
 .../llms/vertex_ai/test_vertex.py             | 25 +++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/litellm/main.py b/litellm/main.py
index 794d30ed384..26f59f4b1c4 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -5206,6 +5206,7 @@ def embedding(  # noqa: PLR0915
                     api_key=api_key,
                     api_base=api_base,
                     client=client,
+                    extra_headers=headers,
                 )
         elif custom_llm_provider == "oobabooga":
             response = oobabooga.embedding(
diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex.py b/tests/test_litellm/llms/vertex_ai/test_vertex.py
index bd12100a88f..f5fb47172c2 100644
--- a/tests/test_litellm/llms/vertex_ai/test_vertex.py
+++ b/tests/test_litellm/llms/vertex_ai/test_vertex.py
@@ -1459,6 +1459,31 @@ def test_aaavertex_embeddings_distances(
         text_embedding = text_response.data[0].embedding
 
 
+def test_vertex_embedding_passes_extra_headers():
+    """
+    Test that extra_headers are forwarded to the vertex embedding handler.
+
+    Relevant issue: https://github.com/BerriAI/litellm/issues/21020
+    """
+    from unittest.mock import MagicMock, patch
+
+    mock_response = MagicMock()
+
+    with patch.object(
+        litellm.main.vertex_embedding, "embedding", return_value=mock_response
+    ) as mock_embedding:
+        litellm.embedding(
+            model="vertex_ai/textembedding-gecko",
+            input=["hello"],
+            extra_headers={"X-Custom-Header": "test-value"},
+        )
+
+        mock_embedding.assert_called_once()
+        call_kwargs = mock_embedding.call_args.kwargs
+        assert "extra_headers" in call_kwargs
+        assert call_kwargs["extra_headers"]["X-Custom-Header"] == "test-value"
+
+
 def test_vertex_parallel_tool_calls_true():
     """
     Test that parallel_tool_calls = True sets the correct tool_config.