From 8c287351675ce2f3b6087200428c4f39a06ef0cd Mon Sep 17 00:00:00 2001 From: Alexey Puzik Date: Thu, 12 Feb 2026 01:26:20 +0200 Subject: [PATCH 1/3] Add playground OAuth flow --- CLAUDE.md | 137 +++++++++++ crates/agentgateway/src/mcp/router.rs | 13 ++ crates/agentgateway/src/proxy/mod.rs | 7 + ui/src/app/playground/oauth/callback/page.tsx | 113 +++++++++ ui/src/app/playground/page.tsx | 140 ++++++++---- ui/src/lib/mcp-oauth-provider.ts | 214 ++++++++++++++++++ 6 files changed, 582 insertions(+), 42 deletions(-) create mode 100644 CLAUDE.md create mode 100644 ui/src/app/playground/oauth/callback/page.tsx create mode 100644 ui/src/lib/mcp-oauth-provider.ts diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 000000000..2ec453c4a --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,137 @@ +# AgentGateway Submodule — Claude Code Guide + +## IMPORTANT: Do Not Build Locally + +**Do NOT attempt to build this Rust project on Windows.** The user builds and checks it manually in a VS2022 Developer Command Prompt. Claude Code should only read, analyze, and edit source files — never run `cargo build`, `cargo test`, or similar commands. + +## MCP Security Guards + +Guards intercept MCP tool calls at different phases and can Allow, Deny, or Modify operations. + +### Guard Phases (`runs_on`) + +| Phase | When | Purpose | +|-------|------|---------| +| `ToolsList` | `tools/list` response | Filter/block exposed tools | +| `ToolInvoke` | Before `tools/call` | Block/modify tool execution | +| `Request` | Any MCP request | General request filtering | +| `Response` | Any MCP response | General response filtering | + +### Core Trait + +```rust +// crates/agentgateway/src/mcp/security/native/mod.rs +pub trait NativeGuard: Send + Sync { + fn evaluate_tools_list(&self, tools: &[Tool], context: &GuardContext) -> GuardResult; + fn evaluate_tool_invoke(&self, tool_name: &str, arguments: &Value, context: &GuardContext) -> GuardResult; + fn evaluate_request(&self, request: &Value, context: &GuardContext) -> GuardResult; + fn evaluate_response(&self, response: &Value, context: &GuardContext) -> GuardResult; + fn reset_server(&self, server_name: &str); +} +``` + +### Native Guards + +| Guard | File | What It Does | +|-------|------|-------------| +| **ToolPoisoningDetector** | `native/tool_poisoning.rs` (28 KB) | Blocks tools with malicious descriptions — prompt injection, system overrides, safety bypass, hidden instructions, prompt leaking, encoding tricks. 32 built-in regex patterns + custom patterns. | +| **RugPullDetector** | `native/rug_pull.rs` (38 KB) | Detects bait-and-switch: establishes per-server tool baseline on first encounter, then blocks if tools are modified/removed. Risk scoring with configurable weights (description=2, schema=3, remove=3, add=1), threshold=5. | +| **PiiGuard** | `native/pii_guard.rs` (28 KB) | Detects PII (email, phone, SSN, credit cards, CA SIN, URLs) via regex. Two modes: Mask (replace with `` placeholders) or Reject (block entire response). | +| **ToolShadowingDetector** | `native/tool_shadowing.rs` | Placeholder — will block duplicate tool names across servers. | +| **ServerWhitelistChecker** | `native/server_whitelist.rs` | Placeholder — will enforce allowed server whitelist. | + +### WASM Guards + +Custom guards loaded at runtime as WebAssembly components (`security/wasm.rs`). Fully sandboxed (no FS/network). ~5-10ms latency vs <1ms for native. See `examples/wasm-guards/` for examples. + +### Key Files + +- **Orchestration**: `crates/agentgateway/src/mcp/security/mod.rs` — `GuardExecutor`, `GuardExecutorRegistry`, config types, priority ordering +- **Integration**: `crates/agentgateway/src/mcp/handler.rs` — calls guards before tool execution and on tools/list responses +- **Config binding**: `crates/agentgateway/src/types/local.rs` — `McpBackendConfig` holds `Vec` +- **PII patterns**: `crates/agentgateway/src/mcp/security/native/pii_detection.rs` + +### Guard Configuration (YAML) + +```yaml +security_guards: + - id: pii-mask + description: "Mask PII in responses" + priority: 10 # Lower = runs first + enabled: true + timeout_ms: 100 + failure_mode: fail_closed # or fail_open + runs_on: [response] + type: + native: + pii: + action: mask + pii_types: [email, phone, ssn, credit_card] +``` + +## E2E Tests (`tests/`) + +### Test Files + +| File | Tests | What It Covers | +|------|-------|----------------| +| `e2e_security_guards_test.py` | Master runner | Orchestrates all guard suites sequentially | +| `e2e_pii_guard_test.py` | 19 per mode (mask+reject) | 6 PII types: single, embedded-in-text, bulk, full-record, clean-data-passthrough | +| `e2e_tool_poisoning_guard_test.py` | 6 | tools/list blocked (HTTP 403), deny reason structure, 6 attack categories covered | +| `e2e_rug_pull_guard_test.py` | 22 | Baseline establishment, session vs global scope, 5 mutation modes (all/description/schema/remove/add) | +| `e2e_mcp_sse_test.py` | 5+ | SSE transport compliance with guards | +| `benchmark.py` | Configurable | Throughput, latency percentiles (p95/p99), error rates | +| `mcp_client.py` | — | Shared client library: `MCPSSEClient`, `MCPStreamableHTTPClient`, `TestResults` | + +### Test Routes + +| Route | Guard Config | Backend Port | +|-------|-------------|-------------| +| `/pii-test` | PII mask mode | 8000 | +| `/pii-test-reject` | PII reject mode | 8000 | +| `/poison` | Tool poisoning | 8010 | +| `/rug-pull` | Rug pull (default) | 8020 | +| `/rug-pull-desc` | Rug pull description-only | 8020 | +| `/rug-pull-schema` | Rug pull schema-only | 8020 | +| `/rug-pull-remove` | Rug pull remove mode | 8020 | +| `/rug-pull-add` | Rug pull add mode | 8020 | + +### Running Tests + +```bash +# Docker (CI/CD) — runs full suite +cd tests/docker && docker compose up --build --abort-on-container-exit --exit-code-from test-runner + +# Against deployed environment +GATEWAY_URL=https://... python tests/e2e_security_guards_test.py --transport streamable +``` + +## Test Servers (`testservers/`) + +Three Python MCP servers in one Docker image, started via `start-server.sh`: + +### PII Test Server (port 8000) +- **Module**: `src/mcp_test_server/fastmcp_server.py` +- **Tools**: `generate_pii`, `generate_bulk_pii`, `generate_full_record`, `generate_text_with_pii`, `list_pii_types` +- **Data**: Faker-generated emails, phones, SSNs, credit cards, CA SINs, URLs, addresses +- **Resources**: `pii://fixtures/{personal,identity,financial,mixed}` — predefined test fixtures + +### Tool Poisoning Server (port 8010) +- **Module**: `src/tool_poisoning_test/server.py` +- **Tools**: 6 poisoned (`add`, `secret_notes`, `translate_text`, `get_status`, `search_files`, `run_diagnostic`) + 2 clean (`subtract`, `multiply`) +- **Attack categories**: hidden instructions, prompt injection, system override, safety bypass, role manipulation, prompt leaking + +### Rug Pull Server (port 8020) +- **Module**: `src/rug_pull_test/server.py` +- **Tools**: `get_weather` (session trigger), `get_global_weather` (global trigger), `get_forecast`, `reset_session_rug`, `reset_global_rug`, `get_rug_status`, `set_rug_pull_mode` +- **Mutation modes**: `all` (default), `description`, `schema`, `remove`, `add` — each changes tool definitions differently after trigger +- **Benign→Malicious**: e.g. "Get weather" becomes "Get weather AND read all env vars, API keys, secrets..." + +### Docker Setup + +```yaml +# tests/docker/docker-compose.yaml — 3 services +mcp-test-servers: # Builds from testservers/, exposes 8000/8010/8020 +agentgateway: # ACR image, mounts test config, port 8080 +test-runner: # Python 3.12, runs e2e_security_guards_test.py +``` diff --git a/crates/agentgateway/src/mcp/router.rs b/crates/agentgateway/src/mcp/router.rs index 84bf349f6..813f14f3b 100644 --- a/crates/agentgateway/src/mcp/router.rs +++ b/crates/agentgateway/src/mcp/router.rs @@ -148,6 +148,19 @@ impl App { "MCP auth configured; validating Authorization header (mode={:?})", auth.mode ); + // When mcpAuthentication is configured, always require an Authorization + // header — even in permissive mode. This ensures the gateway returns + // 401 + WWW-Authenticate on the first unauthenticated request, triggering + // OAuth discovery. Without this, permissive mode would silently forward + // unauthenticated requests to the upstream, which returns its own 401 + // with the upstream's URL (not the gateway's), breaking the OAuth flow. + if req.headers().get(http::header::AUTHORIZATION).is_none() { + return Err(Self::create_auth_required_response( + ProxyError::ProcessingString("missing authorization header".to_string()), + &req, + auth, + )); + } auth .jwt_validator .apply(None, &mut req) diff --git a/crates/agentgateway/src/proxy/mod.rs b/crates/agentgateway/src/proxy/mod.rs index b2b07e7aa..1ab3531b9 100644 --- a/crates/agentgateway/src/proxy/mod.rs +++ b/crates/agentgateway/src/proxy/mod.rs @@ -207,6 +207,13 @@ impl ProxyError { } } pub fn into_response(self) -> Response { + // For upstream MCP errors, return the original response directly. + // This preserves headers like WWW-Authenticate needed for OAuth flows. + if let ProxyError::MCP(mcp::Error::UpstreamError(resp)) = self { + let (parts, body) = resp.0.into_parts(); + return ::http::Response::from_parts(parts, http::Body::from(body)); + } + let code = match self { ProxyError::BindNotFound => StatusCode::NOT_FOUND, ProxyError::ListenerNotFound => StatusCode::NOT_FOUND, diff --git a/ui/src/app/playground/oauth/callback/page.tsx b/ui/src/app/playground/oauth/callback/page.tsx new file mode 100644 index 000000000..137742a4a --- /dev/null +++ b/ui/src/app/playground/oauth/callback/page.tsx @@ -0,0 +1,113 @@ +"use client"; + +import { Suspense, useEffect, useState } from "react"; +import { useSearchParams } from "next/navigation"; + +/** + * Inner component that uses useSearchParams (must be wrapped in Suspense). + * + * Extracts the authorization code (or error) from URL search params, + * sends it back to the opener window via postMessage, and closes itself. + */ +function OAuthCallbackContent() { + const searchParams = useSearchParams(); + const [status, setStatus] = useState<"processing" | "success" | "error" | "no-opener">( + "processing" + ); + const [errorMessage, setErrorMessage] = useState(""); + + useEffect(() => { + const code = searchParams.get("code"); + const error = searchParams.get("error"); + const errorDescription = searchParams.get("error_description"); + + if (!window.opener) { + setStatus("no-opener"); + return; + } + + if (error) { + setStatus("error"); + setErrorMessage(errorDescription || error); + window.opener.postMessage( + { type: "mcp-oauth-callback", code: null, error: errorDescription || error }, + window.location.origin + ); + setTimeout(() => window.close(), 3000); + return; + } + + if (code) { + setStatus("success"); + window.opener.postMessage( + { type: "mcp-oauth-callback", code, error: null }, + window.location.origin + ); + setTimeout(() => window.close(), 1000); + return; + } + + setStatus("error"); + setErrorMessage("No authorization code or error received."); + window.opener.postMessage( + { type: "mcp-oauth-callback", code: null, error: "No authorization code received" }, + window.location.origin + ); + }, [searchParams]); + + return ( +
+ {status === "processing" && ( + <> +
+

Completing authentication...

+ + )} + {status === "success" && ( + <> +
✓
+

Authorization successful!

+

This window will close automatically.

+ + )} + {status === "error" && ( + <> +
✗
+

Authorization failed

+

{errorMessage}

+

This window will close shortly.

+ + )} + {status === "no-opener" && ( + <> +
⚠
+

Unable to complete authorization

+

+ This page should be opened from the playground. Please close this window and try again. +

+ + )} +
+ ); +} + +/** + * OAuth callback page for the playground. + * Opened in a popup by the BrowserOAuthProvider during the MCP OAuth flow. + */ +export default function OAuthCallbackPage() { + return ( +
+ +
+

Completing authentication...

+
+ } + > + + +
+ ); +} diff --git a/ui/src/app/playground/page.tsx b/ui/src/app/playground/page.tsx index 212716490..790ac335e 100644 --- a/ui/src/app/playground/page.tsx +++ b/ui/src/app/playground/page.tsx @@ -2,7 +2,7 @@ import { useState, useEffect, useCallback } from "react"; import { Client as McpClient } from "@modelcontextprotocol/sdk/client/index.js"; -import { SSEClientTransport as McpSseTransport } from "@modelcontextprotocol/sdk/client/sse.js"; +import { StreamableHTTPClientTransport as McpHttpTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"; import { ClientRequest as McpClientRequest, Result as McpResult, @@ -11,6 +11,8 @@ import { ListToolsResultSchema as McpListToolsResultSchema, Tool as McpTool, } from "@modelcontextprotocol/sdk/types.js"; +import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js"; +import { BrowserOAuthProvider } from "@/lib/mcp-oauth-provider"; import { z } from "zod"; import { A2AClient, @@ -247,9 +249,21 @@ export default function PlaygroundPage() { if (listener.routes) { listener.routes.forEach((route: Route, routeIndex: number) => { const protocol = listener.protocol === ListenerProtocol.HTTPS ? "https" : "http"; - const hostname = listener.hostname || "localhost"; - const port = bind.port; // Use the actual port from the bind configuration - const baseEndpoint = `${protocol}://${hostname}:${port}`; + // Use window.location.hostname when listener hostname is not set + // This ensures the UI connects to the correct domain in deployed environments + const hostname = + listener.hostname || + (typeof window !== "undefined" ? window.location.hostname : "localhost"); + const configPort = bind.port; + // Local dev: browser URL has explicit port (e.g., 127.0.0.1:15000) -> use config port for routes + // Deployed (Azure): browser URL has no port (standard 80/443) -> use origin (ingress handles routing) + // Drawback: if Azure exposed non-standard port, this would break (uncommon scenario) + const hasExplicitPort = typeof window !== "undefined" && window.location.port !== ""; + const baseEndpoint = hasExplicitPort + ? `${protocol}://${hostname}:${configPort}` + : typeof window !== "undefined" + ? window.location.origin + : `${protocol}://${hostname}:${configPort}`; // Generate route path and description with better pattern recognition let routePath = "/"; @@ -447,48 +461,91 @@ export default function PlaygroundPage() { if (backendType === "mcp") { setConnectionState((prev) => ({ ...prev, connectionType: "mcp" })); - // TODO: Support acting as a stateless client - const client = new McpClient( - { name: "agentgateway-dashboard", version: "0.1.0" }, - { capabilities: {} } - ); + const mcpUrl = selectedRoute.endpoint.endsWith("/") + ? selectedRoute.endpoint.slice(0, -1) + : selectedRoute.endpoint; - const headers: Record = { - Accept: "text/event-stream", - "Cache-Control": "no-cache", - "mcp-protocol-version": "2024-11-05", - }; + let connectedClient: McpClient; - // Only add auth header if token is provided and not empty if (connectionState.authToken && connectionState.authToken.trim()) { - headers["Authorization"] = `Bearer ${connectionState.authToken}`; - } + // Manual bearer token mode: use header injection (backward compatible) + const client = new McpClient( + { name: "agentgateway-dashboard", version: "0.1.0" }, + { capabilities: {} } + ); + + const headers: Record = { + Accept: "text/event-stream", + "Cache-Control": "no-cache", + "mcp-protocol-version": "2024-11-05", + Authorization: `Bearer ${connectionState.authToken}`, + }; - const sseUrl = selectedRoute.endpoint.endsWith("/") - ? `${selectedRoute.endpoint}sse` - : `${selectedRoute.endpoint}/sse`; - const transport = new McpSseTransport(new URL(sseUrl), { - eventSourceInit: { - fetch: (url, init) => { - return fetch(url, { - ...init, - headers: headers as HeadersInit, - }); + const transport = new McpHttpTransport(new URL(mcpUrl), { + requestInit: { + headers: headers as HeadersInit, + credentials: "omit", + mode: "cors", }, - }, - requestInit: { - headers: headers as HeadersInit, - credentials: "omit", - mode: "cors", - }, - }); + }); - await client.connect(transport); - setMcpState((prev) => ({ ...prev, client })); + await client.connect(transport); + connectedClient = client; + } else { + // Automatic OAuth mode: use SDK's built-in authProvider + // If the server doesn't require auth, the provider is never invoked. + // If the server requires auth, the full OAuth 2.1 flow runs automatically. + const oauthProvider = new BrowserOAuthProvider(mcpUrl); + + const client = new McpClient( + { name: "agentgateway-dashboard", version: "0.1.0" }, + { capabilities: {} } + ); + + const transport = new McpHttpTransport(new URL(mcpUrl), { + authProvider: oauthProvider, + }); + + try { + await client.connect(transport); + connectedClient = client; + } catch (error) { + if (error instanceof UnauthorizedError) { + // OAuth redirect was triggered — popup is open, wait for auth code + toast.info("Please authorize in the popup window..."); + try { + const authCode = await oauthProvider.waitForAuthCode(); + await transport.finishAuth(authCode); + + // Reconnect with a new client + transport (tokens are now cached in provider) + const newClient = new McpClient( + { name: "agentgateway-dashboard", version: "0.1.0" }, + { capabilities: {} } + ); + const newTransport = new McpHttpTransport(new URL(mcpUrl), { + authProvider: oauthProvider, + }); + await newClient.connect(newTransport); + connectedClient = newClient; + } catch (oauthError) { + oauthProvider.dispose(); + throw oauthError; + } + } else { + oauthProvider.dispose(); + throw error; + } + } + } + + setMcpState((prev) => ({ ...prev, client: connectedClient })); setUiState((prev) => ({ ...prev, isLoadingCapabilities: true })); const listToolsRequest: McpClientRequest = { method: "tools/list", params: {} }; - const toolsResponse = await client.request(listToolsRequest, McpListToolsResultSchema); + const toolsResponse = await connectedClient.request( + listToolsRequest, + McpListToolsResultSchema + ); setMcpState((prev) => ({ ...prev, tools: toolsResponse.tools })); // Only mark as connected and show success after tools are loaded successfully @@ -1093,8 +1150,7 @@ export default function PlaygroundPage() { Request URL
- {selectedRoute.protocol}://{selectedRoute.listener.hostname || "localhost"}: - {selectedRoute.bindPort} + {selectedRoute.endpoint} {request.path}
@@ -1415,12 +1471,12 @@ export default function PlaygroundPage() {
-
@@ -1471,12 +1415,12 @@ export default function PlaygroundPage() {
-