fix: authorize provider MCP tools on native Responses WebSockets

docs/plans/provider-executed-mcp.md

@@ -1,6 +1,12 @@
 # Provider-executed MCP registration
 
-Status: registration, permission integration, and native request authorization are approved and pushed through `18cd30435c`. The provider registration UI and generated API schema refresh are verified and approved for commit and push
+Status: registration, permission integration, native request authorization, and registration UI are approved and pushed through `e8fe7788f9`
+
+## URL identity decision, 2026-11-02
+
+Keep the pushed implementation. Provider registrations have unique canonical URLs, enforced by the partial unique index in the original execution-mode migration and config conflict validation. Responses `server_url` and Messages `url` resolve the registration, then its server ID is used for existing object-permission and tool checks. Client labels do not need to match registration names. Existing permission inputs may use IDs, names, or aliases, which are resolved to server IDs for authorization
+
+The uncommitted label-based request lookup, duplicate-URL support, migration removal, and related UI/test changes were discarded. No additional migration or migration-recovery changes are included. The known missing-history adoption gap remains a separate concern, normal migration-based installs and upgrades execute the unique-index SQL
 
 ## Implementation plan
 
@@ -380,3 +386,12 @@ The revision passed 966 focused tests plus 132 catalog/manager tests. The live p
 The latest revision restores the gateway credential-update flow and its existing transaction rules. The shared writer accepts a preparation callback only for provider updates, keeping their stored-row validation under the lock. Neutral connection fields remain accepted without enabling upstream activity. No strict five-field rewrite or scoped-retrieval implementation was added
 
 Validation after this trim passed 967 focused tests, followed by 606 affected tests after the final adjustment. All 24 live proxy/PostgreSQL checks passed, including neutral fields on create/update, unchanged records after rejection, deleted-record errors, and concurrent edits. Provider connections remained zero. Type comparison found no new diagnostics and strict Ruff rule counts did not increase. The revision was reviewed in VS Code and approved for commit and push
+
+
+## WebSocket authorization follow-up
+
+The managed WebSocket path already calls aresponses per turn and inherits gateway discovery and execution permissions. The native OpenAI/Azure relay forwards frames directly and does not expand litellm_proxy gateway references
+
+The local follow-up reuses authorize_provider_mcp_request inside existing turn admission. Native frames are checked after deployment defaults and masking, before token reservation and upstream forwarding. The existing URL registration lookup, server-ID permissions and tool ceilings remain unchanged. Rejections use the existing turn-failure logging and keep the socket open
+
+All 12 new regression cases failed on the previous implementation. The final rerun passes 604 executed tests and 13 real OpenAI WebSocket checks, including the pre-fix bypass, permitted MCP execution, denied frames with zero upstream forwarding, allow/deny/allow on one connection, deployment defaults and overrides, and existing managed gateway execution/denial. Three legacy async endpoint tests were not awaited and are excluded from the evidence. The follow-up is approved for commit and push. Unicode URL normalization remains a separate unresolved finding

litellm/responses/streaming_iterator.py

