fix(image_edit): read vertex_project/location from litellm_params in Imagen get_complete_url

VertexAIImagenImageEditConfig.get_complete_url was resolving vertex_project
and vertex_location only from env vars and global settings, ignoring
litellm_params. Users supplying project/location exclusively via YAML
config would get a ValueError or wrong URL even after auth headers were fixed.

Mirrors the pattern already used by VertexAIGeminiImageEditConfig and
image_generation counterpart (safe_get_vertex_ai_project/location).

Also fixes api_key type hint in MockImageEditConfig (str -> Optional[str])
and adds a test covering get_complete_url credential resolution.

Made-with: Cursor

litellm/llms/azure/image_edit/transformation.py

@@ -14,6 +14,8 @@ class AzureImageEditConfig(OpenAIImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         api_key = (
             api_key

litellm/llms/azure_ai/image_edit/flux2_transformation.py

@@ -65,6 +65,8 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         """
         Validate Azure AI Foundry environment and set up authentication

litellm/llms/azure_ai/image_edit/transformation.py

@@ -25,6 +25,8 @@ class AzureFoundryFluxImageEditConfig(OpenAIImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         """
         Validate Azure AI Foundry environment and set up authentication

litellm/llms/base_llm/image_edit/transformation.py

@@ -67,6 +67,8 @@ class BaseImageEditConfig(ABC):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         return {}
 

litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py

@@ -483,6 +483,8 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         if headers is None:
             headers = {}

litellm/llms/bedrock/image_edit/stability_transformation.py

@@ -372,6 +372,8 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         """
         Validate environment for Bedrock Stability image edit.

litellm/llms/black_forest_labs/image_edit/transformation.py

@@ -123,6 +123,8 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         """
         Validate environment and set up headers for Black Forest Labs.

litellm/llms/custom_httpx/llm_http_handler.py

@@ -5146,6 +5146,8 @@ class BaseLLMHTTPHandler:
             api_key=litellm_params.api_key,
             headers=image_edit_optional_request_params.get("extra_headers", {}) or {},
             model=model,
+            litellm_params=dict(litellm_params),
+            api_base=litellm_params.api_base,
         )
 
         if extra_headers:
@@ -5242,6 +5244,8 @@ class BaseLLMHTTPHandler:
             api_key=litellm_params.api_key,
             headers=image_edit_optional_request_params.get("extra_headers", {}) or {},
             model=model,
+            litellm_params=dict(litellm_params),
+            api_base=litellm_params.api_base,
         )
 
         if extra_headers:

litellm/llms/gemini/image_edit/transformation.py

