diff --git a/.gitignore b/.gitignore index 3329f39ca10..a9003663e79 100644 --- a/.gitignore +++ b/.gitignore @@ -107,6 +107,8 @@ STABILIZATION_TODO.md **/test-results **/playwright-report **/*.storageState.json +**/.dd_session.json +**/litellm-e2e-dd-session.json **/coverage test-config diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 680e0dff67b..7e71f2994b2 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -33,7 +33,7 @@ Every test under `tests/e2e/mcp/` must exercise the proxy against the real Datad - Prefer calling real Datadog tools that prove the product path (e.g. `search_datadog_logs` for list/call and permission denials). Seed a unique marker (`e2e-datadog-mcp-*`) in a chat completion when you need a log the tool can find; dual-read with `dd_logs` from conftest when delivery matters - Delete the MCP server (and any keys) through `resources.defer` the same way every other suite tears down - If a new MCP behavior cannot be covered with Datadog's tool surface, say so in the PR and get agreement before inventing another upstream; the default is always Datadog -- The one standing exception is `test_mcp_chat_completion_oauth_e2e.py`. Datadog authenticates with the static `DD-API-KEY` / `DD-APPLICATION-KEY` headers and exposes no authorize/token dance at all, so it cannot exercise gateway-managed OAuth or per-user token seeding in any form. That test drives a real Linear MCP server instead; it is still a real remote upstream, so the no-mock, no-fixture rule above holds unchanged +- Gateway-managed OAuth (authorization_code / per-user token vault) is not exercisable against Datadog's static-header MCP. When adding chat/responses/messages bridge e2e for that path, use a real remote OAuth MCP server (one shared real OAuth MCP for the suite), still no mocks or local fixtures ## Lay the pattern down in a class @@ -119,14 +119,36 @@ mgmt.. mgmt.key.generate.happy_path (surface=ui) ``` -MCPs - endpoint features with the protocol op as the variant +MCPs - endpoint features with the protocol op as the variant. Feature map and full +vocabulary: `coverage_registry/MCP_FEATURES.md`. `auth_family` is how the *client* +authenticates to LiteLLM; upstream server auth is an assertion (e.g. `upstream_sigv4`) ``` mcp... - operation : list_tools | call_tool | list_resources | read_resource | list_prompts | get_prompt + operation : list_tools | call_tool | list_resources | read_resource | list_resource_templates + | list_prompts | get_prompt | tool_search | tool_call_virtual | test_connection + | sampling | elicitation | chat_completion | responses | messages + | auth | permission | oauth | toolset | health auth_family : none | api_key | bearer | oauth - assertion : succeeds | denied_without_permission + assertion : succeeds | denied_without_permission | access_group_scoped | allowed_tools_scoped + | disallowed_tools_blocked | params_filtered | toolset_scoped + | namespaced_multi_server | partial_on_upstream_fault | progress_forwarded + | cost_logged | concurrent_limit | forwards_extra_headers | resolves_user_env_vars + | auto_executes_tools | stream_auto_executes_tools | semantic_filter_narrows + | enforces_model_access + | upstream_static_auth | upstream_oauth2_client_credentials + | upstream_oauth2_authorization_code | delegate_auth_upstream | oauth_passthrough + | upstream_true_passthrough | upstream_token_exchange | upstream_id_jag + | upstream_oauth_delegate | upstream_sigv4 + | transport_http | transport_sse | transport_stdio | openapi_tools + | returns_401_not_500 | returns_429_on_budget | empty_intersection_denies + | dangling_grant_errors | stable_id_survives_repoint | per_server_route + | sees_newly_added_server | trailing_slash_parity | resource_metadata_public_https + | prefix_separator_honored | byok_not_false_unhealthy + | unregistered_server_blocked | acting_user_oauth_token e.g. mcp.call_tool.oauth.succeeds + mcp.messages.api_key.auto_executes_tools + mcp.chat_completion.api_key.auto_executes_tools ``` Reliability & Performance - behavior features (no route; endpoint is exercised_on) diff --git a/tests/e2e/coverage_registry/MCP_FEATURES.md b/tests/e2e/coverage_registry/MCP_FEATURES.md new file mode 100644 index 00000000000..c86cbbd8949 --- /dev/null +++ b/tests/e2e/coverage_registry/MCP_FEATURES.md @@ -0,0 +1,301 @@ +# LiteLLM MCP gateway feature map + +Inventory of MCP features in the codebase, mapped to the e2e coverage registry +(`mcp.yaml`, plus related `mgmt` / `guardrail` cells). Cross-checked against +customer open-MCP issues (anon export, 2026-08). + +Every customer-noticeable behavior should become a registry cell, then a +`@pytest.mark.covers(...)` test. + +**Primary code:** `litellm/proxy/_experimental/mcp_server/` +**Admin API:** `litellm/proxy/management_endpoints/mcp_management_endpoints.py` +**Types:** `litellm/types/mcp.py`, `litellm/types/mcp_server/` +**LLM bridges:** +- Chat completions: `litellm/responses/mcp/chat_completions_handler.py` (`acompletion_with_mcp`) +- Responses API: `litellm/responses/main.py` (`aresponses_api_with_mcp`), `mcp_streaming_iterator.py` +- **Messages API:** `litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py` (`anthropic_messages_with_mcp`) wired from `messages/handler.py` +- Shared expansion: `litellm/responses/mcp/litellm_proxy_mcp_handler.py` +**Guardrails:** `guardrail_hooks/mcp_*`, `guardrail_translation/` + +Registry grammar (MCP module): + +``` +mcp... + auth_family : none | api_key | bearer | oauth # how the *client* authenticates to LiteLLM +``` + +`auth_family` is **not** the upstream MCP server auth type. Upstream auth +(oauth2, sigv4, basic, …) is expressed in the assertion or operation variant. + +--- + +## 1. Protocol operations (MCP server surface) + +Exposed via MCP protocol handlers in `server.py` and REST mirrors in +`rest_endpoints.py` (`/tools/list`, `/tools/call`, test helpers). Aggregate +route is `/mcp/` (trailing slash matters for some clients); per-server and +toolset routes exist via `dynamic_mcp_route` / `toolset_mcp_route`. + +| Feature | Code | Registry cell(s) | E2E today | +| --- | --- | --- | --- | +| list_tools | `server.py` handle_list_tools; REST `/tools/list` | `mcp.list_tools.{api_key,bearer,oauth,none}.succeeds` | api_key often skipped LIT-5052; oauth bridge e2e TBD on a non-Linear real OAuth MCP | +| call_tool | `server.py` mcp_server_tool_call; REST `/tools/call` | `mcp.call_tool.{api_key,bearer,oauth,none}.succeeds` | same | +| list_tools denied without key scope | permission path | `mcp.list_tools.api_key.denied_without_permission` | yes | +| call_tool denied without permission | rest/manager | `mcp.call_tool.api_key.denied_without_permission` | skipped LIT-5052 | +| list_prompts / get_prompt | `server.py` | `mcp.list_prompts` / `mcp.get_prompt` | missing | +| list_resources / read_resource | `server.py` | `mcp.list_resources` / `mcp.read_resource` | missing (**customer pain:** Atlassian / Claude Code ListResources) | +| list_resource_templates | `server.py` | `mcp.list_resource_templates.api_key.succeeds` | missing | +| Progress on call_tool | `server.py` forward_progress | `mcp.call_tool.api_key.progress_forwarded` | missing | +| Multi-server tool namespace / prefix | short_prefix / `MCP_TOOL_PREFIX_SEPARATOR` | `mcp.list_tools.api_key.namespaced_multi_server` | missing | +| Toolset tool FQ names + separator | toolset UI/DB + prefix | `mcp.toolset.api_key.prefix_separator_honored` | missing (customer PR #34559) | +| Virtual tools: mcp_tool_search / mcp_tool_call | `tool_search.py` | `mcp.tool_search` / `mcp.tool_call_virtual` | missing | +| Partial list on upstream failure | `faults/` | `mcp.list_tools.api_key.partial_on_upstream_fault` | missing | +| Per-server MCP path (`/mcp/{alias}/…`) | dynamic route | `mcp.list_tools.api_key.per_server_route` | missing (customer open) | +| Aggregate `/mcp/` picks up new servers without restart | registry reload | `mcp.list_tools.api_key.sees_newly_added_server` | missing (customer open) | +| Trailing slash parity `/mcp` vs `/mcp/` | server routing | `mcp.list_tools.api_key.trailing_slash_parity` | partial historical fix | +| REST test connection / test list tools | `rest_endpoints.py` | `mcp.test_connection.api_key.succeeds` | missing | +| BYOK health without user token | health + is_byok | `mcp.health.api_key.byok_not_false_unhealthy` | missing (LIT-4896/5136) | + +--- + +## 2. Client auth to the gateway (`auth_family`) + +| Family | Meaning | Registry | +| --- | --- | --- | +| `api_key` | Virtual key / master key | P0 list/call + deny | +| `bearer` | Bearer token (incl. OAuth access token as Bearer) | P1 list/call | +| `oauth` | Interactive OAuth2; gateway-managed per-user tokens | P1 + chat/messages bridges | +| `none` | Anonymous / public / `delegate_auth_to_upstream` | P1 list/call | + +| Related behavior | Code | Registry | Customer issue | +| --- | --- | --- | --- | +| Auth fail → 401 + WWW-Authenticate (not 500) | `extract_mcp_auth_context`, ProxyException mapping | `mcp.auth.api_key.returns_401_not_500` | #2 (PR #31011 claimed) | +| Budget exceeded on /mcp → 429 | same path | `mcp.auth.api_key.returns_429_on_budget` | #2 | +| Pre-emptive 401 for unauthenticated OAuth servers | `server.py` | oauth cells | #15 | +| Stateful session auth contexts + cap | `server.py` session managers | P2 | — | + +--- + +## 3. Upstream server auth (`MCPAuth`) + +| `auth_type` / flag | Behavior | Cell | E2E / issues | +| --- | --- | --- | --- | +| `none` | No upstream auth | generic succeeds | partial | +| static (`api_key`, `bearer_token`, `basic`, `authorization`, `token`) | Inject static headers | `upstream_static_auth` | missing dedicated | +| `oauth2` + `authorization_code` | Per-user vault | oauth cells | e2e TBD on a real OAuth MCP; UI flaky (#15) | +| `oauth2` + `client_credentials` | M2M | `upstream_oauth2_client_credentials` | missing | +| `delegate_auth_to_upstream` | Client PKCE with upstream | `delegate_auth_upstream` | missing | +| `oauth_passthrough` | Proxy metadata + 401 challenges | `oauth_passthrough` | unit | +| `true_passthrough` | Forward client Authorization | `upstream_true_passthrough` | missing | +| `oauth2_token_exchange` / OBO | RFC 8693 / entra_obo | `upstream_token_exchange` | missing | +| `oauth2_id_jag` | ID-JAG | `upstream_id_jag` | missing | +| `oauth_delegate` | Bridge / SSO assertion | `upstream_oauth_delegate` | missing | +| `aws_sigv4` | SigV4 (AgentCore etc.) | `upstream_sigv4` | unit; UI onboard flaky (#18) | +| Session handshake (Databricks-style mint) | **not implemented** as full handshake | product gap | #5 open | +| Resource metadata `https` behind TLS terminator | `oauth_utils` X-Forwarded-Proto / public base | `mcp.oauth.api_key.resource_metadata_public_https` | #16 open | + +--- + +## 4. Transports and MCP spec versions + +| Feature | Values | Notes / issues | +| --- | --- | --- | +| Transport | `sse`, `http`, `stdio` | stdio poorly e2e'd; customer #20 | +| Spec version | 2024-11-05, 2025-03-26, 2025-06-18 | | +| OpenAPI → tools | `openapi_to_mcp_generator.py` | backend exists; **UI missing** (#20) | +| gRPC transport | **not supported** | FR #23 | +| Admin transport allowlist | **not supported** | FR #22 | + +--- + +## 5. Permission and multi-tenant safety + +| Feature | Code | Registry | Customer / e2e | +| --- | --- | --- | --- | +| Key `object_permission.mcp_servers` | manager + auth_mcp | deny cells | partial | +| Empty key∩team (or similar) intersection = **deny-all** | `user_api_key_auth_mcp.py` hierarchy | `mcp.list_tools.api_key.empty_intersection_denies` | **#3 open (A2A same class; MCP too)** | +| MCP access groups | manager + e2e | `access_group_scoped` | list yes; call missing | +| Key access group grants beyond team list | product gap / disputed | `mcp.list_tools.api_key.key_access_group_beyond_team` | #4 open | +| Model-level MCP scoping | **not supported** | FR cell if added | #4 FR | +| `allowed_tools` / `disallowed_tools` | MCPServer | allowed/disallowed cells | missing e2e | +| Toolsets → permission expansion | toolset_db | toolset_scoped | missing; prefix bug #6 | +| `allowed_params` | MCPServer | params_filtered | missing | +| `allow_all_keys` | MCPServer | document + test | missing | +| User-only / exclude service accounts | **FR** | `mcp.call_tool.api_key.user_scoped_only` when built | #8 open | +| `require_key_mcp_access_defined` | auth hierarchy | cell | missing e2e | +| End-user MCP permission guardrail | mcp_end_user_permission | guardrail | missing e2e | +| Team-scoped server list | management | mgmt | missing | +| Gateway allowlist: only registered MCPs reachable | product intent | `mcp.call_tool.api_key.unregistered_server_blocked` | #17 FR | +| Invalid server name / tool prefix validation | utils / UI | `mcp.server.api_key.rejects_invalid_prefix` | #17 | + +### Stable identity (config.yaml servers) + +| Feature | Code | Issue | +| --- | --- | --- | +| `server_id` = hash(name, url, transport, auth_type, alias) | `_generate_stable_server_id` | **#1 OPEN:** rename/repoint URL changes ID; grants silently dangle | +| Stable id OR migrate grants + loud dangling-grant error | needed | cell: `mcp.permission.api_key.dangling_grant_errors` + identity migration | + +--- + +## 6. Admin / management API (`/v1/mcp/...`) + +Registry module often `mgmt`, not `mcp`. + +| Feature | Registry | Customer | +| --- | --- | --- | +| Create / update / delete / list / health | mgmt cells | UI onboard #18; health BYOK #12 | +| Non-admin register + approve/reject | mgmt approve/reject | | +| Temporary / session MCP | gap | | +| User OAuth credential CRUD | mgmt + docs gap | **#9:** deposit API exists, docs + acting-user resolution unclear | +| User env vars CRUD | mgmt | | +| Toolsets CRUD | mgmt | prefix #6 | +| Make public / discover / registry | gap | local registry FR #19 | +| Multi-pod: UI save visible on all processes | reliability / pub-sub | **#7** (v1.96 Redis push claimed) | +| Key update must not wipe MCP toolsets/servers | key management | **#10 OPEN** | +| Remove stale deleted MCP IDs from key (UI) | key UI | **#11 LIT-3278** | +| OpenAPI converter in UI | UI | #20 | +| First-party LiteLLM admin MCP server | **FR product** | #21 | + +--- + +## 7. LLM bridges (MCP tools inside model APIs) + +All three share `LiteLLM_Proxy_MCP_Handler` for `litellm_proxy` / `litellm_proxy/mcp/...` tool references: expand tools under the caller's credentials, optional auto-execute loop. + +| Surface | Entry | Registry | E2E | +| --- | --- | --- | --- | +| **Chat completions** `/v1/chat/completions` | `acompletion_with_mcp` via `main.py` | `mcp.chat_completion.{api_key,oauth}.auto_executes_tools` | e2e TBD on a shared real MCP for api_key and oauth | +| **Responses** `/v1/responses` | `aresponses_api_with_mcp` | `mcp.responses.api_key.auto_executes_tools` | missing | +| **Messages** `/v1/messages` | `anthropic_messages_with_mcp` via experimental pass-through handler | `mcp.messages.{api_key,oauth}.auto_executes_tools` | **was missing from map; no e2e** | +| Semantic tool filter on chat tools | `semantic_tool_filter` + hook | `mcp.chat_completion.api_key.semantic_filter_narrows` | missing | +| Provider-native Anthropic `mcp_servers` tool | `AnthropicMcpServerTool` / beta header | **not gateway**; provider-side MCP | separate from litellm_proxy bridge | +| Playground auto-execution | UI → same bridges | exercises chat/messages paths | customer #12 | + +Claude Code / Desktop often use **messages** + OAuth + resources; that is why resources + messages bridge + OAuth metadata are P0 for gateway maturity. + +--- + +## 8. Sampling and elicitation + +| Feature | Flag / code | Registry | +| --- | --- | --- | +| Sampling createMessage → completion | `allow_sampling`, `sampling_handler.py` | `mcp.sampling.api_key.succeeds` | +| Sampling model access + budget | same | `mcp.sampling.api_key.enforces_model_access` | +| Elicitation relay | `allow_elicitation` | `mcp.elicitation.api_key.succeeds` | + +--- + +## 9. Cost, spend, concurrency, headers + +| Feature | Registry | Customer | +| --- | --- | --- | +| Per-server / per-tool cost on call_tool | `mcp.call_tool.api_key.cost_logged` | | +| list_tools spend log | `mcp.list_tools.api_key.cost_logged` | | +| UI MCP Server Activity: tool invocations not tokens | product/UI | **#13 LIT-4897** (tokens always 0) | +| max_concurrent_requests | `concurrent_limit` | | +| extra_headers / static_headers / user env vars | forward / resolve cells | | +| Timeout per server | reliability-adjacent | | + +--- + +## 10. Guardrails on MCP + +| Feature | Registry | E2E | +| --- | --- | --- | +| Content filter `pre_mcp_call` | `guardrail.litellm_content_filter.pre_mcp_call.blocks` | skipped LIT-5052 | +| MCP security hook | `guardrail.mcp_security.pre_call.blocks` | missing | +| MCP JWT signer | gap | unit | +| End-user permission | gap | missing | + +--- + +## 11. Config load / multi-pod / lifecycle + +| Feature | Notes | Issue | +| --- | --- | --- | +| YAML `mcp_servers` load | manager | #1 identity | +| DB store + reload / Redis pub-sub config sync | proxy config sync | #7 multi-process | +| Temporary server Redis cache | management | | +| Discovery well-known OAuth | discoverable + byok | #15/#16 | +| DCR | gateway_dcr_flow | | +| Operator open servers / allow_all union | manager | | + +--- + +## 12. Customer open-issue → product/registry map + +Source: open MCP issues export (26 rows, anon). Status abbreviated. + +| # | Theme | Product status | Registry / e2e action | +| --- | --- | --- | --- | +| 1 | Stable server_id vs rename/URL; dangling grants silent | OPEN | `mcp.permission.api_key.stable_id_survives_url_change` or migration + `dangling_grant_errors` | +| 2 | /mcp auth 500 vs 401/429 | Claimed fixed #31011 | `returns_401_not_500`, `returns_429_on_budget` — regression e2e | +| 3 | Empty permission intersection allow-all (A2A; same class MCP) | OPEN blocking | `empty_intersection_denies` (MCP + a2a suite) | +| 4 | Model-level MCP scope; key AG beyond team | FR / open | FR cells when designed | +| 5 | Central MCP inherit without per-key; Databricks handshake | OPEN | product; handshake not in codebase | +| 6 | Toolset prefix separator | Claimed #34559 | `toolset.prefix_separator_honored` | +| 7 | Multi-process UI save not sticky | Claimed v1.96 Redis | multi-pod e2e after release | +| 8 | User-only MCP (exclude service accounts) | FR | when built | +| 9 | OBO deposit API + acting user for agents | OPEN docs/behavior | mgmt credential + `call_tool` with acting user | +| 10 | Key budget edit wipes MCP toolset | OPEN | mgmt key update regression | +| 11 | Stale MCP IDs on key UI | LIT-3278 | UI/mgmt e2e | +| 12 | BYOK health false unhealthy | LIT-4896/5136 | `byok_not_false_unhealthy` | +| 13 | MCP activity Total Tokens = 0 | LIT-4897 | UI metric | +| 14 | Per-server routes; aggregate tool list stale | OPEN | `per_server_route`, `sees_newly_added_server` | +| 15 | OAuth E2E UI + token forward + clients | OPEN / churn | oauth + messages/chat bridges | +| 16 | Resource metadata http behind TLS | OPEN | `resource_metadata_public_https` | +| 17 | Gateway allowlist unregistered MCPs; invalid names | FR | allowlist + validation cells | +| 18 | UI register / playground tools empty | OPEN/stale | mgmt + playground | +| 19 | Local/dev MCP registry governance | FR | product | +| 20 | OpenAPI UI; stdio/oauth maturity | OPEN / lost deal | openapi + transports | +| 21 | First-party LiteLLM admin MCP | FR | product | +| 22 | Transport allowlist setting | FR | product | +| 23 | gRPC transport | FR low | product | +| 24 | Resources + Atlassian OAuth via Claude Code | OPEN recurring | **resources + messages + oauth** cells | +| 25–26 | Competitive losses (MCP immaturity) | LOST | treat as quality bar, not single cells | + +--- + +## Coverage snapshot + +| Bucket | In product | In registry | Live e2e | +| --- | --- | --- | --- | +| Core list/call api_key | yes | yes | blocked LIT-5052 | +| Deny / access groups | yes | yes | partial | +| bearer / oauth / none | yes | yes | oauth skipif | +| prompts / **resources** | yes | yes | **none** (Atlassian #24) | +| LLM bridges chat | yes | yes | oauth only | +| LLM bridges **responses** | yes | yes | none | +| LLM bridges **messages** | yes | **added** | **none** | +| Auth status codes | claimed | **added** | need regression | +| Stable id / dangling grants | bug | **added** | need | +| Empty intersection deny | bug | **added** | need | +| Toolsets / multi-pod / BYOK health | partial | partial | thin | +| Model-level scope, user-only mode, gRPC, admin MCP | FR / missing | FR notes only | — | + +Collector after registry expansion: **MCPs ~2/N live** until tests land; denominator is intentionally honest. + +--- + +## Priority for new e2e (reliability / fewer regressions) + +1. **P0** Unskip LIT-5052 (Datadog list/call/deny/guardrail). +2. **P0** `/mcp` auth status codes 401/429 regression (#2). +3. **P0** Empty permission intersection denies (#3). +4. **P0** **messages** bridge auto-execute tools (Claude Code path) + **list_resources**. +5. **P0** OAuth resource metadata public https (#16) + token used on tool call. +6. **P1** Stable server_id / dangling grant loud failure (#1). +7. **P1** allowed/disallowed tools; call_tool access groups. +8. **P1** chat + responses auto-exec with api_key. +9. **P1** Per-server route + sees newly added server (#14). +10. **P1** Multi-pod MCP edit propagates (#7). +11. **P1** Key update does not wipe MCP grants (#10). +12. **P2** Toolset prefix, BYOK health UX, cost/invocation metrics, sampling. + +--- + +## Related docs + +- E2E MCP suite: `tests/e2e/CLAUDE.md` (real Datadog for api_key; separate real OAuth MCP for vault/oauth bridges) +- Registry: `tests/e2e/coverage_registry/mcp.yaml` +- MCP internal note: `litellm/proxy/_experimental/mcp_server/CLAUDE.md` diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index ab644118a47..75a46623908 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -1,12 +1,19 @@ -# MCP module. Grounded in litellm/proxy/_experimental/mcp_server/. See tests/e2e/CLAUDE.md for the grammar. +# MCP module. Grounded in litellm/proxy/_experimental/mcp_server/. +# Feature map: tests/e2e/coverage_registry/MCP_FEATURES.md +# Grammar: mcp... +# auth_family = how the *client* authenticates to LiteLLM (none|api_key|bearer|oauth) +# upstream server auth is encoded in the assertion / operation variant, not auth_family +# See tests/e2e/CLAUDE.md for the full grammar. + +# --- Protocol: list_tools / call_tool by client auth family --- - id: mcp.list_tools.api_key.succeeds module: mcp tier: P0 operation: list_tools auth_family: api_key assertions: [succeeds] - source: "server.py:637" - rationale: Core operation; most common auth path; high usage + source: "server.py:handle_list_tools" + rationale: "Core operation; most common auth path; high usage" - id: mcp.list_tools.api_key.access_group_scoped module: mcp tier: P1 @@ -14,55 +21,55 @@ auth_family: api_key assertions: [access_group_scoped] source: "test_mcp_access_group_e2e.py" - rationale: "A key granted an MCP access group sees the tagged server's tools; a key with a different group does not. Access-group-scoped tool selection at key creation" + rationale: "A key granted an MCP access group sees the tagged server's tools; a key with a different group does not" - id: mcp.list_tools.api_key.denied_without_permission module: mcp tier: P0 operation: list_tools auth_family: api_key assertions: [denied_without_permission] - source: "mcp_server_manager.py:1409" - rationale: Permission guard is high blast-radius; multi-tenant safety + source: "mcp_server_manager.py:get_allowed_mcp_servers" + rationale: "Permission guard is high blast-radius; multi-tenant safety" - id: mcp.call_tool.api_key.succeeds module: mcp tier: P0 operation: call_tool auth_family: api_key assertions: [succeeds] - source: "server.py:849" - rationale: Primary operation; customer-critical; high usage + source: "server.py:mcp_server_tool_call" + rationale: "Primary operation; customer-critical; high usage" - id: mcp.call_tool.api_key.denied_without_permission module: mcp tier: P0 operation: call_tool auth_family: api_key assertions: [denied_without_permission] - source: "rest_endpoints.py:305-386" - rationale: Tool-level permission guard; multi-tenant safety + source: "rest_endpoints.py:call_tool_rest_api" + rationale: "Tool-level permission guard; multi-tenant safety" - id: mcp.list_tools.bearer.succeeds module: mcp tier: P1 operation: list_tools auth_family: bearer assertions: [succeeds] - source: "server.py:662" - rationale: OAuth/bearer token flow; upstream delegation + source: "server.py:handle_list_tools" + rationale: "OAuth/bearer token flow; upstream delegation" - id: mcp.call_tool.bearer.succeeds module: mcp tier: P1 operation: call_tool auth_family: bearer assertions: [succeeds] - source: "server.py:886" - rationale: Bearer token forwarding for tool invocation + source: "server.py:mcp_server_tool_call" + rationale: "Bearer token forwarding for tool invocation" - id: mcp.list_tools.oauth.succeeds module: mcp tier: P1 operation: list_tools auth_family: oauth assertions: [succeeds] - source: "rest_endpoints.py:138-188" - rationale: Interactive OAuth2 flow; live token management + source: "rest_endpoints.py; outbound_credentials/per_user_oauth_store" + rationale: "Interactive OAuth2 flow; live token management" - id: mcp.call_tool.oauth.succeeds module: mcp tier: P1 @@ -70,52 +77,511 @@ auth_family: oauth assertions: [succeeds] source: "db.py user_oauth_credential lookup" - rationale: OAuth2 token passthrough; per-user credential storage + rationale: "OAuth2 token passthrough; per-user credential storage" - id: mcp.list_tools.none.succeeds module: mcp tier: P1 operation: list_tools auth_family: none assertions: [succeeds] - source: "mcp_server_manager.py:1485-1492" - rationale: Public/anonymous servers; delegate_auth_to_upstream + source: "mcp_server_manager.py; delegate_auth_to_upstream" + rationale: "Public/anonymous servers; delegate_auth_to_upstream" - id: mcp.call_tool.none.succeeds module: mcp tier: P1 operation: call_tool auth_family: none assertions: [succeeds] - source: "rest_endpoints.py:305-334" - rationale: No upstream auth required; demo servers + source: "rest_endpoints.py" + rationale: "No LiteLLM key required; demo or upstream-delegated servers" + +# --- Protocol: prompts / resources --- - id: mcp.get_prompt.api_key.succeeds module: mcp tier: P1 operation: get_prompt auth_family: api_key assertions: [succeeds] - source: "server.py:1042" - rationale: Prompt op; same auth stack as tools + source: "server.py:get_prompt" + rationale: "Prompt op; same auth stack as tools" - id: mcp.read_resource.api_key.succeeds module: mcp tier: P1 operation: read_resource auth_family: api_key assertions: [succeeds] - source: "server.py:1177" - rationale: Resource op; same permission model as tools + source: "server.py:read_resource" + rationale: "Resource op; same permission model as tools" - id: mcp.list_prompts.api_key.succeeds module: mcp tier: P2 operation: list_prompts auth_family: api_key assertions: [succeeds] - source: "server.py:993" - rationale: Smoke-level; same auth stack as list_tools + source: "server.py:list_prompts" + rationale: "Smoke-level; same auth stack as list_tools" - id: mcp.list_resources.api_key.succeeds module: mcp tier: P2 operation: list_resources auth_family: api_key assertions: [succeeds] - source: "server.py:1089" - rationale: Smoke; rarely used; same auth model as tools + source: "server.py:list_resources" + rationale: "Smoke; rarely used; same auth model as tools" +- id: mcp.list_resource_templates.api_key.succeeds + module: mcp + tier: P2 + operation: list_resource_templates + auth_family: api_key + assertions: [succeeds] + source: "server.py:list_resource_templates" + rationale: "Resource templates; same permission model as list_resources" + +# --- Multi-tenant tool selection (server + key policy) --- +- id: mcp.call_tool.api_key.access_group_scoped + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [access_group_scoped] + source: "mcp_server_manager.py:get_allowed_mcp_servers" + rationale: "Call path must honor the same MCP access groups as list_tools" +- id: mcp.list_tools.api_key.allowed_tools_scoped + module: mcp + tier: P0 + operation: list_tools + auth_family: api_key + assertions: [allowed_tools_scoped] + source: "types/mcp_server/mcp_server_manager.py:MCPServer.allowed_tools" + rationale: "Server allowed_tools must hide disallowed tools from list_tools" +- id: mcp.call_tool.api_key.allowed_tools_scoped + module: mcp + tier: P0 + operation: call_tool + auth_family: api_key + assertions: [allowed_tools_scoped] + source: "mcp_server_manager.py" + rationale: "Calling a tool outside allowed_tools must be denied" +- id: mcp.call_tool.api_key.disallowed_tools_blocked + module: mcp + tier: P0 + operation: call_tool + auth_family: api_key + assertions: [disallowed_tools_blocked] + source: "types/mcp_server/mcp_server_manager.py:MCPServer.disallowed_tools" + rationale: "Explicit disallowed_tools block must win even if listed upstream" +- id: mcp.call_tool.api_key.params_filtered + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [params_filtered] + source: "types/mcp_server/mcp_server_manager.py:MCPServer.allowed_params" + rationale: "allowed_params strips or rejects tool arguments outside the allow-list" +- id: mcp.call_tool.api_key.toolset_scoped + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [toolset_scoped] + source: "toolset_db.py; resolve_toolset_tool_permissions" + rationale: "Key scoped to a toolset only sees/calls that toolset's tools" + +# --- Namespacing, virtual tools, partial failure --- +- id: mcp.list_tools.api_key.namespaced_multi_server + module: mcp + tier: P1 + operation: list_tools + auth_family: api_key + assertions: [namespaced_multi_server] + source: "mcp_server_manager.py:_assign_unique_short_prefix" + rationale: "Two servers' tools remain distinguishable (prefix/short_prefix) without collision" +- id: mcp.tool_search.api_key.succeeds + module: mcp + tier: P1 + operation: tool_search + auth_family: api_key + assertions: [succeeds] + source: "tool_search.py:handle_mcp_tool_search" + rationale: "Virtual mcp_tool_search tool returns ranked tools when mcp_tool_search_enabled" +- id: mcp.tool_call_virtual.api_key.succeeds + module: mcp + tier: P1 + operation: tool_call_virtual + auth_family: api_key + assertions: [succeeds] + source: "tool_search.py:handle_mcp_tool_call" + rationale: "Virtual mcp_tool_call invokes a real upstream tool by name after search" +- id: mcp.list_tools.api_key.partial_on_upstream_fault + module: mcp + tier: P1 + operation: list_tools + auth_family: api_key + assertions: [partial_on_upstream_fault] + source: "faults/list_outcomes.py" + rationale: "One bad upstream must not blank the entire multi-server tool list" +- id: mcp.call_tool.api_key.progress_forwarded + module: mcp + tier: P2 + operation: call_tool + auth_family: api_key + assertions: [progress_forwarded] + source: "server.py:forward_progress" + rationale: "Upstream progress notifications reach the MCP client during long tool calls" + +# --- Upstream auth variants (client still api_key unless noted) --- +- id: mcp.call_tool.api_key.upstream_static_auth + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [upstream_static_auth] + source: "types/mcp.py:MCPAuth api_key|bearer_token|basic|authorization|token" + rationale: "Gateway injects static upstream credentials so call_tool succeeds" +- id: mcp.call_tool.api_key.upstream_oauth2_client_credentials + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [upstream_oauth2_client_credentials] + source: "MCPServer.has_client_credentials; outbound_credentials/client_credentials.py" + rationale: "M2M oauth2_flow=client_credentials fetches and uses a client token" +- id: mcp.call_tool.api_key.upstream_oauth2_authorization_code + module: mcp + tier: P0 + operation: call_tool + auth_family: api_key + assertions: [upstream_oauth2_authorization_code] + source: "outbound_credentials/per_user_oauth_store.py" + rationale: "Per-user vaulted authorization_code token is used on call_tool after consent" +- id: mcp.list_tools.none.delegate_auth_upstream + module: mcp + tier: P1 + operation: list_tools + auth_family: none + assertions: [delegate_auth_upstream] + source: "MCPServer.delegate_auth_to_upstream" + rationale: "oauth2 + delegate_auth_to_upstream allows anonymous PKCE against upstream" +- id: mcp.call_tool.none.oauth_passthrough + module: mcp + tier: P1 + operation: call_tool + auth_family: none + assertions: [oauth_passthrough] + source: "MCPServer.oauth_passthrough" + rationale: "Non-oauth2 servers with oauth_passthrough proxy metadata and 401 challenges" +- id: mcp.call_tool.api_key.upstream_true_passthrough + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [upstream_true_passthrough] + source: "types/mcp.py:MCPAuth.true_passthrough" + rationale: "Client Authorization is forwarded unchanged to upstream" +- id: mcp.call_tool.api_key.upstream_token_exchange + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [upstream_token_exchange] + source: "types/mcp.py:MCPAuth.oauth2_token_exchange; outbound_credentials/token_exchanger.py" + rationale: "OBO / RFC 8693 (or entra_obo) exchange yields a usable upstream token" +- id: mcp.call_tool.api_key.upstream_id_jag + module: mcp + tier: P2 + operation: call_tool + auth_family: api_key + assertions: [upstream_id_jag] + source: "types/mcp.py:MCPAuth.oauth2_id_jag" + rationale: "ID-JAG two-leg exchange issues a resource token for call_tool" +- id: mcp.call_tool.api_key.upstream_oauth_delegate + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [upstream_oauth_delegate] + source: "types/mcp.py:MCPAuth.oauth_delegate; bridge_token_flow.py" + rationale: "SSO / bridge assertion path obtains upstream credentials" +- id: mcp.call_tool.api_key.upstream_sigv4 + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [upstream_sigv4] + source: "types/mcp.py:MCPAuth.aws_sigv4" + rationale: "SigV4-signed upstream MCP calls (e.g. Bedrock AgentCore)" + +# --- Transports / OpenAPI tools --- +- id: mcp.call_tool.api_key.transport_http + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [transport_http] + source: "types/mcp.py:MCPTransport.http" + rationale: "Streamable HTTP transport reaches a real upstream" +- id: mcp.call_tool.api_key.transport_sse + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [transport_sse] + source: "sse_transport.py; types/mcp.py:MCPTransport.sse" + rationale: "SSE transport list+call works end-to-end" +- id: mcp.call_tool.api_key.transport_stdio + module: mcp + tier: P2 + operation: call_tool + auth_family: api_key + assertions: [transport_stdio] + source: "types/mcp.py:MCPTransport.stdio" + rationale: "Local stdio MCP server (command/args/env) is invocable through the gateway" +- id: mcp.call_tool.api_key.openapi_tools + module: mcp + tier: P2 + operation: call_tool + auth_family: api_key + assertions: [openapi_tools] + source: "openapi_to_mcp_generator.py" + rationale: "OpenAPI spec_path servers expose operations as tools and execute them" + +# --- Cost, concurrency, headers, env --- +- id: mcp.call_tool.api_key.cost_logged + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [cost_logged] + source: "cost_calculator.py:MCPCostCalculator" + rationale: "Spend logs record non-zero (or configured) cost for an MCP tool call" +- id: mcp.list_tools.api_key.cost_logged + module: mcp + tier: P2 + operation: list_tools + auth_family: api_key + assertions: [cost_logged] + source: "server.py:_get_tools_from_mcp_servers log_list_tools_to_spendlogs" + rationale: "list_tools can write a spend row when logging is enabled" +- id: mcp.call_tool.api_key.concurrent_limit + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [concurrent_limit] + source: "mcp_server_manager.py:_get_call_semaphore" + rationale: "max_concurrent_requests queues excess calls instead of unbounded fan-out" +- id: mcp.call_tool.api_key.forwards_extra_headers + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [forwards_extra_headers] + source: "MCPServer.extra_headers" + rationale: "Client headers listed in extra_headers are forwarded upstream" +- id: mcp.call_tool.api_key.resolves_user_env_vars + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [resolves_user_env_vars] + source: "MCPServer.env_vars; management store_mcp_user_env_vars" + rationale: "Per-user env vars interpolate into static_headers for the call" +- id: mcp.test_connection.api_key.succeeds + module: mcp + tier: P2 + operation: test_connection + auth_family: api_key + assertions: [succeeds] + source: "rest_endpoints.py:/test/connection" + rationale: "Pre-save connection test against a candidate server config succeeds" + +# --- Sampling / elicitation --- +- id: mcp.sampling.api_key.succeeds + module: mcp + tier: P1 + operation: sampling + auth_family: api_key + assertions: [succeeds] + source: "sampling_handler.py:handle_sampling_create_message" + rationale: "allow_sampling servers can createMessage through the proxy LLM path" +- id: mcp.sampling.api_key.enforces_model_access + module: mcp + tier: P1 + operation: sampling + auth_family: api_key + assertions: [enforces_model_access] + source: "sampling_handler.py:_check_model_access" + rationale: "Sampling refuses models the key cannot use" +- id: mcp.elicitation.api_key.succeeds + module: mcp + tier: P2 + operation: elicitation + auth_family: api_key + assertions: [succeeds] + source: "elicitation_handler.py:handle_elicitation_request" + rationale: "allow_elicitation relays elicitation requests to the downstream client" + +# --- LLM bridges (chat / responses / messages) --- +# Shared expansion: responses/mcp/litellm_proxy_mcp_handler.py (litellm_proxy tool refs) +- id: mcp.chat_completion.api_key.auto_executes_tools + module: mcp + tier: P0 + operation: chat_completion + auth_family: api_key + assertions: [auto_executes_tools] + source: "responses/mcp/chat_completions_handler.py:acompletion_with_mcp" + rationale: "/chat/completions with MCP tools lists, calls, and returns a final answer" +- id: mcp.chat_completion.oauth.auto_executes_tools + module: mcp + tier: P1 + operation: chat_completion + auth_family: oauth + assertions: [auto_executes_tools] + source: "responses/mcp/chat_completions_handler.py:acompletion_with_mcp" + rationale: "Gateway-managed OAuth: chat completion lists and executes tools with vaulted user token (e2e TBD on a real OAuth MCP)" +- id: mcp.chat_completion.api_key.semantic_filter_narrows + module: mcp + tier: P1 + operation: chat_completion + auth_family: api_key + assertions: [semantic_filter_narrows] + source: "semantic_tool_filter.py; hooks/mcp_semantic_filter" + rationale: "Semantic filter reduces the tool set offered to the model without dropping needed tools" +- id: mcp.responses.api_key.auto_executes_tools + module: mcp + tier: P1 + operation: responses + auth_family: api_key + assertions: [auto_executes_tools] + source: "responses/main.py:aresponses_api_with_mcp; responses/mcp/mcp_streaming_iterator.py" + rationale: "Responses API streams list_tools and call events and auto-executes MCP tools" +- id: mcp.messages.api_key.auto_executes_tools + module: mcp + tier: P0 + operation: messages + auth_family: api_key + assertions: [auto_executes_tools] + source: "llms/anthropic/experimental_pass_through/messages/mcp_handler.py:anthropic_messages_with_mcp" + rationale: "/v1/messages expands litellm_proxy MCP tools, runs the tool_use loop, and returns a final Anthropic message (Claude Code path)" +- id: mcp.messages.oauth.auto_executes_tools + module: mcp + tier: P0 + operation: messages + auth_family: oauth + assertions: [auto_executes_tools] + source: "llms/anthropic/experimental_pass_through/messages/mcp_handler.py:anthropic_messages_with_mcp" + rationale: "Messages bridge uses vaulted per-user OAuth token for list+call under Claude Code / Desktop OAuth clients" +- id: mcp.messages.api_key.stream_auto_executes_tools + module: mcp + tier: P1 + operation: messages + auth_family: api_key + assertions: [stream_auto_executes_tools] + source: "llms/anthropic/experimental_pass_through/messages/mcp_handler.py" + rationale: "Streaming /v1/messages still expands MCP tools and completes the tool loop" + +# --- Gateway reliability / multi-tenant edge cases (customer open issues) --- +- id: mcp.auth.api_key.returns_401_not_500 + module: mcp + tier: P0 + operation: auth + auth_family: api_key + assertions: [returns_401_not_500] + source: "server.py:extract_mcp_auth_context; PR #31011" + rationale: "Malformed/invalid/expired key on /mcp returns 401 with WWW-Authenticate, not a flattened 500" +- id: mcp.auth.api_key.returns_429_on_budget + module: mcp + tier: P0 + operation: auth + auth_family: api_key + assertions: [returns_429_on_budget] + source: "server.py:extract_mcp_auth_context" + rationale: "Over-budget key on /mcp returns 429, matching REST chat semantics" +- id: mcp.list_tools.api_key.empty_intersection_denies + module: mcp + tier: P0 + operation: list_tools + auth_family: api_key + assertions: [empty_intersection_denies] + source: "auth/user_api_key_auth_mcp.py permission hierarchy" + rationale: "Empty key∩team (or similar) MCP permission intersection is deny-all, never allow-all" +- id: mcp.permission.api_key.dangling_grant_errors + module: mcp + tier: P0 + operation: permission + auth_family: api_key + assertions: [dangling_grant_errors] + source: "mcp_server_manager.py:_generate_stable_server_id; object_permission.mcp_servers" + rationale: "Grants pointing at missing server_ids surface a loud error (or migrate), not silent empty tool lists" +- id: mcp.permission.api_key.stable_id_survives_repoint + module: mcp + tier: P1 + operation: permission + auth_family: api_key + assertions: [stable_id_survives_repoint] + source: "mcp_server_manager.py:_generate_stable_server_id" + rationale: "Config.yaml rename or URL repoint keeps grants valid (stable id or automatic migration)" +- id: mcp.list_tools.api_key.per_server_route + module: mcp + tier: P1 + operation: list_tools + auth_family: api_key + assertions: [per_server_route] + source: "server.py dynamic_mcp_route" + rationale: "Per-server MCP URL (e.g. /mcp/{alias}/) lists that server's tools, not only aggregate /mcp/" +- id: mcp.list_tools.api_key.sees_newly_added_server + module: mcp + tier: P1 + operation: list_tools + auth_family: api_key + assertions: [sees_newly_added_server] + source: "mcp_server_manager registry reload / config sync" + rationale: "After admin adds a server, aggregate list_tools includes its tools without process restart" +- id: mcp.list_tools.api_key.trailing_slash_parity + module: mcp + tier: P1 + operation: list_tools + auth_family: api_key + assertions: [trailing_slash_parity] + source: "server.py MCP route registration" + rationale: "/mcp and /mcp/ accept the same authenticated list_tools behavior" +- id: mcp.oauth.api_key.resource_metadata_public_https + module: mcp + tier: P0 + operation: oauth + auth_family: api_key + assertions: [resource_metadata_public_https] + source: "oauth_utils.py X-Forwarded-Proto / public base URL" + rationale: "Behind a TLS terminator, protected-resource metadata advertises https:// so Claude accepts OAuth" +- id: mcp.toolset.api_key.prefix_separator_honored + module: mcp + tier: P1 + operation: toolset + auth_family: api_key + assertions: [prefix_separator_honored] + source: "toolset_db.py; MCP_TOOL_PREFIX_SEPARATOR; PR #34559" + rationale: "Tools stored on a toolset use server{separator}tool with the configured separator and resolve on call" +- id: mcp.health.api_key.byok_not_false_unhealthy + module: mcp + tier: P2 + operation: health + auth_family: api_key + assertions: [byok_not_false_unhealthy] + source: "mcp_management_endpoints.py:health_check_servers; is_byok" + rationale: "BYOK/per-user auth servers are not reported as hard Unhealthy solely because health check has no user token" +- id: mcp.call_tool.api_key.unregistered_server_blocked + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [unregistered_server_blocked] + source: "gateway allowlist intent; user_api_key_auth_mcp" + rationale: "Clients cannot invoke arbitrary upstream MCP URLs that are not registered on the gateway" +- id: mcp.call_tool.api_key.acting_user_oauth_token + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [acting_user_oauth_token] + source: "management store_mcp_oauth_user_credential; per_user_oauth_store" + rationale: "Agent key call resolves a documented acting-user identity to that user's vaulted upstream OAuth token" diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 277478eebaf..120cfeceb1a 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -39,8 +39,6 @@ UI_BASE_URL = os.environ.get("E2E_UI_BASE_URL", PROXY_BASE_URL).rstrip("/") CHEAP_ANTHROPIC_MODEL = os.environ.get("E2E_CHEAP_ANTHROPIC_MODEL", "claude-haiku-4-5") CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5") -LINEAR_MCP_URL = os.environ.get("E2E_LINEAR_MCP_URL", "https://mcp.linear.app/mcp") -LINEAR_STORAGE_STATE = os.environ.get("E2E_LINEAR_STORAGE_STATE", "") # Jaeger query API of the compose stack's OTEL trace destination (the `jaeger` # service in docker-compose.yml maps it to host 16686). Trace-completeness tests diff --git a/tests/e2e/mcp/datadog_mcp.py b/tests/e2e/mcp/datadog_mcp.py index d1ea53a0b3b..2843c73c91b 100644 --- a/tests/e2e/mcp/datadog_mcp.py +++ b/tests/e2e/mcp/datadog_mcp.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +from dataclasses import dataclass from e2e_config import datadog_mcp_url, unique_marker from lifecycle import ResourceManager @@ -11,6 +12,12 @@ from mcp_client import McpClient SEARCH_LOGS_TOOL = "search_datadog_logs" +@dataclass(frozen=True, slots=True) +class DatadogMcpServer: + server_id: str + alias: str + + def _dd_api_key() -> str: return os.environ.get("DD_API_KEY", "").strip() @@ -30,25 +37,31 @@ def assert_dd_mcp_creds() -> None: ) +def _dd_static_headers() -> dict[str, str]: + return { + "DD-API-KEY": _dd_api_key(), + "DD-APPLICATION-KEY": _dd_app_key(), + } + + def register_datadog_mcp( client: McpClient, resources: ResourceManager, *, mcp_access_groups: list[str] | None = None, -) -> str: + allowed_tools: list[str] | None = None, + toolsets: str = "core", +) -> DatadogMcpServer: assert_dd_mcp_creds() name = f"e2e_dd_mcp_{unique_marker()}" server_id = client.register_server( server_name=name, alias=name, - url=datadog_mcp_url(toolsets="core"), + url=datadog_mcp_url(toolsets=toolsets), transport="http", - static_headers={ - "DD-API-KEY": _dd_api_key(), - "DD-APPLICATION-KEY": _dd_app_key(), - }, - allowed_tools=[SEARCH_LOGS_TOOL], + static_headers=_dd_static_headers(), + allowed_tools=allowed_tools if allowed_tools is not None else [SEARCH_LOGS_TOOL], mcp_access_groups=mcp_access_groups, ) resources.defer(lambda: client.delete_server(server_id)) - return server_id + return DatadogMcpServer(server_id=server_id, alias=name) diff --git a/tests/e2e/mcp/dd_oauth.py b/tests/e2e/mcp/dd_oauth.py new file mode 100644 index 00000000000..5e14e2692c2 --- /dev/null +++ b/tests/e2e/mcp/dd_oauth.py @@ -0,0 +1,275 @@ +"""Shared helpers for e2e tests that exercise the Datadog MCP server through +gateway-managed OAuth2 (authorization_code + PKCE). + +Datadog's remote MCP server (mcp.datadoghq.com/v1/mcp) supports OAuth2.1 with +mandatory S256 PKCE. The authorize endpoint (app.datadoghq.com) serves an +interactive consent page, so the browser leg is a headless Chromium primed +with a saved Datadog browser session (E2E_DD_STORAGE_STATE). + +The gateway discovers the OAuth endpoints via /.well-known metadata, so the +server is registered with auth_type=oauth2, oauth2_flow=authorization_code +and no explicit authorize/token URLs. The per-user token is stored via +POST /v1/mcp/server/{server_id}/oauth-user-credential. +""" + +from __future__ import annotations + +import base64 +import hashlib +import os +import re +import secrets +import time +import urllib.parse +from dataclasses import dataclass +from typing import TYPE_CHECKING +from urllib.parse import parse_qsl + +import httpx +import pytest +from pydantic import BaseModel, TypeAdapter + +from e2e_config import REQUEST_TIMEOUT +from e2e_http import AuthHeaders, NoBody, unwrap +from models import McpServerCreateBody, McpServerInfo +from proxy_client import ProxyClient + +if TYPE_CHECKING: + from playwright.async_api import Route + +DD_MCP_URL = "https://mcp.datadoghq.com/v1/mcp" +DD_AUTHORIZE_URL = "https://app.datadoghq.com/oauth2/v1/authorize" +DD_TOKEN_URL = "https://app.datadoghq.com/api/v2/oauth2/token" +DD_REGISTER_URL = "https://app.datadoghq.com/api/v2/oauth2/register" +OAUTH_CLIENT_REDIRECT_URI = "http://127.0.0.1:53682/e2e/callback" +BROWSER_CONSENT_TIMEOUT = 60.0 + + +@dataclass(frozen=True, slots=True) +class PkceChallenge: + verifier: str + challenge: str + state: str + + +@dataclass(frozen=True, slots=True) +class DcrClient: + client_id: str + + +@dataclass(frozen=True, slots=True) +class OAuthToken: + access_token: str + token_type: str + refresh_token: str | None = None + + +@dataclass(frozen=True, slots=True) +class DatadogMcpOAuthServer: + server_id: str + alias: str + + +class OAuthCredentialBody(BaseModel): + access_token: str + refresh_token: str | None = None + expires_in: int | None = None + scopes: list[str] | None = None + + +def assert_dd_oauth_env() -> None: + path = os.environ.get("E2E_DD_STORAGE_STATE", "") + if not path or not os.path.exists(path): + pytest.fail( + "Datadog MCP OAuth e2e requires E2E_DD_STORAGE_STATE to point at a " + "saved Datadog browser session. Capture one with mcp/dd_session_capture.py." + ) + + +def _generate_pkce() -> PkceChallenge: + verifier = secrets.token_urlsafe(64) + challenge = base64.urlsafe_b64encode( + hashlib.sha256(verifier.encode()).digest() + ).rstrip(b"=").decode() + state = secrets.token_urlsafe(32) + return PkceChallenge(verifier=verifier, challenge=challenge, state=state) + + +def _dcr_register() -> DcrClient: + resp = httpx.post( + DD_REGISTER_URL, + json={ + "client_name": "e2e-mcp-dd-oauth", + "redirect_uris": [OAUTH_CLIENT_REDIRECT_URI], + "token_endpoint_auth_method": "none", + "grant_types": ["authorization_code", "refresh_token"], + "response_types": ["code"], + }, + timeout=REQUEST_TIMEOUT, + ) + resp.raise_for_status() + client_id = TypeAdapter(str).validate_python(resp.json()["client_id"]) + return DcrClient(client_id=client_id) + + +async def _browser_authorize( + authorize_url: str, storage_state_path: str +) -> tuple[str, str | None]: + import asyncio + + from playwright.async_api import async_playwright + + captured: dict[str, str] = {} + trail: list[str] = [] + + def _note_request(request: object) -> None: + url = getattr(request, "url", "") + if url.startswith(OAUTH_CLIENT_REDIRECT_URI) and "url" not in captured: + captured["url"] = url + + async def _swallow_redirect(route: "Route") -> None: + await route.fulfill(status=200, content_type="text/plain", body="ok") + + async with async_playwright() as playwright: + browser = await playwright.chromium.launch(headless=True) + context = await browser.new_context(storage_state=storage_state_path) + await context.route( + re.compile(re.escape(OAUTH_CLIENT_REDIRECT_URI) + r".*"), _swallow_redirect + ) + page = await context.new_page() + page.on("request", _note_request) + page.on("framenavigated", lambda frame: trail.append(frame.url.split("?", 1)[0])) + await page.goto(authorize_url, wait_until="domcontentloaded") + deadline = time.monotonic() + BROWSER_CONSENT_TIMEOUT + while "url" not in captured and time.monotonic() < deadline: + try: + await page.wait_for_load_state("networkidle", timeout=8000) + except Exception: + pass + if "url" in captured: + break + control = page.locator( + 'button[name="action"][value="approve"], button:has-text("Authorize"), ' + 'button:has-text("Allow"), button:has-text("@"), a:has-text("@")' + ).first + try: + await control.click(timeout=5000) + except Exception: + await asyncio.sleep(0.5) + final_url = page.url + await browser.close() + + landing = captured.get("url") + assert landing is not None, ( + f"consent flow never reached {OAUTH_CLIENT_REDIRECT_URI}; " + f"final={final_url.split('?', 1)[0]!r}; trail={trail[-6:]}" + ) + params = dict(parse_qsl(httpx.URL(landing).query.decode())) + assert "code" in params, f"client redirect_uri carried no code: {landing}" + return params["code"], params.get("state") + + +def _exchange_code( + code: str, pkce: PkceChallenge, client: DcrClient +) -> OAuthToken: + resp = httpx.post( + DD_TOKEN_URL, + data={ + "grant_type": "authorization_code", + "code": code, + "redirect_uri": OAUTH_CLIENT_REDIRECT_URI, + "client_id": client.client_id, + "code_verifier": pkce.verifier, + }, + timeout=REQUEST_TIMEOUT, + ) + resp.raise_for_status() + data = TypeAdapter(dict[str, object]).validate_python(resp.json()) + access_token = str(data["access_token"]) + token_type = str(data.get("token_type", "Bearer")) + refresh_token_raw = data.get("refresh_token") + refresh_token = str(refresh_token_raw) if refresh_token_raw is not None else None + return OAuthToken( + access_token=access_token, + token_type=token_type, + refresh_token=refresh_token, + ) + + +def fetch_dd_oauth_token(storage_state_path: str) -> OAuthToken: + """Drive the full PKCE dance: DCR, authorize (browser), token exchange.""" + import asyncio + + pkce = _generate_pkce() + dcr = _dcr_register() + + params = { + "response_type": "code", + "client_id": dcr.client_id, + "redirect_uri": OAUTH_CLIENT_REDIRECT_URI, + "code_challenge": pkce.challenge, + "code_challenge_method": "S256", + "state": pkce.state, + } + authorize_url = f"{DD_AUTHORIZE_URL}?{urllib.parse.urlencode(params)}" + code, returned_state = asyncio.run( + _browser_authorize(authorize_url, storage_state_path) + ) + assert returned_state == pkce.state, ( + f"OAuth state mismatch: sent {pkce.state!r}, got {returned_state!r}" + ) + return _exchange_code(code, pkce, dcr) + + +def register_dd_oauth_server( + proxy: ProxyClient, alias: str +) -> DatadogMcpOAuthServer: + """Register the Datadog MCP server with auth_type=oauth2, + oauth2_flow=authorization_code. The gateway discovers the authorize/token + endpoints via /.well-known metadata.""" + resp = unwrap( + proxy.transport.post( + "/v1/mcp/server", + headers=proxy.transport.master, + json=McpServerCreateBody( + alias=alias, + url=DD_MCP_URL, + transport="http", + allow_all_keys=False, + auth_type="oauth2", + oauth2_flow="authorization_code", + ), + response_type=McpServerInfo, + ) + ) + return DatadogMcpOAuthServer(server_id=resp.server_id, alias=alias) + + +def store_dd_oauth_token( + proxy: ProxyClient, + server_id: str, + key: str, + token: OAuthToken, +) -> None: + """Store the OAuth access token in the gateway's per-user credential vault + via POST /v1/mcp/server/{server_id}/oauth-user-credential.""" + unwrap( + proxy.transport.post( + f"/v1/mcp/server/{server_id}/oauth-user-credential", + headers=AuthHeaders(authorization=f"Bearer {key}"), + json=OAuthCredentialBody( + access_token=token.access_token, + refresh_token=token.refresh_token, + ), + response_type=NoBody, + ) + ) + + +def delete_dd_oauth_server(proxy: ProxyClient, server_id: str) -> None: + _ = proxy.transport.delete( + f"/v1/mcp/server/{server_id}", + headers=proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) diff --git a/tests/e2e/mcp/dd_session_capture.py b/tests/e2e/mcp/dd_session_capture.py new file mode 100644 index 00000000000..f60ce7f1bcd --- /dev/null +++ b/tests/e2e/mcp/dd_session_capture.py @@ -0,0 +1,59 @@ +"""Capture a Datadog browser session for the MCP OAuth e2e tests. + +Run this once to log into Datadog and save the browser session outside the +repo (default: $TMPDIR/litellm-e2e-dd-session.json): + + uv run python tests/e2e/mcp/dd_session_capture.py + +Then set the env var and run the OAuth tests: + + export E2E_DD_STORAGE_STATE="$TMPDIR/litellm-e2e-dd-session.json" + uv run pytest tests/e2e/mcp/test_mcp_datadog_oauth_e2e.py -v +""" + +from __future__ import annotations + +import os +import tempfile +from pathlib import Path + +DEFAULT_STATE_PATH = Path(tempfile.gettempdir()) / "litellm-e2e-dd-session.json" + + +def _repo_root() -> Path | None: + here = Path(__file__).resolve() + for parent in here.parents: + if (parent / ".git").exists(): + return parent + return None + + +def capture(state_path: Path) -> None: + state_path = state_path.expanduser().resolve() + repo = _repo_root() + if repo is not None and (state_path == repo or repo in state_path.parents): + raise SystemExit( + f"Refusing to write session state under the repo ({state_path}). " + f"Set E2E_DD_STORAGE_STATE to a path outside the tree " + f"(default: {DEFAULT_STATE_PATH})." + ) + state_path.parent.mkdir(parents=True, exist_ok=True) + + from playwright.sync_api import sync_playwright + + with sync_playwright() as p: + browser = p.chromium.launch(headless=False) + context = browser.new_context() + page = context.new_page() + page.goto("https://app.datadoghq.com/account/login") + print("Log into Datadog in the browser, then press Enter here.") + input() + context.storage_state(path=str(state_path)) + browser.close() + print(f"Session saved to {state_path}") + print(f' export E2E_DD_STORAGE_STATE="{state_path}"') + + +if __name__ == "__main__": + state = Path(os.environ.get("E2E_DD_STORAGE_STATE", str(DEFAULT_STATE_PATH))) + capture(state) diff --git a/tests/e2e/mcp/linear_session_capture.py b/tests/e2e/mcp/linear_session_capture.py deleted file mode 100644 index 1c867e17e22..00000000000 --- a/tests/e2e/mcp/linear_session_capture.py +++ /dev/null @@ -1,58 +0,0 @@ -"""One-time helper to capture a logged-in Linear browser session for the -real-Linear MCP e2e test. - -The real-Linear test drives the genuine gateway-managed authorization_code -dance against ``mcp.linear.app``. The only step that cannot be scripted is -Linear's login (magic link / SSO), so a human authenticates once here and the -resulting session (cookies + local storage) is persisted to disk. The e2e test -then loads that session in a headless Playwright context and clicks Approve on -Linear's consent screen every run, with no human and no login automation. - -Run it with the e2e venv, log into Linear in the window that opens, then return -to the terminal and press Enter: - - LITELLM=~/litellm-mcpe2e - "$LITELLM"/.venv/bin/python "$LITELLM"/tests/e2e/mcp/linear_session_capture.py - -The session is written to ``E2E_LINEAR_STORAGE_STATE`` (default -``~/.litellm-e2e/linear_storage_state.json``), outside the repo. It is a -secret: never commit it. Re-run this whenever Linear expires the session. -""" - -from __future__ import annotations - -import os -from pathlib import Path - -from playwright.sync_api import sync_playwright - -DEFAULT_STATE_PATH = Path.home() / ".litellm-e2e" / "linear_storage_state.json" - - -def capture(state_path: Path) -> None: - """Open a headed browser at Linear, wait for the human to log in, then save - the authenticated session to ``state_path``.""" - state_path.parent.mkdir(parents=True, exist_ok=True) - with sync_playwright() as playwright: - browser = playwright.chromium.launch(headless=False) - context = browser.new_context() - page = context.new_page() - page.goto("https://linear.app/login", wait_until="domcontentloaded") - print("\n" + "=" * 72) - print("Log into Linear in the browser window that just opened.") - print("If Linear emails you a magic link, paste the link into THIS window's") - print("address bar (opening it in your default browser won't capture the") - print("session). Google SSO works too as long as you complete it here.") - print("When your Linear workspace has loaded, come back and press Enter.") - print("=" * 72) - input("Press Enter once you are logged in... ") - page.goto("https://mcp.linear.app/", wait_until="domcontentloaded") - context.storage_state(path=str(state_path)) - browser.close() - print(f"\nSaved Linear session to {state_path}") - print("Point the e2e test at it with:") - print(f' export E2E_LINEAR_STORAGE_STATE="{state_path}"') - - -if __name__ == "__main__": - capture(Path(os.environ.get("E2E_LINEAR_STORAGE_STATE", str(DEFAULT_STATE_PATH)))) diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 73453478e5a..873e14667db 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -19,8 +19,24 @@ from dataclasses import dataclass from pydantic import BaseModel, ConfigDict, Field, RootModel from e2e_config import settle_propagation -from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap -from models import KeyGenerateBody, ObjectPermission +from e2e_http import ( + AnthropicHeaders, + Headers, + NoBody, + Result, + StreamingResponse, + Success, + UnknownApiError, + unwrap, +) +from models import ( + AnthropicMessagesBody, + AnthropicMessagesResponse, + ChatBody, + ChatResponse, + KeyGenerateBody, + ObjectPermission, +) from proxy_client import ProxyClient McpToolArg = str | int | float | bool | list[str] | dict[str, str] @@ -145,6 +161,67 @@ class McpCallToolResponse(BaseModel): return "\n".join(part.text for part in self.content if part.text) +class ResponsesMcpTool(BaseModel): + type: str = "mcp" + server_label: str + server_url: str + require_approval: str = "never" + allowed_tools: list[str] | None = None + + +class ResponsesMcpInputMessage(BaseModel): + role: str = "user" + type: str = "message" + content: str + + +class ResponsesMcpBody(BaseModel): + model: str + input: list[ResponsesMcpInputMessage] + instructions: str | None = None + stream: bool = False + tools: list[ResponsesMcpTool] + + +class ResponsesMcpOutputContent(BaseModel): + type: str | None = None + text: str | None = None + + +class ResponsesMcpOutputItem(BaseModel): + model_config = ConfigDict(extra="allow") + type: str | None = None + content: list[ResponsesMcpOutputContent] = [] + name: str | None = None + arguments: str | None = None + + +class ResponsesMcpResult(BaseModel): + model_config = ConfigDict(extra="allow") + id: str | None = None + status: str | None = None + model: str | None = None + output: list[ResponsesMcpOutputItem] = [] + + @property + def text(self) -> str: + return "".join( + content.text or "" for item in self.output for content in item.content + ) + + @property + def mcp_tools_fetched(self) -> ResponsesMcpOutputItem | None: + return next( + (item for item in self.output if item.type == "mcp_tools_fetched"), None + ) + + @property + def tool_execution_results(self) -> ResponsesMcpOutputItem | None: + return next( + (item for item in self.output if item.type == "tool_execution_results"), None + ) + + @dataclass(frozen=True, slots=True) class McpClient: proxy: ProxyClient @@ -198,17 +275,22 @@ class McpClient: ).root def await_registered(self, server_id: str) -> None: - """Poll /v1/mcp/server until `server_id` is listed. Fails at poll_timeout. + """Poll /v1/mcp/server until `server_id` is listed, then wait out the + propagation budget. Fails at poll_timeout. The DB row exists the moment registration returns, but a data-plane pod answers the listing from a registry it refreshes on a periodic DB sync, so a pod that joined the load balancer after the write reports the server as - absent until its first sync. + absent until its first sync. The poll only proves ONE replica has the row; + settle_propagation is what makes the server safe to call on whichever + replica the next request lands on (mirrors ProxyClient.create_model). """ - deadline = time.monotonic() + self.proxy.poll_timeout + written_at = time.monotonic() + deadline = written_at + self.proxy.poll_timeout while True: registered = frozenset(row.server_id for row in self.registered_servers()) if server_id in registered: + settle_propagation(written_at) return if time.monotonic() >= deadline: raise AssertionError( @@ -374,6 +456,43 @@ class McpClient: response_type=McpCallToolResponse, ) + def chat_with_mcp(self, key: str, body: ChatBody) -> Result[ChatResponse]: + return self.proxy.transport.post( + "/chat/completions", + headers=self.proxy.transport.bearer(key), + json=body, + response_type=ChatResponse, + ) + + def responses_with_mcp( + self, key: str, body: ResponsesMcpBody + ) -> Result[ResponsesMcpResult]: + return self.proxy.transport.post( + "/v1/responses", + headers=self.proxy.transport.bearer(key), + json=body, + response_type=ResponsesMcpResult, + ) + + def messages_with_mcp( + self, key: str, body: AnthropicMessagesBody + ) -> Result[AnthropicMessagesResponse]: + return self.proxy.transport.post( + "/v1/messages", + headers=AnthropicHeaders(authorization=self.proxy.transport.bearer(key).authorization), + json=body, + response_type=AnthropicMessagesResponse, + ) + + def messages_stream_with_mcp( + self, key: str, body: AnthropicMessagesBody + ) -> StreamingResponse: + return self.proxy.transport.stream( + "/v1/messages", + headers=AnthropicHeaders(authorization=self.proxy.transport.bearer(key).authorization), + json=body, + ) + def _is_mcp_not_synced( result: Result[McpCallToolResponse], diff --git a/tests/e2e/mcp/oauth_chat_client.py b/tests/e2e/mcp/oauth_chat_client.py deleted file mode 100644 index 2eaf512cfa5..00000000000 --- a/tests/e2e/mcp/oauth_chat_client.py +++ /dev/null @@ -1,271 +0,0 @@ -"""Client for the mcp chat-completion OAuth e2e suite. - -Registers a gateway-managed OAuth (authorization_code) MCP server, seeds the -per-user upstream token by driving the interactive authorize dance with the -official mcp SDK's OAuthClientProvider (the browser leg is a headless Chromium -primed with a human's saved Linear session), then exercises the server through -/chat/completions, where the gateway lists and executes its tools with the -stored per-user token. - -Management routes (/v1/mcp/server CRUD, /chat/completions) go through the -shared ProxyClient transport. The MCP protocol used to seed the token goes through -the mcp SDK, the same library production MCP hosts run. -""" - -from __future__ import annotations - -import asyncio -import re -import time -from dataclasses import dataclass -from typing import TYPE_CHECKING -from urllib.parse import parse_qsl - -import httpx -import pytest -from mcp import ClientSession -from mcp.client.auth import OAuthClientProvider -from mcp.client.streamable_http import streamable_http_client -from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthToken - -from e2e_config import PROXY_BASE_URL, REQUEST_TIMEOUT -from proxy_client import ProxyClient -from e2e_http import AuthHeaders, NoBody, unwrap -from models import ChatBody, ChatResponse, McpServerCreateBody, McpServerInfo - -if TYPE_CHECKING: - from playwright.async_api import Route - -# Where the "browser" lands at the end of the authorize dance. Nothing listens -# here: the route interceptor short-circuits the final redirect and reads the -# code/state off its query string, exactly like a desktop MCP host intercepting -# its loopback redirect. -OAUTH_CLIENT_REDIRECT_URI = "http://127.0.0.1:53682/e2e/callback" -BROWSER_CONSENT_TIMEOUT = 60.0 - - -def _mcp_url(alias: str) -> str: - return f"{PROXY_BASE_URL}/{alias}/mcp" - - -class InMemoryTokenStorage: - """The mcp SDK's TokenStorage protocol, in memory for one dance: the - DCR-registered client and the gateway tokens minted for it.""" - - def __init__(self) -> None: - self._tokens: OAuthToken | None = None - self._client_info: OAuthClientInformationFull | None = None - - async def get_tokens(self) -> OAuthToken | None: - return self._tokens - - async def set_tokens(self, tokens: OAuthToken) -> None: - self._tokens = tokens - - async def get_client_info(self) -> OAuthClientInformationFull | None: - return self._client_info - - async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: - self._client_info = client_info - - -async def _browser_follow_authorize(start_url: str, storage_state_path: str) -> tuple[str, str | None]: - """Play the browser's role for a real upstream whose authorize endpoint - serves an interactive consent page (Linear). A headless Chromium primed - with a human's saved Linear session opens the gateway authorize URL and - clicks through Linear's consent screens (the mcp.linear.app Approve form, - then the linear.app workspace-selection page), riding the rest of the chain - (Linear -> gateway callback -> host redirect_uri). The final hop is - intercepted and short-circuited, since nothing listens there, and its - code/state are read off the query string.""" - from playwright.async_api import async_playwright - - captured: dict[str, str] = {} # mutable-ok: hand-off from the request listener - trail: list[str] = [] # mutable-ok: navigation diagnostics for a failed dance - - def _note_request(request: object) -> None: - url = getattr(request, "url", "") - if url.startswith(OAUTH_CLIENT_REDIRECT_URI) and "url" not in captured: - captured["url"] = url - - async def _swallow_redirect(route: "Route") -> None: - await route.fulfill(status=200, content_type="text/plain", body="ok") - - async with async_playwright() as playwright: - browser = await playwright.chromium.launch(headless=True) - context = await browser.new_context(storage_state=storage_state_path) - await context.route(re.compile(re.escape(OAUTH_CLIENT_REDIRECT_URI) + r".*"), _swallow_redirect) - page = await context.new_page() - page.on("request", _note_request) - page.on("framenavigated", lambda frame: trail.append(frame.url.split("?", 1)[0])) - await page.goto(start_url, wait_until="domcontentloaded") - deadline = time.monotonic() + BROWSER_CONSENT_TIMEOUT - while "url" not in captured and time.monotonic() < deadline: - try: - await page.wait_for_load_state("networkidle", timeout=8000) - except Exception: # noqa: BLE001 - a busy consent page never idles; fall through and try to advance it - pass - if "url" in captured: - break - control = page.locator( - 'button[name="action"][value="approve"], button:has-text("Authorize"), ' - 'button:has-text("Allow"), button:has-text("@"), a:has-text("@")' - ).first - try: - await control.click(timeout=5000) - except Exception: # noqa: BLE001 - nothing to advance yet; loop and re-check - await asyncio.sleep(0.5) - final_url = page.url - await browser.close() - - landing = captured.get("url") - assert landing is not None, ( - f"consent flow never reached {OAUTH_CLIENT_REDIRECT_URI}; " - f"final={final_url.split('?', 1)[0]!r}; trail={trail[-6:]}" - ) - params = dict(parse_qsl(httpx.URL(landing).query.decode())) - assert "code" in params, f"client redirect_uri carried no code: {landing}" - return params["code"], params.get("state") - - -def _oauth_provider(url: str, storage: InMemoryTokenStorage, storage_state_path: str) -> OAuthClientProvider: - """The SDK's real OAuth machinery (RFC 9728/8414 discovery, RFC 7591 DCR, - PKCE, token exchange) with the browser leg driven by Playwright against the - upstream's consent screen.""" - code_holder: dict[str, str | None] = {} # mutable-ok: hand-off between the two SDK callbacks - - async def redirect_handler(authorize_url: str) -> None: - code, state = await _browser_follow_authorize(authorize_url, storage_state_path) - code_holder["code"] = code - code_holder["state"] = state - - async def callback_handler() -> tuple[str, str | None]: - code = code_holder.get("code") - assert code is not None, "callback_handler ran before the authorize redirect completed" - return code, code_holder.get("state") - - return OAuthClientProvider( - server_url=url, - client_metadata=OAuthClientMetadata.model_validate( - { - "redirect_uris": [OAUTH_CLIENT_REDIRECT_URI], - "token_endpoint_auth_method": "none", - "grant_types": ["authorization_code", "refresh_token"], - "response_types": ["code"], - "client_name": "e2e-mcp-host", - } - ), - storage=storage, - redirect_handler=redirect_handler, - callback_handler=callback_handler, - ) - - -class _HeaderInjectingTransport(httpx.AsyncBaseTransport): - """Adds the caller's LiteLLM key header to every outgoing SDK request - (discovery, DCR, token exchange), so the gateway resolves which user to - store the upstream token for from the key on the token exchange, exactly - like a production MCP host configured with a LiteLLM key header.""" - - def __init__(self, inner: httpx.AsyncBaseTransport, headers: dict[str, str]) -> None: - self._inner = inner - self._headers = headers - - async def handle_async_request(self, request: httpx.Request) -> httpx.Response: - for name, value in self._headers.items(): - if name not in request.headers: - request.headers[name] = value - return await self._inner.handle_async_request(request) - - -def _oauth_http_client(headers: dict[str, str], auth: OAuthClientProvider) -> httpx.AsyncClient: - return httpx.AsyncClient( - headers=headers, - auth=auth, - timeout=httpx.Timeout(REQUEST_TIMEOUT), - follow_redirects=True, - transport=_HeaderInjectingTransport(httpx.AsyncHTTPTransport(), headers), - ) - - -async def _seed_via_dance( - url: str, headers: dict[str, str], storage: InMemoryTokenStorage, storage_state_path: str -) -> tuple[str, ...]: - async with _oauth_http_client(headers, _oauth_provider(url, storage, storage_state_path)) as http_client: - async with streamable_http_client(url, http_client=http_client) as (read, write, _): - async with ClientSession(read, write) as session: - await session.initialize() - listed = await session.list_tools() - return tuple(sorted(tool.name for tool in listed.tools)) - - -@dataclass(frozen=True, slots=True) -class ChatMcpClient: - proxy: ProxyClient - - def create_server(self, body: McpServerCreateBody) -> McpServerInfo: - return unwrap( - self.proxy.transport.post( - "/v1/mcp/server", - headers=self.proxy.transport.master, - json=body, - response_type=McpServerInfo, - ) - ) - - def server_info(self, server_id: str) -> McpServerInfo: - return unwrap( - self.proxy.transport.get( - f"/v1/mcp/server/{server_id}", - headers=self.proxy.transport.master, - params=NoBody(), - response_type=McpServerInfo, - ) - ) - - def delete_server(self, server_id: str) -> None: - _ = self.proxy.transport.delete( - f"/v1/mcp/server/{server_id}", - headers=self.proxy.transport.master, - json=NoBody(), - response_type=NoBody, - ) - - def seed_user_token(self, alias: str, key: str, storage_state_path: str) -> tuple[str, ...]: - """Drive the interactive authorize dance for `key`'s user so the gateway - stores their upstream token, retried to the shared deadline since the - just-created server and key propagate asynchronously. The LiteLLM key - rides x-litellm-api-key so the gateway binds the token to that user. - Returns the upstream tool names the dance listed, proof the token works.""" - headers = {"x-litellm-api-key": f"Bearer {key}"} - storage = InMemoryTokenStorage() - deadline = time.monotonic() + self.proxy.poll_timeout - last_error: Exception | None = None - while time.monotonic() < deadline: - try: - return asyncio.run(_seed_via_dance(_mcp_url(alias), headers, storage, storage_state_path)) - except Exception as exc: # noqa: BLE001 - retried to the deadline; the last error surfaces below - last_error = exc - time.sleep(self.proxy.poll_interval) - pytest.fail( - f"authorize dance for {alias!r} never completed within {self.proxy.poll_timeout}s; " - f"last error: {last_error!r}" - ) - - def chat_with_mcp(self, headers: AuthHeaders, body: ChatBody) -> ChatResponse: - """POST /chat/completions carrying the LiteLLM key in `headers` (either - ingress form) with an MCP server attached in `body.tools`. The gateway - resolves the user from the key and lists/executes the server's tools - with that user's stored upstream token.""" - return unwrap( - self.proxy.transport.post( - "/chat/completions", - headers=headers, - json=body, - response_type=ChatResponse, - ) - ) - - -def build_chat_client(proxy: ProxyClient) -> ChatMcpClient: - return ChatMcpClient(proxy=proxy) diff --git a/tests/e2e/mcp/test_mcp_access_group_e2e.py b/tests/e2e/mcp/test_mcp_access_group_e2e.py index f72b75fd43d..08527cc2c61 100644 --- a/tests/e2e/mcp/test_mcp_access_group_e2e.py +++ b/tests/e2e/mcp/test_mcp_access_group_e2e.py @@ -28,7 +28,7 @@ class TestMcpAccessGroupToolSelection: self, client: McpClient, resources: ResourceManager ) -> None: group = f"e2e-mcp-grp-{unique_marker()}" - server_id = register_datadog_mcp(client, resources, mcp_access_groups=[group]) + server_id = register_datadog_mcp(client, resources, mcp_access_groups=[group]).server_id client.await_registered(server_id) granted = client.generate_key( diff --git a/tests/e2e/mcp/test_mcp_auth_errors_e2e.py b/tests/e2e/mcp/test_mcp_auth_errors_e2e.py new file mode 100644 index 00000000000..cc0f21f07b4 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_auth_errors_e2e.py @@ -0,0 +1,35 @@ +"""Live e2e: the MCP REST endpoints return 401 for an invalid key, not a +flattened 500. + +Before PR #31011 the MCP protocol path flattened auth errors to 500; the REST +path shares the same user_api_key_auth dependency as /chat/completions, so a +401 here proves the gateway's auth error mapping is intact. Budget enforcement +(429) is the same dependency as chat and is covered by the quota_management +suite. +""" + +from __future__ import annotations + +import pytest + +from e2e_http import UnauthorizedError +from lifecycle import ResourceManager +from mcp_client import McpClient + +pytestmark = pytest.mark.e2e + +GARBAGE_KEY = "sk-deadbeef-not-a-real-key" + + +class TestMcpAuthStatusCodes: + @pytest.mark.covers("mcp.auth.api_key.returns_401_not_500") + def test_invalid_key_returns_401_not_500( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + result = client.list_tools(GARBAGE_KEY) + assert isinstance(result, UnauthorizedError), ( + f"invalid key on /mcp-rest/tools/list must return 401, not 500; " + f"got: {result}" + ) diff --git a/tests/e2e/mcp/test_mcp_chat_completion_oauth_e2e.py b/tests/e2e/mcp/test_mcp_chat_completion_oauth_e2e.py deleted file mode 100644 index 01e94f7b86f..00000000000 --- a/tests/e2e/mcp/test_mcp_chat_completion_oauth_e2e.py +++ /dev/null @@ -1,197 +0,0 @@ -"""On-demand e2e: a chat completion drives a gateway-managed OAuth MCP server. - -The real end-user flow for MCP over an OAuth server: a user registers a Linear -authorization_code server, authorizes it once so the gateway stores their -upstream token, then sends a normal /chat/completions request with the Linear -MCP attached. The gateway resolves the user from the LiteLLM key, lists Linear's -tools with the stored per-user token, lets the model call one, executes it -upstream with that token, and returns the answer. This is proven against the -real Linear MCP server (mcp.linear.app) and a real Anthropic model, once per -documented ingress header (x-litellm-api-key and Authorization). - -The authorize dance is seeded through the mcp SDK's OAuthClientProvider; the one -step Linear cannot auto-approve is the human consent, so it is captured once out -of band (mcp/linear_session_capture.py) into a saved browser session and a -headless Chromium clicks Approve every run. The test therefore skips unless -E2E_LINEAR_STORAGE_STATE points at that session, so it never runs on the per-PR -CI path; it is a nightly/on-demand real-server smoke test. - -Fail-before-fix: without the stored per-user token the gateway lists no Linear -tools, so mcp_list_tools comes back empty, nothing is called, and the -assertions fail; a served, called, non-empty Linear tool proves the gateway -pulled and used the user's token. -""" - -from __future__ import annotations - -import os - -import pytest - -from e2e_config import CHEAP_ANTHROPIC_MODEL, LINEAR_MCP_URL, LINEAR_STORAGE_STATE, unique_marker -from e2e_http import AuthHeaders -from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, KeyGenerateBody, McpChatTool, McpServerCreateBody, ObjectPermission -from proxy_client import ProxyClient - -pytest.importorskip("mcp", reason="mcp SDK not installed; run `uv sync --inexact --group e2e-dev`") -pytest.importorskip( - "playwright.async_api", - reason="playwright not installed; run `uv pip install playwright` and `playwright install chromium`", -) - -from oauth_chat_client import ChatMcpClient, build_chat_client # noqa: E402 # imports follow the importorskip guards - -pytestmark = [ - pytest.mark.e2e, - pytest.mark.skipif( - not LINEAR_STORAGE_STATE or not os.path.exists(LINEAR_STORAGE_STATE), - reason="set E2E_LINEAR_STORAGE_STATE to a Linear session captured via mcp/linear_session_capture.py", - ), -] - -# Pinned from a live dance during verification (never guessed); the gateway -# prefixes every upstream tool name with the server alias. list_teams is a -# read-only Linear tool that takes no arguments and returns the caller's teams. -LINEAR_READONLY_TOOL = "list_teams" -LINEAR_PROMPT = "Use the list_teams tool to list my Linear teams, then reply with the name of one of them." - - -@pytest.fixture(scope="session") -def chat_client(proxy: ProxyClient) -> ChatMcpClient: - return build_chat_client(proxy) - - -class TestMcpChatCompletionOauth: - """A scoped internal-user key on a real Linear authorization_code server, - used through /chat/completions once per ingress header: the gateway pulls - the user's stored upstream token, lists and executes Linear's tools during - the completion, and returns the answer.""" - - @pytest.mark.covers("mcp.list_tools.oauth.succeeds") - @pytest.mark.covers("mcp.call_tool.oauth.succeeds") - def test_chat_completion_uses_linear_with_x_litellm_api_key_header( - self, chat_client: ChatMcpClient, resources: ResourceManager - ) -> None: - marker = unique_marker() - alias = f"e2elinear{marker}" - created = chat_client.create_server( - McpServerCreateBody( - alias=alias, - url=LINEAR_MCP_URL, - allow_all_keys=False, - auth_type="oauth2", - oauth2_flow="authorization_code", - ) - ) - resources.defer(lambda: chat_client.delete_server(created.server_id)) - - stored = chat_client.server_info(created.server_id) - assert stored.auth_type == "oauth2" - assert stored.oauth2_flow == "authorization_code" - assert stored.allow_all_keys is False - - key = chat_client.proxy.generate_key( - KeyGenerateBody( - user_id="e2e-test-user", - object_permission=ObjectPermission(mcp_servers=[created.server_id]), - ) - ) - resources.defer(lambda: chat_client.proxy.delete_key(key)) - - seeded = chat_client.seed_user_token(alias, key, LINEAR_STORAGE_STATE) - assert f"{alias}-{LINEAR_READONLY_TOOL}" in seeded, ( - f"the authorize dance listed {seeded}, expected it to include {alias}-{LINEAR_READONLY_TOOL}" - ) - - response = chat_client.chat_with_mcp( - AuthHeaders.model_validate({"x-litellm-api-key": f"Bearer {key}"}), - ChatBody( - model=CHEAP_ANTHROPIC_MODEL, - messages=[ChatMessage(role="user", content=LINEAR_PROMPT)], - tools=[ - McpChatTool( - server_url=f"litellm_proxy/mcp/{alias}", - server_label=alias, - require_approval="never", - ) - ], - ), - ) - - message = response.choices[0].message - assert message is not None and message.content, f"completion returned no answer: {response}" - meta = message.provider_specific_fields - assert meta is not None, f"no MCP metadata on the completion: {response}" - listed = {t.function.name for t in (meta.mcp_list_tools or []) if t.function} - assert f"{alias}-{LINEAR_READONLY_TOOL}" in listed, ( - f"the gateway listed {sorted(listed)}, expected the stored token to surface {alias}-{LINEAR_READONLY_TOOL}" - ) - results = [r for r in (meta.mcp_call_results or []) if r.name == f"{alias}-{LINEAR_READONLY_TOOL}"] - assert results and results[0].result, ( - f"Linear tool {alias}-{LINEAR_READONLY_TOOL} was not executed with a result: {meta.mcp_call_results}" - ) - - @pytest.mark.covers("mcp.list_tools.oauth.succeeds") - @pytest.mark.covers("mcp.call_tool.oauth.succeeds") - def test_chat_completion_uses_linear_with_authorization_bearer_header( - self, chat_client: ChatMcpClient, resources: ResourceManager - ) -> None: - marker = unique_marker() - alias = f"e2elinear{marker}" - created = chat_client.create_server( - McpServerCreateBody( - alias=alias, - url=LINEAR_MCP_URL, - allow_all_keys=False, - auth_type="oauth2", - oauth2_flow="authorization_code", - ) - ) - resources.defer(lambda: chat_client.delete_server(created.server_id)) - - stored = chat_client.server_info(created.server_id) - assert stored.auth_type == "oauth2" - assert stored.oauth2_flow == "authorization_code" - assert stored.allow_all_keys is False - - key = chat_client.proxy.generate_key( - KeyGenerateBody( - user_id="e2e-test-user", - object_permission=ObjectPermission(mcp_servers=[created.server_id]), - ) - ) - resources.defer(lambda: chat_client.proxy.delete_key(key)) - - seeded = chat_client.seed_user_token(alias, key, LINEAR_STORAGE_STATE) - assert f"{alias}-{LINEAR_READONLY_TOOL}" in seeded, ( - f"the authorize dance listed {seeded}, expected it to include {alias}-{LINEAR_READONLY_TOOL}" - ) - - response = chat_client.chat_with_mcp( - AuthHeaders.model_validate({"authorization": f"Bearer {key}"}), - ChatBody( - model=CHEAP_ANTHROPIC_MODEL, - messages=[ChatMessage(role="user", content=LINEAR_PROMPT)], - tools=[ - McpChatTool( - server_url=f"litellm_proxy/mcp/{alias}", - server_label=alias, - require_approval="never", - ) - ], - ), - ) - - message = response.choices[0].message - assert message is not None and message.content, f"completion returned no answer: {response}" - meta = message.provider_specific_fields - assert meta is not None, f"no MCP metadata on the completion: {response}" - listed = {t.function.name for t in (meta.mcp_list_tools or []) if t.function} - assert f"{alias}-{LINEAR_READONLY_TOOL}" in listed, ( - f"the gateway listed {sorted(listed)}, expected the stored token to surface {alias}-{LINEAR_READONLY_TOOL}" - ) - results = [r for r in (meta.mcp_call_results or []) if r.name == f"{alias}-{LINEAR_READONLY_TOOL}"] - assert results and results[0].result, ( - f"Linear tool {alias}-{LINEAR_READONLY_TOOL} was not executed with a result: {meta.mcp_call_results}" - ) diff --git a/tests/e2e/mcp/test_mcp_chat_completions_e2e.py b/tests/e2e/mcp/test_mcp_chat_completions_e2e.py new file mode 100644 index 00000000000..946eb4c0ad3 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_chat_completions_e2e.py @@ -0,0 +1,100 @@ +"""Live e2e: /chat/completions expands a gateway-registered MCP server and +auto-executes its tools in one agentic turn. + +Registers the real Datadog remote MCP server, grants a key access to it, then +sends a /chat/completions request whose ``tools`` array carries an +``{type: "mcp", server_url: "litellm_proxy/mcp/"}`` reference. The gateway +lists the server's tools, feeds them to the model, the model calls +search_datadog_logs, the gateway executes the call upstream, and folds the +result back into a follow-up completion. The response must carry +provider_specific_fields.mcp_list_tools (the gateway listed tools), +mcp_tool_calls (the model called one), and mcp_call_results (the gateway +executed it), proving the full bridge loop ran end to end against a real MCP +server and a real LLM. +""" + +from __future__ import annotations + +import pytest + +from datadog_mcp import assert_dd_mcp_creds, register_datadog_mcp +from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from mcp_client import McpClient +from models import ChatBody, ChatMessage, McpChatTool + +pytestmark = pytest.mark.e2e + + +class TestChatCompletionMcpAutoExecute: + @pytest.mark.covers("mcp.chat_completion.api_key.auto_executes_tools") + def test_chat_completion_lists_calls_and_executes_mcp_tools( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + assert_dd_mcp_creds() + dd = register_datadog_mcp(client, resources) + client.await_registered(dd.server_id) + + marker = f"e2e-mcp-chat-nohit-{unique_marker()}" + key = client.generate_key( + user_id=f"e2e-mcp-chat-{unique_marker()}", + mcp_servers=[dd.server_id], + models=[CHEAP_ANTHROPIC_MODEL], + ) + resources.defer(lambda: client.proxy.delete_key(key)) + + body = ChatBody( + model=CHEAP_ANTHROPIC_MODEL, + messages=[ + ChatMessage( + role="user", + content=( + "Use the search_datadog_logs tool to search for logs " + f"with query '{marker}' from now-30m to now with " + "max_tokens 100. After you get results, reply with ok only." + ), + ) + ], + max_tokens=1024, + tools=[ + McpChatTool( + type="mcp", + server_url=f"litellm_proxy/mcp/{dd.alias}", + server_label="datadog", + require_approval="never", + ) + ], + ) + response = unwrap(client.chat_with_mcp(key, body)) + + assert response.choices, f"chat completion returned no choices: {response}" + message = response.choices[0].message + assert message is not None, f"choice had no message: {response}" + + psf = message.provider_specific_fields + assert psf is not None, ( + "provider_specific_fields missing; the gateway did not attach MCP metadata " + f"(mcp_list_tools / mcp_tool_calls / mcp_call_results): {message}" + ) + + assert psf.mcp_list_tools, ( + "mcp_list_tools is empty; the gateway never listed the Datadog server's tools " + "through the chat bridge" + ) + assert psf.mcp_tool_calls, ( + "mcp_tool_calls is empty; the model did not call any MCP tool " + "(it may not have seen the expanded tools)" + ) + assert psf.mcp_call_results, ( + "mcp_call_results is empty; the gateway did not execute the tool call upstream" + ) + + result_text = next( + (r.result for r in psf.mcp_call_results if r.result), None + ) + assert result_text, ( + "mcp_call_results has no result text; the tool call returned nothing" + ) diff --git a/tests/e2e/mcp/test_mcp_datadog_e2e.py b/tests/e2e/mcp/test_mcp_datadog_e2e.py index 138f654272d..03ed26f18d5 100644 --- a/tests/e2e/mcp/test_mcp_datadog_e2e.py +++ b/tests/e2e/mcp/test_mcp_datadog_e2e.py @@ -49,16 +49,6 @@ def _seed_completion(proxy: ProxyClient, *, key: str, marker: str) -> None: class TestDatadogMcpRoundTrip: - @pytest.mark.skip( - reason=( - "LIT-5052: this test sends a `telemetry` argument that Datadog's " - "search_datadog_logs tool now rejects, so every tool call fails validation with " - "'unexpected additional properties [\"telemetry\"]' before the round-trip " - "assertion is reached. `telemetry` was never a documented Datadog parameter; the " - "test relied on the server ignoring unknown properties. Unskip once the argument " - "is dropped." - ) - ) @pytest.mark.covers("mcp.list_tools.api_key.succeeds", "mcp.call_tool.api_key.succeeds") def test_search_logs_finds_seeded_completion( self, @@ -69,13 +59,13 @@ class TestDatadogMcpRoundTrip: assert_dd_mcp_creds() _assert_datadog_logger_active(client.proxy) - server_id = register_datadog_mcp(client, resources) - client.await_registered(server_id) + dd = register_datadog_mcp(client, resources) + client.await_registered(dd.server_id) marker = f"{MARKER_PREFIX}{unique_marker()}" key = client.generate_key( user_id=f"e2e-dd-mcp-{unique_marker()}", - mcp_servers=[server_id], + mcp_servers=[dd.server_id], models=[CHEAP_ANTHROPIC_MODEL], ) resources.defer(lambda: client.proxy.delete_key(key)) @@ -88,19 +78,16 @@ class TestDatadogMcpRoundTrip: "within the poll deadline; MCP search would have nothing to find" ) - tool_name = client.await_tool(key, server_id, SEARCH_LOGS_TOOL) + tool_name = client.await_tool(key, dd.server_id, SEARCH_LOGS_TOOL) call = client.await_call_tool( key, - server_id=server_id, + server_id=dd.server_id, name=tool_name, arguments={ "query": marker, "from": DD_SEARCH_FROM, "to": "now", "max_tokens": 5000, - "telemetry": { - "intent": "e2e assert seeded litellm completion log is searchable via MCP" - }, }, ) assert call.is_error is not True, f"search_datadog_logs errored: {call}" diff --git a/tests/e2e/mcp/test_mcp_datadog_oauth_e2e.py b/tests/e2e/mcp/test_mcp_datadog_oauth_e2e.py new file mode 100644 index 00000000000..57d2f17cd9a --- /dev/null +++ b/tests/e2e/mcp/test_mcp_datadog_oauth_e2e.py @@ -0,0 +1,159 @@ +"""Live e2e: the Datadog MCP server through gateway-managed OAuth2 (PKCE). + +Registers the Datadog MCP server with auth_type=oauth2, +oauth2_flow=authorization_code. The gateway discovers the OAuth endpoints via +/.well-known metadata. A real PKCE authorize dance (DCR, browser consent, +token exchange) produces an access token, which is stored in the gateway's +per-user credential vault. Then the key lists tools, calls one, and drives a +chat completion through the MCP bridge, all using the stored per-user token. + +Requires E2E_DD_STORAGE_STATE pointing at a saved Datadog browser session. +""" + +from __future__ import annotations + +import os + +import pytest + +from dd_oauth import ( + assert_dd_oauth_env, + delete_dd_oauth_server, + fetch_dd_oauth_token, + register_dd_oauth_server, + store_dd_oauth_token, +) +from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from mcp_client import McpClient +from models import ChatBody, ChatMessage, KeyGenerateBody, McpChatTool, ObjectPermission + +pytestmark = [ + pytest.mark.e2e, + pytest.mark.skipif( + not os.environ.get("E2E_DD_STORAGE_STATE"), + reason="set E2E_DD_STORAGE_STATE to a Datadog session captured via mcp/dd_session_capture.py", + ), +] + +pytest.importorskip("mcp", reason="mcp SDK not installed") +pytest.importorskip("playwright.async_api", reason="playwright not installed") + +SEARCH_LOGS_TOOL = "search_datadog_logs" + + +class TestDatadogMcpOAuth: + @pytest.mark.covers( + "mcp.list_tools.oauth.succeeds", + "mcp.call_tool.oauth.succeeds", + ) + def test_oauth_list_and_call_tools( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + assert_dd_oauth_env() + + marker = unique_marker() + alias = f"e2e_dd_oauth_{marker}" + dd = register_dd_oauth_server(client.proxy, alias) + resources.defer(lambda: delete_dd_oauth_server(client.proxy, dd.server_id)) + client.await_registered(dd.server_id) + + key = client.proxy.generate_key( + KeyGenerateBody( + models=[CHEAP_ANTHROPIC_MODEL], + user_id=f"e2e-dd-oauth-{marker}", + object_permission=ObjectPermission(mcp_servers=[dd.server_id]), + ) + ) + resources.defer(lambda: client.proxy.delete_key(key)) + + token = fetch_dd_oauth_token(os.environ["E2E_DD_STORAGE_STATE"]) + store_dd_oauth_token(client.proxy, dd.server_id, key, token) + + tool_name = client.await_tool(key, dd.server_id, SEARCH_LOGS_TOOL) + assert tool_name, f"OAuth token did not surface any tools for server {dd.server_id}" + + call = client.await_call_tool( + key, + server_id=dd.server_id, + name=tool_name, + arguments={ + "query": "service:litellm", + "from": "now-30m", + "to": "now", + "max_tokens": 500, + }, + ) + assert call.is_error is not True, f"search_datadog_logs errored via OAuth: {call}" + assert call.all_text, f"OAuth tool call returned empty text: {call}" + + @pytest.mark.covers("mcp.chat_completion.oauth.auto_executes_tools") + def test_oauth_chat_completion_auto_executes_tools( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + assert_dd_oauth_env() + + marker = unique_marker() + alias = f"e2e_dd_oauth_chat_{marker}" + dd = register_dd_oauth_server(client.proxy, alias) + resources.defer(lambda: delete_dd_oauth_server(client.proxy, dd.server_id)) + client.await_registered(dd.server_id) + + key = client.proxy.generate_key( + KeyGenerateBody( + models=[CHEAP_ANTHROPIC_MODEL], + user_id=f"e2e-dd-oauth-chat-{marker}", + object_permission=ObjectPermission(mcp_servers=[dd.server_id]), + ) + ) + resources.defer(lambda: client.proxy.delete_key(key)) + + token = fetch_dd_oauth_token(os.environ["E2E_DD_STORAGE_STATE"]) + store_dd_oauth_token(client.proxy, dd.server_id, key, token) + + response = unwrap( + client.chat_with_mcp( + key, + ChatBody( + model=CHEAP_ANTHROPIC_MODEL, + messages=[ + ChatMessage( + role="user", + content=( + "Use the search_datadog_logs tool to search for logs " + f"with query 'e2e-mcp-oauth-nohit-{marker}' from now-30m to now with " + "max_tokens 100. After you get results, reply with ok only." + ), + ) + ], + max_tokens=1024, + tools=[ + McpChatTool( + type="mcp", + server_url=f"litellm_proxy/mcp/{dd.alias}", + server_label="datadog", + require_approval="never", + ) + ], + ), + ) + ) + + assert response.choices, f"chat completion returned no choices: {response}" + message = response.choices[0].message + assert message is not None, f"choice had no message: {response}" + psf = message.provider_specific_fields + assert psf is not None, ( + f"provider_specific_fields missing; the gateway did not attach MCP metadata: {message}" + ) + assert psf.mcp_list_tools, ( + f"mcp_list_tools is empty; the gateway never listed tools via the stored OAuth token: {psf}" + ) + assert psf.mcp_call_results, ( + f"mcp_call_results is empty; the gateway did not execute any tool via OAuth: {psf}" + ) diff --git a/tests/e2e/mcp/test_mcp_guardrail_e2e.py b/tests/e2e/mcp/test_mcp_guardrail_e2e.py deleted file mode 100644 index 60a349ddc5e..00000000000 --- a/tests/e2e/mcp/test_mcp_guardrail_e2e.py +++ /dev/null @@ -1,178 +0,0 @@ -"""Live e2e: a guardrail on the MCP tool-call path blocks banned content in the -tool arguments before the call reaches the upstream MCP server. - -A general litellm_content_filter guardrail is configured with mode=pre_mcp_call -(the event type the proxy rewrites pre_call to for a call_mcp_tool) and default_on -(per-key/request guardrail selection is dropped from the synthetic MCP request the -hook sees, so default_on is how it attaches to tools/call). The banned keyword is -unique per run, so default_on only ever intercepts this test's own banned call. - -Against the real Datadog MCP server, calling search_datadog_logs with the banned -keyword in the query is blocked with HTTP 400 attributed to the pre_mcp_call hook, -and the tool never runs; the same guardrail lets a clean query through to Datadog. -This is the enforced half (the block) plus the pass-through half in one spec. -""" - -from __future__ import annotations - -import time -from collections.abc import Callable - -import pytest - -from datadog_mcp import SEARCH_LOGS_TOOL, assert_dd_mcp_creds, register_datadog_mcp -from e2e_config import DD_SEARCH_FROM, unique_marker -from e2e_http import Result, Success, UnknownApiError -from lifecycle import ResourceManager -from mcp_client import McpCallToolResponse, McpClient, McpToolArguments - -pytestmark = pytest.mark.e2e - -# Stage runs several data-plane pods behind the shared key, and each picks up a -# newly registered guardrail or MCP server only on its next periodic DB sync (~30s in -# proxy_server.py). Every pod is guaranteed to have refreshed only once a full sync -# interval has elapsed since the later of those two writes; before then a banned call -# routed to a lagging pod passes through as legitimate in-flight propagation, not a leak. -FULL_SYNC_SECONDS = 40.0 -POST_SYNC_VERIFICATION_CALLS = 4 - - -def _poll_until_blocked( - search: Callable[[str], Result[McpCallToolResponse]], banned_keyword: str, client: McpClient -) -> Result[McpCallToolResponse]: - """Retry a banned tool call until the guardrail blocks it (400) or the deadline - passes, returning the last result. Absorbs the control-plane -> data-plane - guardrail-sync delay so the check waits for enforcement instead of racing it.""" - deadline = time.monotonic() + client.proxy.poll_timeout - last: Result[McpCallToolResponse] = search(f"tell me about {banned_keyword}") - while time.monotonic() < deadline: - if isinstance(last, UnknownApiError) and last.status_code == 400: - return last - time.sleep(client.proxy.poll_interval) - last = search(f"tell me about {banned_keyword}") - return last - - -def _pod_lacks_mcp_server(result: Result[McpCallToolResponse]) -> bool: - """True when the pod that served the call answered as though the MCP server or its - tool does not exist (500 "Tool ... not found"), i.e. its MCP registry has not synced - yet and the request never reached the guardrail at all.""" - if not isinstance(result, UnknownApiError) or result.status_code != 500: - return False - body = result.body.lower() - return "not found" in body and ("tool" in body or "server" in body) - - -def _search_on_synced_pod( - search: Callable[[str], Result[McpCallToolResponse]], query: str, client: McpClient -) -> Result[McpCallToolResponse]: - """Issue `query`, retrying to the poll deadline only while the serving pod does not - know the MCP server yet. Every other outcome, guardrail block or pass-through, comes - back untouched so the caller's assertion still decides it.""" - deadline = time.monotonic() + client.proxy.poll_timeout - last = search(query) - while _pod_lacks_mcp_server(last) and time.monotonic() < deadline: - time.sleep(client.proxy.poll_interval) - last = search(query) - return last - - -class TestMcpToolCallGuardrail: - @pytest.mark.skip( - reason=( - "LIT-5052: the control call sends a `telemetry` argument that Datadog's " - "search_datadog_logs tool now rejects, so the clean-argument half of this test " - "errors with 'unexpected additional properties [\"telemetry\"]' and the guardrail " - "block it exists to prove is never exercised. `telemetry` was never a documented " - "Datadog parameter; the test relied on the server ignoring unknown properties. " - "Unskip once the argument is dropped." - ) - ) - @pytest.mark.covers( - "guardrail.litellm_content_filter.pre_mcp_call.blocks", - exercised_on=["mcp_operations"], - ) - def test_content_filter_blocks_banned_keyword_in_tool_args( - self, client: McpClient, resources: ResourceManager - ) -> None: - assert_dd_mcp_creds() - marker = unique_marker() - banned_keyword = f"e2eblocked{marker}" - - guardrail_id = client.register_mcp_content_filter( - name=f"e2e-mcp-cf-{marker}", blocked_keyword=banned_keyword - ) - guardrail_created_at = time.monotonic() - resources.defer(lambda: client.delete_guardrail(guardrail_id)) - - server_id = register_datadog_mcp(client, resources) - server_registered_at = time.monotonic() - key = client.generate_key(user_id=f"e2e-mcp-guard-{marker}", mcp_servers=[server_id]) - resources.defer(lambda: client.proxy.delete_key(key)) - - tool_name = client.await_tool(key, server_id, SEARCH_LOGS_TOOL) - - def search(query: str) -> Result[McpCallToolResponse]: - arguments: McpToolArguments = { - "query": query, - "from": DD_SEARCH_FROM, - "to": "now", - "max_tokens": 500, - "telemetry": {"intent": "e2e mcp guardrail check"}, - } - return client.call_tool(key, server_id=server_id, name=tool_name, arguments=arguments) - - # Registering the guardrail is a control-plane write; the data-plane worker - # that serves tools/call picks it up on its next guardrail sync, so an - # immediate call can race the propagation and slip through. Poll the banned - # call to the deadline and require a block, so the check proves enforcement - # rather than catching a pre-sync pass-through. The keyword is unique per - # run, so this only ever intercepts this test's own call. - blocked = _poll_until_blocked(search, banned_keyword, client) - match blocked: - case UnknownApiError(status_code=400, body=body): - assert banned_keyword in body or "content blocked" in body.lower(), ( - f"the block must name the content-filter reason, got: {body[:300]}" - ) - assert "pre_mcp_call" in body, ( - f"the block must be attributed to the MCP tool-call hook (pre_mcp_call), got: {body[:300]}" - ) - case _: - pytest.fail( - "content_filter never blocked the banned keyword on the MCP tool call within " - f"{client.proxy.poll_timeout}s (the guardrail or the MCP server never synced to " - f"the data plane); last result: {blocked}" - ) - - # The block above only proves the one pod that served it has synced; another - # pod could still lack the guardrail and let the banned call reach Datadog. - # Wait out the full sync interval from the later of the guardrail create and the - # MCP server registration (each syncs on its own clock, so the earlier write's - # deadline can elapse while a pod still lacks the other) so every pod has - # refreshed from the DB, then require the banned call to stay blocked across - # several attempts. A pass-through now is a genuine partial-propagation leak, not - # a race. Client load balancing still can't guarantee every pod is hit, so this - # samples several worker selections rather than proving all pods synced. - sync_remaining = max(guardrail_created_at, server_registered_at) + FULL_SYNC_SECONDS - time.monotonic() - if sync_remaining > 0: - time.sleep(sync_remaining) - for attempt in range(1, POST_SYNC_VERIFICATION_CALLS + 1): - reblocked = _search_on_synced_pod(search, f"still about {banned_keyword} #{attempt}", client) - assert isinstance(reblocked, UnknownApiError) and reblocked.status_code == 400, ( - "after the sync interval every data-plane pod must block the banned keyword, but " - f"attempt {attempt} of {POST_SYNC_VERIFICATION_CALLS} was not blocked (a pod still " - f"lacks the guardrail, or never synced the MCP server): {reblocked}" - ) - if attempt < POST_SYNC_VERIFICATION_CALLS: - time.sleep(client.proxy.poll_interval) - - allowed = _search_on_synced_pod(search, f"e2e-clean-{marker}", client) - match allowed: - case Success(data=result): - assert result.is_error is not True, ( - f"a clean MCP tool call must reach the server and not error, got: {result}" - ) - case _: - pytest.fail( - f"a clean MCP tool call must pass the guardrail and reach the server; got {allowed}" - ) diff --git a/tests/e2e/mcp/test_mcp_key_access_e2e.py b/tests/e2e/mcp/test_mcp_key_access_e2e.py index 788a0a3f45c..429f0840e4b 100644 --- a/tests/e2e/mcp/test_mcp_key_access_e2e.py +++ b/tests/e2e/mcp/test_mcp_key_access_e2e.py @@ -37,7 +37,7 @@ class TestMcpKeyWithoutAccessIsDenied: client: McpClient, resources: ResourceManager, ) -> None: - server_id = register_datadog_mcp(client, resources) + server_id = register_datadog_mcp(client, resources).server_id client.await_registered(server_id) permitted_key = _key(client, resources, mcp_servers=[server_id]) @@ -51,23 +51,13 @@ class TestMcpKeyWithoutAccessIsDenied: f"boundary: {denied_tools}" ) - @pytest.mark.skip( - reason=( - "LIT-5052: the control call proving a granted key CAN invoke the tool sends a " - "`telemetry` argument that Datadog's search_datadog_logs tool now rejects, so it " - "errors with 'unexpected additional properties [\"telemetry\"]' and the denial " - "assertion is never reached. `telemetry` was never a documented Datadog " - "parameter; the test relied on the server ignoring unknown properties. Unskip " - "once the argument is dropped." - ) - ) @pytest.mark.covers("mcp.call_tool.api_key.denied_without_permission") def test_call_tool_denied_without_permission( self, client: McpClient, resources: ResourceManager, ) -> None: - server_id = register_datadog_mcp(client, resources) + server_id = register_datadog_mcp(client, resources).server_id client.await_registered(server_id) permitted_key = _key(client, resources, mcp_servers=[server_id]) @@ -80,7 +70,6 @@ class TestMcpKeyWithoutAccessIsDenied: "from": DD_SEARCH_FROM, "to": "now", "max_tokens": 1000, - "telemetry": {"intent": "e2e control call proving granted key can invoke Datadog MCP"}, } permitted_call = client.await_call_tool( permitted_key, server_id=server_id, name=tool_name, arguments=search_args diff --git a/tests/e2e/mcp/test_mcp_messages_e2e.py b/tests/e2e/mcp/test_mcp_messages_e2e.py new file mode 100644 index 00000000000..0426c740068 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_messages_e2e.py @@ -0,0 +1,119 @@ +"""Live e2e: /v1/messages expands a gateway-registered MCP server and +auto-executes its tools, running the tool_use loop internally and returning a +final Anthropic response (the Claude Code path). + +Registers the real Datadog remote MCP server, grants a key access to it, then +sends a /v1/messages request whose ``tools`` array carries an +``{type: "mcp", server_url: "litellm_proxy/mcp/"}`` reference. The +gateway intercepts the litellm_proxy reference (which Anthropic cannot reach), +expands it into native Anthropic custom tools under the caller's credentials, +runs the tool_use loop (model calls search_datadog_logs, gateway executes it, +feeds the result back as a tool_result), and returns the final answer. The +response has no MCP-specific metadata; the proof is the final text answer, +meaning the loop completed and the model used the tool result. + +The streaming variant exercises the same loop but with stream=True, where the +gateway runs the tool_use loop non-streaming internally and then fakes a stream +of the final answer as Anthropic SSE events. +""" + +from __future__ import annotations + +import pytest + +from datadog_mcp import assert_dd_mcp_creds, register_datadog_mcp +from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker +from e2e_http import require_successful_call, unwrap +from lifecycle import ResourceManager +from mcp_client import McpClient +from models import AnthropicMcpTool, AnthropicMessagesBody, ChatMessage + +pytestmark = pytest.mark.e2e + +def _tool_prompt(marker: str) -> str: + return ( + "Use the search_datadog_logs tool to search for logs " + f"with query '{marker}' from now-30m to now with " + "max_tokens 100. After you get results, reply with ok only." + ) + + +def _messages_body(model: str, alias: str, marker: str) -> AnthropicMessagesBody: + return AnthropicMessagesBody( + model=model, + max_tokens=1024, + messages=[ChatMessage(role="user", content=_tool_prompt(marker))], + tools=[ + AnthropicMcpTool( + server_label="datadog", + server_url=f"litellm_proxy/mcp/{alias}", + require_approval="never", + ) + ], + ) + + +class TestMessagesMcpAutoExecute: + @pytest.mark.covers("mcp.messages.api_key.auto_executes_tools") + def test_messages_runs_tool_loop_and_returns_final_answer( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + assert_dd_mcp_creds() + dd = register_datadog_mcp(client, resources) + client.await_registered(dd.server_id) + + marker = f"e2e-mcp-msg-nohit-{unique_marker()}" + key = client.generate_key( + user_id=f"e2e-mcp-msg-{unique_marker()}", + mcp_servers=[dd.server_id], + models=[CHEAP_ANTHROPIC_MODEL], + ) + resources.defer(lambda: client.proxy.delete_key(key)) + + response = unwrap( + client.messages_with_mcp( + key, _messages_body(CHEAP_ANTHROPIC_MODEL, dd.alias, marker) + ) + ) + + assert response.content, f"/v1/messages returned no content blocks: {response}" + text = "".join(block.text or "" for block in response.content) + assert text.strip(), ( + f"/v1/messages returned no text after the MCP tool loop; the gateway " + f"may not have completed the tool_use loop: {response}" + ) + + @pytest.mark.covers("mcp.messages.api_key.stream_auto_executes_tools") + def test_messages_stream_runs_tool_loop_and_returns_final_answer( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + assert_dd_mcp_creds() + dd = register_datadog_mcp(client, resources) + client.await_registered(dd.server_id) + + marker = f"e2e-mcp-msg-stream-nohit-{unique_marker()}" + key = client.generate_key( + user_id=f"e2e-mcp-msg-stream-{unique_marker()}", + mcp_servers=[dd.server_id], + models=[CHEAP_ANTHROPIC_MODEL], + ) + resources.defer(lambda: client.proxy.delete_key(key)) + + body = _messages_body(CHEAP_ANTHROPIC_MODEL, dd.alias, marker) + body.stream = True + + result = client.messages_stream_with_mcp(key, body) + require_successful_call(result) + assert result.is_streaming, f"response was not streamed: {result.headers}" + assert not result.stream_error, f"stream errored: {result.stream_error}" + assert result.stream_events, "stream produced no SSE events" + assert any("content_block_delta" in event for event in result.stream_events), ( + "stream carried no content deltas" + ) + assert any("message_stop" in event for event in result.stream_events), ( + "stream never reached message_stop" + ) diff --git a/tests/e2e/mcp/test_mcp_responses_e2e.py b/tests/e2e/mcp/test_mcp_responses_e2e.py new file mode 100644 index 00000000000..234bfb7ec11 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_responses_e2e.py @@ -0,0 +1,84 @@ +"""Live e2e: /v1/responses expands a gateway-registered MCP server and +auto-executes its tools, surfacing the results as response output items. + +Registers the real Datadog remote MCP server, grants a key access to it, then +sends a /v1/responses request whose ``tools`` array carries an +``{type: "mcp", server_url: "litellm_proxy/mcp/"}`` reference. The +gateway lists the server's tools, feeds them to the model in Responses API +format, the model calls search_datadog_logs, the gateway executes the call +upstream, and appends ``mcp_tools_fetched`` and ``tool_execution_results`` +output items to the response. Their presence proves the full Responses API +bridge loop ran end to end against a real MCP server and a real LLM. +""" + +from __future__ import annotations + +import pytest + +from datadog_mcp import assert_dd_mcp_creds, register_datadog_mcp +from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from mcp_client import McpClient, ResponsesMcpBody, ResponsesMcpInputMessage, ResponsesMcpTool + +pytestmark = pytest.mark.e2e + + +class TestResponsesMcpAutoExecute: + @pytest.mark.covers("mcp.responses.api_key.auto_executes_tools") + def test_responses_lists_calls_and_executes_mcp_tools( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + assert_dd_mcp_creds() + dd = register_datadog_mcp(client, resources) + client.await_registered(dd.server_id) + + marker = f"e2e-mcp-resp-nohit-{unique_marker()}" + key = client.generate_key( + user_id=f"e2e-mcp-resp-{unique_marker()}", + mcp_servers=[dd.server_id], + models=[CHEAP_ANTHROPIC_MODEL], + ) + resources.defer(lambda: client.proxy.delete_key(key)) + + body = ResponsesMcpBody( + model=CHEAP_ANTHROPIC_MODEL, + input=[ + ResponsesMcpInputMessage( + content=( + "Use the search_datadog_logs tool to search for logs " + f"with query '{marker}' from now-30m to now with " + "max_tokens 100. After you get results, reply with ok only." + ) + ) + ], + instructions="You are a helpful assistant.", + tools=[ + ResponsesMcpTool( + server_label="datadog", + server_url=f"litellm_proxy/mcp/{dd.alias}", + require_approval="never", + ) + ], + ) + result = unwrap(client.responses_with_mcp(key, body)) + + fetched = result.mcp_tools_fetched + assert fetched is not None, ( + "response.output has no mcp_tools_fetched item; the gateway did not list " + f"the Datadog server's tools through the responses bridge: {result.output}" + ) + assert fetched.content, ( + "mcp_tools_fetched item has no content; the tool list was empty" + ) + + executed = result.tool_execution_results + assert executed is not None, ( + "response.output has no tool_execution_results item; the gateway did not " + f"execute any MCP tool through the responses bridge: {result.output}" + ) + assert executed.content, ( + "tool_execution_results item has no content; the tool call returned nothing" + ) diff --git a/tests/e2e/mcp/test_mcp_server_config_e2e.py b/tests/e2e/mcp/test_mcp_server_config_e2e.py new file mode 100644 index 00000000000..fad4a27827a --- /dev/null +++ b/tests/e2e/mcp/test_mcp_server_config_e2e.py @@ -0,0 +1,134 @@ +"""Live e2e: the gateway brokers Datadog's static-header HTTP transport and +multi-server tool namespacing. + +Three cells in one spec: +1. upstream_static_auth: a server registered with static_headers (Datadog's + DD-API-KEY / DD-APPLICATION-KEY) injects them on every upstream call, so + search_datadog_logs succeeds. The tool result must be non-empty, proving + the static credentials reached Datadog. +2. transport_http: the Datadog server uses the streamable HTTP transport, and + a successful tools/list + tools/call round-trip proves that transport path + works end to end. This is the same call as upstream_static_auth but asserts + the transport-specific cell. +3. namespaced_multi_server: two registered servers' tools remain + distinguishable on the aggregate tools/list (each tool carries its own + mcp_info.server_id), so a multi-server tenant never sees tools collide. +""" + +from __future__ import annotations + +import pytest + +from datadog_mcp import SEARCH_LOGS_TOOL, assert_dd_mcp_creds, register_datadog_mcp +from e2e_config import DD_SEARCH_FROM, unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from mcp_client import McpClient, McpToolArguments + +pytestmark = pytest.mark.e2e + + +def _search_args(query: str) -> McpToolArguments: + return { + "query": query, + "from": DD_SEARCH_FROM, + "to": "now", + "max_tokens": 500, + } + + +class TestUpstreamStaticAuthAndTransport: + @pytest.mark.covers( + "mcp.call_tool.api_key.upstream_static_auth", + "mcp.call_tool.api_key.transport_http", + ) + def test_static_header_http_transport_call_succeeds( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + assert_dd_mcp_creds() + dd = register_datadog_mcp(client, resources) + client.await_registered(dd.server_id) + + key = client.generate_key( + user_id=f"e2e-mcp-static-{unique_marker()}", + mcp_servers=[dd.server_id], + ) + resources.defer(lambda: client.proxy.delete_key(key)) + + tool_name = client.await_tool(key, dd.server_id, SEARCH_LOGS_TOOL) + call = client.await_call_tool( + key, + server_id=dd.server_id, + name=tool_name, + arguments=_search_args("service:litellm"), + ) + assert call.is_error is not True, ( + f"search_datadog_logs errored with static headers + http transport: {call}" + ) + assert call.all_text, ( + "search_datadog_logs returned empty text; the static DD-API-KEY / " + "DD-APPLICATION-KEY headers may not have reached the upstream" + ) + + +class TestNamespacedMultiServer: + @pytest.mark.covers("mcp.list_tools.api_key.namespaced_multi_server") + def test_two_servers_tools_remain_distinguishable( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + assert_dd_mcp_creds() + dd_a = register_datadog_mcp(client, resources) + dd_b = register_datadog_mcp(client, resources) + client.await_registered(dd_a.server_id) + client.await_registered(dd_b.server_id) + + key_a = client.generate_key( + user_id=f"e2e-mcp-ns-a-{unique_marker()}", + mcp_servers=[dd_a.server_id], + ) + resources.defer(lambda: client.proxy.delete_key(key_a)) + key_b = client.generate_key( + user_id=f"e2e-mcp-ns-b-{unique_marker()}", + mcp_servers=[dd_b.server_id], + ) + resources.defer(lambda: client.proxy.delete_key(key_b)) + + _ = client.await_tool(key_a, dd_a.server_id, SEARCH_LOGS_TOOL) + _ = client.await_tool(key_b, dd_b.server_id, SEARCH_LOGS_TOOL) + + key_both = client.generate_key( + user_id=f"e2e-mcp-ns-both-{unique_marker()}", + mcp_servers=[dd_a.server_id, dd_b.server_id], + ) + resources.defer(lambda: client.proxy.delete_key(key_both)) + + tools = unwrap(client.list_tools(key_both)) + a_tools = tools.tool_names_for_server(dd_a.server_id) + b_tools = tools.tool_names_for_server(dd_b.server_id) + + assert a_tools, ( + f"server A's tools are missing from the aggregate list; " + f"the multi-server namespace collapsed: {tools.tools}" + ) + assert b_tools, ( + f"server B's tools are missing from the aggregate list; " + f"the multi-server namespace collapsed: {tools.tools}" + ) + + a_entries = tuple( + t for t in tools.tools + if t.mcp_info and t.mcp_info.server_id == dd_a.server_id + ) + b_entries = tuple( + t for t in tools.tools + if t.mcp_info and t.mcp_info.server_id == dd_b.server_id + ) + assert len(a_entries) == len(b_tools) and len(b_entries) == len(b_tools), ( + f"each server's tools must carry its own mcp_info.server_id so a " + f"multi-server tenant can tell them apart; " + f"A entries={a_entries}, B entries={b_entries}" + ) diff --git a/tests/e2e/mcp/test_mcp_tool_filtering_e2e.py b/tests/e2e/mcp/test_mcp_tool_filtering_e2e.py new file mode 100644 index 00000000000..3b08478f566 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_tool_filtering_e2e.py @@ -0,0 +1,100 @@ +"""Live e2e: the gateway's allowed_tools filtering on the Datadog MCP server. + +A server registered with a narrow allowed_tools list hides every other tool +from tools/list, and tools/call on a hidden tool is blocked (403 or 404, +depending on whether the tool was in the gateway's resolved tool map). + +Note: disallowed_tools and allowed_params are config-only fields today, they +have no DB column in the Prisma schema and are silently dropped on the +management API path. Those cells are product gaps, not test gaps. + +All against the real Datadog remote MCP server. +""" + +from __future__ import annotations + +import pytest + +from datadog_mcp import SEARCH_LOGS_TOOL, assert_dd_mcp_creds, register_datadog_mcp +from e2e_config import DD_SEARCH_FROM, unique_marker +from e2e_http import UnknownApiError, unwrap +from lifecycle import ResourceManager +from mcp_client import McpClient, McpToolArguments + +pytestmark = pytest.mark.e2e + +BOGUS_TOOL = "nonexistent_e2e_tool" + + +def _search_args(query: str) -> McpToolArguments: + return { + "query": query, + "from": DD_SEARCH_FROM, + "to": "now", + "max_tokens": 500, + } + + +def _key( + client: McpClient, resources: ResourceManager, server_id: str, label: str +) -> str: + key = client.generate_key( + user_id=f"e2e-mcp-{label}-{unique_marker()}", + mcp_servers=[server_id], + ) + resources.defer(lambda: client.proxy.delete_key(key)) + return key + + +class TestAllowedToolsScoping: + @pytest.mark.covers("mcp.list_tools.api_key.allowed_tools_scoped") + def test_list_tools_hides_non_allowed_tools( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + assert_dd_mcp_creds() + dd = register_datadog_mcp(client, resources, allowed_tools=[BOGUS_TOOL]) + client.await_registered(dd.server_id) + + key = _key(client, resources, dd.server_id, "allow-list") + + tools = unwrap(client.list_tools(key)).tool_names_for_server(dd.server_id) + assert SEARCH_LOGS_TOOL not in tools, ( + f"search_datadog_logs must be hidden by the allowed_tools filter " + f"(only {BOGUS_TOOL!r} is allowed), but it appeared in tools/list: {tools}" + ) + + @pytest.mark.covers("mcp.call_tool.api_key.allowed_tools_scoped") + def test_call_tool_denied_outside_allowed_tools( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + assert_dd_mcp_creds() + dd_all = register_datadog_mcp(client, resources) + client.await_registered(dd_all.server_id) + dd_narrow = register_datadog_mcp(client, resources, allowed_tools=[BOGUS_TOOL]) + client.await_registered(dd_narrow.server_id) + + key_all = _key(client, resources, dd_all.server_id, "allow-call-control") + _ = client.await_tool(key_all, dd_all.server_id, SEARCH_LOGS_TOOL) + + key_narrow = _key(client, resources, dd_narrow.server_id, "allow-call-narrow") + + result = client.call_tool( + key_narrow, + server_id=dd_narrow.server_id, + name=SEARCH_LOGS_TOOL, + arguments=_search_args("service:litellm"), + ) + match result: + case UnknownApiError(status_code=403): + pass + case UnknownApiError(status_code=404): + pass + case _: + pytest.fail( + f"calling a tool outside allowed_tools must be blocked (403 or 404), " + f"got: {result}" + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 8305b42e9ad..a3ffd2142d4 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -381,7 +381,15 @@ class AnthropicCustomTool(BaseModel): input_schema: ToolInputSchema -type AnthropicTool = AnthropicToolSearchTool | AnthropicWebSearchTool | AnthropicCustomTool +class AnthropicMcpTool(BaseModel): + type: Literal["mcp"] = "mcp" + server_label: str + server_url: str + require_approval: str = "never" + allowed_tools: list[str] | None = None + + +type AnthropicTool = AnthropicToolSearchTool | AnthropicWebSearchTool | AnthropicCustomTool | AnthropicMcpTool class AnthropicMessagesBody(BaseModel):