@@ -2159,13 +2159,26 @@ ResponsesWsTurnAdmission: TypeAlias = Callable[[], Awaitable[Mapping[str, object
 
 
 async def _admit_frame_turn(
-    turn_admission: ResponsesWsTurnAdmission | None, raw_message: str
+    turn_admission: ResponsesWsTurnAdmission | None,
+    raw_message: str,
+    *,
+    user_api_key_auth: UserAPIKeyAuth | None = None,
 ) -> _AdmittedTurn | _RefusedTurn:
     frame: Final = _json_object_or_none(raw_message)
-    if turn_admission is None or frame is None or frame.get("type") != "response.create":
+    if frame is None or frame.get("type") != "response.create":
         return _AdmittedTurn()
     try:
-        return _AdmittedTurn(await turn_admission())
+        if user_api_key_auth is not None:
+            from litellm.responses.mcp.request_context import authorize_provider_mcp_request
+
+            nested: Final = frame.get("response")
+            body: Final = nested if _is_json_object(nested) else frame
+            await authorize_provider_mcp_request(
+                body,
+                {"litellm_metadata": {"user_api_key_auth": user_api_key_auth}},
+                provider_bound=True,
+            )
+        return _AdmittedTurn(await turn_admission()) if turn_admission is not None else _AdmittedTurn()
     except RateLimitError as refusal:
         return _RefusedTurn(refusal)
     except Exception as e:
@@ -2180,7 +2193,14 @@ def _turn_refusal_frame(refusal: RateLimitError) -> str:
 
 
 def _turn_admission_error_frame(failure: Exception) -> str:
-    return json.dumps({"type": "error", "status": 500, "error": {"type": "server_error", "message": str(failure)}})
+    _, error_type, error_code = _error_event_fields(failure)
+    return json.dumps(
+        {
+            "type": "error",
+            "status": _status_code_for_error_fields(error_type, error_code),
+            "error": {"type": error_type or "server_error", "message": str(failure)},
+        }
+    )
 
 
 def _turn_rejection_frame(rejection: Exception) -> str:
@@ -2595,12 +2615,12 @@ class ResponsesWebSocketStreaming:
         guarded: Final = await self._guard_frame(message)
         if guarded is None:
             return
-        admitted: Final = await self._admit_or_reject_turn(guarded, message, turn_admission)
-        if admitted is None:
-            return
         if guarded.pii_tokens:
             self._remember_pii_tokens(guarded.pii_tokens)
         masked: Final = await self._mask_response_create(guarded.frame)
+        admitted: Final = await self._admit_or_reject_turn(guarded, masked, turn_admission)
+        if admitted is None:
+            return
         self._record_client_frame(masked, guarded.guardrail_information, admitted.reservation)
         await self.backend_ws.send(masked)
 
@@ -3003,7 +3023,7 @@ class ResponsesWebSocketStreaming:
     async def _admit_or_reject_turn(
         self, guarded: GuardedResponsesWsFrame, message: str, turn_admission: ResponsesWsTurnAdmission | None
     ) -> _AdmittedTurn | None:
-        verdict: Final = await _admit_frame_turn(turn_admission, message)
+        verdict: Final = await _admit_frame_turn(turn_admission, message, user_api_key_auth=self.user_api_key_dict)
         match verdict:
             case _AdmittedTurn():
                 return verdict

tests/unit/responses/test_responses_websocket_all_providers.py

@@ -55,6 +55,7 @@ if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging
     from litellm.responses.streaming_iterator import ResponsesWsTurnAdmission
     from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck
+    from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
 
 
 class TestResponsesAPIWebSocketSupport:
@@ -3062,6 +3063,120 @@ def _make_streaming(**kwargs):
     return ResponsesWebSocketStreaming(**kwargs)
 
 
+@pytest.fixture
+async def provider_mcp_auth(monkeypatch: pytest.MonkeyPatch) -> "UserAPIKeyAuth":
+    from litellm.proxy import proxy_server
+    from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth
+    from litellm.proxy.mcp_server import mcp_server_manager
+
+    manager: Final = mcp_server_manager.MCPServerManager()
+    monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
+    monkeypatch.setattr(proxy_server, "prisma_client", None)
+    await manager.load_servers_from_config(
+        {
+            "docs": {
+                "server_id": "native-id",
+                "execution_mode": "provider",
+                "url": "https://native.example/mcp",
+                "allowed_tools": ["search", "fetch"],
+            }
+        }
+    )
+    return UserAPIKeyAuth(
+        object_permission=LiteLLM_ObjectPermissionTable(
+            object_permission_id="key-grant",
+            mcp_servers=["native-id"],
+            mcp_tool_permissions={"native-id": ["search"]},
+        )
+    )
+
+
+def _native_mcp_tools(tool: str = "search", url: str = "https://native.example/mcp") -> list[dict[str, object]]:
+    return [{"type": "mcp", "server_label": "client_label", "server_url": url, "allowed_tools": [tool]}]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("nested", [False, True])
+async def test_native_websocket_checks_mcp_permissions_on_first_and_later_frames(
+    provider_mcp_auth: "UserAPIKeyAuth", nested: bool
+) -> None:
+    allowed: Final = {"input": "search docs", "tools": _native_mcp_tools()}
+    denied: Final = {**allowed, "tools": _native_mcp_tools("fetch")}
+    frames: Final = tuple(
+        json.dumps({"type": "response.create", **({"response": body} if nested else body)})
+        for body in (denied, allowed, denied, allowed)
+    )
+    client: Final = _ClientSocket(*frames[1:])
+    backend: Final = AsyncMock()
+    handler: Final = _make_streaming(
+        websocket=client, backend_ws=backend, user_api_key_dict=provider_mcp_auth, first_message=frames[0]
+    )
+
+    await handler.client_to_backend()
+
+    assert [call.args[0] for call in backend.send.await_args_list] == [frames[1], frames[3]]
+    assert [(frame["status"], frame["error"]["type"]) for frame in client.sent] == [
+        (403, "tool_access_denied"), (403, "tool_access_denied")
+    ]
+    assert client.close_code is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("source", ["client", "deployment_default", "deployment_override"])
+@pytest.mark.parametrize("denial", ["url", "tool", "unbounded"])
+async def test_native_websocket_authorizes_effective_mcp_tools_before_turn_reservation(
+    provider_mcp_auth: "UserAPIKeyAuth", source: str, denial: str
+) -> None:
+    from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults
+
+    selected: Final = {
+        "url": _native_mcp_tools(url="https://unregistered.example/mcp"),
+        "tool": _native_mcp_tools("fetch"),
+        "unbounded": [{"type": "mcp", "server_label": "client_label", "server_url": "https://native.example/mcp"}],
+    }[denial]
+    body: Final = {
+        "type": "response.create",
+        "input": "search docs",
+        **({"tools": selected} if source == "client" else {}),
+        **({"tools": _native_mcp_tools()} if source == "deployment_override" else {}),
+    }
+    client: Final = _ClientSocket()
+    backend: Final = AsyncMock()
+    reserve: Final = AsyncMock(return_value={})
+    handler: Final = _make_streaming(
+        websocket=client,
+        backend_ws=backend,
+        user_api_key_dict=provider_mcp_auth,
+        request_defaults=ResponsesWebSocketRequestDefaults(
+            fill_missing={"tools": selected} if source == "deployment_default" else {},
+            overrides={"tools": selected} if source == "deployment_override" else {},
+        ),
+    )
+
+    await handler._forward_guarded_frame(json.dumps(body), reserve)
+
+    backend.send.assert_not_awaited()
+    reserve.assert_not_awaited()
+    assert [(frame["status"], frame["error"]["type"]) for frame in client.sent] == [(403, "tool_access_denied")]
+
+
+@pytest.mark.asyncio
+async def test_native_websocket_rechecks_registration_on_an_open_connection(provider_mcp_auth: "UserAPIKeyAuth") -> None:
+    from litellm.proxy.mcp_server import mcp_server_manager
+
+    frame: Final = json.dumps({"type": "response.create", "input": "search docs", "tools": _native_mcp_tools()})
+    client: Final = _ClientSocket()
+    backend: Final = AsyncMock()
+    handler: Final = _make_streaming(websocket=client, backend_ws=backend, user_api_key_dict=provider_mcp_auth)
+
+    await handler._forward_guarded_frame(frame)
+    await mcp_server_manager.global_mcp_server_manager.load_servers_from_config({})
+    await handler._forward_guarded_frame(frame)
+
+    backend.send.assert_awaited_once_with(frame)
+    assert [(event["status"], event["error"]["type"]) for event in client.sent] == [(403, "tool_access_denied")]
+
+
 class TestNativeWebSocketGuardrailMasking:
     """Exercises the input/output PII masking hooks on ResponsesWebSocketStreaming."""