test(mcp): cover scoped gateway JWT admission and rejection

tests/e2e/coverage_registry/mcp.yaml

@@ -127,3 +127,93 @@
   assertions: [toolset_scoped]
   source: "user_api_key_auth_mcp.py:2137"
   rationale: "A key granted a toolset lists exactly the toolset's tools: the rest of the server's catalog stays hidden and every stored name resolves"
+
+- id: mcp.list_tools.bearer.jwt_valid_scoped
+  module: mcp
+  tier: P0
+  operation: list_tools
+  auth_family: bearer
+  assertions: [valid_scoped]
+  source: "test_mcp_jwt_auth_e2e.py"
+  rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials"
+
+- id: mcp.call_tool.bearer.jwt_valid_scoped
+  module: mcp
+  tier: P0
+  operation: call_tool
+  auth_family: bearer
+  assertions: [valid_scoped]
+  source: "test_mcp_jwt_auth_e2e.py"
+  rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials"
+
+- id: mcp.list_tools.bearer.jwt_invalid_signature_denied
+  module: mcp
+  tier: P0
+  operation: list_tools
+  auth_family: bearer
+  assertions: [invalid_signature_denied]
+  source: "test_mcp_jwt_auth_e2e.py"
+  rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials"
+
+- id: mcp.call_tool.bearer.jwt_invalid_signature_denied
+  module: mcp
+  tier: P0
+  operation: call_tool
+  auth_family: bearer
+  assertions: [invalid_signature_denied]
+  source: "test_mcp_jwt_auth_e2e.py"
+  rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials"
+
+- id: mcp.list_tools.bearer.jwt_expired_denied
+  module: mcp
+  tier: P0
+  operation: list_tools
+  auth_family: bearer
+  assertions: [expired_denied]
+  source: "test_mcp_jwt_auth_e2e.py"
+  rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials"
+
+- id: mcp.call_tool.bearer.jwt_expired_denied
+  module: mcp
+  tier: P0
+  operation: call_tool
+  auth_family: bearer
+  assertions: [expired_denied]
+  source: "test_mcp_jwt_auth_e2e.py"
+  rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials"
+
+- id: mcp.list_tools.bearer.jwt_inactive_user_denied
+  module: mcp
+  tier: P0
+  operation: list_tools
+  auth_family: bearer
+  assertions: [inactive_user_denied]
+  source: "test_mcp_jwt_auth_e2e.py"
+  rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials"
+
+- id: mcp.call_tool.bearer.jwt_inactive_user_denied
+  module: mcp
+  tier: P0
+  operation: call_tool
+  auth_family: bearer
+  assertions: [inactive_user_denied]
+  source: "test_mcp_jwt_auth_e2e.py"
+  rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials"
+
+- id: mcp.list_tools.bearer.jwt_header_precedence
+  module: mcp
+  tier: P0
+  operation: list_tools
+  auth_family: bearer
+  assertions: [header_precedence]
+  source: "test_mcp_jwt_auth_e2e.py"
+  rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials"
+
+- id: mcp.call_tool.bearer.jwt_header_precedence
+  module: mcp
+  tier: P0
+  operation: call_tool
+  auth_family: bearer
+  assertions: [header_precedence]
+  source: "test_mcp_jwt_auth_e2e.py"
+  rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials"

tests/e2e/mcp/mcp_client.py

@@ -17,12 +17,11 @@ from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import Literal
 
-from pydantic import BaseModel, ConfigDict, Field, RootModel
-
 from e2e_config import settle_propagation
-from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap
+from e2e_http import AuthHeaders, Headers, NoBody, Result, Success, UnknownApiError, unwrap
 from models import KeyGenerateBody, McpServerListResponse, McpServerRow, ObjectPermission
 from proxy_client import ProxyClient
+from pydantic import BaseModel, ConfigDict, Field, RootModel
 
 McpToolArg = str | int | float | bool | list[str] | dict[str, str]
 McpToolArguments = Mapping[str, McpToolArg]
@@ -254,10 +253,10 @@ class McpClient:
             )
         )
 
-    def list_tools(self, key: str) -> Result[McpToolsListResponse]:
+    def list_tools(self, key: str, *, headers: AuthHeaders | None = None) -> Result[McpToolsListResponse]:
         return self.proxy.transport.get(
             "/mcp-rest/tools/list",
-            headers=ApiKeyHeaders(x_litellm_api_key=key),
+            headers=headers if headers is not None else ApiKeyHeaders(x_litellm_api_key=key),
             params=NoBody(),
             response_type=McpToolsListResponse,
         )
@@ -393,10 +392,11 @@ class McpClient:
         server_id: str,
         name: str,
         arguments: McpToolArguments,
