From d12b5a13373cd99144beb1b8b9077140e850e58b Mon Sep 17 00:00:00 2001 From: Wen-Hao Lue Date: Wed, 30 Sep 2026 14:12:23 -0700 Subject: [PATCH 1/2] Bring agent_run, web search, and the HTTP runtime up to hosted parity agent_run advertises run ID patterns that compile as unicode-mode regexes (strict providers rejected the old tool list), declares an MCP outputSchema and returns structuredContent, accepts the ultra effort and the polymarket and macrobond providers, is marked openWorldHint, names only registered sibling tools in its description, explains 402 as not retryable, is on by default for callers with their own key, and hands its run ID back after a 45s call window. The Exa Agent skill documents the new result shape. web_search_exa takes an optional objective, forwarded to /search, and web_search_advanced_exa is marked openWorldHint. EXA_API_BASE_URL overrides the Exa API origin the tools call. The HTTP handler validates Accept media ranges (406), redacts ?exaApiKey= from debug logs, and prefers a ?exaApiKey= URL key over a Bearer JWT, retrying a rejected key once on the JWT through McpConfig.apiKeyFallback. The container runtime keeps idle SSE streams alive, drains in-flight requests on SIGTERM/SIGINT without waiting out kept-alive connections, survives unhandled rejections, and answers malformed requests with 400. Tests are split into unit and integration suites; the integration suite drives the runtime server against a fake Exa API and runs in npm run ci. Co-Authored-By: Claude Opus 5.5 (1M context) --- README.md | 4 +- api/mcp.ts | 165 ++++++--- api/well-known-mcp-config.ts | 2 +- env.example | 12 +- package.json | 3 +- skills/exa-agent/SKILL.md | 16 +- src/mcp-handler.ts | 14 +- src/runtime-server.ts | 108 +----- src/runtime.ts | 243 +++++++++++++ src/toolRegistry.ts | 2 +- src/tools/agentRun.ts | 186 ++++++++-- src/tools/config.ts | 15 +- src/tools/webSearch.ts | 23 +- src/tools/webSearchAdvanced.ts | 2 +- src/types.ts | 8 +- src/utils/agentErrorHandler.ts | 9 +- src/utils/apiKeyFallback.ts | 160 +++++++++ src/utils/mcpAccept.ts | 75 ++++ src/utils/requestTarget.ts | 25 ++ src/utils/response.ts | 14 + tests/helpers/advertisedTools.ts | 43 +++ tests/helpers/fakeMcpServer.ts | 20 ++ tests/integration/accept.test.ts | 44 +++ tests/integration/apiKeyFallback.test.ts | 154 ++++++++ tests/integration/gracefulShutdown.test.ts | 92 +++++ tests/integration/harness.ts | 321 +++++++++++++++++ tests/integration/webSearch.test.ts | 144 ++++++++ tests/unit/api/mcp.test.ts | 173 +++++++++ tests/unit/contract.test.ts | 11 + tests/unit/mcp-handler.test.ts | 120 ++++++- tests/unit/runtime.test.ts | 158 +++++++++ tests/unit/toolSchemas.test.ts | 103 ++++++ tests/unit/tools/agentProgress.test.ts | 168 ++++++++- tests/unit/tools/agentRun.test.ts | 386 ++++++++++++++++++++- tests/unit/tools/config.test.ts | 18 +- tests/unit/tools/webSearch.test.ts | 59 ++++ tests/unit/utils/agentErrorHandler.test.ts | 40 +++ tests/unit/utils/apiKeyFallback.test.ts | 118 +++++++ tests/unit/utils/mcpAccept.test.ts | 41 +++ tests/unit/utils/requestTarget.test.ts | 27 ++ tsconfig.test.json | 8 +- vitest.config.ts | 3 + vitest.integration.config.ts | 22 ++ 43 files changed, 3118 insertions(+), 241 deletions(-) create mode 100644 src/runtime.ts create mode 100644 src/utils/apiKeyFallback.ts create mode 100644 src/utils/mcpAccept.ts create mode 100644 src/utils/requestTarget.ts create mode 100644 tests/helpers/advertisedTools.ts create mode 100644 tests/integration/accept.test.ts create mode 100644 tests/integration/apiKeyFallback.test.ts create mode 100644 tests/integration/gracefulShutdown.test.ts create mode 100644 tests/integration/harness.ts create mode 100644 tests/integration/webSearch.test.ts create mode 100644 tests/unit/runtime.test.ts create mode 100644 tests/unit/toolSchemas.test.ts create mode 100644 tests/unit/utils/agentErrorHandler.test.ts create mode 100644 tests/unit/utils/apiKeyFallback.test.ts create mode 100644 tests/unit/utils/mcpAccept.test.ts create mode 100644 tests/unit/utils/requestTarget.test.ts create mode 100644 vitest.integration.config.ts diff --git a/README.md b/README.md index 8455797e..3662a4f4 100644 --- a/README.md +++ b/README.md @@ -95,12 +95,12 @@ Exa MCP works with most other clients, point them at `https://mcp.exa.ai/mcp`. | --- | --- | | `web_search_exa` | Search the web for any topic and get clean, ready-to-use content | | `web_fetch_exa` | Read a webpage's full content as clean markdown from one or more URLs | +| `agent_run` | Run an [Exa Agent](https://docs.exa.ai/reference/agent-api-guide) for multi-step research, list-building, enrichment, and structured output (authenticated connections only) | ### Optional Tools (enable via the `tools` parameter) | Tool | Description | | --- | --- | -| `agent_run` | Run an [Exa Agent](https://docs.exa.ai/reference/agent-api-guide) for multi-step research, list-building, enrichment, and structured output | | `web_search_advanced_exa` | Advanced search with filters, domains, dates, highlights, summaries, and subpage crawling | Enable tools by appending them to the MCP URL (this will replace the defaults, so include all you want): @@ -110,7 +110,7 @@ https://mcp.exa.ai/mcp?tools=web_search_advanced_exa https://mcp.exa.ai/mcp?tools=web_search_exa,web_fetch_exa,agent_run ``` -Exa Agent requires authentication (OAuth or an [API key](https://dashboard.exa.ai/api-keys)). +Exa Agent requires authentication (OAuth or an [API key](https://dashboard.exa.ai/api-keys)) and is on by default once you connect with either. ## Agent Skills diff --git a/api/mcp.ts b/api/mcp.ts index b16d7f47..a4ee0e5f 100644 --- a/api/mcp.ts +++ b/api/mcp.ts @@ -6,6 +6,7 @@ import { DEFAULT_MCP_MAX_DURATION_SECONDS, parsePositiveInteger } from "../src/t import type { Ratelimit } from "@upstash/ratelimit"; import type { Redis } from "@upstash/redis"; import { isJwtToken, verifyOAuthToken } from "../src/utils/auth.js"; +import { acceptsMcpResponses } from "../src/utils/mcpAccept.js"; import { expandToolSelection, requiresUserProvidedApiKey, @@ -378,7 +379,10 @@ async function checkRateLimits( * - DEBUG: Enable debug logging (true/false) * - ENABLED_TOOLS: Comma-separated list of tools to enable * - * Priority: x-api-key header > Authorization header > URL query parameter > environment variable. + * Priority: x-api-key header > plain Bearer API key > ?exaApiKey= > Bearer JWT > environment variable. + * An API key always outranks a JWT: MCP clients run OAuth discovery and attach a + * JWT even when the connector URL already carries the org's key. If upstream + * rejects that URL key, the JWT is verified and the tool call retried on it. * * ARCHITECTURE NOTE: * The mcp-handler library creates a single server instance and doesn't pass @@ -401,10 +405,14 @@ function getBearerToken(request: Request): string | undefined { return undefined; } -/** - * Extract configuration from request headers, URL, or environment variables. - * Priority: header > query parameter > environment variable. - */ +/** API key from `?exaApiKey=`, or undefined when absent, empty, or the URL is unparseable. */ +function getApiKeyFromUrl(request: Request): string | undefined { + try { + return new URL(request.url).searchParams.get("exaApiKey") || undefined; + } catch { + return undefined; + } +} interface RequestConfig { exaApiKey?: string; @@ -419,13 +427,56 @@ interface RequestConfig { oauthAccessToken?: string; /** True when a Bearer token was a JWT but failed OAuth verification (expired, bad sig, wrong issuer/audience). */ invalidOAuthJwt: boolean; + /** + * Unverified Bearer JWT that lost to a `?exaApiKey=` URL key. Verified only if + * upstream rejects that key, in which case the tool call is retried on it. + */ + fallbackJwt?: string; + apiKeyFallback?: () => Promise; agentCallWindowMs?: number; mcpMaxDurationSeconds?: number; } +/** + * Build the once-per-request fallback the tool-call wrapper invokes when + * upstream rejects the URL API key: verify the JWT the request also carried + * and, if it is good, switch `config` over to it so the retried call runs as + * the OAuth identity. + */ +function createApiKeyFallback(config: RequestConfig): (() => Promise) | undefined { + const jwt = config.fallbackJwt; + if (!jwt) { + return undefined; + } + let attempted = false; + return async () => { + if (attempted) { + return false; + } + attempted = true; + const claims = await verifyOAuthToken(jwt); + if (!claims) { + console.error("[EXA-MCP] URL API key rejected upstream; fallback JWT invalid"); + return false; + } + config.exaApiKey = undefined; + config.oauthAccessToken = jwt; + config.userProvidedApiKey = true; + config.authMethod = "oauth"; + if (config.debug) { + console.log("[EXA-MCP] URL API key rejected upstream; retrying as oauth"); + } + return true; + }; +} + /** * Extract configuration from request headers, URL, or environment variables. - * Priority: x-api-key header > OAuth JWT > plain Bearer API key > query parameter > environment variable. + * Priority: x-api-key header > plain Bearer API key > ?exaApiKey= > Bearer JWT (OAuth) > + * environment variable. A JWT is only consulted when no API key was supplied, so a key in + * the connector URL is never overridden by a token the client obtained through OAuth + * discovery. That JWT is kept as `fallbackJwt`, to be verified only if upstream rejects the + * URL key (see createApiKeyFallback). */ async function getConfigFromRequest(request: Request): Promise { let exaApiKey = process.env.EXA_API_KEY; @@ -436,41 +487,47 @@ async function getConfigFromRequest(request: Request): Promise { let defaultSearchType: "auto" | "fast" | "instant" | undefined; let oauthAccessToken: string | undefined; let invalidOAuthJwt = false; + let fallbackJwt: string | undefined; let agentCallWindowMs: number | undefined; - // 1. Check x-api-key header (highest priority) const xApiKey = request.headers.get("x-api-key"); + const bearerToken = xApiKey ? undefined : getBearerToken(request); + const apiKeyFromUrl = xApiKey ? undefined : getApiKeyFromUrl(request); + // A bearer JWT yields to a key in the URL; a plain bearer API key does not. + const bearerIsJwt = bearerToken !== undefined && isJwtToken(bearerToken); + if (xApiKey) { + // 1. x-api-key header (highest priority) exaApiKey = xApiKey; userProvidedApiKey = true; authMethod = "api_key"; - } - - // 2. Check Authorization: Bearer header (fallback when no x-api-key) - if (!xApiKey) { - const bearerToken = getBearerToken(request); - if (bearerToken) { - // Distinguish JWT (OAuth) from plain API key - if (isJwtToken(bearerToken)) { - const claims = await verifyOAuthToken(bearerToken); - if (claims) { - oauthAccessToken = bearerToken; - exaApiKey = undefined; - userProvidedApiKey = true; - authMethod = "oauth"; - } else { - // JWT verification failed — flag so the caller can return 401 with - // a WWW-Authenticate challenge instead of silently falling through to - // the env API key or free tier. - invalidOAuthJwt = true; - console.error("[EXA-MCP] Invalid OAuth JWT token"); - } - } else { - // Plain API key in Bearer header - exaApiKey = bearerToken; - userProvidedApiKey = true; - authMethod = "api_key"; - } + } else if (bearerToken && !bearerIsJwt) { + // 2. Plain API key in Authorization: Bearer header + exaApiKey = bearerToken; + userProvidedApiKey = true; + authMethod = "api_key"; + } else if (apiKeyFromUrl) { + // 3. ?exaApiKey=YOUR_KEY (backwards compatible; the connector's configured key). + // A bearer JWT is not verified here: the key is the credential the connector + // owner configured. It is kept only as the fallback for a rejected key. + exaApiKey = apiKeyFromUrl; + userProvidedApiKey = true; + authMethod = "api_key"; + fallbackJwt = bearerToken; + } else if (bearerToken) { + // 4. Bearer JWT from Exa OAuth. + const claims = await verifyOAuthToken(bearerToken); + if (claims) { + oauthAccessToken = bearerToken; + exaApiKey = undefined; + userProvidedApiKey = true; + authMethod = "oauth"; + } else { + // JWT verification failed — flag so the caller can return 401 with + // a WWW-Authenticate challenge instead of silently falling through to + // the env API key or free tier. + invalidOAuthJwt = true; + console.error("[EXA-MCP] Invalid OAuth JWT token"); } } @@ -478,16 +535,6 @@ async function getConfigFromRequest(request: Request): Promise { const parsedUrl = new URL(request.url); const params = parsedUrl.searchParams; - // 3. Check ?exaApiKey=YOUR_KEY (fallback for backwards compat, only if no header) - if (!xApiKey && !getBearerToken(request) && params.has("exaApiKey")) { - const keyFromUrl = params.get("exaApiKey"); - if (keyFromUrl) { - exaApiKey = keyFromUrl; - userProvidedApiKey = true; - authMethod = "api_key"; - } - } - // Support ?tools=tool1,tool2 if (params.has("tools")) { const toolsParam = params.get("tools"); @@ -553,6 +600,7 @@ async function getConfigFromRequest(request: Request): Promise { defaultSearchType, oauthAccessToken, invalidOAuthJwt, + fallbackJwt, mcpMaxDurationSeconds: parsePositiveInteger(process.env.MCP_MAX_DURATION_SECONDS), agentCallWindowMs: agentCallWindowMs ?? parsePositiveInteger(process.env.AGENT_CALL_WINDOW_MS), }; @@ -682,6 +730,22 @@ async function processRequest( options?: { forceOAuth?: boolean; resourcePath?: string }, ): Promise { const debug = process.env.DEBUG === "true"; + if (request.method === "POST" && !acceptsMcpResponses(request.headers.get("accept"))) { + return withCors( + Response.json( + { + jsonrpc: "2.0", + error: { + code: -32000, + message: + "Not Acceptable: Client must accept both application/json and text/event-stream", + }, + id: null, + }, + { status: 406 }, + ), + ); + } const body = request.method === "POST" ? await request.clone().text() : undefined; if (body !== undefined && declaresJsonBody(request) && !isParsableJson(body)) { return createParseErrorResponse(); @@ -746,9 +810,15 @@ async function processRequest( stored: storedMcpClient, userAgent, }); + config.apiKeyFallback = createApiKeyFallback(config); if (config.debug) { - console.log(`[EXA-MCP] Request URL: ${request.url}`); + // Redact the API key from the logged URL; the request is only sanitized below. + const debugUrl = new URL(request.url); + if (debugUrl.searchParams.has("exaApiKey")) { + debugUrl.searchParams.set("exaApiKey", "[REDACTED]"); + } + console.log(`[EXA-MCP] Request URL: ${debugUrl}`); console.log(`[EXA-MCP] Enabled tools: ${config.enabledTools?.join(", ") || "default"}`); console.log(`[EXA-MCP] Auth method: ${config.authMethod}`); console.log( @@ -815,6 +885,11 @@ async function processRequest( // access without exposing it to downstream request instrumentation. url.searchParams.delete("exaApiKey"); const sanitizedHeaders = new Headers(request.headers); + if (request.method === "POST") { + // Validated above; the transport's own Accept check is a substring match, + // so hand it the canonical form. + sanitizedHeaders.set("accept", "application/json, text/event-stream"); + } sanitizedHeaders.delete("x-api-key"); sanitizedHeaders.delete("authorization"); request = new Request(url.toString(), { diff --git a/api/well-known-mcp-config.ts b/api/well-known-mcp-config.ts index b2d0eee3..0e36534e 100644 --- a/api/well-known-mcp-config.ts +++ b/api/well-known-mcp-config.ts @@ -22,7 +22,7 @@ const configSchema = { type: "string", title: "Enabled Tools", description: - "Comma-separated tools. When set, replaces defaults (web_search_exa, web_fetch_exa). agent_run requires OAuth or an API key.", + "Comma-separated tools. When set, replaces defaults (web_search_exa, web_fetch_exa, plus agent_run with OAuth or an API key). agent_run requires OAuth or an API key.", examples: [ "web_search_advanced_exa", "web_search_exa,web_fetch_exa,agent_run", diff --git a/env.example b/env.example index f7a0de77..345c8c20 100644 --- a/env.example +++ b/env.example @@ -2,14 +2,18 @@ EXA_API_KEY= # Comma-separated tools. When set, replaces the defaults. -# Defaults: web_search_exa, web_fetch_exa -# Optional: agent_run, web_search_advanced_exa +# Defaults: web_search_exa, web_fetch_exa, plus agent_run when EXA_API_KEY is set +# Optional: web_search_advanced_exa # Alias: agent_tools → agent_run # ENABLED_TOOLS=web_search_exa,web_fetch_exa,agent_run # DEBUG=true # DEFAULT_SEARCH_TYPE=auto # auto | fast | instant -# agent_run timing (optional) +# agent_run timing (optional): each call hands back the run ID for resumption +# once the call window passes. # MCP_MAX_DURATION_SECONDS=800 -# AGENT_CALL_WINDOW_MS=750000 +# AGENT_CALL_WINDOW_MS=45000 + +# Exa API origin override (optional; defaults to https://api.exa.ai) +# EXA_API_BASE_URL=https://api.exa.ai diff --git a/package.json b/package.json index 848af2a1..6cafd567 100644 --- a/package.json +++ b/package.json @@ -45,9 +45,10 @@ "format": "prettier --write \"**/*.{js,jsx,ts,tsx}\"", "format:check": "prettier --check \"**/*.{js,jsx,ts,tsx}\"", "test": "vitest run", + "test:integration": "vitest run --config vitest.integration.config.ts", "test:watch": "vitest", "test:coverage": "vitest run --coverage", - "ci": "npm run format:check && npm run typecheck && npm run test", + "ci": "npm run format:check && npm run typecheck && npm run test && npm run test:integration", "prepare": "npm run build", "watch": "./node_modules/.bin/tsc --watch", "dev": "tsx src/stdio-cli.ts", diff --git a/skills/exa-agent/SKILL.md b/skills/exa-agent/SKILL.md index 2654082c..dde8e6f0 100644 --- a/skills/exa-agent/SKILL.md +++ b/skills/exa-agent/SKILL.md @@ -1,6 +1,6 @@ --- name: Exa Agent -description: "Use Exa Agent for multi-step web research, list-building, enrichment, structured output, run continuation, and coverage validation. Exa Agent can access additional data providers: fiber, financial_datasets, similarweb, baselayer, affiliate, particle, and jinko." +description: "Use Exa Agent for multi-step web research, list-building, enrichment, structured output, run continuation, and coverage validation. Exa Agent can access additional data providers: fiber, financial_datasets, similarweb, baselayer, affiliate, particle, jinko, polymarket, and macrobond." --- # Exa Agent Research @@ -23,7 +23,9 @@ Use only the currently usable self-serve providers: - `baselayer`: US business verification, officers, registrations, and KYB - `affiliate`: product catalog search, pricing, brands, and merchant links - `particle`: podcast transcript search with speaker attribution and timestamps -- `jinko`: travel destination discovery ranked by fare +- `jinko`: flight and hotel search with real-time pricing, plus travel destination discovery +- `polymarket`: prediction market odds, price history, order books, and trader positions +- `macrobond`: macroeconomic and financial time series, entity metadata, and release calendars Do not suggest request-only providers unless the user explicitly says their Exa account already has them enabled. @@ -143,6 +145,7 @@ Example with Exa Connect: 1. Run the agent - Call `agent_run`. - Omit `effort` to use the tool's `low` default. Choose `auto` or a higher effort only when the user asks for more depth or the task clearly requires it. + - `ultra` is the highest-effort tier; use it only when the user explicitly wants exhaustive results. - Include `outputSchema` for structured work. - Use `input.data` for known rows. - Use `input.exclusion` for entities already returned or disallowed. @@ -150,11 +153,13 @@ Example with Exa Connect: - Name the provider-specific data you want in both the query and the schema so Agent uses the provider instead of falling back to web search. - Save the returned `id` when a later continuation may use `previousRunId`. - If the response has `status: "running"`, call `agent_run` again with that `runId` until `outputReady` is true. This continuation is available for retained runs that outlive one MCP call. - - Zero Data Retention (ZDR) teams: new runs always stream, and output is only available on that live stream (not via `runId` resumption). The MCP call window is ~750 seconds; if a ZDR run cannot finish in one call, retry with lower effort or split the task. `previousRunId` is not available on ZDR. + - Zero Data Retention (ZDR) teams: new runs always stream, and output is only available on that live stream (not via `runId` resumption). If a ZDR run cannot finish within one MCP call window (45 seconds by default), retry with lower effort or split the task. `previousRunId` is not available on ZDR. 2. Read the result + - `agent_run` declares an MCP `outputSchema`; every completed/running/failed/cancelled result carries the same object as `structuredContent` and as a JSON text block. Read `structuredContent` when your client exposes it; otherwise parse the text. + - Result fields: `success`, `id`, `status` (`completed` | `running` | `failed` | `cancelled`), `outputReady`, `output` (completed runs only; may be null), optional `stopReason`, `usage`, `costDollars`, `error` (failed), `message` (running). - Wait until `outputReady` is true (or status is failed/cancelled). - - Read both `output.structured` and `output.grounding`. + - Read `output.text`, `output.structured`, and `output.grounding`. `structured` follows your `outputSchema`; `grounding` carries per-field citations and confidence. - Do not assume results are exhaustive just because the run completed. 3. Validate coverage @@ -247,7 +252,8 @@ If a run fails to start: If the run fails: -- Explain the failure from the returned terminal status. +- Explain the failure from the returned terminal status (`status: "failed"`, `stopReason`, and `error` when present). +- A text-only error without `structuredContent` means the tool could not report a run result; unless the text names an `agent_run_...` ID, the upstream state is unknown, so do not retry blindly. - Create a corrected follow-up/new run only if the correction is clear. If the run objective/schema is wrong, abort the streaming call. The server will attempt to cancel the upstream run; you will then need to create a new run with the corrected objective/schema. diff --git a/src/mcp-handler.ts b/src/mcp-handler.ts index 3ed8f8be..7109893e 100644 --- a/src/mcp-handler.ts +++ b/src/mcp-handler.ts @@ -15,6 +15,7 @@ import { } from "./toolRegistry.js"; import { log } from "./utils/logger.js"; import { loadAgentSkillContent } from "./utils/agentSkill.js"; +import { withApiKeyFallback } from "./utils/apiKeyFallback.js"; import type { McpAnalytics } from "./analytics.js"; export interface McpConfig { @@ -40,6 +41,12 @@ export interface McpConfig { * the package ships no tracking backend. See {@link McpAnalytics}. */ analytics?: McpAnalytics; + /** + * Invoked when the Exa API rejects this request's API key (HTTP 401/403). + * Return true after switching this config to another credential (for + * example by setting `oauthAccessToken`) to retry the tool call once on it. + */ + apiKeyFallback?: () => Promise; } /** @@ -50,7 +57,11 @@ export interface McpConfig { * @param config - Configuration object with API key and tool settings */ export function initializeMcpServer(server: any, config: McpConfig = {}) { + const baseServer = server; try { + if (config.apiKeyFallback) { + server = withApiKeyFallback(server, config.apiKeyFallback); + } if (config.debug) { log("Initializing Exa MCP Server in debug mode"); if (config.enabledTools) { @@ -102,6 +113,7 @@ export function initializeMcpServer(server: any, config: McpConfig = {}) { agentCallWindowMs: config.agentCallWindowMs, mcpMaxDurationSeconds: config.mcpMaxDurationSeconds, }), + siblingTools: registeredTools, }); registeredTools.push("agent_run"); } @@ -232,7 +244,7 @@ export function initializeMcpServer(server: any, config: McpConfig = {}) { // Apply the embedder's analytics server wrapper, if any. The server object // might be wrapped by a transport handler, so pass the underlying server. if (config.analytics?.wrapServer) { - const underlyingServer = (server as any).server || server; + const underlyingServer = (baseServer as any).server || baseServer; try { config.analytics.wrapServer(underlyingServer); diff --git a/src/runtime-server.ts b/src/runtime-server.ts index c5fca4e4..a71494f7 100644 --- a/src/runtime-server.ts +++ b/src/runtime-server.ts @@ -1,107 +1,3 @@ -import { createServer, type IncomingMessage, type ServerResponse } from "node:http"; -import { Readable } from "node:stream"; -import { DELETE, GET, OPTIONS, POST } from "../api/mcp.js"; +import { startRuntimeServer } from "./runtime.js"; -const HOST = process.env.HOST || "0.0.0.0"; -const PORT = Number(process.env.PORT || 8000); - -type Handler = (request: Request) => Promise | Response; - -async function toRequest(request: IncomingMessage, body: Buffer): Promise { - const headers = new Headers(); - for (const [name, value] of Object.entries(request.headers)) { - if (Array.isArray(value)) { - headers.set(name, value.join(", ")); - } else if (value !== undefined) { - headers.set(name, value); - } - } - - const url = new URL(request.url || "/", `http://${HOST}:${PORT}`); - const hasInboundApiKey = - headers.has("x-api-key") || headers.has("authorization") || url.searchParams.has("exaApiKey"); - if (!hasInboundApiKey && process.env.EXA_API_KEY) { - headers.set("x-api-key", process.env.EXA_API_KEY); - } - - return new Request(url, { - method: request.method, - headers, - body: - request.method === "GET" || request.method === "HEAD" || body.length === 0 - ? undefined - : body.toString(), - }); -} - -async function readBody(request: IncomingMessage): Promise { - const chunks: Buffer[] = []; - for await (const chunk of request) { - chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk)); - } - return Buffer.concat(chunks); -} - -async function writeResponse(response: Response, target: ServerResponse): Promise { - target.statusCode = response.status; - target.statusMessage = response.statusText; - response.headers.forEach((value, key) => target.setHeader(key, value)); - - if (!response.body) { - target.end(); - return; - } - - await new Promise((resolve, reject) => { - Readable.fromWeb(response.body as import("node:stream/web").ReadableStream) - .on("error", reject) - .pipe(target) - .on("finish", resolve) - .on("error", reject); - }); -} - -function handlerForMethod(method: string | undefined): Handler | undefined { - switch (method) { - case "GET": - return GET; - case "POST": - return POST; - case "DELETE": - return DELETE; - case "OPTIONS": - return OPTIONS; - default: - return undefined; - } -} - -const server = createServer(async (request, response) => { - try { - if (request.method === "GET" && request.url?.split("?")[0] === "/ping") { - response.writeHead(200, { "Content-Type": "text/plain; charset=utf-8" }); - response.end("ok\n"); - return; - } - - const handler = handlerForMethod(request.method); - if (!handler) { - response.writeHead(405, { "Content-Type": "application/json" }); - response.end(JSON.stringify({ error: "Method not allowed" })); - return; - } - - const webRequest = await toRequest(request, await readBody(request)); - await writeResponse(await handler(webRequest), response); - } catch (error) { - console.error("[EXA-MCP] Runtime request failed:", error); - if (!response.headersSent) { - response.writeHead(500, { "Content-Type": "application/json" }); - } - response.end(JSON.stringify({ error: "Internal server error" })); - } -}); - -server.listen(PORT, HOST, () => { - console.log(`[EXA-MCP] AgentCore Runtime server listening on http://${HOST}:${PORT}`); -}); +startRuntimeServer(); diff --git a/src/runtime.ts b/src/runtime.ts new file mode 100644 index 00000000..1a427284 --- /dev/null +++ b/src/runtime.ts @@ -0,0 +1,243 @@ +/** + * Node HTTP server for the container runtime (see Dockerfile). + * + * Hosts the fetch-style MCP handlers from api/mcp.ts behind node:http. The + * process entrypoint is src/runtime-server.ts. + */ +import { createServer, type IncomingMessage, type Server, type ServerResponse } from "node:http"; +import { DELETE, GET, OPTIONS, POST } from "../api/mcp.js"; +import { parseRequestTarget } from "./utils/requestTarget.js"; + +type Handler = (request: Request) => Promise | Response; + +/** + * SSE responses idle between events (a long tool call sends nothing until it + * completes unless the client passed a progressToken), and proxies in front of + * the runtime close responses that stay idle too long. Comment lines are part + * of the SSE grammar and ignored by clients, so an idle-triggered comment keeps + * the connection alive without touching the event stream. + */ +export const SSE_KEEPALIVE_MS = 25_000; +const SSE_KEEPALIVE_COMMENT = ": keepalive\n\n"; +const NEWLINE = 0x0a; + +const DEFAULT_SHUTDOWN_TIMEOUT_SECS = 10; + +export interface RuntimeServerOptions { + /** Idle time after which an event stream gets a keepalive comment. */ + sseKeepaliveMs?: number; +} + +/** The client hung up (or the connection broke) before the request body arrived. */ +class ClientBodyAbortError extends Error {} + +function handlerForMethod(method: string | undefined): Handler | undefined { + switch (method) { + case "GET": + return GET; + case "POST": + return POST; + case "DELETE": + return DELETE; + case "OPTIONS": + return OPTIONS; + default: + return undefined; + } +} + +function sendJsonError(res: ServerResponse, status: number, error: string): void { + if (res.headersSent) { + // The response is already partly written; cut it off rather than append + // an error body that would make it look complete. + res.destroy(); + return; + } + res.writeHead(status, { "Content-Type": "application/json" }); + res.end(JSON.stringify({ error })); +} + +// The body is buffered rather than streamed: MCP request bodies are small +// JSON-RPC payloads, and a buffered body keeps a mid-body client abort as a +// distinguishable 400 instead of a stream error surfacing later as a 500. +async function toWebRequest(req: IncomingMessage, url: URL): Promise { + const headers = new Headers(); + for (const [name, value] of Object.entries(req.headers)) { + if (value === undefined) continue; + if (Array.isArray(value)) { + for (const item of value) headers.append(name, item); + } else { + headers.set(name, value); + } + } + + // The container's EXA_API_KEY is the operator's own key, not a shared + // free-tier key: requests that bring no credential run on it. + const hasInboundApiKey = + headers.has("x-api-key") || headers.has("authorization") || url.searchParams.has("exaApiKey"); + if (!hasInboundApiKey && process.env.EXA_API_KEY) { + headers.set("x-api-key", process.env.EXA_API_KEY); + } + + const method = req.method ?? "GET"; + let body: string | undefined; + if (method !== "GET" && method !== "HEAD") { + const chunks: Buffer[] = []; + try { + for await (const chunk of req) { + chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk)); + } + } catch (error) { + throw new ClientBodyAbortError("Client aborted while sending the request body", { + cause: error, + }); + } + if (chunks.length > 0) body = Buffer.concat(chunks).toString(); + } + + return new Request(url, { method, headers, body }); +} + +/** + * Write a web `Response` to a Node response. An event stream that goes + * `keepaliveMs` without output gets a keepalive comment, written only where + * an event ended so it never splits a partially written event. + */ +export async function writeWebResponse( + res: ServerResponse, + response: Response, + keepaliveMs = SSE_KEEPALIVE_MS, +): Promise { + const headers: Record = {}; + response.headers.forEach((value, key) => { + const existing = headers[key]; + if (existing === undefined) { + headers[key] = value; + } else { + headers[key] = Array.isArray(existing) ? [...existing, value] : [existing, value]; + } + }); + if (response.statusText) res.statusMessage = response.statusText; + res.writeHead(response.status, headers); + + const isEventStream = (response.headers.get("content-type") ?? "").includes("text/event-stream"); + let keepaliveTimer: NodeJS.Timeout | undefined; + // The last two bytes written; an event ends with a blank line ("\n\n"). + let tail = [NEWLINE, NEWLINE]; + if (isEventStream && response.body) { + keepaliveTimer = setInterval(() => { + if (!res.writableEnded && tail[0] === NEWLINE && tail[1] === NEWLINE) { + res.write(SSE_KEEPALIVE_COMMENT); + } + }, keepaliveMs); + keepaliveTimer.unref(); + } + + try { + if (response.body) { + for await (const chunk of response.body) { + keepaliveTimer?.refresh(); + res.write(chunk); + if (chunk.length >= 2) { + tail = [chunk[chunk.length - 2], chunk[chunk.length - 1]]; + } else if (chunk.length === 1) { + tail = [tail[1], chunk[0]]; + } + } + } + } finally { + if (keepaliveTimer) clearInterval(keepaliveTimer); + } + res.end(); +} + +/** Create the runtime HTTP server without binding it. */ +export function createRuntimeServer(options: RuntimeServerOptions = {}): Server { + const server = createServer((req, res) => { + // Once close() has stopped the listener, connections are not reused: a + // kept-alive connection would take new requests the shutdown then cuts + // off, and would hold the drain open until its keep-alive timeout. + res.on("finish", () => { + if (!server.listening) server.closeIdleConnections(); + }); + + const url = parseRequestTarget(req.url, req.headers.host); + // An unparseable target is a malformed request, not a server fault. + if (url === null) { + sendJsonError(res, 400, "Bad request target"); + return; + } + + if (req.method === "GET" && url.pathname === "/ping") { + res.writeHead(200, { "Content-Type": "text/plain; charset=utf-8" }); + res.end("ok\n"); + return; + } + + const handler = handlerForMethod(req.method); + if (!handler) { + sendJsonError(res, 405, "Method not allowed"); + return; + } + + toWebRequest(req, url) + .then((request) => handler(request)) + .then((response) => { + if (!server.listening) res.setHeader("Connection", "close"); + return writeWebResponse(res, response, options.sseKeepaliveMs); + }) + .catch((error: unknown) => { + if (error instanceof ClientBodyAbortError) { + sendJsonError(res, 400, "Client closed request"); + return; + } + console.error("[EXA-MCP] Runtime request failed:", error); + sendJsonError(res, 500, "Internal server error"); + }); + }); + return server; +} + +function shutdownTimeoutMs(): number { + const seconds = Number(process.env.SHUTDOWN_TIMEOUT_SECS || DEFAULT_SHUTDOWN_TIMEOUT_SECS); + return (Number.isFinite(seconds) && seconds > 0 ? seconds : DEFAULT_SHUTDOWN_TIMEOUT_SECS) * 1000; +} + +/** + * Serve on HOST:PORT until SIGTERM/SIGINT, then stop accepting connections, + * let in-flight requests finish, and exit (forced after SHUTDOWN_TIMEOUT_SECS, + * or at once on a second signal). + */ +export function startRuntimeServer(): Server { + // Backstop for rejections raised outside any request's promise chain, most + // often by a dependency that starts work without awaiting it. Node's default + // is to exit, which drops every in-flight request over one bad one; log + // instead and keep serving. + process.on("unhandledRejection", (reason) => { + console.error("[EXA-MCP] Unhandled promise rejection:", reason); + }); + + const host = process.env.HOST || "0.0.0.0"; + const port = Number(process.env.PORT || 8000); + const server = createRuntimeServer(); + server.listen(port, host, () => { + console.log(`[EXA-MCP] AgentCore Runtime server listening on http://${host}:${port}`); + }); + + const timeoutMs = shutdownTimeoutMs(); + let shuttingDown = false; + const shutdown = (signal: NodeJS.Signals): void => { + if (shuttingDown) { + console.log(`[EXA-MCP] Received ${signal} again, exiting without waiting for the drain`); + process.exit(0); + } + shuttingDown = true; + console.log(`[EXA-MCP] Received ${signal}, shutting down`); + server.close(() => process.exit(0)); + setTimeout(() => process.exit(0), timeoutMs).unref(); + }; + process.on("SIGTERM", shutdown); + process.on("SIGINT", shutdown); + + return server; +} diff --git a/src/toolRegistry.ts b/src/toolRegistry.ts index 7b312330..11f84f63 100644 --- a/src/toolRegistry.ts +++ b/src/toolRegistry.ts @@ -32,7 +32,7 @@ export const TOOL_REGISTRY = { name: "Run Exa Agent", description: "Run an Exa Agent for multi-step research, list-building, enrichment, or structured output. Returns the final output or a run ID for retained-run continuation.", - enabled: false, + enabled: true, group: "agent", requiresUserProvidedApiKey: true, }, diff --git a/src/tools/agentRun.ts b/src/tools/agentRun.ts index 49b86aae..d6984aba 100644 --- a/src/tools/agentRun.ts +++ b/src/tools/agentRun.ts @@ -8,7 +8,7 @@ import type { AgentEffort, AgentRunInput, ToolContent } from "../types.js"; import { formatAgentToolError } from "../utils/agentErrorHandler.js"; import { delay } from "../utils/errorHandler.js"; import { createRequestLogger } from "../utils/logger.js"; -import { jsonContent } from "../utils/response.js"; +import { structuredContent } from "../utils/response.js"; import { AgentProgressBridge, DEFAULT_PROGRESS_THROTTLE_MS, @@ -17,8 +17,11 @@ import { } from "./agentProgress.js"; import { createExaClient } from "./config.js"; -const effortSchema = z.enum(["minimal", "low", "medium", "high", "xhigh", "auto"]); +const effortSchema = z.enum(["minimal", "low", "medium", "high", "xhigh", "ultra", "auto"]); const recordSchema = () => z.record(z.unknown()); + +const agentRunIdSchema = () => + z.string().regex(/^agent_run_/, { message: 'Must start with "agent_run_"' }); const dataSourceProviderSchema = z.enum([ "fiber", "financial_datasets", @@ -27,8 +30,25 @@ const dataSourceProviderSchema = z.enum([ "affiliate", "particle", "jinko", + "polymarket", + "macrobond", ]); +/** + * Tool description shown to the client model. The cheaper alternatives are named only + * when registered in the same session, so an agent-only `?tools=agent_run` server never + * points the model at tools it cannot call. + */ +export function buildAgentRunDescription(siblingTools: Iterable = []): string { + const registered = new Set(siblingTools); + const searchTools = ["web_search_exa", "web_search_advanced_exa"].filter((t) => + registered.has(t), + ); + const searchAlt = searchTools.length > 0 ? ` (use ${searchTools.join(" or ")})` : ""; + const fetchAlt = registered.has("web_fetch_exa") ? " (use web_fetch_exa)" : ""; + return `Start or resume an Exa Agent run for multi-step research, list-building, or enrichment. Use it instead of running many searches yourself when a task needs multiple searches, cross-source verification, or a structured table of results. Long-running: returns a run ID; resume with runId when the tool reports the run is still running. Not for a single question or lookup${searchAlt} or reading a known URL${fetchAlt}. An interrupted tool call is not an explicit cancellation request.`; +} + export const agentRunInputShape = { query: z .string() @@ -37,9 +57,7 @@ export const agentRunInputShape = { .describe( "Natural-language research or enrichment objective. Provide query or runId, not both.", ), - runId: z - .string() - .startsWith("agent_run_") + runId: agentRunIdSchema() .optional() .describe( "agent_run_... ID returned by an earlier call. Use it to check or continue waiting for the same run; do not start a duplicate run.", @@ -67,22 +85,66 @@ export const agentRunInputShape = { .max(5) .optional() .describe("Optional Exa Connect providers to enable for this run."), - previousRunId: z - .string() - .startsWith("agent_run_") + previousRunId: agentRunIdSchema() .optional() .describe("Completed prior agent_run_... ID to use as context for a new run."), effort: effortSchema .optional() - .describe("Agent effort: minimal, low, medium, high, xhigh, or auto. Defaults to low."), + .describe("Agent effort: minimal, low, medium, high, xhigh, ultra, or auto. Defaults to low."), +}; + +/** + * `string | null` that serializes as `anyOf: [{type: "string"}, {type: "null"}]`. + * zod-to-json-schema collapses unions of check-free primitives into + * `type: ["string", "null"]`, which several MCP clients misread as a single + * type; the no-op `minLength: 0` keeps the branches separate. + */ +function nullableString(description: string) { + return z.union([z.string().min(0), z.null()]).describe(description); +} + +/** + * MCP `outputSchema` for agent_run (distinct from the Agent API `outputSchema` + * request field above). One flat object covers every reportable status rather + * than a per-status `oneOf`, because some clients validate `structuredContent` + * even on `isError` results and not all JSON Schema validators handle + * discriminators. Field names and semantics follow the Agent API run object; + * `output.structured` and `output.grounding` stay loose because they vary per + * run and per caller-supplied schema. + */ +export const agentRunOutputShape = { + success: z + .boolean() + .describe("True for completed and running runs; false for failed and cancelled."), + id: z.string().describe("agent_run_... ID. Pass as runId to resume a running run."), + status: z.enum(["completed", "running", "failed", "cancelled"]), + outputReady: z.boolean().describe("True only when status is completed."), + output: z + .object({ + text: nullableString("Prose answer.").optional(), + structured: z.unknown().describe("Output matching the caller-supplied outputSchema, if any."), + grounding: z.unknown().describe("Per-field citations and confidence."), + }) + .passthrough() + .nullable() + .optional(), + stopReason: nullableString( + "Why the run stopped: schema_satisfied, budget_reached, time_limit_reached, stopped, error, or cancelled.", + ).optional(), + usage: z.record(z.unknown()).optional(), + costDollars: z.record(z.unknown()).optional(), + error: z.record(z.unknown()).optional(), + message: z.string().optional().describe("Next-step guidance for running runs."), }; -// Default MCP function duration (seconds) and headroom before the platform kills -// the invocation. The call window is derived from these unless overridden. +const agentRunOutputSchema = z.object(agentRunOutputShape); + +// Default MCP function duration (seconds) and the headroom kept before the platform +// kills the invocation; together they cap the call window, which otherwise defaults +// to DEFAULT_CALL_WINDOW_MS so the run ID comes back before clients time out. export const DEFAULT_MCP_MAX_DURATION_SECONDS = 800; export const CALL_WINDOW_HEADROOM_MS = 50_000; -export const DEFAULT_CALL_WINDOW_MS = - DEFAULT_MCP_MAX_DURATION_SECONDS * 1000 - CALL_WINDOW_HEADROOM_MS; +export const DEFAULT_CALL_WINDOW_MS = 45_000; export const DEFAULT_HEARTBEAT_MS = 15_000; export const DEFAULT_POLL_INTERVAL_MS = 4_000; export const DEFAULT_PROGRESS_TIMEOUT_MS = 2_000; @@ -101,7 +163,7 @@ export function resolveAgentCallWindowMs(options?: { const ceiling = Number.isFinite(maxDurationSeconds) && maxDurationSeconds > 0 ? Math.max(1, maxDurationSeconds * 1000 - CALL_WINDOW_HEADROOM_MS) - : DEFAULT_CALL_WINDOW_MS; + : DEFAULT_MCP_MAX_DURATION_SECONDS * 1000 - CALL_WINDOW_HEADROOM_MS; if ( options?.agentCallWindowMs != null && @@ -111,7 +173,7 @@ export function resolveAgentCallWindowMs(options?: { return Math.min(Math.trunc(options.agentCallWindowMs), ceiling); } - return ceiling; + return Math.min(DEFAULT_CALL_WINDOW_MS, ceiling); } const TERMINAL_EVENTS = new Map([ @@ -336,6 +398,7 @@ export async function streamAgentRun(params: { const event = next.value.value; eventCount += 1; + if (runId == null) { const eventData = typeof event.data === "object" && event.data !== null && !Array.isArray(event.data) @@ -440,33 +503,81 @@ function outcomePayload(outcome: RunOutcome): Record { return outcome.terminalEvent?.data ?? {}; } +/** + * Validate an outgoing payload against the declared MCP output schema before it + * leaves the server. A mismatch means the upstream run shape drifted from the + * Agent API contract; surface it as an error rather than emit `structuredContent` + * that clients will reject. + */ +function agentRunResult(payload: Record): ToolContent { + const parsed = agentRunOutputSchema.safeParse(payload); + if (!parsed.success) { + throw new Error( + `agent_run result for ${String(payload.id)} did not match the declared output schema: ${parsed.error.issues + .map((issue) => `${issue.path.join(".")}: ${issue.message}`) + .join("; ")}`, + ); + } + return structuredContent(parsed.data); +} + +function validRunId(value: unknown): string | undefined { + return typeof value === "string" && value.length > 0 ? value : undefined; +} + +function missingRunIdResult(status: "completed" | "cancelled" | "failed" | "running"): ToolContent { + return { + content: [ + { + type: "text", + text: `agent_run error: the upstream run reported status=${status} without a run ID, so the result cannot be resumed or referenced.`, + }, + ], + isError: true, + }; +} + function outcomeToToolContent(outcome: RunOutcome): ToolContent { const run = outcomePayload(outcome); - const id = typeof run.id === "string" ? run.id : outcome.runId; + const id = validRunId(run.id) ?? validRunId(outcome.runId); + const stopReason = run.stopReason != null ? { stopReason: run.stopReason } : {}; switch (outcome.status) { case "completed": - return jsonContent({ + if (id == null) return missingRunIdResult(outcome.status); + return agentRunResult({ success: true, id, status: "completed", outputReady: true, output: run.output ?? null, + ...stopReason, ...(run.usage != null ? { usage: run.usage } : {}), ...(run.costDollars != null ? { costDollars: run.costDollars } : {}), }); case "cancelled": - return jsonContent({ success: false, id, status: "cancelled", outputReady: false }); - case "failed": + if (id == null) return missingRunIdResult(outcome.status); + return agentRunResult({ + success: false, + id, + status: "cancelled", + outputReady: false, + ...stopReason, + }); + case "failed": { + if (id == null) return missingRunIdResult(outcome.status); + const result = agentRunResult({ + success: false, + id, + status: "failed", + outputReady: false, + ...stopReason, + ...(run.error != null ? { error: run.error } : {}), + }); return { + ...result, content: [ - ...jsonContent({ - success: false, - id, - status: "failed", - outputReady: false, - ...(run.error != null ? { error: run.error } : {}), - }).content, + ...result.content, { type: "text", text: "The Agent run failed. Inspect the error above and verify any outputSchema before retrying.", @@ -474,8 +585,10 @@ function outcomeToToolContent(outcome: RunOutcome): ToolContent { ], isError: true, }; + } case "running": - return jsonContent({ + if (id == null) return missingRunIdResult(outcome.status); + return agentRunResult({ success: true, id, status: "running", @@ -524,6 +637,7 @@ export type AgentRunToolOptions = { pollIntervalMs?: number; progressTimeoutMs?: number; clientFactory?: (config: AgentRunConfig | undefined) => AgentRunClient; + siblingTools?: Iterable; }; function defaultClientFactory(config: AgentRunConfig | undefined): AgentRunClient { @@ -541,11 +655,19 @@ export function registerAgentRunTool( config?: AgentRunConfig, options?: AgentRunToolOptions, ): void { - server.tool( + server.registerTool( "agent_run", - "Start or resume an Exa Agent run; runs may take several minutes. Retain the returned run ID and resume with runId when the tool reports the run is still running. An interrupted tool call is not an explicit cancellation request.", - agentRunInputShape, - { readOnlyHint: true, destructiveHint: false, idempotentHint: false }, + { + description: buildAgentRunDescription(options?.siblingTools), + inputSchema: agentRunInputShape, + outputSchema: agentRunOutputShape, + annotations: { + readOnlyHint: true, + destructiveHint: false, + idempotentHint: false, + openWorldHint: true, + }, + }, async ( { query, runId, systemPrompt, outputSchema, input, dataSources, previousRunId, effort }, extra: StreamToolExtra, @@ -615,6 +737,8 @@ export function registerAgentRunTool( ...(systemPrompt != null ? { systemPrompt } : {}), ...(outputSchema != null ? { outputSchema } : {}), ...(input != null ? { input } : {}), + // Under non-strict compilation zod infers `provider` as + // optional; normalize so both modes typecheck. ...(dataSources != null ? { dataSources: dataSources.flatMap(({ provider }) => diff --git a/src/tools/config.ts b/src/tools/config.ts index 220a8c56..6480e50e 100644 --- a/src/tools/config.ts +++ b/src/tools/config.ts @@ -53,15 +53,26 @@ export function createExaClient(config?: Record, tool?: string) return exa; } +/** + * Construct an exa-js client for `apiKey`. `EXA_API_BASE_URL` overrides the + * Exa API origin the tools call (the integration suite points it at a local + * double); unset, exa-js keeps its default of https://api.exa.ai. + */ +function newExaClient(apiKey: string): Exa { + // exa-js appends endpoint paths ("/search") to the origin as given. + const baseUrl = process.env.EXA_API_BASE_URL?.replace(/\/+$/, ""); + return baseUrl ? new Exa(apiKey, baseUrl) : new Exa(apiKey); +} + function createBaseExaClient(config?: Record) { const oauthAccessToken = config?.oauthAccessToken; if (typeof oauthAccessToken === "string" && oauthAccessToken.length > 0) { - const exa = new Exa("oauth"); + const exa = newExaClient("oauth"); (exa as unknown as { headers: Headers }).headers.delete("x-api-key"); return exa; } const exaApiKey = config?.exaApiKey; - return new Exa( + return newExaClient( typeof exaApiKey === "string" && exaApiKey.length > 0 ? exaApiKey : process.env.EXA_API_KEY || "", diff --git a/src/tools/webSearch.ts b/src/tools/webSearch.ts index e8d0504e..6ffdde9f 100644 --- a/src/tools/webSearch.ts +++ b/src/tools/webSearch.ts @@ -8,6 +8,15 @@ import { sanitizeSearchResponse } from "../utils/exaResponseSanitizer.js"; import { lenientString, lenientOptionalNumber } from "./validation.js"; import type { McpAnalytics } from "../analytics.js"; + +// The Exa /search API caps `objective` at this length; enforcing it here keeps +// schema-valid tool calls from failing upstream validation. +const MAX_OBJECTIVE_LENGTH = 4096; + +// Tool-parameter wording the Exa API reference recommends for `objective`. +const OBJECTIVE_DESCRIPTION = + "Goal for this search turn; say which documents should rank first, which should be excluded, and what specific facts or figures to pull from them."; + type WebSearchConfig = { exaApiKey?: string; userProvidedApiKey?: boolean; @@ -36,6 +45,17 @@ export function registerWebSearchTool(server: McpServer, config?: WebSearchConfi numResults: lenientOptionalNumber().describe( "Number of search results to return (default: 10).", ), + objective: z + .preprocess( + // Clients that send unused optional parameters as null or "" search + // without an objective instead of failing the call. + (value) => + value === null || (typeof value === "string" && value.trim() === "") + ? undefined + : value, + z.string().trim().min(1).max(MAX_OBJECTIVE_LENGTH).optional(), + ) + .describe(OBJECTIVE_DESCRIPTION), }, { readOnlyHint: true, @@ -43,7 +63,7 @@ export function registerWebSearchTool(server: McpServer, config?: WebSearchConfi openWorldHint: true, idempotentHint: true, }, - async ({ query, numResults }) => { + async ({ query, numResults, objective }) => { const toolId = "web_search_exa"; const logger = createRequestLogger(toolId); @@ -73,6 +93,7 @@ export function registerWebSearchTool(server: McpServer, config?: WebSearchConfi type: config?.defaultSearchType || "auto", numResults: numResults || API_CONFIG.DEFAULT_NUM_RESULTS, ...(category && { category }), + ...(objective && { objective }), contents: { highlights: true, }, diff --git a/src/tools/webSearchAdvanced.ts b/src/tools/webSearchAdvanced.ts index 95de8c0f..c51dc7c3 100644 --- a/src/tools/webSearchAdvanced.ts +++ b/src/tools/webSearchAdvanced.ts @@ -132,7 +132,7 @@ Returns: Search results with optional highlights, summaries, and subpage content { readOnlyHint: true, destructiveHint: false, - openWorldHint: false, + openWorldHint: true, idempotentHint: true, }, async (params) => { diff --git a/src/types.ts b/src/types.ts index b8e0899a..b0845ad2 100644 --- a/src/types.ts +++ b/src/types.ts @@ -18,6 +18,7 @@ export interface ExaSearchRequest { numResults?: number; additionalQueries?: string[]; outputSchema?: Record; + objective?: string; contents: { text?: | { @@ -277,7 +278,7 @@ export interface ExaCodeResponse { traces?: any; } -export type AgentEffort = "minimal" | "low" | "medium" | "high" | "xhigh" | "auto"; +export type AgentEffort = "minimal" | "low" | "medium" | "high" | "xhigh" | "ultra" | "auto"; export type AgentDataSourceProvider = | "fiber" | "financial_datasets" @@ -285,7 +286,9 @@ export type AgentDataSourceProvider = | "baselayer" | "affiliate" | "particle" - | "jinko"; + | "jinko" + | "polymarket" + | "macrobond"; export type AgentRunInput = { query: string; @@ -304,5 +307,6 @@ export type AgentRunInput = { export type ToolContent = { content: Array<{ type: "text"; text: string }>; + structuredContent?: Record; isError?: true; }; diff --git a/src/utils/agentErrorHandler.ts b/src/utils/agentErrorHandler.ts index 92f8099e..47303837 100644 --- a/src/utils/agentErrorHandler.ts +++ b/src/utils/agentErrorHandler.ts @@ -2,11 +2,13 @@ import { ExaError } from "exa-js"; import type { ToolContent } from "../types.js"; import { EXA_API_KEYS_URL } from "./errorHandler.js"; +export const EXA_BILLING_URL = "https://dashboard.exa.ai/billing"; + export function formatAgentToolError(error: unknown, toolName: string): ToolContent { if (isExaError(error)) { const status = error.statusCode; const apiMessage = error.message; - const guidance = guidanceForStatus(status); + const guidance = guidanceForStatus(status, toolName); return { content: [ { @@ -40,10 +42,13 @@ function isExaError(error: unknown): error is ExaError { ); } -function guidanceForStatus(status: number | "unknown"): string { +function guidanceForStatus(status: number | "unknown", toolName: string): string { if (status === 400) { return "Check the run body and outputSchema. Use a top-level object schema, bound arrays with maxItems when possible, and use input.data for known rows."; } + if (status === 402) { + return `Not retryable: the account is out of credits or over its spending budget, so every ${toolName} call will be rejected until that changes. Do not call ${toolName} again; tell the user to add credits or raise the budget at ${EXA_BILLING_URL}.`; + } if (status === 401 || status === 403) { return `Authenticate with an Exa API key. API keys are available at ${EXA_API_KEYS_URL}.`; } diff --git a/src/utils/apiKeyFallback.ts b/src/utils/apiKeyFallback.ts new file mode 100644 index 00000000..63a10741 --- /dev/null +++ b/src/utils/apiKeyFallback.ts @@ -0,0 +1,160 @@ +/** + * Retry a tool call on the request's fallback credential when upstream rejects + * its API key. + * + * An HTTP entrypoint can see two credentials on one request: an API key (a + * connector's `?exaApiKey=`) and the Bearer JWT an OAuth-discovering client + * attaches alongside it. The key wins (api/mcp.ts). When the Exa API rejects + * that key, the JWT may still be the user's working credential: `fallback` + * (built per request and passed as `McpConfig.apiKeyFallback`) verifies it and + * switches the request config over, and the call runs once more on it. Any + * other failure returns as-is. + */ + +export type ToolHandler = (...args: unknown[]) => Promise; + +/** An MCP server (McpServer or mcp-handler's) as far as tool registration goes. */ +export interface ToolRegistrar { + tool: (...args: never[]) => unknown; + registerTool?: (...args: never[]) => unknown; +} + +const TOOL_REGISTRATION_METHODS = new Set(["tool", "registerTool"]); + +/** + * Wrap an MCP server so every `server.tool(name, ..., handler)` and + * `server.registerTool(name, config, handler)` registration passes its handler + * (the trailing function argument) through `wrap`. All other properties pass + * through. + */ +export function wrapToolRegistrations( + server: T, + wrap: (tool: string, handler: ToolHandler) => ToolHandler, +): T { + return new Proxy(server, { + get(target, prop, receiver) { + const value = Reflect.get(target, prop, receiver); + if (!TOOL_REGISTRATION_METHODS.has(prop) || typeof value !== "function") { + return typeof value === "function" ? value.bind(target) : value; + } + return (...args: unknown[]) => { + const name = args[0]; + const last = args[args.length - 1]; + if (typeof name === "string" && typeof last === "function") { + args[args.length - 1] = wrap(name, last as ToolHandler); + } + return (value as (...a: unknown[]) => unknown).apply(target, args); + }; + }, + }); +} + +/** Whether a tool result is an MCP tool error (`isError: true`). */ +function isToolError(result: unknown): boolean { + return ( + typeof result === "object" && + result !== null && + (result as { isError?: unknown }).isError === true + ); +} + +/** The text of a tool result's first content block, or "" when it has none. */ +function firstText(result: unknown): string { + const content = (result as { content?: unknown }).content; + const first = Array.isArray(content) ? (content[0] as { text?: unknown } | undefined) : undefined; + return typeof first?.text === "string" ? first.text : ""; +} + +function numericStatus(value: unknown): number | undefined { + if (typeof value !== "object" || value === null) return undefined; + const candidate = value as { + status?: unknown; + statusCode?: unknown; + code?: unknown; + response?: { status?: unknown }; + }; + for (const status of [ + candidate.status, + candidate.statusCode, + candidate.code, + candidate.response?.status, + ]) { + if (typeof status === "number") return status; + } + return undefined; +} + +/** + * The status formatToolError and formatAgentToolError put at the start of a + * rejected call's error text: " error (): ". + */ +const TOOL_ERROR_STATUS = /^\S+ error \((\d{3})\):/; + +/** Whether a tool outcome (returned result or thrown error) is an upstream API-key rejection. */ +export function isInvalidApiKeyOutcome(errorOrResult: unknown): boolean { + if (isToolError(errorOrResult)) { + // A result with structuredContent reports an Agent run's own outcome (a + // failed run), and the run was created, so this call's key worked. Only + // the leading status of an error text says the API rejected this call; + // a status or "unauthorized" further in is detail from elsewhere (a + // failed run's error, a fetched page's URL). + if ((errorOrResult as { structuredContent?: unknown }).structuredContent !== undefined) { + return false; + } + const status = firstText(errorOrResult).match(TOOL_ERROR_STATUS)?.[1]; + return status === "401" || status === "403"; + } + + if (!(errorOrResult instanceof Error)) return false; + const code = (errorOrResult as { code?: unknown }).code; + const codeText = typeof code === "string" ? code : ""; + if ( + errorOrResult.name === "AbortError" || + /abort|timed?\s*out|timeout/i.test(errorOrResult.message) || + /timeout|aborted/i.test(codeText) + ) { + return false; + } + if ( + /invalid[_\s-]*(input|argument)|validation|zod/i.test(`${codeText} ${errorOrResult.message}`) + ) { + return false; + } + const status = numericStatus(errorOrResult); + return status === 401 || status === 403; +} + +/** + * Wrap `handler` so an invalid-API-key outcome triggers `fallback()` and, if + * that switched credentials, a single retry with the same arguments. + */ +export function withApiKeyFallbackHandler( + handler: ToolHandler, + fallback: () => Promise, +): ToolHandler { + return async (...args: unknown[]): Promise => { + let outcome: unknown; + try { + outcome = await handler(...args); + } catch (error) { + if (isInvalidApiKeyOutcome(error) && (await fallback())) { + return handler(...args); + } + throw error; + } + if (isInvalidApiKeyOutcome(outcome) && (await fallback())) { + return handler(...args); + } + return outcome; + }; +} + +/** Apply `withApiKeyFallbackHandler` to every tool registered on `server`. */ +export function withApiKeyFallback( + server: T, + fallback: () => Promise, +): T { + return wrapToolRegistrations(server, (_tool, handler) => + withApiKeyFallbackHandler(handler, fallback), + ); +} diff --git a/src/utils/mcpAccept.ts b/src/utils/mcpAccept.ts new file mode 100644 index 00000000..5a2e7805 --- /dev/null +++ b/src/utils/mcpAccept.ts @@ -0,0 +1,75 @@ +/** Parse the explicit response media types required by MCP Streamable HTTP. */ +const TOKEN = /^[!#$%&'*+.^_`|~0-9A-Za-z-]+$/; +const QUOTED_VALUE = /^"(?:[\t !#-\[\]-~\x80-\xff]|\\[\t !-~\x80-\xff])*"$/; + +/** Split an HTTP list without treating delimiters inside quoted strings as separators. */ +function splitHeader(value: string, separator: string): string[] | undefined { + const parts: string[] = []; + let start = 0; + let quoted = false; + let escaped = false; + for (let index = 0; index < value.length; index++) { + const character = value[index]; + if (escaped) { + escaped = false; + } else if (quoted && character === "\\") { + escaped = true; + } else if (character === '"') { + quoted = !quoted; + } else if (!quoted && character === separator) { + parts.push(value.slice(start, index).trim()); + start = index + 1; + } + } + if (quoted || escaped) return undefined; + parts.push(value.slice(start).trim()); + return parts; +} + +/** The media type and quality of one media range, or undefined when the range is malformed. */ +function parseMediaRange(range: string): { mediaType: string; quality: number } | undefined { + const parts = splitHeader(range, ";"); + if (!parts) return undefined; + const [mediaType, ...parameters] = parts; + const components = mediaType.split("/"); + if (components.length !== 2 || !components.every((part) => TOKEN.test(part))) return undefined; + let quality = 1; + let hasQuality = false; + for (const parameter of parameters) { + // RFC 9110 allows empty parameters (`type/subtype;`). + if (!parameter) continue; + const equals = parameter.indexOf("="); + if (equals < 1) return undefined; + const name = parameter.slice(0, equals).trim(); + const parameterValue = parameter.slice(equals + 1).trim(); + if (!TOKEN.test(name) || (!TOKEN.test(parameterValue) && !QUOTED_VALUE.test(parameterValue))) { + return undefined; + } + if (name.toLowerCase() === "q") { + if (hasQuality || !/^(?:0(?:\.\d{0,3})?|1(?:\.0{0,3})?)$/.test(parameterValue)) { + return undefined; + } + quality = Number(parameterValue); + hasQuality = true; + } + } + return { mediaType: mediaType.toLowerCase(), quality }; +} + +/** + * Require both concrete MCP response types, each in a well-formed media range + * with nonzero quality. Other malformed ranges in the list are ignored rather + * than failing the whole header. + */ +export function acceptsMcpResponses(value: string | null): boolean { + if (!value || /[\x00-\x08\x0a-\x1f\x7f]/.test(value)) return false; + const ranges = splitHeader(value, ","); + if (!ranges) return false; + const accepted = new Set(); + for (const range of ranges) { + if (!range) continue; + const parsed = parseMediaRange(range); + if (parsed && parsed.quality > 0) accepted.add(parsed.mediaType); + } + return accepted.has("application/json") && accepted.has("text/event-stream"); +} diff --git a/src/utils/requestTarget.ts b/src/utils/requestTarget.ts new file mode 100644 index 00000000..4f151a1e --- /dev/null +++ b/src/utils/requestTarget.ts @@ -0,0 +1,25 @@ +/** + * Parsing of the raw HTTP request target (`IncomingMessage.url`). + * + * The request target and the Host header are both fully client-controlled and + * are not guaranteed to form a valid URL, so parsing them is a validation step + * that reports failure rather than an operation that always succeeds. + */ + +/** + * Parse a request target against its Host header into an absolute URL. + * + * Returns null when the pair cannot form a URL — most notably a target of + * `//`, which WHATWG parsing reads as a protocol-relative URL with an empty + * host and rejects, and a Host header that is not a valid authority. + */ +export function parseRequestTarget( + rawTarget: string | undefined, + hostHeader: string | undefined, +): URL | null { + try { + return new URL(rawTarget ?? "/", `http://${hostHeader ?? "localhost"}`); + } catch { + return null; + } +} diff --git a/src/utils/response.ts b/src/utils/response.ts index c553d0b2..058d4458 100644 --- a/src/utils/response.ts +++ b/src/utils/response.ts @@ -11,6 +11,20 @@ export function jsonContent(value: unknown): ToolContent { }; } +/** + * Result for a tool that declares an MCP `outputSchema`: the object goes out as + * `structuredContent` and, per the spec's back-compat guidance, as a serialized + * JSON text block. The text block is compact (no indentation) since clients + * forward exactly one of the two channels to the model and the text is the + * fallback for those that ignore `structuredContent`. + */ +export function structuredContent(value: Record): ToolContent { + return { + content: [{ type: "text", text: JSON.stringify(value) }], + structuredContent: value, + }; +} + export function clampInteger( value: number | undefined, fallback: number, diff --git a/tests/helpers/advertisedTools.ts b/tests/helpers/advertisedTools.ts new file mode 100644 index 00000000..301c44fe --- /dev/null +++ b/tests/helpers/advertisedTools.ts @@ -0,0 +1,43 @@ +import { Client } from "@modelcontextprotocol/sdk/client/index.js"; +import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js"; +import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import type { Tool } from "@modelcontextprotocol/sdk/types.js"; + +/** + * An MCP client connected in memory to a real McpServer set up by `register`, + * so requests pass through the SDK's own schema conversion and validation. + * Close the client when done. + */ +export async function connectInMemory(register: (server: McpServer) => void): Promise { + const server = new McpServer({ name: "exa-mcp-test", version: "0.0.0" }); + register(server); + const [serverTransport, clientTransport] = InMemoryTransport.createLinkedPair(); + await server.connect(serverTransport); + const client = new Client({ name: "exa-mcp-test-client", version: "0.0.0" }); + await client.connect(clientTransport); + return client; +} + +/** + * The tools a real McpServer advertises on `tools/list` after `register` runs, + * with input/output schemas converted to JSON Schema exactly as clients see + * them. + */ +export async function listAdvertisedTools(register: (server: McpServer) => void): Promise { + const client = await connectInMemory(register); + try { + return (await client.listTools()).tools; + } finally { + await client.close(); + } +} + +/** Every `pattern` string in a JSON Schema, keyed by its JSON path. */ +export function collectPatterns(schema: unknown, path = "$"): Array<[string, string]> { + if (schema === null || typeof schema !== "object") return []; + return Object.entries(schema).flatMap(([key, value]) => + key === "pattern" && typeof value === "string" + ? [[`${path}.pattern`, value] as [string, string]] + : collectPatterns(value, `${path}.${key}`), + ); +} diff --git a/tests/helpers/fakeMcpServer.ts b/tests/helpers/fakeMcpServer.ts index 3dd8e436..137c73cc 100644 --- a/tests/helpers/fakeMcpServer.ts +++ b/tests/helpers/fakeMcpServer.ts @@ -2,10 +2,19 @@ export interface RegisteredTool { name: string; description?: string; inputSchema?: unknown; + outputSchema?: unknown; annotations?: unknown; handler: (args: any, extra?: unknown) => unknown | Promise; } +interface RegisterToolConfig { + title?: string; + description?: string; + inputSchema?: unknown; + outputSchema?: unknown; + annotations?: unknown; +} + export interface RegisteredPrompt { name: string; description?: string; @@ -46,6 +55,17 @@ export class FakeMcpServer { }); } + registerTool(name: string, config: RegisterToolConfig, handler: RegisteredTool["handler"]): void { + this.tools.push({ + name, + description: config.description, + inputSchema: config.inputSchema, + outputSchema: config.outputSchema, + annotations: config.annotations, + handler, + }); + } + prompt(...args: unknown[]): void { const [name] = args; const handler = getLastArg(args); diff --git a/tests/integration/accept.test.ts b/tests/integration/accept.test.ts new file mode 100644 index 00000000..a0dd5078 --- /dev/null +++ b/tests/integration/accept.test.ts @@ -0,0 +1,44 @@ +/** Accept negotiation through the runtime server's real HTTP transport. */ +import { beforeAll, describe, expect, it } from "vitest"; +import { startInProcessServer, type ServedMcp } from "./harness.js"; + +let served: ServedMcp; + +beforeAll(async () => { + served = await startInProcessServer({ EXA_API_KEY: "" }); +}); + +describe("MCP Accept negotiation", () => { + it.each([ + "application/json;text/event-stream", + "application/json.text/event-stream", + 'application/json;note="text/event-stream"', + "application/json, text/event-stream;q=0", + "application/json;q=0, text/event-stream", + "application/json, text/event-stream;q=invalid", + "application/json", + "text/event-stream", + "", + ])("rejects %j before dispatching the request", async (accept) => { + const response = await fetch(`${served.url}/mcp`, { + method: "POST", + headers: { "Content-Type": "application/json", Accept: accept }, + body: JSON.stringify({ jsonrpc: "2.0", id: 1, method: "tools/list" }), + }); + expect(response.status).toBe(406); + expect(await response.json()).toMatchObject({ error: { code: -32000 } }); + }); + + it.each([ + "application/json, text/event-stream", + 'APPLICATION/JSON;note="comma,semi;colon", TEXT/EVENT-STREAM;q=0.5', + ])("serves tools/list with %j", async (accept) => { + const response = await fetch(`${served.url}/mcp`, { + method: "POST", + headers: { "Content-Type": "application/json", Accept: accept }, + body: JSON.stringify({ jsonrpc: "2.0", id: 1, method: "tools/list" }), + }); + expect(response.status).toBe(200); + expect(await response.text()).toContain('"tools":'); + }); +}); diff --git a/tests/integration/apiKeyFallback.test.ts b/tests/integration/apiKeyFallback.test.ts new file mode 100644 index 00000000..33ba49e1 --- /dev/null +++ b/tests/integration/apiKeyFallback.test.ts @@ -0,0 +1,154 @@ +/** + * The `?exaApiKey=` + Bearer JWT connector shape through the served MCP + * endpoint (what an OAuth-discovering client sends once an admin has put the + * org's key in the connector URL): + * + * - a working URL key answers on its own; the JWT is never verified; + * - a URL key the Exa API rejects makes the server verify the JWT and run the + * same tool call again on it; + * - a URL key the Exa API rejects with no usable JWT stays a tool error. + * + * The authorization server is a local issuer: `OAUTH_ISSUER` points the + * server's JWKS fetch at a keypair this file generates, so the JWT it mints + * verifies exactly as an auth.exa.ai token does. + */ +import { createServer, type Server } from "node:http"; +import type { AddressInfo } from "node:net"; +import { afterAll, beforeAll, beforeEach, describe, expect, it } from "vitest"; +import type { Client } from "@modelcontextprotocol/sdk/client/index.js"; +import * as jose from "jose"; +import { + connectMcpClient, + FakeExaApi, + searchResults, + startInProcessServer, + type ServedMcp, +} from "./harness.js"; + +const AUDIENCE = "https://mcp.exa.ai"; +const ORG_KEY = "org-url-key"; +const REVOKED_ORG_KEY = "revoked-org-url-key"; + +let exaApi: FakeExaApi; +let issuer: Server; +let issuerUrl: string; +let signingKey: CryptoKey; +let served: ServedMcp; +const clients: Client[] = []; + +beforeAll(async () => { + exaApi = await FakeExaApi.start(); + + const { publicKey, privateKey } = await jose.generateKeyPair("RS256"); + signingKey = privateKey; + const jwk = { ...(await jose.exportJWK(publicKey)), kid: "integration", alg: "RS256" }; + issuer = createServer((req, res) => { + if (req.url === "/api/oauth/jwks") { + res.writeHead(200, { "Content-Type": "application/json" }); + res.end(JSON.stringify({ keys: [jwk] })); + return; + } + res.writeHead(404); + res.end(); + }); + await new Promise((resolve) => issuer.listen(0, "127.0.0.1", resolve)); + issuerUrl = `http://127.0.0.1:${(issuer.address() as AddressInfo).port}`; + + served = await startInProcessServer({ + EXA_API_BASE_URL: exaApi.url, + EXA_API_KEY: "", + OAUTH_ISSUER: issuerUrl, + OAUTH_AUDIENCE: AUDIENCE, + }); +}); + +afterAll(async () => { + await Promise.all(clients.map((client) => client.close())); + await exaApi.close(); + await new Promise((resolve) => issuer.close(() => resolve())); +}); + +beforeEach(() => { + // The Exa API double: the org's live key and the OAuth identity's Bearer + // token both search; the revoked key is rejected the way api.exa.ai does. + exaApi.respondWith((call) => { + const apiKey = call.headers["x-api-key"]; + const authorization = call.headers["authorization"]; + if (apiKey === ORG_KEY || (authorization ?? "").startsWith("Bearer ey")) { + return { status: 200, body: searchResults("https://exa.ai/") }; + } + return { status: 401, body: { error: "Invalid API key" } }; + }); +}); + +/** A JWT the server accepts as an Exa OAuth access token for `sub`, valid until `expiresIn`. */ +async function accessToken(sub: string, expiresIn = "5m"): Promise { + return new jose.SignJWT({ "exa:team_id": "team-oauth", scope: "mcp:tools" }) + .setProtectedHeader({ alg: "RS256", kid: "integration" }) + .setIssuer(issuerUrl) + .setAudience(AUDIENCE) + .setSubject(sub) + .setIssuedAt() + .setExpirationTime(expiresIn) + .sign(signingKey); +} + +async function connectorClient(urlKey: string, jwt: string | undefined): Promise { + const headers: Record = jwt ? { Authorization: `Bearer ${jwt}` } : {}; + const connected = await connectMcpClient(served, headers, { exaApiKey: urlKey }); + clients.push(connected); + return connected; +} + +async function search(mcp: Client) { + return mcp.callTool({ name: "web_search_exa", arguments: { query: "exa" } }); +} + +describe("?exaApiKey= alongside a Bearer JWT", () => { + it("searches on a working URL key and never touches the JWT", async () => { + const mcp = await connectorClient(ORG_KEY, await accessToken("user-1")); + + const result = await search(mcp); + + expect(result.isError).toBeFalsy(); + const calls = exaApi.callsTo("/search"); + expect(calls).toHaveLength(1); + expect(calls[0].headers["x-api-key"]).toBe(ORG_KEY); + expect(calls[0].headers["authorization"]).toBeUndefined(); + }); + + it("retries a search on the JWT when the URL key is rejected upstream", async () => { + const jwt = await accessToken("user-2"); + const mcp = await connectorClient(REVOKED_ORG_KEY, jwt); + + const result = await search(mcp); + + expect(result.isError).toBeFalsy(); + const calls = exaApi.callsTo("/search"); + expect(calls).toHaveLength(2); + expect(calls[0].headers["x-api-key"]).toBe(REVOKED_ORG_KEY); + expect(calls[0].headers["authorization"]).toBeUndefined(); + expect(calls[1].headers["authorization"]).toBe(`Bearer ${jwt}`); + expect(calls[1].headers["x-api-key"]).toBeUndefined(); + expect(calls[1].body).toEqual(calls[0].body); + }); + + it("keeps the rejection when the JWT alongside the rejected key is expired", async () => { + const jwt = await accessToken("user-3", "-5m"); + const mcp = await connectorClient(REVOKED_ORG_KEY, jwt); + + const result = await search(mcp); + + expect(result.isError).toBe(true); + expect(exaApi.callsTo("/search")).toHaveLength(1); + }); + + it("keeps the rejection when the URL key is the only credential", async () => { + const mcp = await connectorClient(REVOKED_ORG_KEY, undefined); + + const result = await search(mcp); + + expect(result.isError).toBe(true); + expect(exaApi.callsTo("/search")).toHaveLength(1); + }); +}); diff --git a/tests/integration/gracefulShutdown.test.ts b/tests/integration/gracefulShutdown.test.ts new file mode 100644 index 00000000..cd954991 --- /dev/null +++ b/tests/integration/gracefulShutdown.test.ts @@ -0,0 +1,92 @@ +/** + * Graceful shutdown of the runtime entrypoint: on SIGTERM/SIGINT the server + * stops accepting connections, finishes the tool call already in flight, and + * exits cleanly well inside its drain timeout — without waiting out kept-alive + * connections. A second signal exits without waiting for the drain. + */ +import { afterAll, afterEach, beforeAll, describe, expect, it } from "vitest"; +import { + FakeExaApi, + postMcp, + searchResults, + ServerProcess, + waitFor, + webSearchCall, +} from "./harness.js"; + +/** Far above the test's own wait, so a clean exit proves drain rather than the force-exit timer. */ +const DRAIN_TIMEOUT_SECS = 60; +const UPSTREAM_DELAY_MS = 1_500; +/** Well under the 5s keep-alive timeout an idle client connection would otherwise hold the drain for. */ +const EXIT_AFTER_DRAIN_MS = 2_500; + +let exaApi: FakeExaApi; +let server: ServerProcess | undefined; + +beforeAll(async () => { + exaApi = await FakeExaApi.start(); +}); + +afterEach(() => server?.stop()); + +afterAll(async () => { + await exaApi.close(); +}); + +describe("graceful shutdown", () => { + it.each(["SIGTERM", "SIGINT"] as const)( + "drains the in-flight tool call on %s, then exits 0", + async (signal) => { + exaApi.respondWith(() => ({ + status: 200, + body: searchResults("https://exa.ai/"), + delayMs: UPSTREAM_DELAY_MS, + })); + server = await ServerProcess.start({ + EXA_API_BASE_URL: exaApi.url, + EXA_API_KEY: "", + SHUTDOWN_TIMEOUT_SECS: String(DRAIN_TIMEOUT_SECS), + }); + const inFlight = postMcp(server, webSearchCall("exa"), { "x-api-key": "caller-key" }); + await waitFor(() => exaApi.callsTo("/search").length === 1, "the upstream search call"); + + server.signal(signal); + await waitFor( + () => server!.logs.includes(`Received ${signal}, shutting down`), + "the shutdown log line", + ); + + await expect(fetch(`${server.url}/ping`)).rejects.toThrow(); + const reply = await inFlight; + const drainedAt = Date.now(); + expect(reply.status).toBe(200); + expect(JSON.stringify(reply.message?.result)).toContain("URL: https://exa.ai/"); + const exit = await server.exited; + expect(exit).toEqual({ code: 0, signal: null }); + expect(Date.now() - drainedAt).toBeLessThan(EXIT_AFTER_DRAIN_MS); + }, + ); + + it("exits at once on a second signal instead of waiting for the drain", async () => { + exaApi.respondWith(() => ({ + status: 200, + body: searchResults("https://exa.ai/"), + delayMs: DRAIN_TIMEOUT_SECS * 1000, + })); + server = await ServerProcess.start({ + EXA_API_BASE_URL: exaApi.url, + EXA_API_KEY: "", + SHUTDOWN_TIMEOUT_SECS: String(DRAIN_TIMEOUT_SECS), + }); + const inFlight = postMcp(server, webSearchCall("exa"), { "x-api-key": "caller-key" }); + await waitFor(() => exaApi.callsTo("/search").length === 1, "the upstream search call"); + + server.signal("SIGTERM"); + await waitFor(() => server!.logs.includes("Received SIGTERM, shutting down"), "shutdown"); + server.signal("SIGTERM"); + + expect(await server.exited).toEqual({ code: 0, signal: null }); + expect(server.logs).toContain("Received SIGTERM again, exiting without waiting for the drain"); + await expect(inFlight).rejects.toThrow(); + }); +}); diff --git a/tests/integration/harness.ts b/tests/integration/harness.ts new file mode 100644 index 00000000..60753f33 --- /dev/null +++ b/tests/integration/harness.ts @@ -0,0 +1,321 @@ +/** + * Integration harness for exa-mcp-server. + * + * Every suite drives the container runtime's HTTP server (src/runtime.ts) over + * loopback, exactly as an MCP client reaches a deployed instance, with the Exa + * API replaced by `FakeExaApi` — an HTTP double that `EXA_API_BASE_URL` points + * the tools at. Each test scripts the replies it needs and inspects the calls + * the server made. + * + * `startInProcessServer` serves from the test process; `ServerProcess` runs the + * entrypoint (src/runtime-server.ts) as its own Node process for behavior that + * ends the process (signal handling). + */ +import { spawn, type ChildProcess } from "node:child_process"; +import { + createServer, + type IncomingHttpHeaders, + type IncomingMessage, + type Server, + type ServerResponse, +} from "node:http"; +import { createServer as createNetServer, type AddressInfo } from "node:net"; +import { fileURLToPath } from "node:url"; +import { Client } from "@modelcontextprotocol/sdk/client/index.js"; +import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"; +import type { RuntimeServerOptions } from "../../src/runtime.js"; + +/** Repository root (tests/integration/..). */ +export const PROJECT_ROOT = fileURLToPath(new URL("../..", import.meta.url)); + +/** The TCP port a listening server is bound to. */ +function listeningPort(server: { address(): AddressInfo | string | null }): number { + const address = server.address(); + if (address === null || typeof address === "string") { + throw new Error(`expected a listening TCP server, got address ${String(address)}`); + } + return address.port; +} + +/** Bind `server` to an ephemeral loopback port. */ +async function listenOnLoopback(server: { + once(event: "error", listener: (error: Error) => void): unknown; + listen(port: number, host: string, callback: () => void): unknown; +}): Promise { + await new Promise((resolveListen, reject) => { + server.once("error", reject); + server.listen(0, "127.0.0.1", resolveListen); + }); +} + +/** Reserve an ephemeral loopback port and release it for the caller to bind. */ +export async function freePort(): Promise { + const server = createNetServer(); + await listenOnLoopback(server); + const port = listeningPort(server); + await new Promise((resolveClose) => server.close(() => resolveClose())); + return port; +} + +/** Poll `check` until it returns true, failing with `what` after `timeoutMs`. */ +export async function waitFor( + check: () => Promise | boolean, + what: string, + timeoutMs = 15_000, +): Promise { + const deadline = Date.now() + timeoutMs; + let lastError: unknown; + while (Date.now() < deadline) { + try { + if (await check()) return; + } catch (error) { + lastError = error; + } + await new Promise((r) => setTimeout(r, 50)); + } + throw new Error(`Timed out after ${timeoutMs}ms waiting for ${what}`, { cause: lastError }); +} + +/** One request the server sent to the Exa API double. */ +export interface ExaApiCall { + method: string; + path: string; + headers: IncomingHttpHeaders; + body: unknown; +} + +/** The double's answer to one call. */ +export interface ExaApiReply { + status: number; + body: unknown; + /** Hold the response this long before answering. */ + delayMs?: number; +} + +type ExaApiHandler = (call: ExaApiCall) => ExaApiReply; + +/** + * HTTP double of the Exa API that records calls and answers with scripted + * replies. A call the double itself cannot serve (unparsable body, throwing + * handler) is answered 500 and recorded in `failures`, so the server sees an + * upstream error and the test can tell the two apart. + */ +export class FakeExaApi { + readonly calls: ExaApiCall[] = []; + readonly failures: unknown[] = []; + private readonly server: Server; + private handler: ExaApiHandler = (call) => ({ + status: 500, + body: { error: `FakeExaApi: no reply scripted for ${call.method} ${call.path}` }, + }); + + private constructor() { + this.server = createServer((req, res) => { + this.serve(req, res).catch((error: unknown) => { + this.failures.push(error); + if (!res.headersSent) res.writeHead(500, { "Content-Type": "application/json" }); + res.end(JSON.stringify({ error: `FakeExaApi failed: ${String(error)}` })); + }); + }); + } + + static async start(): Promise { + const api = new FakeExaApi(); + await listenOnLoopback(api.server); + return api; + } + + get url(): string { + return `http://127.0.0.1:${listeningPort(this.server)}`; + } + + /** Script the reply for every following call and forget earlier calls. */ + respondWith(handler: ExaApiHandler): void { + this.calls.length = 0; + this.handler = handler; + } + + /** Calls made to `path`, in arrival order. */ + callsTo(path: string): ExaApiCall[] { + return this.calls.filter((call) => call.path === path); + } + + async close(): Promise { + this.server.closeAllConnections(); + await new Promise((resolveClose) => this.server.close(() => resolveClose())); + } + + private async serve(req: IncomingMessage, res: ServerResponse): Promise { + const chunks: Buffer[] = []; + for await (const chunk of req) { + chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(String(chunk))); + } + const raw = Buffer.concat(chunks).toString("utf8"); + const call: ExaApiCall = { + method: req.method ?? "GET", + path: new URL(req.url ?? "/", "http://fake").pathname, + headers: req.headers, + body: raw ? JSON.parse(raw) : undefined, + }; + this.calls.push(call); + const reply = this.handler(call); + if (reply.delayMs) await new Promise((r) => setTimeout(r, reply.delayMs)); + res.writeHead(reply.status, { "Content-Type": "application/json" }); + res.end(JSON.stringify(reply.body)); + } +} + +/** Where a started server serves MCP. */ +export interface ServedMcp { + url: string; +} + +/** + * Serve src/runtime.ts inside the test process with `env` applied. The env is + * set before the server's module graph loads, since some settings (the OAuth + * issuer) are read at import; call it once per file (vitest runs each file in + * its own process). + */ +export async function startInProcessServer( + env: Record, + options?: RuntimeServerOptions, +): Promise { + Object.assign(process.env, env); + const { createRuntimeServer } = await import("../../src/runtime.js"); + const server = createRuntimeServer(options); + await listenOnLoopback(server); + return { url: `http://127.0.0.1:${listeningPort(server)}` }; +} + +/** How a `ServerProcess` ended. */ +export interface ProcessExit { + code: number | null; + signal: NodeJS.Signals | null; +} + +/** The runtime entrypoint (src/runtime-server.ts) running as its own Node process. */ +export class ServerProcess implements ServedMcp { + private output = ""; + + private constructor( + private readonly child: ChildProcess, + readonly url: string, + readonly exited: Promise, + ) { + child.stdout?.on("data", (chunk: Buffer) => (this.output += chunk.toString())); + child.stderr?.on("data", (chunk: Buffer) => (this.output += chunk.toString())); + } + + static async start(env: Record): Promise { + const port = await freePort(); + const child = spawn(process.execPath, ["--import", "tsx", "src/runtime-server.ts"], { + cwd: PROJECT_ROOT, + env: { ...process.env, HOST: "127.0.0.1", PORT: String(port), ...env }, + stdio: ["ignore", "pipe", "pipe"], + }); + const exited = new Promise((resolveExit) => + child.once("exit", (code, signal) => resolveExit({ code, signal })), + ); + const served = new ServerProcess(child, `http://127.0.0.1:${port}`, exited); + await waitFor(async () => (await fetch(`${served.url}/ping`)).ok, `${served.url}/ping`); + return served; + } + + /** Everything the process wrote to stdout and stderr so far. */ + get logs(): string { + return this.output; + } + + signal(signal: NodeJS.Signals): void { + this.child.kill(signal); + } + + /** Kill the process if it is still running. */ + stop(): void { + if (this.child.exitCode === null && this.child.signalCode === null) { + this.child.kill("SIGKILL"); + } + } +} + +/** + * Connect an MCP SDK client to the served `/mcp` endpoint, sending `headers` + * on every request. `query` is appended to the endpoint URL, the way a + * connector configured with `?exaApiKey=` reaches the server. + */ +export async function connectMcpClient( + served: ServedMcp, + headers: Record = {}, + query: Record = {}, +): Promise { + const client = new Client({ name: "exa-mcp-integration", version: "1.0.0" }); + const endpoint = new URL("/mcp", served.url); + endpoint.search = new URLSearchParams(query).toString(); + const transport = new StreamableHTTPClientTransport(endpoint, { + requestInit: { headers }, + }); + await client.connect(transport); + return client; +} + +/** A raw JSON-RPC exchange with `/mcp`: the HTTP status, body, and decoded message. */ +export interface McpHttpReply { + status: number; + headers: Headers; + text: string; + message: { result?: unknown; error?: { code: number; message: string } } | undefined; +} + +/** + * POST one JSON-RPC message to `/mcp` without an MCP client, for assertions + * on the HTTP layer. Decodes both JSON and single-event SSE replies. `path` + * may carry the endpoint's query parameters (`/mcp?tools=...`). + */ +export async function postMcp( + served: ServedMcp, + message: Record, + headers: Record = {}, + path = "/mcp", +): Promise { + const response = await fetch(new URL(path, served.url), { + method: "POST", + headers: { + "Content-Type": "application/json", + Accept: "application/json, text/event-stream", + ...headers, + }, + body: JSON.stringify({ jsonrpc: "2.0", ...message }), + }); + const text = await response.text(); + const payload = (response.headers.get("content-type") ?? "").includes("text/event-stream") + ? text + .split("\n") + .find((line) => line.startsWith("data: ")) + ?.slice("data: ".length) + : text; + return { + status: response.status, + headers: response.headers, + text, + message: payload ? JSON.parse(payload) : undefined, + }; +} + +/** A `tools/call` JSON-RPC message for `web_search_exa`. */ +export function webSearchCall(query: string, id = 1): Record { + return { id, method: "tools/call", params: { name: "web_search_exa", arguments: { query } } }; +} + +/** A minimal successful Exa `/search` response body. */ +export function searchResults(...urls: string[]): Record { + return { + requestId: "req-integration", + results: urls.map((url, i) => ({ + id: url, + url, + title: `Result ${i + 1}`, + highlights: [`highlight for ${url}`], + })), + costDollars: { total: 0.005 }, + }; +} diff --git a/tests/integration/webSearch.test.ts b/tests/integration/webSearch.test.ts new file mode 100644 index 00000000..60cd8d53 --- /dev/null +++ b/tests/integration/webSearch.test.ts @@ -0,0 +1,144 @@ +/** + * web_search_exa through the served MCP endpoint: what an MCP client gets back + * when the Exa API answers, answers slowly, fails transiently, or rejects the + * key — and which key the server spends on the caller's behalf. + */ +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import type { Client } from "@modelcontextprotocol/sdk/client/index.js"; +import { + connectMcpClient, + FakeExaApi, + postMcp, + searchResults, + startInProcessServer, + webSearchCall, + type ServedMcp, +} from "./harness.js"; + +const CONTAINER_KEY = "container-env-key"; +const SSE_KEEPALIVE_MS = 50; + +let exaApi: FakeExaApi; +let served: ServedMcp; +const clients: Client[] = []; + +beforeAll(async () => { + exaApi = await FakeExaApi.start(); + served = await startInProcessServer( + { EXA_API_BASE_URL: exaApi.url, EXA_API_KEY: CONTAINER_KEY }, + { sseKeepaliveMs: SSE_KEEPALIVE_MS }, + ); +}); + +afterAll(async () => { + await Promise.all(clients.map((client) => client.close())); + await exaApi.close(); +}); + +async function client(headers: Record = {}): Promise { + const connected = await connectMcpClient(served, headers); + clients.push(connected); + return connected; +} + +/** The text of a tool result's single content block. */ +function resultText(result: Awaited>): string { + const content = result.content as Array<{ type: string; text: string }>; + expect(content).toHaveLength(1); + return content[0].text; +} + +describe("web_search_exa over MCP", () => { + it("advertises search, fetch, and agent_run to a caller with their own key", async () => { + const mcp = await client({ "x-api-key": "caller-key" }); + + const { tools } = await mcp.listTools(); + + expect(tools.map((tool) => tool.name)).toEqual([ + "web_search_exa", + "web_fetch_exa", + "agent_run", + ]); + }); + + it("searches with the caller's own API key and renders the results", async () => { + exaApi.respondWith(() => ({ + status: 200, + body: searchResults("https://exa.ai/", "https://docs.exa.ai/"), + })); + const mcp = await client({ "x-api-key": "caller-key" }); + + const result = await mcp.callTool({ + name: "web_search_exa", + arguments: { + query: "neural search engines", + numResults: 2, + objective: "Rank engineering blog posts first", + }, + }); + + expect(result.isError).toBeFalsy(); + expect(resultText(result)).toContain("Title: Result 1\nURL: https://exa.ai/"); + expect(resultText(result)).toContain("URL: https://docs.exa.ai/"); + const [call] = exaApi.callsTo("/search"); + expect(call.headers["x-api-key"]).toBe("caller-key"); + expect(call.headers["x-exa-integration"]).toBe("web-search-mcp"); + expect(call.body).toMatchObject({ + query: "neural search engines", + numResults: 2, + objective: "Rank engineering blog posts first", + }); + }); + + it("runs credential-less callers on the container's EXA_API_KEY", async () => { + exaApi.respondWith(() => ({ status: 200, body: searchResults("https://exa.ai/") })); + const mcp = await client(); + + const result = await mcp.callTool({ name: "web_search_exa", arguments: { query: "exa" } }); + + expect(result.isError).toBeFalsy(); + expect(exaApi.callsTo("/search")[0].headers["x-api-key"]).toBe(CONTAINER_KEY); + }); + + it("retries transient Exa API failures before answering", async () => { + let attempts = 0; + exaApi.respondWith(() => + ++attempts <= 2 + ? { status: 503, body: { error: "Service Unavailable" } } + : { status: 200, body: searchResults("https://exa.ai/") }, + ); + const mcp = await client({ "x-api-key": "caller-key" }); + + const result = await mcp.callTool({ name: "web_search_exa", arguments: { query: "exa" } }); + + expect(result.isError).toBeFalsy(); + expect(resultText(result)).toContain("URL: https://exa.ai/"); + expect(exaApi.callsTo("/search")).toHaveLength(3); + }); + + it("reports a rejected API key as a tool error without retrying", async () => { + exaApi.respondWith(() => ({ status: 401, body: { error: "Invalid API key" } })); + const mcp = await client({ "x-api-key": "revoked-key" }); + + const result = await mcp.callTool({ name: "web_search_exa", arguments: { query: "exa" } }); + + expect(result.isError).toBe(true); + expect(resultText(result)).toContain("web_search_exa error (401): Invalid API key"); + expect(exaApi.callsTo("/search")).toHaveLength(1); + }); + + it("keeps a slow tool call's event stream alive with SSE comments", async () => { + exaApi.respondWith(() => ({ + status: 200, + body: searchResults("https://exa.ai/"), + delayMs: SSE_KEEPALIVE_MS * 6, + })); + + const reply = await postMcp(served, webSearchCall("exa"), { "x-api-key": "caller-key" }); + + expect(reply.status).toBe(200); + expect(reply.headers.get("content-type")).toContain("text/event-stream"); + expect(reply.text).toContain(": keepalive\n\n"); + expect(JSON.stringify(reply.message?.result)).toContain("URL: https://exa.ai/"); + }); +}); diff --git a/tests/unit/api/mcp.test.ts b/tests/unit/api/mcp.test.ts index 2c5329b2..c05c3bce 100644 --- a/tests/unit/api/mcp.test.ts +++ b/tests/unit/api/mcp.test.ts @@ -109,6 +109,9 @@ function expectMcpCorsHeaders(response: Response) { } async function callHandleRequest(request: Request, options?: { forceOAuth?: boolean }) { + if (!request.headers.has("accept")) { + request.headers.set("accept", "application/json, text/event-stream"); + } const { handleRequest } = await import("../../../api/mcp.js"); const response = await handleRequest(request, options); const config = initializeMcpServerMock.mock.calls.at(-1)?.[1]; @@ -460,6 +463,108 @@ describe("api/mcp handler", () => { expect(initializeMcpServerMock).not.toHaveBeenCalled(); }); + it("prefers an exaApiKey query parameter over a valid OAuth JWT", async () => { + verifyOAuthTokenMock.mockResolvedValue({ + sub: "user-1", + "exa:team_id": "team-1", + scope: "mcp:tools", + }); + + const { response, config, forwardedRequest } = await callHandleRequest( + new Request("https://mcp.exa.ai/mcp?exaApiKey=query-key", { + headers: { + authorization: "Bearer jwt-token", + }, + }), + ); + + expect(response.status).toBe(200); + expect(verifyOAuthTokenMock).not.toHaveBeenCalled(); + expect(config).toMatchObject({ + exaApiKey: "query-key", + userProvidedApiKey: true, + authMethod: "api_key", + fallbackJwt: "jwt-token", + }); + expect(config.oauthAccessToken).toBeUndefined(); + expect(forwardedRequest?.headers.get("authorization")).toBeNull(); + expect(new URL(forwardedRequest?.url ?? "").searchParams.has("exaApiKey")).toBe(false); + }); + + it("switches the config to OAuth when the URL key is rejected and the JWT verifies", async () => { + verifyOAuthTokenMock.mockResolvedValue({ + sub: "user-1", + "exa:team_id": "team-1", + scope: "mcp:tools", + }); + + const { config } = await callHandleRequest( + new Request("https://mcp.exa.ai/mcp?exaApiKey=query-key", { + headers: { + authorization: "Bearer jwt-token", + }, + }), + ); + + await expect(config.apiKeyFallback()).resolves.toBe(true); + expect(verifyOAuthTokenMock).toHaveBeenCalledWith("jwt-token"); + expect(config).toMatchObject({ + exaApiKey: undefined, + oauthAccessToken: "jwt-token", + userProvidedApiKey: true, + authMethod: "oauth", + }); + + // One verification per request, however many tool calls fail. + await expect(config.apiKeyFallback()).resolves.toBe(false); + expect(verifyOAuthTokenMock).toHaveBeenCalledTimes(1); + }); + + it("keeps the URL key when it is rejected and the JWT is invalid", async () => { + verifyOAuthTokenMock.mockResolvedValue(null); + + const { response, config } = await callHandleRequest( + new Request("https://mcp.exa.ai/mcp?exaApiKey=query-key", { + headers: { + authorization: "Bearer invalid-jwt", + }, + }), + ); + + expect(response.status).toBe(200); + expect(verifyOAuthTokenMock).not.toHaveBeenCalled(); + + await expect(config.apiKeyFallback()).resolves.toBe(false); + expect(verifyOAuthTokenMock).toHaveBeenCalledWith("invalid-jwt"); + expect(config).toMatchObject({ exaApiKey: "query-key", authMethod: "api_key" }); + expect(config.oauthAccessToken).toBeUndefined(); + }); + + it("offers no fallback when the URL key is the only credential", async () => { + const { config } = await callHandleRequest( + new Request("https://mcp.exa.ai/mcp?exaApiKey=query-key"), + ); + + expect(config.fallbackJwt).toBeUndefined(); + expect(config.apiKeyFallback).toBeUndefined(); + }); + + it("ignores an empty exaApiKey query parameter and falls through to the JWT", async () => { + verifyOAuthTokenMock.mockResolvedValue(null); + + const { response, config } = await callHandleRequest( + new Request("https://mcp.exa.ai/mcp?exaApiKey=", { + headers: { + authorization: "Bearer invalid-jwt", + }, + }), + ); + + expect(verifyOAuthTokenMock).toHaveBeenCalledWith("invalid-jwt"); + expect(response.status).toBe(401); + expect(config).toBeUndefined(); + }); + it("uses exaApiKey query parameters when no key header is present", async () => { const { config, forwardedRequest } = await callHandleRequest( new Request("https://mcp.exa.ai/mcp?exaApiKey=query-key"), @@ -726,6 +831,74 @@ describe("api/mcp handler", () => { expectMcpCorsHeaders(response); }); + it.each([ + ["no Accept header", undefined], + ["application/json only", "application/json"], + ["text/event-stream only", "text/event-stream"], + ["a zero-quality media range", "application/json, text/event-stream;q=0"], + ["a malformed media range", "application/json;text/event-stream"], + ])("rejects POST requests with %s as 406 Not Acceptable", async (_label, accept) => { + const { handleRequest } = await import("../../../api/mcp.js"); + const headers = new Headers({ "Content-Type": "application/json" }); + if (accept !== undefined) headers.set("accept", accept); + + const response = await handleRequest( + new Request("https://mcp.exa.ai/mcp", { + method: "POST", + headers, + body: JSON.stringify({ jsonrpc: "2.0", id: 1, method: "tools/list" }), + }), + ); + + expect(response.status).toBe(406); + await expect(response.json()).resolves.toMatchObject({ + jsonrpc: "2.0", + error: { code: -32000 }, + id: null, + }); + expectMcpCorsHeaders(response); + expect(createMcpHandlerMock).not.toHaveBeenCalled(); + }); + + it("hands the MCP transport a canonical Accept header once the media ranges validate", async () => { + const { response, forwardedRequest } = await callHandleRequest( + new Request("https://mcp.exa.ai/mcp", { + method: "POST", + headers: { + Accept: 'APPLICATION/JSON;note="comma,semi;colon", TEXT/EVENT-STREAM;q=0.5', + "Content-Type": "application/json", + }, + body: JSON.stringify({ jsonrpc: "2.0", id: 1, method: "tools/list" }), + }), + ); + + expect(response.status).toBe(200); + expect(forwardedRequest?.headers.get("accept")).toBe("application/json, text/event-stream"); + }); + + it("does not require an Accept header on GET requests", async () => { + const { handleRequest } = await import("../../../api/mcp.js"); + + const response = await handleRequest(new Request("https://mcp.exa.ai/mcp")); + + expect(response.status).toBe(200); + expect(createMcpHandlerMock).toHaveBeenCalled(); + }); + + it("redacts the URL API key from the debug request log", async () => { + await callHandleRequest( + new Request( + "https://mcp.exa.ai/mcp?debug=true&exaApiKey=secret-url-key&tools=web_search_exa", + ), + ); + + const logged = vi.mocked(console.log).mock.calls.map((args) => args.join(" ")); + expect(logged).toContain( + "[EXA-MCP] Request URL: https://mcp.exa.ai/mcp?debug=true&exaApiKey=%5BREDACTED%5D&tools=web_search_exa", + ); + expect(logged.join("\n")).not.toContain("secret-url-key"); + }); + it("rejects unparsable JSON bodies before they reach the MCP handler", async () => { const { response } = await callHandleRequest( new Request("https://mcp.exa.ai/mcp", { diff --git a/tests/unit/contract.test.ts b/tests/unit/contract.test.ts index ba6645c1..3c63c093 100644 --- a/tests/unit/contract.test.ts +++ b/tests/unit/contract.test.ts @@ -40,6 +40,17 @@ describe("public API contract", () => { expect(server.tools.map((tool) => tool.name)).toEqual(["web_search_exa", "web_fetch_exa"]); }); + it("adds agent_run to the defaults when the caller brings their own credential", () => { + const server = new FakeMcpServer(); + initializeMcpServer(server, { exaApiKey: "embedder-key", userProvidedApiKey: true }); + + expect(server.tools.map((tool) => tool.name)).toEqual([ + "web_search_exa", + "web_fetch_exa", + "agent_run", + ]); + }); + it("exposes only the supported v4 tool IDs", () => { expect(publicApi.AVAILABLE_TOOL_IDS).toEqual([ "web_search_exa", diff --git a/tests/unit/mcp-handler.test.ts b/tests/unit/mcp-handler.test.ts index 39485420..deee0e5a 100644 --- a/tests/unit/mcp-handler.test.ts +++ b/tests/unit/mcp-handler.test.ts @@ -1,7 +1,28 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; -import { initializeMcpServer } from "../../src/mcp-handler.js"; +import { ExaError } from "exa-js"; +import { initializeMcpServer, type McpConfig } from "../../src/mcp-handler.js"; +import { searchResponse } from "../fixtures/exaResponses.js"; import { FakeMcpServer } from "../helpers/fakeMcpServer.js"; +const { ExaMock, exaConstructorMock, requestMock } = vi.hoisted(() => { + const requestMock = vi.fn(); + const exaConstructorMock = vi.fn(); + class ExaMock { + headers = new Headers(); + request = requestMock; + + constructor(...args: unknown[]) { + exaConstructorMock(...args); + } + } + return { ExaMock, exaConstructorMock, requestMock }; +}); + +vi.mock("exa-js", async (importOriginal) => ({ + ...(await importOriginal()), + Exa: ExaMock, +})); + describe("initializeMcpServer", () => { beforeEach(() => { vi.clearAllMocks(); @@ -38,6 +59,103 @@ describe("initializeMcpServer", () => { ); }); + it("registers agent_run by default when the user provided an API key", async () => { + const server = new FakeMcpServer(); + + initializeMcpServer(server, { userProvidedApiKey: true }); + + expect(server.tools.map((tool) => tool.name)).toEqual([ + "web_search_exa", + "web_fetch_exa", + "agent_run", + ]); + expect(server.prompts.map((prompt) => prompt.name)).toEqual([ + "web_search_help", + "agent_research_help", + ]); + expect(server.resources.map((resource) => resource.name)).toEqual([ + "tools_list", + "agent_research_guide", + "agent_schema_templates", + ]); + + const resourceResult = await server.resources[0].handler(); + const toolsList = JSON.parse((resourceResult as any).contents[0].text); + expect(toolsList).toEqual( + expect.arrayContaining([ + expect.objectContaining({ id: "web_search_exa", enabled: true }), + expect.objectContaining({ id: "web_fetch_exa", enabled: true }), + expect.objectContaining({ id: "agent_run", enabled: true }), + expect.objectContaining({ id: "web_search_advanced_exa", enabled: false }), + ]), + ); + }); + + it("does not register agent_run by default for free-tier callers", () => { + const server = new FakeMcpServer(); + + initializeMcpServer(server, { userProvidedApiKey: false }); + + expect(server.tools.map((tool) => tool.name)).toEqual(["web_search_exa", "web_fetch_exa"]); + expect(server.prompts.map((prompt) => prompt.name)).toEqual(["web_search_help"]); + }); + + it("names only the registered search and fetch tools in the agent_run description", () => { + const withSearchOnly = new FakeMcpServer(); + initializeMcpServer(withSearchOnly, { + enabledTools: ["web_search_exa", "agent_run"], + userProvidedApiKey: true, + }); + const description = withSearchOnly.getTool("agent_run").description; + expect(description).toContain("(use web_search_exa)"); + expect(description).not.toContain("web_fetch_exa"); + + const agentOnly = new FakeMcpServer(); + initializeMcpServer(agentOnly, { enabledTools: ["agent_run"], userProvidedApiKey: true }); + expect(agentOnly.getTool("agent_run").description).not.toMatch(/web_(search|fetch)/); + }); + + it("retries a tool call on the fallback credential after an upstream API-key rejection", async () => { + const server = new FakeMcpServer(); + const config: McpConfig = { + exaApiKey: "revoked-url-key", + userProvidedApiKey: true, + apiKeyFallback: vi.fn(async () => { + config.exaApiKey = undefined; + config.oauthAccessToken = "jwt-token"; + return true; + }), + }; + requestMock + .mockRejectedValueOnce(new ExaError("x-api-key header is invalid", 401)) + .mockResolvedValueOnce(searchResponse); + + initializeMcpServer(server, config); + const result = await server.getTool("web_search_exa").handler({ query: "exa" }); + + expect(result).not.toMatchObject({ isError: true }); + expect(config.apiKeyFallback).toHaveBeenCalledTimes(1); + expect(exaConstructorMock.mock.calls.map(([apiKey]) => apiKey)).toEqual([ + "revoked-url-key", + "oauth", + ]); + expect(requestMock).toHaveBeenCalledTimes(2); + expect(requestMock.mock.calls[1][4]).toMatchObject({ Authorization: "Bearer jwt-token" }); + }); + + it("keeps the analytics wrapServer hook on the underlying server when a fallback is set", () => { + const wrapServer = vi.fn(); + const server = new FakeMcpServer(); + + initializeMcpServer(server, { + analytics: { wrapServer }, + apiKeyFallback: async () => false, + }); + + expect(wrapServer).toHaveBeenCalledWith(server.server); + expect(server.tools.map((tool) => tool.name)).toEqual(["web_search_exa", "web_fetch_exa"]); + }); + it("registers only supported tools from an explicit selection", () => { const server = new FakeMcpServer(); diff --git a/tests/unit/runtime.test.ts b/tests/unit/runtime.test.ts new file mode 100644 index 00000000..4202a92c --- /dev/null +++ b/tests/unit/runtime.test.ts @@ -0,0 +1,158 @@ +import { connect, type AddressInfo } from "node:net"; +import type { Server } from "node:http"; +import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; + +const { handlerMock } = vi.hoisted(() => ({ handlerMock: vi.fn() })); + +vi.mock("../../api/mcp.js", () => ({ + GET: handlerMock, + POST: handlerMock, + DELETE: handlerMock, + OPTIONS: handlerMock, +})); + +const { createRuntimeServer } = await import("../../src/runtime.js"); + +const KEEPALIVE_MS = 20; + +let server: Server; +let baseUrl: string; + +beforeAll(async () => { + server = createRuntimeServer({ sseKeepaliveMs: KEEPALIVE_MS }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + baseUrl = `http://127.0.0.1:${(server.address() as AddressInfo).port}`; +}); + +afterAll(async () => { + server.closeAllConnections(); + await new Promise((resolve) => server.close(() => resolve())); +}); + +beforeEach(() => { + handlerMock.mockReset(); + delete process.env.EXA_API_KEY; + vi.spyOn(console, "error").mockImplementation(() => undefined); +}); + +const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); + +/** A response body the test writes chunk by chunk, pausing between writes. */ +function scriptedBody(steps: Array): ReadableStream { + const encoder = new TextEncoder(); + return new ReadableStream({ + async start(controller) { + for (const step of steps) { + if (typeof step === "number") await sleep(step); + else controller.enqueue(encoder.encode(step)); + } + controller.close(); + }, + }); +} + +/** Send a raw HTTP request line (bypassing fetch's URL normalization) and return the status line. */ +async function rawStatusLine(requestLine: string): Promise { + const { port } = server.address() as AddressInfo; + return new Promise((resolve, reject) => { + const socket = connect(port, "127.0.0.1", () => { + socket.write(`${requestLine}\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n`); + }); + let received = ""; + socket.on("data", (data) => (received += data.toString())); + socket.on("end", () => resolve(received.split("\r\n")[0])); + socket.on("error", reject); + }); +} + +describe("runtime server", () => { + it("answers /ping without reaching the MCP handler", async () => { + const response = await fetch(`${baseUrl}/ping`); + + expect(response.status).toBe(200); + await expect(response.text()).resolves.toBe("ok\n"); + expect(handlerMock).not.toHaveBeenCalled(); + }); + + it("rejects an unparseable request target with 400", async () => { + await expect(rawStatusLine("GET // HTTP/1.1")).resolves.toBe("HTTP/1.1 400 Bad Request"); + expect(handlerMock).not.toHaveBeenCalled(); + }); + + it("rejects unsupported methods with 405", async () => { + const response = await fetch(`${baseUrl}/mcp`, { method: "PUT" }); + + expect(response.status).toBe(405); + await expect(response.json()).resolves.toEqual({ error: "Method not allowed" }); + expect(handlerMock).not.toHaveBeenCalled(); + }); + + it("forwards the request and runs credential-less callers on the container's EXA_API_KEY", async () => { + process.env.EXA_API_KEY = "container-key"; + handlerMock.mockResolvedValue(new Response("{}", { status: 200 })); + + await fetch(`${baseUrl}/mcp?tools=web_search_exa`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: '{"jsonrpc":"2.0","id":1,"method":"tools/list"}', + }); + await fetch(`${baseUrl}/mcp`, { + method: "POST", + headers: { Authorization: "Bearer caller-key" }, + body: "{}", + }); + + const [anonymous, withCredential] = handlerMock.mock.calls.map( + ([request]) => request as Request, + ); + expect(anonymous.method).toBe("POST"); + expect(new URL(anonymous.url).pathname).toBe("/mcp"); + expect(new URL(anonymous.url).searchParams.get("tools")).toBe("web_search_exa"); + expect(anonymous.headers.get("x-api-key")).toBe("container-key"); + await expect(anonymous.text()).resolves.toBe('{"jsonrpc":"2.0","id":1,"method":"tools/list"}'); + expect(withCredential.headers.get("x-api-key")).toBeNull(); + expect(withCredential.headers.get("authorization")).toBe("Bearer caller-key"); + }); + + it("answers 500 when the MCP handler throws", async () => { + handlerMock.mockRejectedValue(new Error("transport failed")); + + const response = await fetch(`${baseUrl}/mcp`); + + expect(response.status).toBe(500); + await expect(response.json()).resolves.toEqual({ error: "Internal server error" }); + }); + + it("keeps an idle event stream alive with comments written only between events", async () => { + handlerMock.mockResolvedValue( + new Response( + scriptedBody([ + "event: message\ndata: 1\n\n", + KEEPALIVE_MS * 4, + "event: message\n", + KEEPALIVE_MS * 4, + "data: 2\n\n", + ]), + { headers: { "Content-Type": "text/event-stream" } }, + ), + ); + + const text = await (await fetch(`${baseUrl}/mcp`)).text(); + + expect(text).toMatch( + /^event: message\ndata: 1\n\n(: keepalive\n\n)+event: message\ndata: 2\n\n$/, + ); + }); + + it("never adds keepalives to responses that are not event streams", async () => { + handlerMock.mockResolvedValue( + new Response(scriptedBody(['{"a":', KEEPALIVE_MS * 4, "1}"]), { + headers: { "Content-Type": "application/json" }, + }), + ); + + const response = await fetch(`${baseUrl}/mcp`); + + await expect(response.text()).resolves.toBe('{"a":1}'); + }); +}); diff --git a/tests/unit/toolSchemas.test.ts b/tests/unit/toolSchemas.test.ts new file mode 100644 index 00000000..e73bda4c --- /dev/null +++ b/tests/unit/toolSchemas.test.ts @@ -0,0 +1,103 @@ +/** + * Tool definitions as clients see them on `tools/list`. + * + * JSON Schema defines `pattern` as an ECMA-262 regex, and strict model + * providers compile it in unicode mode, where an escape of a non-syntax + * character (`\_`) is an error. One such pattern makes the provider reject + * every request that carries the tool list, so every advertised pattern must + * compile with the `u` flag. + */ +import { beforeAll, describe, expect, it } from "vitest"; +import type { Tool } from "@modelcontextprotocol/sdk/types.js"; +import { initializeMcpServer } from "../../src/mcp-handler.js"; +import { AVAILABLE_TOOL_IDS } from "../../src/toolRegistry.js"; +import { collectPatterns, listAdvertisedTools } from "../helpers/advertisedTools.js"; + +/** The compile error of `pattern` as a unicode-mode ECMA-262 regex, or null when it compiles. */ +function unicodeRegexError(pattern: string): string | null { + try { + new RegExp(pattern, "u"); + return null; + } catch (error) { + return error instanceof Error ? error.message : String(error); + } +} + +let tools: Tool[]; + +beforeAll(async () => { + tools = await listAdvertisedTools((server) => + initializeMcpServer(server, { + exaApiKey: "test-key", + userProvidedApiKey: true, + enabledTools: [...AVAILABLE_TOOL_IDS], + }), + ); +}); + +function advertised(name: string): Tool { + const tool = tools.find((candidate) => candidate.name === name); + if (!tool) throw new Error(`${name} is not advertised`); + return tool; +} + +describe("tools/list", () => { + it("advertises every available tool", () => { + expect(tools.map((tool) => tool.name).sort()).toEqual([...AVAILABLE_TOOL_IDS].sort()); + }); + + it("carries only patterns that compile as unicode-mode regexes", () => { + const invalid = tools.flatMap((tool) => + [ + ...collectPatterns(tool.inputSchema, `${tool.name}.inputSchema`), + ...collectPatterns(tool.outputSchema, `${tool.name}.outputSchema`), + ].flatMap(([path, pattern]) => { + const error = unicodeRegexError(pattern); + return error === null ? [] : [`${path} ${JSON.stringify(pattern)}: ${error}`]; + }), + ); + expect(invalid).toEqual([]); + }); + + it("constrains agent_run IDs to the agent_run_ prefix", () => { + expect(collectPatterns(advertised("agent_run").inputSchema)).toEqual([ + ["$.properties.runId.pattern", "^agent_run_"], + ["$.properties.previousRunId.pattern", "^agent_run_"], + ]); + }); + + it("marks every tool read-only and open-world", () => { + for (const tool of tools) { + expect(tool.annotations, tool.name).toMatchObject({ + readOnlyHint: true, + destructiveHint: false, + openWorldHint: true, + }); + } + }); + + it("advertises web_search_exa's objective as an optional bounded string", () => { + const inputSchema = advertised("web_search_exa").inputSchema as { + required?: string[]; + properties: Record; + }; + expect(inputSchema.required).toEqual(["query"]); + expect(inputSchema.properties.objective).toEqual({ + type: "string", + minLength: 1, + maxLength: 4096, + description: + "Goal for this search turn; say which documents should rank first, which should be excluded, and what specific facts or figures to pull from them.", + }); + }); + + it("declares an output schema for agent_run only", () => { + expect(tools.filter((tool) => tool.outputSchema).map((tool) => tool.name)).toEqual([ + "agent_run", + ]); + expect(advertised("agent_run").outputSchema).toMatchObject({ + type: "object", + required: expect.arrayContaining(["success", "id", "status", "outputReady"]), + }); + }); +}); diff --git a/tests/unit/tools/agentProgress.test.ts b/tests/unit/tools/agentProgress.test.ts index caa3266c..23c85989 100644 --- a/tests/unit/tools/agentProgress.test.ts +++ b/tests/unit/tools/agentProgress.test.ts @@ -8,36 +8,170 @@ import { summarizeAgentProgress, } from "../../../src/tools/agentProgress.js"; -function event(name: string, data: Record = {}): AgentEvent { +function event(name: string, data: Record): AgentEvent { return { event: name, data }; } describe("agent progress state", () => { - it("tracks meaningful events for heartbeats", () => { - const state = createAgentProgressState(null, 1_000); - ingestAgentEvent(state, event("agent_run.created", { id: "agent_run_1" }), 1_100); + it("correlates nested tool and source activity to the parent call", () => { + const state = createAgentProgressState("agent_run_1", 1_000); ingestAgentEvent( state, - event("agent_run.source.added", { source: { url: "https://example.com" } }), - 1_200, + event("agent_run.output_item.added", { + item: { + id: "item_1", + type: "function_call", + call_id: "call_parent", + name: "search", + status: "in_progress", + }, + }), + 1_001, ); + ingestAgentEvent( + state, + event("agent_run.output_item.done", { + item: { + id: "item_2", + type: "function_call", + call_id: "call_nested", + parent_call_id: "call_parent", + name: "search", + status: "completed", + metadata: { sourceCount: 2 }, + }, + }), + 1_002, + ); + ingestAgentEvent( + state, + event("agent_run.source.added", { + source: { url: "https://example.com", callId: "call_parent" }, + }), + 1_003, + ); + + expect(state.sourcesByCallId.get("call_parent")).toBe(1); + expect(state.toolsByCallId.get("call_parent")).toMatchObject({ + name: "search", + started: true, + finished: true, + }); + expect(state.toolsByCallId.has("call_nested")).toBe(false); + }); + + it("uses native search trace tool names and call IDs", () => { + const state = createAgentProgressState("agent_run_1", 1_000); + ingestAgentEvent( + state, + event("agent_run.search_trace", { + tool: "contents", + callId: "call_contents", + text: "Reading the relevant pages", + }), + 1_001, + ); + + expect(state.toolsByCallId.get("call_contents")).toMatchObject({ + name: "contents", + started: true, + finished: false, + }); + expect(state.latestNarrationByCallId.get("call_contents")).toBe("Reading the relevant pages"); + }); + + it("ignores unknown events without changing meaningful activity", () => { + const state = createAgentProgressState("agent_run_1", 1_000); + const update = ingestAgentEvent(state, event("agent_run.future", { status: "new" }), 2_000); + expect(update.kind).toBe("ignored"); + expect(state.lastMeaningfulAt).toBe(1_000); + }); + + it("builds bounded aggregate summaries and heartbeat messages", () => { + const state = createAgentProgressState("agent_run_1", 1_000); + ingestAgentEvent( + state, + event("agent_run.search_trace", { + tool: "search", + callId: "call_1", + text: "Looking for recent funding announcements", + }), + 2_000, + ); + for (let i = 0; i < 27; i += 1) { + ingestAgentEvent( + state, + event("agent_run.source.added", { + source: { url: `https://example.com/${i}`, callId: "call_1" }, + }), + 2_001 + i, + ); + } - expect(state.runId).toBe("agent_run_1"); - expect(summarizeAgentProgress(state)).toContain("1 source"); - expect(heartbeatMessage(state, 3_200)).toContain("elapsed 2s"); + expect(summarizeAgentProgress(state)).toContain("27 sources"); + expect(summarizeAgentProgress(state).length).toBeLessThanOrEqual(200); + expect(heartbeatMessage(state, 48_000)).toContain("elapsed 47s"); + expect(heartbeatMessage(state, 48_000)).toContain("Looking for recent funding"); }); }); describe("AgentProgressBridge", () => { - it("coalesces rapid activity and flushes the latest state on cleanup", async () => { - const emit = vi.fn(async () => {}); - const bridge = new AgentProgressBridge({ emit, throttleMs: 1_000, now: () => 1_000 }); + it("emits leading and trailing updates for a burst", async () => { + vi.useFakeTimers(); + try { + const messages: string[] = []; + const bridge = new AgentProgressBridge({ + emit: async (message) => { + messages.push(message); + }, + }); - await bridge.handle(event("agent_run.created", { id: "agent_run_1" })); - await bridge.handle(event("agent_run.search_trace", { text: "Searching" })); - await bridge.cleanup(); + await bridge.handle( + event("agent_run.source.added", { + source: { url: "https://example.com/1", callId: "call_1" }, + }), + ); + await bridge.handle( + event("agent_run.source.added", { + source: { url: "https://example.com/2", callId: "call_1" }, + }), + ); + await bridge.handle( + event("agent_run.source.added", { + source: { url: "https://example.com/3", callId: "call_1" }, + }), + ); + + expect(messages).toHaveLength(1); + expect(messages[0]).toContain("1 source"); + await vi.advanceTimersByTimeAsync(1_000); + expect(messages).toHaveLength(2); + expect(messages[1]).toContain("3 sources"); + await bridge.cleanup(); + } finally { + vi.useRealTimers(); + } + }); - expect(emit).toHaveBeenCalledTimes(2); - expect((emit.mock.calls as unknown[][])[1]?.[0]).toContain("Searching"); + it("flushes coalesced state before an immediate milestone", async () => { + const messages: string[] = []; + const bridge = new AgentProgressBridge({ + emit: async (message) => { + messages.push(message); + }, + now: (() => { + let time = 1_000; + return () => time++; + })(), + }); + + await bridge.handle( + event("agent_run.source.added", { source: { url: "https://example.com/1" } }), + ); + await bridge.handle(event("agent_run.started", { status: "running" })); + expect(messages).toHaveLength(2); + expect(messages[0]).toContain("1 source"); + expect(messages[1]).toContain("started"); + await bridge.cleanup(); }); }); diff --git a/tests/unit/tools/agentRun.test.ts b/tests/unit/tools/agentRun.test.ts index b7a1d30d..22fd175d 100644 --- a/tests/unit/tools/agentRun.test.ts +++ b/tests/unit/tools/agentRun.test.ts @@ -4,6 +4,7 @@ import type { AgentEvent, AgentRun } from "exa-js"; import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import { agentRunInputShape, + agentRunOutputShape, DEFAULT_CALL_WINDOW_MS, formatProgressMessage, pollAgentRun, @@ -13,6 +14,7 @@ import { type AgentRunClient, } from "../../../src/tools/agentRun.js"; import type { AgentRunInput } from "../../../src/types.js"; +import { connectInMemory, listAdvertisedTools } from "../../helpers/advertisedTools.js"; import { FakeMcpServer } from "../../helpers/fakeMcpServer.js"; function event(name: string, data: Record = {}): AgentEvent { @@ -65,12 +67,12 @@ function setup( config?: { exaApiKey?: string; oauthAccessToken?: string }; callWindowMs?: number; heartbeatMs?: number; - progressThrottleMs?: number; pollIntervalMs?: number; progressTimeoutMs?: number; progressToken?: string | number; signal?: AbortSignal; sendNotification?: (notification: Notification) => void | Promise; + siblingTools?: string[]; } = {}, ) { const fake = new FakeMcpServer(); @@ -85,9 +87,9 @@ function setup( clientFactory: () => client, callWindowMs: options.callWindowMs, heartbeatMs: options.heartbeatMs ?? 10_000, - progressThrottleMs: options.progressThrottleMs, pollIntervalMs: options.pollIntervalMs, progressTimeoutMs: options.progressTimeoutMs, + siblingTools: options.siblingTools ?? ["web_search_exa", "web_fetch_exa"], }); const notifications: Notification[] = []; @@ -105,6 +107,7 @@ function setup( const invoke = (args: Record, extraArg: unknown = extra) => tool.handler(args, extraArg) as Promise<{ content: Array<{ type: "text"; text: string }>; + structuredContent?: Record; isError?: true; }>; @@ -119,17 +122,28 @@ function setup( }; } +/** agent_run's `outputSchema` as a real McpServer advertises it in tools/list. */ +async function advertisedOutputSchema(): Promise { + const [tool] = await listAdvertisedTools((server) => + registerAgentRunTool(server, { exaApiKey: "test-key" }), + ); + return tool.outputSchema; +} + function payload(result: { content: Array<{ text: string }> }): Record { return JSON.parse(result.content[0].text) as Record; } describe("resolveAgentCallWindowMs", () => { - it("defaults to max duration minus headroom", () => { + it("defaults every client to a 45-second window", () => { expect(resolveAgentCallWindowMs()).toBe(DEFAULT_CALL_WINDOW_MS); - expect(resolveAgentCallWindowMs({ mcpMaxDurationSeconds: 600 })).toBe(550_000); }); - it("honors AGENT_CALL_WINDOW_MS when within the ceiling", () => { + it("clamps the default window to a shorter platform ceiling", () => { + expect(resolveAgentCallWindowMs({ mcpMaxDurationSeconds: 60 })).toBe(10_000); + }); + + it("honors AGENT_CALL_WINDOW_MS above the default window", () => { expect(resolveAgentCallWindowMs({ agentCallWindowMs: 300_000 })).toBe(300_000); }); @@ -139,7 +153,7 @@ describe("resolveAgentCallWindowMs", () => { agentCallWindowMs: 900_000, mcpMaxDurationSeconds: 800, }), - ).toBe(DEFAULT_CALL_WINDOW_MS); + ).toBe(750_000); }); }); @@ -236,6 +250,26 @@ describe("streamAgentRun", () => { expect(client.cancelRun).not.toHaveBeenCalled(); }); + it("cannot hand off a run when the stream has not yielded an ID by the boundary", async () => { + const client: AgentRunClient = { + createStream: vi.fn(async () => new Promise>(() => {})), + getRun: vi.fn(), + cancelRun: vi.fn(), + }; + + const outcome = await streamAgentRun({ + client, + runInput: { query: "test", effort: "low" }, + callWindowMs: 20, + }); + + expect(outcome).toMatchObject({ + status: "unrecoverable_stream", + runId: null, + handoffReason: "stream_window_exceeded", + }); + }); + it("turns clean EOF and read errors with a known ID into recoverable handoffs", async () => { for (const events of [ streamOf([event("agent_run.created", { id: "agent_run_1" })]), @@ -278,7 +312,10 @@ describe("streamAgentRun", () => { signal: controller.signal, }); - expect(outcome).toMatchObject({ status: "running", runId: "agent_run_1" }); + expect(outcome).toMatchObject({ + status: "running", + runId: "agent_run_1", + }); expect(client.cancelRun).not.toHaveBeenCalled(); }); @@ -397,16 +434,92 @@ describe("agent_run tool", () => { const { invoke, notifications } = setup({ progressToken: "progress-1" }); await invoke({ query: "test" }); - expect(notifications.length).toBeGreaterThanOrEqual(2); - expect( - notifications.every((notification) => notification.method === "notifications/progress"), - ).toBe(true); - expect(notifications.map((notification) => notification.params.progress)).toEqual( - notifications.map((_, index) => index + 1), - ); + expect(notifications).toHaveLength(2); + expect(notifications.map((notification) => notification.method)).toEqual([ + "notifications/progress", + "notifications/progress", + ]); + expect(notifications.map((notification) => notification.params.progress)).toEqual([1, 2]); expect( notifications.every((notification) => notification.params.progressToken === "progress-1"), ).toBe(true); + expect(notifications.every((notification) => !("total" in notification.params))).toBe(true); + }); + + it("bridges richer events and ignores unknown events until terminal", async () => { + const { invoke, notifications } = setup({ + progressToken: "progress-1", + events: streamOf([ + event("agent_run.created", { id: "agent_run_1", status: "queued" }), + event("agent_run.output_item.added", { + item: { + id: "item_1", + type: "function_call", + call_id: "call_search", + name: "search", + status: "in_progress", + }, + }), + event("agent_run.future", { detail: "ignored" }), + event("agent_run.source.added", { + source: { url: "https://example.com", callId: "call_search" }, + }), + event("agent_run.search_trace", { + tool: "search", + callId: "call_search", + text: "Finding relevant results", + }), + event("agent_run.completed", completedRun()), + ]), + }); + + const result = await invoke({ query: "test" }); + expect(result.isError).toBeUndefined(); + expect(notifications.map((notification) => notification.params.message)).toEqual([ + "run agent_run_1 queued", + expect.stringContaining("1 tool call"), + expect.stringContaining("Finding relevant results"), + ]); + expect(notifications.every((notification) => !("total" in notification.params))).toBe(true); + }); + + it("uses the heartbeat only after meaningful progress becomes idle", async () => { + vi.useFakeTimers(); + try { + const controller = new AbortController(); + const pending = setup({ + progressToken: "progress-1", + signal: controller.signal, + events: hangingStream([event("agent_run.created", { id: "agent_run_1" })]), + heartbeatMs: 10, + callWindowMs: 1_000, + }); + const resultPromise = pending.invoke({ query: "test" }); + await Promise.resolve(); + expect( + pending.notifications.some((notification) => { + const message = notification.params.message; + return typeof message === "string" && message.includes("still working"); + }), + ).toBe(false); + await vi.advanceTimersByTimeAsync(10); + expect( + pending.notifications.filter((notification) => { + const message = notification.params.message; + return typeof message === "string" && message.includes("still working"); + }), + ).toHaveLength(1); + expect( + pending.notifications.some((notification) => { + const message = notification.params.message; + return typeof message === "string" && message.includes("still working"); + }), + ).toBe(true); + controller.abort(); + await resultPromise; + } finally { + vi.useRealTimers(); + } }); it("returns a non-error run-ID handoff without cancelling at the stream boundary", async () => { @@ -567,6 +680,219 @@ describe("agent_run tool", () => { }); }); + describe("structuredContent", () => { + const outputSchema = z.object(agentRunOutputShape); + + function expectStructured(result: { + content: Array<{ text: string }>; + structuredContent?: Record; + }): Record { + expect(result.structuredContent).toBeDefined(); + expect(outputSchema.safeParse(result.structuredContent).success).toBe(true); + expect(result.content[0].text).toBe(JSON.stringify(result.structuredContent)); + expect(result.content[0].text).not.toContain("\n"); + return result.structuredContent ?? {}; + } + + it("registers the MCP outputSchema alongside the input schema", () => { + const { tool } = setup(); + expect(tool.outputSchema).toBe(agentRunOutputShape); + expect(tool.inputSchema).toBe(agentRunInputShape); + }); + + it("serializes nullable strings as anyOf branches, not a type array", async () => { + const json = await advertisedOutputSchema(); + expect(json).toMatchObject({ + properties: { + stopReason: { anyOf: [{ type: "string" }, { type: "null" }] }, + output: { + anyOf: [ + { properties: { text: { anyOf: [{ type: "string" }, { type: "null" }] } } }, + { type: "null" }, + ], + }, + }, + }); + }); + + it("emits completed runs with nullable output, extra output fields, and stopReason", async () => { + const output = { + text: "done", + structured: { companies: [{ name: "Exa" }] }, + grounding: [{ field: "companies[0].name", confidence: "high" }], + files: null, + }; + const { invoke } = setup({ + events: streamOf([ + event("agent_run.created", { id: "agent_run_1" }), + event("agent_run.completed", { + id: "agent_run_1", + status: "completed", + output, + stopReason: "schema_satisfied", + usage: { searches: 1 }, + costDollars: { total: 0.012 }, + }), + ]), + }); + const result = await invoke({ query: "test" }); + expect(result.isError).toBeUndefined(); + expect(expectStructured(result)).toEqual({ + success: true, + id: "agent_run_1", + status: "completed", + outputReady: true, + output, + stopReason: "schema_satisfied", + usage: { searches: 1 }, + costDollars: { total: 0.012 }, + }); + + const nullOutput = setup({ + events: streamOf([ + event("agent_run.created", { id: "agent_run_1" }), + event("agent_run.completed", { id: "agent_run_1", output: null }), + ]), + }); + expect(expectStructured(await nullOutput.invoke({ query: "test" }))).toMatchObject({ + status: "completed", + output: null, + }); + }); + + it("names the run in the error when its result does not match the output schema", async () => { + const { invoke } = setup({ + events: streamOf([ + event("agent_run.created", { id: "agent_run_1" }), + event("agent_run.completed", { id: "agent_run_1", output: { text: 42 } }), + ]), + }); + + const result = await invoke({ query: "test" }); + + expect(result.isError).toBe(true); + expect(result.structuredContent).toBeUndefined(); + expect(result.content[0].text).toContain( + "agent_run result for agent_run_1 did not match the declared output schema", + ); + }); + + it("emits running handoffs", async () => { + const { invoke } = setup({ + events: hangingStream([event("agent_run.created", { id: "agent_run_1" })]), + callWindowMs: 20, + }); + expect(expectStructured(await invoke({ query: "test" }))).toMatchObject({ + success: true, + id: "agent_run_1", + status: "running", + outputReady: false, + }); + }); + + it("emits cancelled runs", async () => { + const { invoke } = setup({ + events: streamOf([ + event("agent_run.created", { id: "agent_run_1" }), + event("agent_run.cancelled", { id: "agent_run_1", stopReason: "cancelled" }), + ]), + }); + const result = await invoke({ query: "test" }); + expect(result.isError).toBeUndefined(); + expect(expectStructured(result)).toEqual({ + success: false, + id: "agent_run_1", + status: "cancelled", + outputReady: false, + stopReason: "cancelled", + }); + }); + + it("emits failed runs as isError with schema-conformant structuredContent", async () => { + const { invoke } = setup({ + events: streamOf([ + event("agent_run.created", { id: "agent_run_1" }), + event("agent_run.failed", { + id: "agent_run_1", + stopReason: "error", + error: { message: "bad schema" }, + }), + ]), + }); + const result = await invoke({ query: "test" }); + expect(result.isError).toBe(true); + expect(expectStructured(result)).toEqual({ + success: false, + id: "agent_run_1", + status: "failed", + outputReady: false, + stopReason: "error", + error: { message: "bad schema" }, + }); + expect(result.content[1].text).toContain("The Agent run failed"); + }); + + it("falls back to a text-only error when the terminal payload has no usable ID", async () => { + for (const id of [undefined, "", 42]) { + const { invoke } = setup({ + events: streamOf([event("agent_run.completed", { id, output: null })]), + }); + const result = await invoke({ query: "test" }); + expect(result.isError).toBe(true); + expect(result.structuredContent).toBeUndefined(); + expect(result.content[0].text).toContain("without a run ID"); + } + }); + + it("passes the MCP SDK's own output validation for every reportable status", async () => { + const streams: Record AsyncIterable> = { + completed: () => streamOf(COMPLETED_EVENTS), + running: () => hangingStream([event("agent_run.created", { id: "agent_run_1" })]), + cancelled: () => + streamOf([ + event("agent_run.created", { id: "agent_run_1" }), + event("agent_run.cancelled", { id: "agent_run_1", stopReason: "cancelled" }), + ]), + failed: () => + streamOf([ + event("agent_run.created", { id: "agent_run_1" }), + event("agent_run.failed", { id: "agent_run_1", error: { message: "bad schema" } }), + ]), + }; + + for (const [status, stream] of Object.entries(streams)) { + const client = await connectInMemory((server) => + registerAgentRunTool( + server, + { exaApiKey: "test-key" }, + { + callWindowMs: 50, + clientFactory: () => ({ + createStream: async () => stream(), + getRun: vi.fn(), + cancelRun: vi.fn(), + }), + }, + ), + ); + try { + const result = await client.callTool({ name: "agent_run", arguments: { query: "test" } }); + expect(result.structuredContent, status).toMatchObject({ id: "agent_run_1", status }); + expect(result.isError ?? false, status).toBe(status === "failed"); + } finally { + await client.close(); + } + } + }); + + it("never emits structuredContent for pre-ID stream failures", async () => { + const { invoke } = setup({ events: failsBeforeId() }); + const result = await invoke({ query: "test" }); + expect(result.isError).toBe(true); + expect(result.structuredContent).toBeUndefined(); + }); + }); + it("rejects invalid argument combinations and resume-only create options", async () => { const { invoke, createStream, getRun } = setup(); @@ -592,10 +918,38 @@ describe("agent_run tool", () => { it("registers the streaming handoff contract and non-idempotent annotations", () => { const { tool } = setup(); - expect(tool.description).toContain("runs may take several minutes"); + expect(tool.description).toContain("Long-running: returns a run ID; resume with runId"); + expect(tool.description).toContain("Use it instead of running many searches yourself"); expect(tool.description).toContain( "interrupted tool call is not an explicit cancellation request", ); - expect(tool.annotations).toMatchObject({ readOnlyHint: true, idempotentHint: false }); + expect(tool.description).toContain("use web_search_exa"); + expect(tool.description).toContain("use web_fetch_exa"); + expect(tool.annotations).toMatchObject({ + readOnlyHint: true, + destructiveHint: false, + idempotentHint: false, + openWorldHint: true, + }); + }); + + it("only names alternative tools that are registered in the session", () => { + const searchOnly = setup({ siblingTools: ["web_search_exa"] }).tool.description; + expect(searchOnly).toContain("use web_search_exa"); + expect(searchOnly).not.toContain("web_fetch_exa"); + + const advancedOnly = setup({ siblingTools: ["web_search_advanced_exa"] }).tool.description; + expect(advancedOnly).toContain("(use web_search_advanced_exa)"); + expect(advancedOnly).not.toContain("use web_search_exa"); + + const bothSearch = setup({ + siblingTools: ["web_search_exa", "web_search_advanced_exa"], + }).tool.description; + expect(bothSearch).toContain("(use web_search_exa or web_search_advanced_exa)"); + + const agentOnly = setup({ siblingTools: [] }).tool.description; + expect(agentOnly).not.toContain("web_search_exa"); + expect(agentOnly).not.toContain("web_fetch_exa"); + expect(agentOnly).toContain("Not for a single question"); }); }); diff --git a/tests/unit/tools/config.test.ts b/tests/unit/tools/config.test.ts index e30dfe1c..b8c9535b 100644 --- a/tests/unit/tools/config.test.ts +++ b/tests/unit/tools/config.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it } from "vitest"; +import { afterEach, describe, expect, it } from "vitest"; import { createExaClient, integrationHeaders } from "../../../src/tools/config.js"; function clientHeaders(client: unknown): Headers { @@ -142,6 +142,10 @@ describe("integrationHeaders", () => { }); describe("createExaClient", () => { + afterEach(() => { + delete process.env.EXA_API_BASE_URL; + }); + it("sets API key auth and Agent integration headers", () => { const exa = createExaClient( { @@ -158,6 +162,18 @@ describe("createExaClient", () => { expect(headers.get("x-exa-mcp-session-id")).toBe("session-123"); }); + it("calls the Exa API origin named by EXA_API_BASE_URL, or api.exa.ai by default", () => { + expect((createExaClient({ exaApiKey: "k" }) as unknown as { baseURL: string }).baseURL).toBe( + "https://api.exa.ai", + ); + + process.env.EXA_API_BASE_URL = "http://127.0.0.1:9999/"; + for (const config of [{ exaApiKey: "k" }, { oauthAccessToken: "jwt-token" }]) { + const exa = createExaClient(config) as unknown as { baseURL: string }; + expect(exa.baseURL).toBe("http://127.0.0.1:9999"); + } + }); + it("uses OAuth without falling back to x-api-key", () => { const exa = createExaClient({ oauthAccessToken: "jwt-token" }, "agent-mcp"); const headers = clientHeaders(exa); diff --git a/tests/unit/tools/webSearch.test.ts b/tests/unit/tools/webSearch.test.ts index 16df5b9b..ea9a138d 100644 --- a/tests/unit/tools/webSearch.test.ts +++ b/tests/unit/tools/webSearch.test.ts @@ -1,4 +1,5 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; +import { z } from "zod"; import { emptySearchResponse, searchResponse } from "../../fixtures/exaResponses.js"; import { FakeMcpServer } from "../../helpers/fakeMcpServer.js"; @@ -122,6 +123,64 @@ describe("registerWebSearchTool", () => { ); }); + it("forwards an objective to /search", async () => { + const { registerWebSearchTool } = await import("../../../src/tools/webSearch.js"); + const server = new FakeMcpServer(); + requestMock.mockResolvedValue(searchResponse); + + registerWebSearchTool(server as any, { exaApiKey: "test-key" }); + + await server.getTool("web_search_exa").handler({ + query: "AI breakthroughs", + objective: "Compile a briefing on recent AI research", + }); + + expect(requestMock).toHaveBeenCalledWith( + "/search", + "POST", + expect.objectContaining({ + query: "AI breakthroughs", + objective: "Compile a briefing on recent AI research", + }), + undefined, + { "x-exa-integration": "web-search-mcp" }, + ); + }); + + it("leaves the objective out of the request when the caller omits it", async () => { + const { registerWebSearchTool } = await import("../../../src/tools/webSearch.js"); + const server = new FakeMcpServer(); + requestMock.mockResolvedValue(searchResponse); + + registerWebSearchTool(server as any, { exaApiKey: "test-key" }); + await server.getTool("web_search_exa").handler({ query: "AI breakthroughs" }); + + expect(requestMock.mock.calls[0][2]).not.toHaveProperty("objective"); + }); + + it("validates objectives the way /search does and drops unused ones", async () => { + const { registerWebSearchTool } = await import("../../../src/tools/webSearch.js"); + const server = new FakeMcpServer(); + + registerWebSearchTool(server as any); + + const args = z.object( + server.getTool("web_search_exa").inputSchema as Record, + ); + expect(args.parse({ query: "q" })).toEqual({ query: "q" }); + expect(args.parse({ query: "q", objective: " Rank primary sources first " })).toEqual({ + query: "q", + objective: "Rank primary sources first", + }); + expect(args.safeParse({ query: "q", objective: "x".repeat(4096) }).success).toBe(true); + expect(args.safeParse({ query: "q", objective: "x".repeat(4097) }).success).toBe(false); + expect(args.safeParse({ query: "q", objective: 42 }).success).toBe(false); + // Unused optional values some clients send search without an objective. + for (const unused of [null, "", " "]) { + expect(args.parse({ query: "q", objective: unused })).toEqual({ query: "q" }); + } + }); + it("returns a friendly message when Exa has no results", async () => { const { registerWebSearchTool } = await import("../../../src/tools/webSearch.js"); const server = new FakeMcpServer(); diff --git a/tests/unit/utils/agentErrorHandler.test.ts b/tests/unit/utils/agentErrorHandler.test.ts new file mode 100644 index 00000000..68c11824 --- /dev/null +++ b/tests/unit/utils/agentErrorHandler.test.ts @@ -0,0 +1,40 @@ +import { ExaError } from "exa-js"; +import { describe, expect, it } from "vitest"; +import { EXA_BILLING_URL, formatAgentToolError } from "../../../src/utils/agentErrorHandler.js"; + +describe("formatAgentToolError", () => { + it("marks 402 as not retryable and points at billing", () => { + const result = formatAgentToolError( + new ExaError("You have exceeded your credits limit.", 402), + "agent_run", + ); + + expect(result.isError).toBe(true); + const text = result.content[0].text; + expect(text).toContain("agent_run error (402): You have exceeded your credits limit."); + expect(text).toContain("Not retryable"); + expect(text).toContain("Do not call agent_run again"); + expect(text).toContain(EXA_BILLING_URL); + }); + + it("keeps the API message and adds no guidance for statuses without one", () => { + const result = formatAgentToolError(new ExaError("Something went wrong.", 500), "agent_run"); + + expect(result).toEqual({ + content: [{ type: "text", text: "agent_run error (500): Something went wrong." }], + isError: true, + }); + }); + + it("formats non-Exa errors with the tool name", () => { + const result = formatAgentToolError( + new Error("Provide exactly one of query or runId."), + "agent_run", + ); + + expect(result).toEqual({ + content: [{ type: "text", text: "agent_run error: Provide exactly one of query or runId." }], + isError: true, + }); + }); +}); diff --git a/tests/unit/utils/apiKeyFallback.test.ts b/tests/unit/utils/apiKeyFallback.test.ts new file mode 100644 index 00000000..d66ade24 --- /dev/null +++ b/tests/unit/utils/apiKeyFallback.test.ts @@ -0,0 +1,118 @@ +import { describe, expect, it, vi } from "vitest"; + +import { withApiKeyFallback } from "../../../src/utils/apiKeyFallback.js"; +import { FakeMcpServer } from "../../helpers/fakeMcpServer.js"; + +const invalidKeyResult = { + content: [{ type: "text", text: "web_search_exa error (401): x-api-key header is invalid" }], + isError: true, +}; + +const okResult = { content: [{ type: "text", text: "results" }] }; + +describe("withApiKeyFallback", () => { + it("retries once on the fallback credential when the API key is rejected", async () => { + const handler = vi.fn().mockResolvedValueOnce(invalidKeyResult).mockResolvedValueOnce(okResult); + const fallback = vi.fn().mockResolvedValue(true); + const server = withApiKeyFallback(new FakeMcpServer(), fallback); + + server.tool("web_search_exa", "desc", {}, handler); + const result = await server.getTool("web_search_exa").handler({ query: "q" }, { extra: 1 }); + + expect(result).toBe(okResult); + expect(fallback).toHaveBeenCalledTimes(1); + expect(handler).toHaveBeenCalledTimes(2); + expect(handler).toHaveBeenNthCalledWith(2, { query: "q" }, { extra: 1 }); + }); + + it("returns the rejection unchanged when no other credential is usable", async () => { + const handler = vi.fn().mockResolvedValue(invalidKeyResult); + const fallback = vi.fn().mockResolvedValue(false); + const server = withApiKeyFallback(new FakeMcpServer(), fallback); + + server.registerTool("web_search_exa", { description: "desc" }, handler); + const result = await server.getTool("web_search_exa").handler({ query: "q" }); + + expect(result).toBe(invalidKeyResult); + expect(fallback).toHaveBeenCalledTimes(1); + expect(handler).toHaveBeenCalledTimes(1); + }); + + it("leaves other tool errors alone", async () => { + const upstream5xx = { + content: [{ type: "text", text: "web_search_exa error (503): unavailable" }], + isError: true, + }; + const handler = vi.fn().mockResolvedValue(upstream5xx); + const fallback = vi.fn().mockResolvedValue(true); + const server = withApiKeyFallback(new FakeMcpServer(), fallback); + + server.tool("web_search_exa", "desc", {}, handler); + const result = await server.getTool("web_search_exa").handler({ query: "q" }); + + expect(result).toBe(upstream5xx); + expect(fallback).not.toHaveBeenCalled(); + expect(handler).toHaveBeenCalledTimes(1); + }); + + it.each([ + [ + "a failed Agent run whose own error carries an auth status", + { + content: [ + { + type: "text", + text: '{"success":false,"id":"agent_run_1","status":"failed","error":{"message":"fiber provider error (403): Forbidden"}}', + }, + { type: "text", text: "The Agent run failed." }, + ], + structuredContent: { success: false, id: "agent_run_1", status: "failed" }, + isError: true, + }, + ], + [ + "a fetch failure for a URL that mentions unauthorized", + { + content: [ + { + type: "text", + text: "Error fetching URL(s): https://example.com/unauthorized: CRAWL_NOT_FOUND", + }, + ], + isError: true, + }, + ], + [ + "an error that only quotes a 401 further in", + { + content: [{ type: "text", text: "web_search_exa error: upstream said error (401): nope" }], + isError: true, + }, + ], + ])("never falls back for %s", async (_label, result) => { + const handler = vi.fn().mockResolvedValue(result); + const fallback = vi.fn().mockResolvedValue(true); + const server = withApiKeyFallback(new FakeMcpServer(), fallback); + + server.tool("agent_run", "desc", {}, handler); + await expect(server.getTool("agent_run").handler({ query: "q" })).resolves.toBe(result); + + expect(fallback).not.toHaveBeenCalled(); + expect(handler).toHaveBeenCalledTimes(1); + }); + + it("retries a thrown API-key rejection once, then lets the retry's outcome through", async () => { + const unauthorized = Object.assign(new Error("Unauthorized"), { statusCode: 401 }); + const handler = vi.fn().mockRejectedValue(unauthorized); + const fallback = vi.fn().mockResolvedValue(true); + const server = withApiKeyFallback(new FakeMcpServer(), fallback); + + server.tool("web_search_exa", "desc", {}, handler); + + await expect(server.getTool("web_search_exa").handler({ query: "q" })).rejects.toBe( + unauthorized, + ); + expect(handler).toHaveBeenCalledTimes(2); + expect(fallback).toHaveBeenCalledTimes(1); + }); +}); diff --git a/tests/unit/utils/mcpAccept.test.ts b/tests/unit/utils/mcpAccept.test.ts new file mode 100644 index 00000000..8309245a --- /dev/null +++ b/tests/unit/utils/mcpAccept.test.ts @@ -0,0 +1,41 @@ +import { describe, expect, it } from "vitest"; +import { acceptsMcpResponses } from "../../../src/utils/mcpAccept.js"; + +describe("acceptsMcpResponses", () => { + it.each([ + "application/json, text/event-stream", + "text/event-stream,application/json", + "APPLICATION/JSON, Text/Event-Stream", + "application/json;q=0.9, text/event-stream;q=1.0", + 'application/json;note="a, b; c", text/event-stream', + 'application/json;note="escaped \\" quote", text/event-stream', + "application/json, text/event-stream, */*;q=0.1", + "application/json,, text/event-stream", + "application/json, text/event-stream;", + "application/json;;charset=utf-8, text/event-stream", + "application/json, text/event-stream, garbage", + "application/json, text/event-stream, text/html;q=invalid", + ])("accepts %j", (accept) => { + expect(acceptsMcpResponses(accept)).toBe(true); + }); + + it.each([ + [null], + [""], + ["*/*"], + ["application/*, text/*"], + ["application/json"], + ["text/event-stream"], + ["application/json, text/event-stream;q=0"], + ["application/json;q=0.000, text/event-stream"], + ["application/json;q=1.5, text/event-stream"], + ["application/json;q=0.5;q=0.5, text/event-stream"], + ["application/json;charset, text/event-stream"], + ['application/json;note="unterminated, text/event-stream'], + ["application/json\n, text/event-stream"], + ["application/json;text/event-stream"], + ["application/jsonx, text/event-streamy"], + ])("rejects %j", (accept) => { + expect(acceptsMcpResponses(accept)).toBe(false); + }); +}); diff --git a/tests/unit/utils/requestTarget.test.ts b/tests/unit/utils/requestTarget.test.ts new file mode 100644 index 00000000..1c84b5e5 --- /dev/null +++ b/tests/unit/utils/requestTarget.test.ts @@ -0,0 +1,27 @@ +import { describe, expect, it } from "vitest"; +import { parseRequestTarget } from "../../../src/utils/requestTarget.js"; + +describe("parseRequestTarget", () => { + it("resolves a target against the Host header", () => { + const url = parseRequestTarget("/mcp?x=1", "mcp.exa.ai"); + expect(url?.host).toBe("mcp.exa.ai"); + expect(url?.pathname).toBe("/mcp"); + expect(url?.search).toBe("?x=1"); + }); + + it("defaults a missing target and Host header", () => { + expect(parseRequestTarget(undefined, undefined)?.href).toBe("http://localhost/"); + }); + + it("returns null for a protocol-relative target with an empty host", () => { + expect(parseRequestTarget("//", "mcp.exa.ai")).toBeNull(); + }); + + it("returns null for an unparseable Host header", () => { + expect(parseRequestTarget("/mcp", "not a host")).toBeNull(); + }); + + it("keeps a protocol-relative target that carries a host", () => { + expect(parseRequestTarget("//example.com/mcp", "mcp.exa.ai")?.host).toBe("example.com"); + }); +}); diff --git a/tsconfig.test.json b/tsconfig.test.json index a178d770..3ccc205e 100644 --- a/tsconfig.test.json +++ b/tsconfig.test.json @@ -1,5 +1,11 @@ { "extends": "./tsconfig.json", - "include": ["src/**/*", "api/**/*", "tests/**/*", "vitest.config.ts"], + "include": [ + "src/**/*", + "api/**/*", + "tests/**/*", + "vitest.config.ts", + "vitest.integration.config.ts" + ], "exclude": ["node_modules", "build", "dist", "coverage"] } diff --git a/vitest.config.ts b/vitest.config.ts index 9be4e449..f364f365 100644 --- a/vitest.config.ts +++ b/vitest.config.ts @@ -11,6 +11,9 @@ export default defineConfig({ }, environment: "node", include: ["tests/**/*.test.ts"], + // The integration suite binds sockets and spawns processes; it runs + // through vitest.integration.config.ts. + exclude: ["tests/integration/**", "node_modules/**"], restoreMocks: true, }, }); diff --git a/vitest.integration.config.ts b/vitest.integration.config.ts new file mode 100644 index 00000000..4082c133 --- /dev/null +++ b/vitest.integration.config.ts @@ -0,0 +1,22 @@ +import { defineConfig } from "vitest/config"; + +// Integration suite (tests/integration): drives the runtime HTTP server against +// a local Exa API double. Coverage spans the api/ handlers the server hosts. +export default defineConfig({ + test: { + clearMocks: true, + coverage: { + provider: "v8", + reporter: ["text", "lcov"], + include: ["src/**/*.ts", "api/**/*.ts"], + exclude: ["src/**/*.d.ts"], + }, + environment: "node", + include: ["tests/integration/**/*.test.ts"], + // Each file boots its own server in its own process. + pool: "forks", + restoreMocks: true, + testTimeout: 60_000, + hookTimeout: 60_000, + }, +}); From 4102b944702cba7598df4c38e3c6a0d004074684 Mon Sep 17 00:00:00 2001 From: Wen-Hao Lue Date: Wed, 30 Sep 2026 14:19:16 -0700 Subject: [PATCH 2/2] Advertise web_search_exa's objective as required Models see objective in the tool schema's required list, so they state a goal for every search, while calls that omit it (or send null or "") still search without one. advertisedRequired decouples the advertised required list from what the schema accepts. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/tools/validation.ts | 30 +++++++++++++++++++++++++++++ src/tools/webSearch.ts | 14 +++++++------- tests/unit/toolSchemas.test.ts | 4 ++-- tests/unit/tools/validation.test.ts | 30 +++++++++++++++++++++++++++++ tests/unit/tools/webSearch.test.ts | 25 ++++++++++++++++++++++++ 5 files changed, 94 insertions(+), 9 deletions(-) diff --git a/src/tools/validation.ts b/src/tools/validation.ts index 1da1fcfc..23cb5483 100644 --- a/src/tools/validation.ts +++ b/src/tools/validation.ts @@ -11,6 +11,36 @@ export function lenientString() { ); } +/** + * ZodOptional whose `isOptional()` probe always reports false. + * + * zod-to-json-schema (what the MCP SDK uses to advertise tool schemas) puts a + * property in the JSON-schema `required` list iff `isOptional()` is false, + * and `isOptional()` is implemented as "does `undefined` parse?". This + * subclass decouples the two so a field can be advertised as required while + * still parsing a missing value. + */ +class AdvertisedRequiredOptional extends z.ZodOptional { + isOptional(): boolean { + return false; + } +} + +/** + * Advertise `inner` as a required parameter in the tool's JSON schema while + * accepting a missing value at call time; any value that is present is parsed + * by `inner`. Models see the field as required; programmatic callers that omit + * it keep working. + */ +export function advertisedRequired( + inner: T, +): AdvertisedRequiredOptional { + return new AdvertisedRequiredOptional({ + innerType: inner, + typeName: z.ZodFirstPartyTypeKind.ZodOptional, + }); +} + /** * Lenient optional number: uses z.coerce for string-to-number conversion, * falls back to undefined on invalid input instead of throwing a validation error. diff --git a/src/tools/webSearch.ts b/src/tools/webSearch.ts index 6ffdde9f..2cb5f7a7 100644 --- a/src/tools/webSearch.ts +++ b/src/tools/webSearch.ts @@ -5,7 +5,7 @@ import { ExaSearchRequest, ExaSearchResponse } from "../types.js"; import { createRequestLogger } from "../utils/logger.js"; import { retryWithBackoff, formatToolError, withTimeout } from "../utils/errorHandler.js"; import { sanitizeSearchResponse } from "../utils/exaResponseSanitizer.js"; -import { lenientString, lenientOptionalNumber } from "./validation.js"; +import { advertisedRequired, lenientString, lenientOptionalNumber } from "./validation.js"; import type { McpAnalytics } from "../analytics.js"; @@ -45,17 +45,17 @@ export function registerWebSearchTool(server: McpServer, config?: WebSearchConfi numResults: lenientOptionalNumber().describe( "Number of search results to return (default: 10).", ), - objective: z - .preprocess( - // Clients that send unused optional parameters as null or "" search - // without an objective instead of failing the call. + // Advertised as required so models state a goal for every search; calls + // that omit it (or send null or "") still search, without an objective. + objective: advertisedRequired( + z.preprocess( (value) => value === null || (typeof value === "string" && value.trim() === "") ? undefined : value, z.string().trim().min(1).max(MAX_OBJECTIVE_LENGTH).optional(), - ) - .describe(OBJECTIVE_DESCRIPTION), + ), + ).describe(OBJECTIVE_DESCRIPTION), }, { readOnlyHint: true, diff --git a/tests/unit/toolSchemas.test.ts b/tests/unit/toolSchemas.test.ts index e73bda4c..99ad0ea0 100644 --- a/tests/unit/toolSchemas.test.ts +++ b/tests/unit/toolSchemas.test.ts @@ -76,12 +76,12 @@ describe("tools/list", () => { } }); - it("advertises web_search_exa's objective as an optional bounded string", () => { + it("advertises web_search_exa's objective as a required bounded string", () => { const inputSchema = advertised("web_search_exa").inputSchema as { required?: string[]; properties: Record; }; - expect(inputSchema.required).toEqual(["query"]); + expect(inputSchema.required).toEqual(["query", "objective"]); expect(inputSchema.properties.objective).toEqual({ type: "string", minLength: 1, diff --git a/tests/unit/tools/validation.test.ts b/tests/unit/tools/validation.test.ts index f1718ad9..26343a17 100644 --- a/tests/unit/tools/validation.test.ts +++ b/tests/unit/tools/validation.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it, vi } from "vitest"; import { z } from "zod"; import { zodToJsonSchema } from "zod-to-json-schema"; import { + advertisedRequired, lenientOptionalBoolean, lenientOptionalNumber, lenientOptionalPositiveNumber, @@ -162,6 +163,35 @@ describe("web_fetch_exa urls preprocess", () => { }); }); +describe("advertisedRequired", () => { + const fields = { goal: advertisedRequired(z.string().min(1)).describe("The goal.") }; + + it("lists the field as required in the advertised JSON Schema", () => { + const json = zodToJsonSchema(z.object(fields)) as { + required?: string[]; + properties: Record; + }; + + expect(json.required).toEqual(["goal"]); + expect(json.properties.goal).toEqual({ + type: "string", + minLength: 1, + description: "The goal.", + }); + }); + + it("still accepts a call that omits the field and validates one that sends it", () => { + const schema = z.object(fields); + + expect(schema.parse({})).toEqual({}); + expect(schema.parse({ goal: "rank primary sources" })).toEqual({ + goal: "rank primary sources", + }); + expect(schema.safeParse({ goal: "" }).success).toBe(false); + expect(schema.safeParse({ goal: null }).success).toBe(false); + }); +}); + describe("advertised JSON Schema retains usable type info", () => { it("each lenient helper produces a JSON Schema with a type field", () => { const fields = { diff --git a/tests/unit/tools/webSearch.test.ts b/tests/unit/tools/webSearch.test.ts index ea9a138d..50636f2f 100644 --- a/tests/unit/tools/webSearch.test.ts +++ b/tests/unit/tools/webSearch.test.ts @@ -1,6 +1,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { z } from "zod"; import { emptySearchResponse, searchResponse } from "../../fixtures/exaResponses.js"; +import { connectInMemory } from "../../helpers/advertisedTools.js"; import { FakeMcpServer } from "../../helpers/fakeMcpServer.js"; const { ExaMock, exaConstructorMock, requestMock } = vi.hoisted(() => { @@ -181,6 +182,30 @@ describe("registerWebSearchTool", () => { } }); + it("advertises the objective as required but still searches when a caller omits it", async () => { + const { registerWebSearchTool } = await import("../../../src/tools/webSearch.js"); + requestMock.mockResolvedValue(searchResponse); + const client = await connectInMemory((server) => + registerWebSearchTool(server, { exaApiKey: "test-key" }), + ); + + try { + const [tool] = (await client.listTools()).tools; + expect(tool.inputSchema.required).toEqual(["query", "objective"]); + + for (const args of [{ query: "AI breakthroughs" }, { query: "AI", objective: null }]) { + const result = await client.callTool({ name: "web_search_exa", arguments: args }); + expect(result.isError).toBeFalsy(); + } + for (const [, , body] of requestMock.mock.calls) { + expect(body).not.toHaveProperty("objective"); + } + expect(requestMock).toHaveBeenCalledTimes(2); + } finally { + await client.close(); + } + }); + it("returns a friendly message when Exa has no results", async () => { const { registerWebSearchTool } = await import("../../../src/tools/webSearch.js"); const server = new FakeMcpServer();