@@ -54,6 +54,8 @@ class GeminiImageEditConfig(BaseImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         final_api_key: Optional[str] = api_key or get_secret_str("GEMINI_API_KEY")
         if not final_api_key:

litellm/llms/litellm_proxy/image_edit/transformation.py

@@ -8,7 +8,12 @@ class LiteLLMProxyImageEditConfig(OpenAIImageEditConfig):
     """Configuration for image edit requests routed through LiteLLM Proxy."""
 
     def validate_environment(
-        self, headers: dict, model: str, api_key: Optional[str] = None
+        self,
+        headers: dict,
+        model: str,
+        api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         api_key = api_key or get_secret_str("LITELLM_PROXY_API_KEY")
         headers.update({"Authorization": f"Bearer {api_key}"})

litellm/llms/openai/image_edit/transformation.py

@@ -165,6 +165,8 @@ class OpenAIImageEditConfig(BaseImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         api_key = (
             api_key

litellm/llms/openrouter/image_edit/transformation.py

@@ -116,6 +116,8 @@ class OpenRouterImageEditConfig(BaseImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         api_key = api_key or litellm.api_key or get_secret_str("OPENROUTER_API_KEY")
         if not api_key:

litellm/llms/recraft/image_edit/transformation.py

@@ -81,6 +81,8 @@ class RecraftImageEditConfig(BaseImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         final_api_key: Optional[str] = api_key or get_secret_str("RECRAFT_API_KEY")
         if not final_api_key:

litellm/llms/stability/image_edit/transformations.py

@@ -149,6 +149,8 @@ class StabilityImageEditConfig(BaseImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         """
         Validate environment and set up headers for Stability AI.

litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py

@@ -103,10 +103,24 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         headers = headers or {}
-        vertex_project = self._resolve_vertex_project()
+        litellm_params = litellm_params or {}
-        vertex_credentials = self._resolve_vertex_credentials()
+
+        _api_base = litellm_params.get("api_base") or api_base
+        if _api_base is not None:
+            return headers
+
+        vertex_project = (
+            self.safe_get_vertex_ai_project(litellm_params)
+            or self._resolve_vertex_project()
+        )
+        vertex_credentials = (
+            self.safe_get_vertex_ai_credentials(litellm_params)
+            or self._resolve_vertex_credentials()
+        )
         access_token, _ = self._ensure_access_token(
             credentials=vertex_credentials,
             project_id=vertex_project,
@@ -123,8 +137,14 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
         """
         Get the complete URL for Vertex AI Imagen predict API
         """
-        vertex_project = self._resolve_vertex_project()
+        vertex_project = (
-        vertex_location = self._resolve_vertex_location()
+            self.safe_get_vertex_ai_project(litellm_params)
+            or self._resolve_vertex_project()
+        )
+        vertex_location = (
+            self.safe_get_vertex_ai_location(litellm_params)
+            or self._resolve_vertex_location()
+        )
 
         if not vertex_project or not vertex_location:
             raise ValueError(

litellm/proxy/_experimental/mcp_server/server.py

@@ -2594,13 +2594,19 @@ if MCP_AVAILABLE:
                     server_name, client_ip=_client_ip
                 )
                 if server and server.auth_type == MCPAuth.oauth2 and not oauth2_headers:
-                    # For servers that store per-user tokens server-side, skip the
+                    # For per-user OAuth servers, only skip the pre-emptive 401 when
-                    # pre-emptive 401 — the call_tool / list_tools dispatch will look
+                    # a stored token actually exists for this user+server pair.
-                    # up the stored token from Redis / DB and only fail at the MCP
+                    # If no stored token exists, fail fast with 401 so clients can
-                    # protocol level if none is found, giving the client a proper
+                    # kick off PKCE/interactive OAuth flow immediately.
-                    # tool-execution error rather than an HTTP 401.
                     if server.needs_user_oauth_token:
-                        continue
+                        stored_oauth_headers = (
+                            await _get_user_oauth_extra_headers_from_db(
+                                server=server,
+                                user_api_key_auth=user_api_key_auth,
+                            )
+                        )
+                        if stored_oauth_headers:
+                            continue
 
                     request = StarletteRequest(scope)
                     base_url = get_request_base_url(request)

tests/test_litellm/images/test_image_edit_utils.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -28,7 +28,12 @@ class MockImageEditConfig(BaseImageEditConfig):
         return "https://example.com/api"
 
     def validate_environment(
-        self, headers: dict, model: str, api_key: str = None
+        self,
+        headers: dict,
+        model: str,
+        api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         return headers
 
@@ -261,3 +266,141 @@ class TestImageEditCustomPricing:
     def test_custom_pricing_not_detected_without_model_info(self):
         litellm_params = {"litellm_call_id": "test-call-id"}
         assert use_custom_pricing_for_model(litellm_params) is False
+
+
+class TestImageEditHandlerCredentialsForwarding:
+    """
+    Regression tests for Vertex AI image_edit credentials bug.
+
+    image_edit handler must forward litellm_params to validate_environment,
+    so that credentials passed via YAML config (vertex_ai_project,
+    vertex_ai_credentials, etc.) reach the auth layer instead of falling
+    through to Application Default Credentials.
+    """
+
+    def test_vertex_gemini_image_edit_reads_credentials_from_litellm_params(self):
+        """
+        VertexAIGeminiImageEditConfig.validate_environment should read
+        vertex_ai_project/vertex_ai_credentials from litellm_params first.
+        """
+        from litellm.llms.vertex_ai.image_edit.vertex_gemini_transformation import (
+            VertexAIGeminiImageEditConfig,
+        )
+
+        config = VertexAIGeminiImageEditConfig()
+
+        litellm_params = {
+            "vertex_ai_project": "test-project-from-params",
+            "vertex_ai_credentials": "/path/to/creds.json",
+        }
+
+        with patch.object(
+            config, "_ensure_access_token", return_value=("token", "project")
+        ) as mock_ensure:
+            config.validate_environment(
+                headers={},
+                model="test-model",
+                litellm_params=litellm_params,
+            )
+
+            mock_ensure.assert_called_once()
+            call_kwargs = mock_ensure.call_args[1]
+
+            assert call_kwargs["credentials"] == "/path/to/creds.json"
+            assert call_kwargs["project_id"] == "test-project-from-params"
+
+    def test_vertex_imagen_image_edit_reads_credentials_from_litellm_params(self):
+        """
+        VertexAIImagenImageEditConfig.validate_environment should read
+        vertex_ai_project/vertex_ai_credentials from litellm_params first.
+        """
+        from litellm.llms.vertex_ai.image_edit.vertex_imagen_transformation import (
+            VertexAIImagenImageEditConfig,
+        )
+
+        config = VertexAIImagenImageEditConfig()
+
+        litellm_params = {
+            "vertex_ai_project": "test-project-from-params",
+            "vertex_ai_credentials": "/path/to/creds.json",
+        }
+
+        with patch.object(
+            config, "_ensure_access_token", return_value=("token", "project")
+        ) as mock_ensure:
+            config.validate_environment(
+                headers={},
+                model="test-model",
+                litellm_params=litellm_params,
+            )
+
+            mock_ensure.assert_called_once()
+            call_kwargs = mock_ensure.call_args[1]
+
+            assert call_kwargs["credentials"] == "/path/to/creds.json"
+            assert call_kwargs["project_id"] == "test-project-from-params"
+
+    def test_vertex_imagen_get_complete_url_reads_project_and_location_from_litellm_params(
+        self,
+    ):
+        """
+        VertexAIImagenImageEditConfig.get_complete_url should read
+        vertex_ai_project and vertex_ai_location from litellm_params,
+        not only from env vars / global settings.
+        """
+        from litellm.llms.vertex_ai.image_edit.vertex_imagen_transformation import (
+            VertexAIImagenImageEditConfig,
+        )
+
+        config = VertexAIImagenImageEditConfig()
+
+        litellm_params = {
+            "vertex_ai_project": "param-project",
+            "vertex_ai_location": "us-east1",
+        }
+
+        url = config.get_complete_url(
+            model="vertex_ai/imagegeneration@002",
+            api_base=None,
+            litellm_params=litellm_params,
+        )
+
+        assert "param-project" in url
+        assert "us-east1" in url
+
+    def test_validate_environment_signature_includes_litellm_params(self):
+        """
+        All image_edit config validate_environment methods should accept
+        litellm_params to allow credentials to be forwarded from the handler.
+        """
+        import inspect
+
+        from litellm.llms.vertex_ai.image_edit.vertex_gemini_transformation import (
+            VertexAIGeminiImageEditConfig,
+        )
+        from litellm.llms.vertex_ai.image_edit.vertex_imagen_transformation import (
+            VertexAIImagenImageEditConfig,
+        )
+        from litellm.llms.openai.image_edit.transformation import (
+            OpenAIImageEditConfig,
+        )
+
+        configs = [
+            VertexAIGeminiImageEditConfig(),
+            VertexAIImagenImageEditConfig(),
+            OpenAIImageEditConfig(),
+            MockImageEditConfig(),
+        ]
+
+        for config in configs:
+            sig = inspect.signature(config.validate_environment)
+            params = list(sig.parameters.keys())
+
+            assert "litellm_params" in params, (
+                f"{config.__class__.__name__}.validate_environment "
+                "missing litellm_params parameter"
+            )
+            assert "api_base" in params, (
+                f"{config.__class__.__name__}.validate_environment "
+                "missing api_base parameter"
+            )

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py

@@ -10,6 +10,9 @@ they may send a stale `mcp-session-id` header. This test verifies that:
 import pytest
 from unittest.mock import AsyncMock, MagicMock, patch
 
+from fastapi import HTTPException
+from litellm.types.mcp import MCPAuth
+
 
 class TestHandleStaleMcpSession:
     """Unit tests for the _handle_stale_mcp_session helper."""
@@ -438,3 +441,129 @@ async def test_no_mcp_session_id_header_works_normally():
     header_names = [k for k, v in captured_scope.get("headers", [])]
     assert b"mcp-session-id" not in header_names
     assert b"content-type" in header_names
+
+
+@pytest.mark.asyncio
+async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
+    """
+    Per-user OAuth server with no stored token should fail fast with 401 +
+    WWW-Authenticate so PKCE can start.
+    """
+    try:
+        from litellm.proxy._experimental.mcp_server.server import (
+            handle_streamable_http_mcp,
+            session_manager,
+        )
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    scope = {
+        "type": "http",
+        "method": "POST",
+        "path": "/mcp",
+        "headers": [
+            (b"content-type", b"application/json"),
+        ],
+    }
+    receive = AsyncMock()
+    send = AsyncMock()
+    user_auth = MagicMock()
+    user_auth.user_id = "test-user-id"
+    oauth_server = MagicMock()
+    oauth_server.auth_type = MCPAuth.oauth2
+    oauth_server.needs_user_oauth_token = True
+
+    with patch(
+        "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
+        new_callable=AsyncMock,
+        return_value=(user_auth, None, ["repro_oauth_server"], None, None, None),
+    ), patch(
+        "litellm.proxy._experimental.mcp_server.server.set_auth_context",
+    ), patch(
+        "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED",
+        True,
+    ), patch(
+        "litellm.proxy._experimental.mcp_server.server._handle_stale_mcp_session",
+        new_callable=AsyncMock,
+        return_value=False,
+    ), patch(
+        "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
+        new_callable=AsyncMock,
+        return_value=None,
+    ) as mock_get_stored_token, patch(
+        "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
+        return_value=oauth_server,
+    ), patch.object(
+        session_manager,
+        "handle_request",
+        new_callable=AsyncMock,
+    ) as mock_handle_request:
+        with pytest.raises(HTTPException) as exc_info:
+            await handle_streamable_http_mcp(scope, receive, send)
+
+    exc = exc_info.value
+    assert exc.status_code == 401
+    assert "www-authenticate" in exc.headers
+    assert mock_get_stored_token.await_count == 1
+    assert mock_handle_request.await_count == 0
+
+
+@pytest.mark.asyncio
+async def test_per_user_oauth_with_stored_token_skips_preemptive_401():
+    """
+    Per-user OAuth server with an existing stored token should skip pre-emptive
+    401 and continue to session manager request handling.
+    """
+    try:
+        from litellm.proxy._experimental.mcp_server.server import (
+            handle_streamable_http_mcp,
+            session_manager,
+        )
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    scope = {
+        "type": "http",
+        "method": "POST",
+        "path": "/mcp",
+        "headers": [
+            (b"content-type", b"application/json"),
+        ],
+    }
+    receive = AsyncMock()
+    send = AsyncMock()
+    user_auth = MagicMock()
+    user_auth.user_id = "test-user-id"
+    oauth_server = MagicMock()
+    oauth_server.auth_type = MCPAuth.oauth2
+    oauth_server.needs_user_oauth_token = True
+
+    with patch(
+        "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
+        new_callable=AsyncMock,
+        return_value=(user_auth, None, ["repro_oauth_server"], None, None, None),
+    ), patch(
+        "litellm.proxy._experimental.mcp_server.server.set_auth_context",
+    ), patch(
+        "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED",
+        True,
+    ), patch(
+        "litellm.proxy._experimental.mcp_server.server._handle_stale_mcp_session",
+        new_callable=AsyncMock,
+        return_value=False,
+    ), patch(
+        "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
+        new_callable=AsyncMock,
+        return_value={"Authorization": "Bearer cached-token"},
+    ) as mock_get_stored_token, patch(
+        "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
+        return_value=oauth_server,
+    ), patch.object(
+        session_manager,
+        "handle_request",
+        new_callable=AsyncMock,
+    ) as mock_handle_request:
+        await handle_streamable_http_mcp(scope, receive, send)
+
+    assert mock_get_stored_token.await_count == 1
+    assert mock_handle_request.await_count == 1