+        headers: AuthHeaders | None = None,
     ) -> Result[McpCallToolResponse]:
         return self.proxy.transport.post(
             "/mcp-rest/tools/call",
-            headers=ApiKeyHeaders(x_litellm_api_key=key),
+            headers=headers if headers is not None else ApiKeyHeaders(x_litellm_api_key=key),
             json=McpCallToolBody(
                 name=name, arguments=dict(arguments), server_id=server_id
             ),

tests/e2e/mcp/test_mcp_jwt_auth_e2e.py

@@ -0,0 +1,178 @@
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+from typing import Final, Literal
+
+import pytest
+from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp
+from e2e_config import DD_SEARCH_FROM, unique_marker
+from e2e_http import AuthHeaders, UnauthorizedError, unwrap
+from idp import SHORT_LIVED_CLIENT_ID, Identity, Keycloak, token_claims
+from lifecycle import ResourceManager
+from management.management_client import build_client as build_management_client
+from mcp_client import McpClient
+from models import KeyGenerateBody, ObjectPermission, UserScimMetadata, UserUpdateBody
+
+pytestmark = pytest.mark.e2e
+
+
+@dataclass(frozen=True, slots=True)
+class GrantedMcp:
+    server_id: str
+    other_server_id: str
+    tool: str
+    token: str = field(repr=False)
+
+
+def _allowed(client: McpClient, access: GrantedMcp, headers: AuthHeaders) -> None:
+    listing: Final = unwrap(client.list_tools(access.token, headers=headers))
+    assert listing.tool_names_for_server(access.server_id) == frozenset((access.tool,))
+    assert listing.tool_names_for_server(access.other_server_id) == frozenset()
+    result: Final = unwrap(
+        client.call_tool(
+            access.token,
+            server_id=access.server_id,
+            name=access.tool,
+            arguments={"query": f"lit4506-{unique_marker()}", "from": DD_SEARCH_FROM, "to": "now", "max_tokens": 1000},
+            headers=headers,
+        )
+    )
+    assert result.is_error is False, result
+    assert result.all_text.strip(), "The successful control must return a real tool result"
+
+
+def _denied(client: McpClient, access: GrantedMcp, headers: AuthHeaders, reason: str) -> None:
+    listing: Final = client.list_tools(access.token, headers=headers)
+    assert isinstance(listing, UnauthorizedError), listing
+    assert reason in listing.body.lower(), listing.body
+    execution: Final = client.call_tool(
+        access.token,
+        server_id=access.server_id,
+        name=access.tool,
+        arguments={"query": f"lit4506-{unique_marker()}", "from": DD_SEARCH_FROM, "to": "now", "max_tokens": 1000},
+        headers=headers,
+    )
+    assert isinstance(execution, UnauthorizedError), execution
+    assert reason in execution.body.lower(), execution.body
+
+
+@pytest.fixture
+def granted_mcp(client: McpClient, resources: ResourceManager, jwt_identity: Identity, idp: Keycloak) -> GrantedMcp:
+    server_id: Final = register_datadog_mcp(client, resources)
+    other_server_id: Final = register_datadog_mcp(client, resources)
+    client.await_registered(server_id)
+    client.await_registered(other_server_id)
+    management: Final = build_management_client(client.proxy)
+    management.update_user(
+        UserUpdateBody(
+            user_id=jwt_identity.user_id,
+            user_role="internal_user",
+            object_permission=ObjectPermission(mcp_servers=[server_id]),
+        )
+    )
+    stored: Final = management.user_info(jwt_identity.user_id)
+    assert stored.user_info.user_id == jwt_identity.user_id
+    assert stored.user_info.user_role == "internal_user"
+    token: Final = idp.access_token(jwt_identity)
+    assert token_claims(token).sub == jwt_identity.user_id
+    tool: Final = client.await_tool(token, server_id, SEARCH_LOGS_TOOL)
+    return GrantedMcp(server_id, other_server_id, tool, token)
+
+
+class TestMcpGatewayJwt:
+    @pytest.mark.covers("mcp.list_tools.bearer.jwt_valid_scoped", "mcp.call_tool.bearer.jwt_valid_scoped")
+    @pytest.mark.parametrize("header", ("authorization", "x_litellm_api_key"))
+    def test_valid_scoped_jwt_lists_and_calls(
+        self, client: McpClient, granted_mcp: GrantedMcp, header: Literal["authorization", "x_litellm_api_key"]
+    ) -> None:
+        headers: Final = (
+            AuthHeaders(authorization=f"Bearer {granted_mcp.token}")
+            if header == "authorization"
+            else AuthHeaders.model_validate({"x-litellm-api-key": granted_mcp.token})
+        )
+        _allowed(client, granted_mcp, headers)
+
+    @pytest.mark.covers(
+        "mcp.list_tools.bearer.jwt_invalid_signature_denied", "mcp.call_tool.bearer.jwt_invalid_signature_denied"
+    )
+    def test_tampered_signature_denies_both_operations(self, client: McpClient, granted_mcp: GrantedMcp) -> None:
+        valid: Final = AuthHeaders(authorization=f"Bearer {granted_mcp.token}")
+        _allowed(client, granted_mcp, valid)
+        header, payload, signature = granted_mcp.token.split(".")
+        flipped: Final = "A" if signature[10] != "A" else "B"
+        tampered: Final = f"{header}.{payload}.{signature[:10]}{flipped}{signature[11:]}"
+        _denied(client, granted_mcp, AuthHeaders(authorization=f"Bearer {tampered}"), "signature")
+        _allowed(client, granted_mcp, valid)
+
+    @pytest.mark.covers("mcp.list_tools.bearer.jwt_expired_denied", "mcp.call_tool.bearer.jwt_expired_denied")
+    def test_expired_signed_jwt_denies_both_operations(
+        self, client: McpClient, granted_mcp: GrantedMcp, jwt_identity: Identity, idp: Keycloak
+    ) -> None:
+        valid: Final = AuthHeaders(authorization=f"Bearer {granted_mcp.token}")
+        _allowed(client, granted_mcp, valid)
+        expiring: Final = idp.access_token(jwt_identity, client_id=SHORT_LIVED_CLIENT_ID)
+        claims: Final = token_claims(expiring)
+        assert claims.sub == jwt_identity.user_id
+        delay: Final = claims.exp - time.time() + 1
+        assert delay <= 5, "Short-lived token configuration or IdP clock drifted"
+        time.sleep(max(0, delay))
+        _denied(client, granted_mcp, AuthHeaders(authorization=f"Bearer {expiring}"), "expired")
+        _allowed(client, granted_mcp, valid)
+
+    @pytest.mark.covers(
+        "mcp.list_tools.bearer.jwt_inactive_user_denied", "mcp.call_tool.bearer.jwt_inactive_user_denied"
+    )
+    def test_deactivated_user_cannot_reuse_warm_jwt(
+        self, client: McpClient, granted_mcp: GrantedMcp, jwt_identity: Identity
+    ) -> None:
+        headers: Final = AuthHeaders(authorization=f"Bearer {granted_mcp.token}")
+        _allowed(client, granted_mcp, headers)
+        management: Final = build_management_client(client.proxy)
+        management.update_user(
+            UserUpdateBody(
+                user_id=jwt_identity.user_id, user_role="internal_user", metadata=UserScimMetadata(scim_active=False)
+            )
+        )
+        stored: Final = management.user_info(jwt_identity.user_id)
+        assert stored.user_info.metadata is not None and stored.user_info.metadata.scim_active is False
+        _denied(client, granted_mcp, headers, "deactivated")
+        management.update_user(
+            UserUpdateBody(
+                user_id=jwt_identity.user_id, user_role="internal_user", metadata=UserScimMetadata(scim_active=True)
+            )
+        )
+        _allowed(client, granted_mcp, headers)
+
+    @pytest.mark.covers("mcp.list_tools.bearer.jwt_header_precedence", "mcp.call_tool.bearer.jwt_header_precedence")
+    def test_explicit_gateway_header_wins_without_fallback(
+        self, client: McpClient, granted_mcp: GrantedMcp, resources: ResourceManager
+    ) -> None:
+        _allowed(
+            client,
+            granted_mcp,
+            AuthHeaders.model_validate(
+                {"x-litellm-api-key": granted_mcp.token, "authorization": "Bearer invalid-secondary-token"}
+            ),
+        )
+        _denied(
+            client,
+            granted_mcp,
+            AuthHeaders.model_validate(
+                {"x-litellm-api-key": "sk-invalid-primary-token", "authorization": f"Bearer {granted_mcp.token}"}
+            ),
+            "key",
+        )
+        sibling_key: Final = client.proxy.generate_key(
+            KeyGenerateBody(object_permission=ObjectPermission(mcp_servers=[granted_mcp.other_server_id]))
+        )
+        resources.defer(lambda: client.proxy.delete_key(sibling_key))
+        sibling_tool: Final = client.await_tool(sibling_key, granted_mcp.other_server_id, SEARCH_LOGS_TOOL)
+        sibling: Final = GrantedMcp(granted_mcp.other_server_id, granted_mcp.server_id, sibling_tool, sibling_key)
+        _allowed(
+            client,
+            sibling,
+            AuthHeaders.model_validate(
+                {"x-litellm-api-key": sibling_key, "authorization": f"Bearer {granted_mcp.token}"}
+            ),
+        )

tests/e2e/models.py

@@ -1229,9 +1229,15 @@ class UserNewResponse(BaseModel):
     user_id: str
 
 
+class UserScimMetadata(BaseModel):
+    scim_active: bool | None = None
+
+
 class UserUpdateBody(BaseModel):
     user_id: str
     user_role: UserRole
+    object_permission: ObjectPermission | None = None
+    metadata: UserScimMetadata | None = None
 
 
 class UserInfoParams(BaseModel):
@@ -1242,6 +1248,7 @@ class UserData(BaseModel):
     user_id: str | None = None
     user_email: str | None = None
     user_role: str | None = None
+    metadata: UserScimMetadata | None = None
 
 
 class UserInfoResponse(BaseModel):