From 7bb79ebbbba88a503851617d053b13d8fd9228bb Mon Sep 17 00:00:00 2001 From: Felix Weinberger <3823880+felixweinberger@users.noreply.github.com> Date: Tue, 3 Feb 2026 20:35:02 +0000 Subject: [PATCH] refactor: replace RequestHandlerExtra with structured context types (#1467) --- CLAUDE.md | 38 ++- docs/migration-SKILL.md | 38 ++- docs/migration.md | 79 ++++- examples/server/src/elicitationUrlExample.ts | 10 +- .../server/src/jsonResponseStreamableHttp.ts | 26 +- .../src/simpleStatelessStreamableHttp.ts | 10 +- examples/server/src/simpleStreamableHttp.ts | 58 +--- examples/server/src/simpleTaskInteractive.ts | 10 +- examples/server/src/ssePollingExample.ts | 48 +-- packages/client/src/client/client.ts | 32 +- .../core/src/experimental/tasks/interfaces.ts | 21 +- packages/core/src/shared/protocol.ts | 298 ++++++++++------- packages/core/test/shared/protocol.test.ts | 304 ++++++++++++++---- .../shared/protocolTransportHandling.test.ts | 12 +- .../node/test/streamableHttp.test.ts | 62 ++-- .../src/experimental/tasks/interfaces.ts | 8 +- .../src/experimental/tasks/mcpServer.ts | 12 +- packages/server/src/server/mcp.ts | 134 +++----- packages/server/src/server/server.ts | 49 ++- test/conformance/src/everythingServer.ts | 48 +-- test/integration/test/client/client.test.ts | 152 ++++----- test/integration/test/server.test.ts | 122 +++---- test/integration/test/server/mcp.test.ts | 299 ++++++++++++----- test/integration/test/taskLifecycle.test.ts | 116 +++---- .../integration/test/taskResumability.test.ts | 8 +- 25 files changed, 1224 insertions(+), 770 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index cad93e59..67601a31 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -144,7 +144,8 @@ When a request arrives from the remote side: 2. **`Protocol.connect()`** routes to `_onrequest()`, `_onresponse()`, or `_onnotification()` 3. **`Protocol._onrequest()`**: - Looks up handler in `_requestHandlers` map (keyed by method name) - - Creates `RequestHandlerExtra` with `signal`, `sessionId`, `sendNotification`, `sendRequest` + - Creates `BaseContext` with `signal`, `sessionId`, `sendNotification`, `sendRequest`, etc. + - Calls `buildContext()` to let subclasses enrich the context (e.g., Server adds `requestInfo`) - Invokes handler, sends JSON-RPC response back via transport 4. **Handler** was registered via `setRequestHandler('method', handler)` @@ -152,29 +153,42 @@ When a request arrives from the remote side: ```typescript // In Client (for server→client requests like sampling, elicitation) -client.setRequestHandler('sampling/createMessage', async (request, extra) => { +client.setRequestHandler('sampling/createMessage', async (request, ctx) => { // Handle sampling request from server return { role: "assistant", content: {...}, model: "..." }; }); // In Server (for client→server requests like tools/call) -server.setRequestHandler('tools/call', async (request, extra) => { +server.setRequestHandler('tools/call', async (request, ctx) => { // Handle tool call from client return { content: [...] }; }); ``` -### Request Handler Extra +### Request Handler Context -The `extra` parameter in handlers (`RequestHandlerExtra`) provides: +The `ctx` parameter in handlers provides a structured context: -- `signal`: AbortSignal for cancellation -- `sessionId`: Transport session identifier -- `authInfo`: Validated auth token info (if authenticated) -- `requestId`: JSON-RPC message ID -- `sendNotification(notification)`: Send related notification back -- `sendRequest(request, schema)`: Send related request (for bidirectional flows) -- `taskStore`: Task storage interface (if tasks enabled) +**`BaseContext`** (common to both Server and Client), fields organized into nested groups: + +- `sessionId?`: Transport session identifier +- `mcpReq`: Request-level concerns + - `id`: JSON-RPC message ID + - `method`: Request method string (e.g., 'tools/call') + - `_meta?`: Request metadata + - `signal`: AbortSignal for cancellation + - `send(request, schema, options?)`: Send related request (for bidirectional flows) + - `notify(notification)`: Send related notification back +- `http?`: HTTP transport info (undefined for stdio) + - `authInfo?`: Validated auth token info +- `task?`: Task context (`{ id?, store, requestedTtl? }`) when task storage is configured + +**`ServerContext`** extends `BaseContext.mcpReq` and `BaseContext.http?` via type intersection: + +- `mcpReq` adds: `log(level, data, logger?)`, `elicitInput(params, options?)`, `requestSampling(params, options?)` +- `http?` adds: `req?` (HTTP request info), `closeSSE?`, `closeStandaloneSSE?` + +**`ClientContext`** is currently identical to `BaseContext`. ### Capability Checking diff --git a/docs/migration-SKILL.md b/docs/migration-SKILL.md index 4e100ef2..ec793aa7 100644 --- a/docs/migration-SKILL.md +++ b/docs/migration-SKILL.md @@ -297,7 +297,7 @@ extra.requestInfo?.headers['mcp-session-id'] // v2: Headers object, .get() access headers: new Headers({ 'Authorization': 'Bearer token' }) -extra.requestInfo?.headers.get('mcp-session-id') +ctx.http?.req?.headers.get('mcp-session-id') ``` ## 8. Removed Server Features @@ -366,11 +366,41 @@ Schema to method string mapping: Request/notification params remain fully typed. Remove unused schema imports after migration. -## 10. Client Behavioral Changes +## 10. Request Handler Context Types + +`RequestHandlerExtra` → structured context types with nested groups. Rename `extra` → `ctx` in all handler callbacks. + +| v1 | v2 | +|----|-----| +| `RequestHandlerExtra` | `ServerContext` (server) / `ClientContext` (client) / `BaseContext` (base) | +| `extra` (param name) | `ctx` | +| `extra.signal` | `ctx.mcpReq.signal` | +| `extra.requestId` | `ctx.mcpReq.id` | +| `extra._meta` | `ctx.mcpReq._meta` | +| `extra.sendRequest(...)` | `ctx.mcpReq.send(...)` | +| `extra.sendNotification(...)` | `ctx.mcpReq.notify(...)` | +| `extra.authInfo` | `ctx.http?.authInfo` | +| `extra.sessionId` | `ctx.sessionId` | +| `extra.requestInfo` | `ctx.http?.req` (only `ServerContext`) | +| `extra.closeSSEStream` | `ctx.http?.closeSSE` (only `ServerContext`) | +| `extra.closeStandaloneSSEStream` | `ctx.http?.closeStandaloneSSE` (only `ServerContext`) | +| `extra.taskStore` | `ctx.task?.store` | +| `extra.taskId` | `ctx.task?.id` | +| `extra.taskRequestedTtl` | `ctx.task?.requestedTtl` | + +`ServerContext` convenience methods (new in v2, no v1 equivalent): + +| Method | Description | Replaces | +|--------|-------------|----------| +| `ctx.mcpReq.log(level, data, logger?)` | Send log notification (respects client's level filter) | `server.sendLoggingMessage(...)` from within handler | +| `ctx.mcpReq.elicitInput(params, options?)` | Elicit user input (form or URL) | `server.elicitInput(...)` from within handler | +| `ctx.mcpReq.requestSampling(params, options?)` | Request LLM sampling from client | `server.createMessage(...)` from within handler | + +## 11. Client Behavioral Changes `Client.listPrompts()`, `listResources()`, `listResourceTemplates()`, `listTools()` now return empty results when the server lacks the corresponding capability (instead of sending the request). Set `enforceStrictCapabilities: true` in `ClientOptions` to throw an error instead. -## 11. Runtime-Specific JSON Schema Validators (Enhancement) +## 12. Runtime-Specific JSON Schema Validators (Enhancement) The SDK now auto-selects the appropriate JSON Schema validator based on runtime: - Node.js → `AjvJsonSchemaValidator` (no change from v1) @@ -390,7 +420,7 @@ new McpServer({ name: 'server', version: '1.0.0' }, {}); Access validators via `_shims` export: `import { DefaultJsonSchemaValidator } from '@modelcontextprotocol/server/_shims';` -## 12. Migration Steps (apply in this order) +## 13. Migration Steps (apply in this order) 1. Update `package.json`: `npm uninstall @modelcontextprotocol/sdk`, install the appropriate v2 packages 2. Replace all imports from `@modelcontextprotocol/sdk/...` using the import mapping tables (sections 3-4), including `StreamableHTTPServerTransport` → `NodeStreamableHTTPServerTransport` diff --git a/docs/migration.md b/docs/migration.md index 8093c3ca..233d8bb2 100644 --- a/docs/migration.md +++ b/docs/migration.md @@ -156,7 +156,7 @@ const transport = new StreamableHTTPClientTransport(url, { }); // Reading headers in a request handler -const sessionId = extra.requestInfo?.headers.get('mcp-session-id'); +const sessionId = ctx.http?.req?.headers.get('mcp-session-id'); ``` ### `McpServer.tool()`, `.prompt()`, `.resource()` removed @@ -381,6 +381,83 @@ import { JSONRPCError, ResourceReference, isJSONRPCError } from '@modelcontextpr import { JSONRPCErrorResponse, ResourceTemplateReference, isJSONRPCErrorResponse } from '@modelcontextprotocol/core'; ``` +### Request handler context types + +The `RequestHandlerExtra` type has been replaced with a structured context type hierarchy using nested groups: + +| v1 | v2 | +|----|-----| +| `RequestHandlerExtra` (flat, all fields) | `ServerContext` (server handlers) or `ClientContext` (client handlers) | +| `extra` parameter name | `ctx` parameter name | +| `extra.signal` | `ctx.mcpReq.signal` | +| `extra.requestId` | `ctx.mcpReq.id` | +| `extra._meta` | `ctx.mcpReq._meta` | +| `extra.sendRequest(...)` | `ctx.mcpReq.send(...)` | +| `extra.sendNotification(...)` | `ctx.mcpReq.notify(...)` | +| `extra.authInfo` | `ctx.http?.authInfo` | +| `extra.requestInfo` | `ctx.http?.req` (only on `ServerContext`) | +| `extra.closeSSEStream` | `ctx.http?.closeSSE` (only on `ServerContext`) | +| `extra.closeStandaloneSSEStream` | `ctx.http?.closeStandaloneSSE` (only on `ServerContext`) | +| `extra.sessionId` | `ctx.sessionId` | +| `extra.taskStore` | `ctx.task?.store` | +| `extra.taskId` | `ctx.task?.id` | +| `extra.taskRequestedTtl` | `ctx.task?.requestedTtl` | + +**Before (v1):** + +```typescript +server.setRequestHandler(CallToolRequestSchema, async (request, extra) => { + const headers = extra.requestInfo?.headers; + const taskStore = extra.taskStore; + await extra.sendNotification({ method: 'notifications/progress', params: { progressToken: 'abc', progress: 50, total: 100 } }); + return { content: [{ type: 'text', text: 'result' }] }; +}); +``` + +**After (v2):** + +```typescript +server.setRequestHandler('tools/call', async (request, ctx) => { + const headers = ctx.http?.req?.headers; + const taskStore = ctx.task?.store; + await ctx.mcpReq.notify({ method: 'notifications/progress', params: { progressToken: 'abc', progress: 50, total: 100 } }); + return { content: [{ type: 'text', text: 'result' }] }; +}); +``` + +Context fields are organized into 4 groups: + +- **`mcpReq`** — request-level concerns: `id`, `method`, `_meta`, `signal`, `send()`, `notify()`, plus server-only `log()`, `elicitInput()`, and `requestSampling()` +- **`http?`** — HTTP transport concerns (undefined for stdio): `authInfo`, plus server-only `req`, `closeSSE`, `closeStandaloneSSE` +- **`task?`** — task lifecycle: `id`, `store`, `requestedTtl` + +`BaseContext` is the common base type shared by both `ServerContext` and `ClientContext`. `ServerContext` extends each group with server-specific additions via type intersection. + +`ServerContext` also provides convenience methods for common server→client operations: + +```typescript +server.setRequestHandler('tools/call', async (request, ctx) => { + // Send a log message (respects client's log level filter) + await ctx.mcpReq.log('info', 'Processing tool call', 'my-logger'); + + // Request client to sample an LLM + const samplingResult = await ctx.mcpReq.requestSampling({ + messages: [{ role: 'user', content: { type: 'text', text: 'Hello' } }], + maxTokens: 100, + }); + + // Elicit user input via a form + const elicitResult = await ctx.mcpReq.elicitInput({ + message: 'Please provide details', + requestedSchema: { type: 'object', properties: { name: { type: 'string' } } }, + }); + + return { content: [{ type: 'text', text: 'done' }] }; +}); +``` + +These replace the pattern of calling `server.sendLoggingMessage()`, `server.createMessage()`, and `server.elicitInput()` from within handlers. + ### Error hierarchy refactoring The SDK now distinguishes between two types of errors: diff --git a/examples/server/src/elicitationUrlExample.ts b/examples/server/src/elicitationUrlExample.ts index 7f22f9b3..c38dd75e 100644 --- a/examples/server/src/elicitationUrlExample.ts +++ b/examples/server/src/elicitationUrlExample.ts @@ -46,12 +46,12 @@ const getServer = () => { cartId: z.string().describe('The ID of the cart to confirm') }) }, - async ({ cartId }, extra): Promise => { + async ({ cartId }, ctx): Promise => { /* In a real world scenario, there would be some logic here to check if the user has the provided cartId. For the purposes of this example, we'll throw an error (-> elicits the client to open a URL to confirm payment) */ - const sessionId = extra.sessionId; + const sessionId = ctx.sessionId; if (!sessionId) { throw new Error('Expected a Session ID'); } @@ -79,15 +79,15 @@ const getServer = () => { param1: z.string().describe('First parameter') }) }, - async (_, extra): Promise => { + async (_, ctx): Promise => { /* In a real world scenario, there would be some logic here to check if we already have a valid access token for the user. - Auth info (with a subject or `sub` claim) can be typically be found in `extra.authInfo`. + Auth info (with a subject or `sub` claim) can be typically be found in `ctx.http?.authInfo`. If we do, we can just return the result of the tool call. If we don't, we can throw an ElicitationRequiredError to request the user to authenticate. For the purposes of this example, we'll throw an error (-> elicits the client to open a URL to authenticate). */ - const sessionId = extra.sessionId; + const sessionId = ctx.sessionId; if (!sessionId) { throw new Error('Expected a Session ID'); } diff --git a/examples/server/src/jsonResponseStreamableHttp.ts b/examples/server/src/jsonResponseStreamableHttp.ts index 2e041e88..7a3aad67 100644 --- a/examples/server/src/jsonResponseStreamableHttp.ts +++ b/examples/server/src/jsonResponseStreamableHttp.ts @@ -51,36 +51,18 @@ const getServer = () => { name: z.string().describe('Name to greet') }) }, - async ({ name }, extra): Promise => { + async ({ name }, ctx): Promise => { const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms)); - await server.sendLoggingMessage( - { - level: 'debug', - data: `Starting multi-greet for ${name}` - }, - extra.sessionId - ); + await ctx.mcpReq.log('debug', `Starting multi-greet for ${name}`); await sleep(1000); // Wait 1 second before first greeting - await server.sendLoggingMessage( - { - level: 'info', - data: `Sending first greeting to ${name}` - }, - extra.sessionId - ); + await ctx.mcpReq.log('info', `Sending first greeting to ${name}`); await sleep(1000); // Wait another second before second greeting - await server.sendLoggingMessage( - { - level: 'info', - data: `Sending second greeting to ${name}` - }, - extra.sessionId - ); + await ctx.mcpReq.log('info', `Sending second greeting to ${name}`); return { content: [ diff --git a/examples/server/src/simpleStatelessStreamableHttp.ts b/examples/server/src/simpleStatelessStreamableHttp.ts index 57d60c0a..2b4f0363 100644 --- a/examples/server/src/simpleStatelessStreamableHttp.ts +++ b/examples/server/src/simpleStatelessStreamableHttp.ts @@ -49,20 +49,14 @@ const getServer = () => { count: z.number().describe('Number of notifications to send (0 for 100)').default(10) }) }, - async ({ interval, count }, extra): Promise => { + async ({ interval, count }, ctx): Promise => { const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms)); let counter = 0; while (count === 0 || counter < count) { counter++; try { - await server.sendLoggingMessage( - { - level: 'info', - data: `Periodic notification #${counter} at ${new Date().toISOString()}` - }, - extra.sessionId - ); + await ctx.mcpReq.log('info', `Periodic notification #${counter} at ${new Date().toISOString()}`); } catch (error) { console.error('Error sending notification:', error); } diff --git a/examples/server/src/simpleStreamableHttp.ts b/examples/server/src/simpleStreamableHttp.ts index 7a2b22d4..dc86b17b 100644 --- a/examples/server/src/simpleStreamableHttp.ts +++ b/examples/server/src/simpleStreamableHttp.ts @@ -88,36 +88,18 @@ const getServer = () => { openWorldHint: false } }, - async ({ name }, extra): Promise => { + async ({ name }, ctx): Promise => { const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms)); - await server.sendLoggingMessage( - { - level: 'debug', - data: `Starting multi-greet for ${name}` - }, - extra.sessionId - ); + await ctx.mcpReq.log('debug', `Starting multi-greet for ${name}`); await sleep(1000); // Wait 1 second before first greeting - await server.sendLoggingMessage( - { - level: 'info', - data: `Sending first greeting to ${name}` - }, - extra.sessionId - ); + await ctx.mcpReq.log('info', `Sending first greeting to ${name}`); await sleep(1000); // Wait another second before second greeting - await server.sendLoggingMessage( - { - level: 'info', - data: `Sending second greeting to ${name}` - }, - extra.sessionId - ); + await ctx.mcpReq.log('info', `Sending second greeting to ${name}`); return { content: [ @@ -139,7 +121,7 @@ const getServer = () => { infoType: z.enum(['contact', 'preferences', 'feedback']).describe('Type of information to collect') }) }, - async ({ infoType }, extra): Promise => { + async ({ infoType }, ctx): Promise => { let message: string; let requestedSchema: { type: 'object'; @@ -238,8 +220,8 @@ const getServer = () => { } try { - // Use sendRequest through the extra parameter to elicit input - const result = await extra.sendRequest( + // Use sendRequest through the ctx parameter to elicit input + const result = await ctx.mcpReq.send( { method: 'elicitation/create', params: { @@ -327,20 +309,14 @@ const getServer = () => { count: z.number().describe('Number of notifications to send (0 for 100)').default(50) }) }, - async ({ interval, count }, extra): Promise => { + async ({ interval, count }, ctx): Promise => { const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms)); let counter = 0; while (count === 0 || counter < count) { counter++; try { - await server.sendLoggingMessage( - { - level: 'info', - data: `Periodic notification #${counter} at ${new Date().toISOString()}` - }, - extra.sessionId - ); + await ctx.mcpReq.log('info', `Periodic notification #${counter} at ${new Date().toISOString()}`); } catch (error) { console.error('Error sending notification:', error); } @@ -484,16 +460,16 @@ const getServer = () => { }) }, { - async createTask({ duration }, { taskStore, taskRequestedTtl }) { + async createTask({ duration }, ctx) { // Create the task - const task = await taskStore.createTask({ - ttl: taskRequestedTtl + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); // Simulate out-of-band work (async () => { await new Promise(resolve => setTimeout(resolve, duration)); - await taskStore.storeTaskResult(task.taskId, 'completed', { + await ctx.task.store.storeTaskResult(task.taskId, 'completed', { content: [ { type: 'text', @@ -508,11 +484,11 @@ const getServer = () => { task }; }, - async getTask(_args, { taskId, taskStore }) { - return await taskStore.getTask(taskId); + async getTask(_args, ctx) { + return await ctx.task.store.getTask(ctx.task.id); }, - async getTaskResult(_args, { taskId, taskStore }) { - const result = await taskStore.getTaskResult(taskId); + async getTaskResult(_args, ctx) { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as CallToolResult; } } diff --git a/examples/server/src/simpleTaskInteractive.ts b/examples/server/src/simpleTaskInteractive.ts index 3a3fe80c..9092926d 100644 --- a/examples/server/src/simpleTaskInteractive.ts +++ b/examples/server/src/simpleTaskInteractive.ts @@ -507,7 +507,7 @@ const createServer = (): Server => { }); // Handle tool calls - server.setRequestHandler('tools/call', async (request, extra): Promise => { + server.setRequestHandler('tools/call', async (request, ctx): Promise => { const { name, arguments: args } = request.params; const taskParams = (request.params._meta?.task || request.params.task) as { ttl?: number; pollInterval?: number } | undefined; @@ -522,7 +522,7 @@ const createServer = (): Server => { pollInterval: taskParams.pollInterval ?? 1000 }; - const task = await taskStore.createTask(taskOptions, extra.requestId, request, extra.sessionId); + const task = await taskStore.createTask(taskOptions, ctx.mcpReq.id, request, ctx.sessionId); console.log(`\n[Server] ${name} called, task created: ${task.taskId}`); @@ -600,7 +600,7 @@ const createServer = (): Server => { activeTaskExecutions.set(task.taskId, { promise: taskExecution, server, - sessionId: extra.sessionId ?? '' + sessionId: ctx.sessionId ?? '' }); return { task }; @@ -617,10 +617,10 @@ const createServer = (): Server => { }); // Handle tasks/result - server.setRequestHandler('tasks/result', async (request, extra): Promise => { + server.setRequestHandler('tasks/result', async (request, ctx): Promise => { const { taskId } = request.params; console.log(`[Server] tasks/result called for task ${taskId}`); - return taskResultHandler.handle(taskId, server, extra.sessionId ?? ''); + return taskResultHandler.handle(taskId, server, ctx.sessionId ?? ''); }); return server; diff --git a/examples/server/src/ssePollingExample.ts b/examples/server/src/ssePollingExample.ts index 2416d2ec..3897b392 100644 --- a/examples/server/src/ssePollingExample.ts +++ b/examples/server/src/ssePollingExample.ts @@ -7,7 +7,7 @@ * Key features: * - Configures `retryInterval` to tell clients how long to wait before reconnecting * - Uses `eventStore` to persist events for replay after reconnection - * - Uses `extra.closeSSEStream()` callback to gracefully disconnect clients mid-operation + * - Uses `ctx.http?.closeSSE()` callback to gracefully disconnect clients mid-operation * * Run with: pnpm tsx src/ssePollingExample.ts * Test with: curl or the MCP Inspector @@ -40,60 +40,36 @@ server.registerTool( { description: 'A long-running task that sends progress updates. Server will disconnect mid-task to demonstrate polling.' }, - async (extra): Promise => { + async (ctx): Promise => { const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms)); - console.log(`[${extra.sessionId}] Starting long-task...`); + console.log(`[${ctx.sessionId}] Starting long-task...`); // Send first progress notification - await server.sendLoggingMessage( - { - level: 'info', - data: 'Progress: 25% - Starting work...' - }, - extra.sessionId - ); + await ctx.mcpReq.log('info', 'Progress: 25% - Starting work...'); await sleep(1000); // Send second progress notification - await server.sendLoggingMessage( - { - level: 'info', - data: 'Progress: 50% - Halfway there...' - }, - extra.sessionId - ); + await ctx.mcpReq.log('info', 'Progress: 50% - Halfway there...'); await sleep(1000); // Server decides to disconnect the client to free resources // Client will reconnect via GET with Last-Event-ID after the transport's retryInterval - // Use extra.closeSSEStream callback - available when eventStore is configured - if (extra.closeSSEStream) { - console.log(`[${extra.sessionId}] Closing SSE stream to trigger client polling...`); - extra.closeSSEStream(); + // Use ctx.http?.closeSSE callback - available when eventStore is configured + if (ctx.http?.closeSSE) { + console.log(`[${ctx.sessionId}] Closing SSE stream to trigger client polling...`); + ctx.http?.closeSSE(); } // Continue processing while client is disconnected // Events are stored in eventStore and will be replayed on reconnect await sleep(500); - await server.sendLoggingMessage( - { - level: 'info', - data: 'Progress: 75% - Almost done (sent while client disconnected)...' - }, - extra.sessionId - ); + await ctx.mcpReq.log('info', 'Progress: 75% - Almost done (sent while client disconnected)...'); await sleep(500); - await server.sendLoggingMessage( - { - level: 'info', - data: 'Progress: 100% - Complete!' - }, - extra.sessionId - ); + await ctx.mcpReq.log('info', 'Progress: 100% - Complete!'); - console.log(`[${extra.sessionId}] Task complete`); + console.log(`[${ctx.sessionId}] Task complete`); return { content: [ diff --git a/packages/client/src/client/client.ts b/packages/client/src/client/client.ts index 9e015987..95519a7b 100644 --- a/packages/client/src/client/client.ts +++ b/packages/client/src/client/client.ts @@ -1,7 +1,9 @@ import { DefaultJsonSchemaValidator } from '@modelcontextprotocol/client/_shims'; import type { + BaseContext, CallToolRequest, ClientCapabilities, + ClientContext, ClientNotification, ClientRequest, ClientResult, @@ -19,12 +21,12 @@ import type { ListResourceTemplatesRequest, ListToolsRequest, LoggingLevel, + MessageExtraInfo, Notification, NotificationMethod, ProtocolOptions, ReadResourceRequest, Request, - RequestHandlerExtra, RequestMethod, RequestOptions, RequestTypeMap, @@ -215,7 +217,12 @@ export class Client< RequestT extends Request = Request, NotificationT extends Notification = Notification, ResultT extends Result = Result -> extends Protocol { +> extends Protocol< + ClientRequest | RequestT, + ClientNotification | NotificationT, + ClientResult | ResultT, + ClientContext +> { private _serverCapabilities?: ServerCapabilities; private _serverVersion?: Implementation; private _capabilities: ClientCapabilities; @@ -247,6 +254,13 @@ export class Client< } } + protected override buildContext( + ctx: BaseContext, + _transportInfo?: MessageExtraInfo + ): ClientContext { + return ctx; + } + /** * Set up handlers for list changed notifications based on config and server capabilities. * This should only be called after initialization when server capabilities are known. @@ -312,13 +326,13 @@ export class Client< method: M, handler: ( request: RequestTypeMap[M], - extra: RequestHandlerExtra + ctx: ClientContext ) => ClientResult | ResultT | Promise ): void { if (method === 'elicitation/create') { const wrappedHandler = async ( request: RequestTypeMap[M], - extra: RequestHandlerExtra + ctx: ClientContext ): Promise => { const validatedRequest = parseSchema(ElicitRequestSchema, request); if (!validatedRequest.success) { @@ -340,7 +354,7 @@ export class Client< throw new ProtocolError(ProtocolErrorCode.InvalidParams, 'Client does not support URL-mode elicitation requests'); } - const result = await Promise.resolve(handler(request, extra)); + const result = await Promise.resolve(handler(request, ctx)); // When task creation is requested, validate and return CreateTaskResult if (params.task) { @@ -385,13 +399,13 @@ export class Client< }; // Install the wrapped handler - return super.setRequestHandler(method, wrappedHandler as unknown as typeof handler); + return super.setRequestHandler(method, wrappedHandler); } if (method === 'sampling/createMessage') { const wrappedHandler = async ( request: RequestTypeMap[M], - extra: RequestHandlerExtra + ctx: ClientContext ): Promise => { const validatedRequest = parseSchema(CreateMessageRequestSchema, request); if (!validatedRequest.success) { @@ -402,7 +416,7 @@ export class Client< const { params } = validatedRequest.data; - const result = await Promise.resolve(handler(request, extra)); + const result = await Promise.resolve(handler(request, ctx)); // When task creation is requested, validate and return CreateTaskResult if (params.task) { @@ -431,7 +445,7 @@ export class Client< }; // Install the wrapped handler - return super.setRequestHandler(method, wrappedHandler as unknown as typeof handler); + return super.setRequestHandler(method, wrappedHandler); } // Other handlers use default behavior diff --git a/packages/core/src/experimental/tasks/interfaces.ts b/packages/core/src/experimental/tasks/interfaces.ts index c1901d70..57a24639 100644 --- a/packages/core/src/experimental/tasks/interfaces.ts +++ b/packages/core/src/experimental/tasks/interfaces.ts @@ -3,7 +3,7 @@ * WARNING: These APIs are experimental and may change without notice. */ -import type { RequestHandlerExtra, RequestTaskStore } from '../../shared/protocol.js'; +import type { RequestTaskStore, ServerContext } from '../../shared/protocol.js'; import type { JSONRPCErrorResponse, JSONRPCNotification, @@ -12,8 +12,6 @@ import type { Request, RequestId, Result, - ServerNotification, - ServerRequest, Task, ToolExecution } from '../../types/types.js'; @@ -23,21 +21,20 @@ import type { // ============================================================================ /** - * Extended handler extra with task store for task creation. + * Server context with guaranteed task store for task creation. * @experimental */ -export interface CreateTaskRequestHandlerExtra extends RequestHandlerExtra { - taskStore: RequestTaskStore; -} +export type CreateTaskServerContext = ServerContext & { + task: { store: RequestTaskStore; requestedTtl?: number | null }; +}; /** - * Extended handler extra with task ID and store for task operations. + * Server context with guaranteed task ID and store for task operations. * @experimental */ -export interface TaskRequestHandlerExtra extends RequestHandlerExtra { - taskId: string; - taskStore: RequestTaskStore; -} +export type TaskServerContext = ServerContext & { + task: { id: string; store: RequestTaskStore; requestedTtl?: number | null }; +}; /** * Task-specific execution configuration. diff --git a/packages/core/src/shared/protocol.ts b/packages/core/src/shared/protocol.ts index 7aaab459..fadf5222 100644 --- a/packages/core/src/shared/protocol.ts +++ b/packages/core/src/shared/protocol.ts @@ -5,6 +5,14 @@ import type { AuthInfo, CancelledNotification, ClientCapabilities, + ClientNotification, + ClientRequest, + CreateMessageRequest, + CreateMessageResult, + CreateMessageResultWithTools, + ElicitRequestFormParams, + ElicitRequestURLParams, + ElicitResult, GetTaskPayloadRequest, GetTaskRequest, GetTaskResult, @@ -13,6 +21,7 @@ import type { JSONRPCRequest, JSONRPCResponse, JSONRPCResultResponse, + LoggingLevel, MessageExtraInfo, Notification, NotificationMethod, @@ -28,6 +37,8 @@ import type { RequestTypeMap, Result, ServerCapabilities, + ServerNotification, + ServerRequest, Task, TaskCreationParams, TaskStatusNotification @@ -244,75 +255,134 @@ export interface RequestTaskStore { } /** - * Extra data given to request handlers. + * Task context provided to request handlers when task storage is configured. */ -export type RequestHandlerExtra = { - /** - * An abort signal used to communicate if the request was cancelled from the sender's side. - */ - signal: AbortSignal; - - /** - * Information about a validated access token, provided to request handlers. - */ - authInfo?: AuthInfo; +export type TaskContext = { + id?: string; + store: RequestTaskStore; + requestedTtl?: number | null; +}; +/** + * Base context provided to all request handlers. + */ +export type BaseContext = { /** * The session ID from the transport, if available. */ sessionId?: string; /** - * Metadata from the original request. + * Information about the MCP request being handled. */ - _meta?: RequestMeta; + mcpReq: { + /** + * The JSON-RPC ID of the request being handled. + */ + id: RequestId; + + /** + * The method name of the request (e.g., 'tools/call', 'ping'). + */ + method: string; + + /** + * Metadata from the original request. + */ + _meta?: RequestMeta; + + /** + * An abort signal used to communicate if the request was cancelled from the sender's side. + */ + signal: AbortSignal; + + /** + * Sends a request that relates to the current request being handled. + * + * This is used by certain transports to correctly associate related messages. + */ + send: (request: SendRequestT, resultSchema: U, options?: TaskRequestOptions) => Promise>; + + /** + * Sends a notification that relates to the current request being handled. + * + * This is used by certain transports to correctly associate related messages. + */ + notify: (notification: SendNotificationT) => Promise; + }; /** - * The JSON-RPC ID of the request being handled. - * This can be useful for tracking or logging purposes. + * HTTP transport information, only available when using an HTTP-based transport. */ - requestId: RequestId; - - taskId?: string; - - taskStore?: RequestTaskStore; - - taskRequestedTtl?: number | null; + http?: { + /** + * Information about a validated access token, provided to request handlers. + */ + authInfo?: AuthInfo; + }; /** - * The original HTTP request. + * Task context, available when task storage is configured. */ - requestInfo?: RequestInfo; - - /** - * Sends a notification that relates to the current request being handled. - * - * This is used by certain transports to correctly associate related messages. - */ - sendNotification: (notification: SendNotificationT) => Promise; - - /** - * Sends a request that relates to the current request being handled. - * - * This is used by certain transports to correctly associate related messages. - */ - sendRequest: (request: SendRequestT, resultSchema: U, options?: TaskRequestOptions) => Promise>; - - /** - * Closes the SSE stream for this request, triggering client reconnection. - * Only available when using a StreamableHTTPServerTransport with eventStore configured. - * Use this to implement polling behavior during long-running operations. - */ - closeSSEStream?: () => void; - - /** - * Closes the standalone GET SSE stream, triggering client reconnection. - * Only available when using a StreamableHTTPServerTransport with eventStore configured. - * Use this to implement polling behavior for server-initiated notifications. - */ - closeStandaloneSSEStream?: () => void; + task?: TaskContext; }; +/** + * Context provided to server-side request handlers, extending BaseContext with server-specific fields. + */ +export type ServerContext< + SendRequestT extends Request = ServerRequest, + SendNotificationT extends Notification = ServerNotification +> = BaseContext & { + mcpReq: { + /** + * Send a log message notification to the client. + * Respects the client's log level filter set via logging/setLevel. + */ + log: (level: LoggingLevel, data: unknown, logger?: string) => Promise; + + /** + * Send an elicitation request to the client, requesting user input. + */ + elicitInput: (params: ElicitRequestFormParams | ElicitRequestURLParams, options?: RequestOptions) => Promise; + + /** + * Request LLM sampling from the client. + */ + requestSampling: ( + params: CreateMessageRequest['params'], + options?: RequestOptions + ) => Promise; + }; + + http?: { + /** + * The original HTTP request information. + */ + req?: RequestInfo; + + /** + * Closes the SSE stream for this request, triggering client reconnection. + * Only available when using a StreamableHTTPServerTransport with eventStore configured. + */ + closeSSE?: () => void; + + /** + * Closes the standalone GET SSE stream, triggering client reconnection. + * Only available when using a StreamableHTTPServerTransport with eventStore configured. + */ + closeStandaloneSSE?: () => void; + }; +}; + +/** + * Context provided to client-side request handlers. + */ +export type ClientContext< + SendRequestT extends Request = ClientRequest, + SendNotificationT extends Notification = ClientNotification +> = BaseContext; + /** * Information about a request's timeout state */ @@ -329,13 +399,15 @@ type TimeoutInfo = { * Implements MCP protocol framing on top of a pluggable transport, including * features like request/response linking, notifications, and progress. */ -export abstract class Protocol { +export abstract class Protocol< + SendRequestT extends Request, + SendNotificationT extends Notification, + SendResultT extends Result, + ContextT extends BaseContext +> { private _transport?: Transport; private _requestMessageId = 0; - private _requestHandlers: Map< - string, - (request: JSONRPCRequest, extra: RequestHandlerExtra) => Promise - > = new Map(); + private _requestHandlers: Map Promise> = new Map(); private _requestHandlerAbortControllers: Map = new Map(); private _notificationHandlers: Map Promise> = new Map(); private _responseHandlers: Map void> = new Map(); @@ -370,7 +442,7 @@ export abstract class Protocol) => Promise; + fallbackRequestHandler?: (request: JSONRPCRequest, ctx: ContextT) => Promise; /** * A handler to invoke for any notification types that do not have their own handler installed. @@ -398,8 +470,8 @@ export abstract class Protocol { - const task = await this._taskStore!.getTask(request.params.taskId, extra.sessionId); + this.setRequestHandler('tasks/get', async (request, ctx) => { + const task = await this._taskStore!.getTask(request.params.taskId, ctx.sessionId); if (!task) { throw new ProtocolError(ProtocolErrorCode.InvalidParams, 'Failed to retrieve task: Task not found'); } @@ -411,14 +483,14 @@ export abstract class Protocol { + this.setRequestHandler('tasks/result', async (request, ctx) => { const handleTaskResult = async (): Promise => { const taskId = request.params.taskId; // Deliver queued messages if (this._taskMessageQueue) { let queuedMessage: QueuedMessage | undefined; - while ((queuedMessage = await this._taskMessageQueue.dequeue(taskId, extra.sessionId))) { + while ((queuedMessage = await this._taskMessageQueue.dequeue(taskId, ctx.sessionId))) { // Handle response and error messages by routing them to the appropriate resolver if (queuedMessage.type === 'response' || queuedMessage.type === 'error') { const message = queuedMessage.message; @@ -456,12 +528,12 @@ export abstract class Protocol { + this.setRequestHandler('tasks/list', async (request, ctx) => { try { - const { tasks, nextCursor } = await this._taskStore!.listTasks(request.params?.cursor, extra.sessionId); + const { tasks, nextCursor } = await this._taskStore!.listTasks(request.params?.cursor, ctx.sessionId); return { tasks, nextCursor, @@ -514,10 +586,10 @@ export abstract class Protocol { + this.setRequestHandler('tasks/cancel', async (request, ctx) => { try { // Get the current task to check if it's in a terminal state, in case the implementation is not atomic - const task = await this._taskStore!.getTask(request.params.taskId, extra.sessionId); + const task = await this._taskStore!.getTask(request.params.taskId, ctx.sessionId); if (!task) { throw new ProtocolError(ProtocolErrorCode.InvalidParams, `Task not found: ${request.params.taskId}`); @@ -532,12 +604,12 @@ export abstract class Protocol, transportInfo?: MessageExtraInfo): ContextT; + private async _oncancel(notification: CancelledNotification): Promise { if (!notification.params.requestId) { return; @@ -734,43 +812,40 @@ export abstract class Protocol = { - signal: abortController.signal, + const task: TaskContext | undefined = taskStore + ? { id: relatedTaskId, store: taskStore, requestedTtl: taskCreationParams?.ttl } + : undefined; + + const baseCtx: BaseContext = { sessionId: capturedTransport?.sessionId, - _meta: request.params?._meta, - sendNotification: async notification => { - // Include related-task metadata if this request is part of a task - const notificationOptions: NotificationOptions = { relatedRequestId: request.id }; - if (relatedTaskId) { - notificationOptions.relatedTask = { taskId: relatedTaskId }; + mcpReq: { + id: request.id, + method: request.method, + _meta: request.params?._meta, + signal: abortController.signal, + send: async (r, resultSchema, options?) => { + const requestOptions: RequestOptions = { ...options, relatedRequestId: request.id }; + if (relatedTaskId && !requestOptions.relatedTask) { + requestOptions.relatedTask = { taskId: relatedTaskId }; + } + const effectiveTaskId = requestOptions.relatedTask?.taskId ?? relatedTaskId; + if (effectiveTaskId && taskStore) { + await taskStore.updateTaskStatus(effectiveTaskId, 'input_required'); + } + return await this.request(r, resultSchema, requestOptions); + }, + notify: async notification => { + const notificationOptions: NotificationOptions = { relatedRequestId: request.id }; + if (relatedTaskId) { + notificationOptions.relatedTask = { taskId: relatedTaskId }; + } + await this.notification(notification, notificationOptions); } - await this.notification(notification, notificationOptions); }, - sendRequest: async (r, resultSchema, options?) => { - // Include related-task metadata if this request is part of a task - const requestOptions: RequestOptions = { ...options, relatedRequestId: request.id }; - if (relatedTaskId && !requestOptions.relatedTask) { - requestOptions.relatedTask = { taskId: relatedTaskId }; - } - - // Set task status to input_required when sending a request within a task context - // Use the taskId from options (explicit) or fall back to relatedTaskId (inherited) - const effectiveTaskId = requestOptions.relatedTask?.taskId ?? relatedTaskId; - if (effectiveTaskId && taskStore) { - await taskStore.updateTaskStatus(effectiveTaskId, 'input_required'); - } - - return await this.request(r, resultSchema, requestOptions); - }, - authInfo: extra?.authInfo, - requestId: request.id, - requestInfo: extra?.requestInfo, - taskId: relatedTaskId, - taskStore: taskStore, - taskRequestedTtl: taskCreationParams?.ttl, - closeSSEStream: extra?.closeSSEStream, - closeStandaloneSSEStream: extra?.closeStandaloneSSEStream + http: extra?.authInfo ? { authInfo: extra.authInfo } : undefined, + task }; + const ctx = this.buildContext(baseCtx, extra); // Starting with Promise.resolve() puts any synchronous errors into the monad as well. Promise.resolve() @@ -781,7 +856,7 @@ export abstract class Protocol handler(request, fullExtra)) + .then(() => handler(request, ctx)) .then( async result => { if (abortController.signal.aborted) { @@ -1418,17 +1493,14 @@ export abstract class Protocol( method: M, - handler: ( - request: RequestTypeMap[M], - extra: RequestHandlerExtra - ) => SendResultT | Promise + handler: (request: RequestTypeMap[M], ctx: ContextT) => SendResultT | Promise ): void { this.assertRequestHandlerCapability(method); const schema = getRequestSchema(method); - this._requestHandlers.set(method, (request, extra) => { - const parsed = schema.parse(request); - return Promise.resolve(handler(parsed, extra)); + this._requestHandlers.set(method, (request, ctx) => { + const parsed = schema.parse(request) as RequestTypeMap[M]; + return Promise.resolve(handler(parsed, ctx)); }); } diff --git a/packages/core/test/shared/protocol.test.ts b/packages/core/test/shared/protocol.test.ts index 535af4ed..67502bfc 100644 --- a/packages/core/test/shared/protocol.test.ts +++ b/packages/core/test/shared/protocol.test.ts @@ -11,6 +11,7 @@ import type { TaskStore } from '../../src/experimental/tasks/interfaces.js'; import { InMemoryTaskMessageQueue } from '../../src/experimental/tasks/stores/inMemory.js'; +import type { BaseContext } from '../../src/shared/protocol.js'; import { mergeCapabilities, Protocol } from '../../src/shared/protocol.js'; import type { ErrorMessage, ResponseMessage } from '../../src/shared/responseMessage.js'; import { toArrayAsync } from '../../src/shared/responseMessage.js'; @@ -149,18 +150,21 @@ function assertQueuedRequest(o?: QueuedMessage): asserts o is QueuedRequest { } describe('protocol tests', () => { - let protocol: Protocol; + let protocol: Protocol>; let transport: MockTransport; let sendSpy: MockInstance; beforeEach(() => { transport = new MockTransport(); sendSpy = vi.spyOn(transport, 'send'); - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })(); }); @@ -630,11 +634,14 @@ describe('protocol tests', () => { it('should NOT debounce a notification that has parameters', async () => { // ARRANGE - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ debouncedNotificationMethods: ['test/debounced_with_params'] }); await protocol.connect(transport); @@ -653,11 +660,14 @@ describe('protocol tests', () => { it('should NOT debounce a notification that has a relatedRequestId', async () => { // ARRANGE - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ debouncedNotificationMethods: ['test/debounced_with_options'] }); await protocol.connect(transport); @@ -674,11 +684,14 @@ describe('protocol tests', () => { it('should clear pending debounced notifications on connection close', async () => { // ARRANGE - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ debouncedNotificationMethods: ['test/debounced'] }); await protocol.connect(transport); @@ -700,11 +713,14 @@ describe('protocol tests', () => { it('should debounce multiple synchronous calls when params property is omitted', async () => { // ARRANGE - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ debouncedNotificationMethods: ['test/debounced'] }); await protocol.connect(transport); @@ -729,11 +745,14 @@ describe('protocol tests', () => { it('should debounce calls when params is explicitly undefined', async () => { // ARRANGE - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ debouncedNotificationMethods: ['test/debounced'] }); await protocol.connect(transport); @@ -756,11 +775,14 @@ describe('protocol tests', () => { it('should send non-debounced notifications immediately and multiple times', async () => { // ARRANGE - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ debouncedNotificationMethods: ['test/debounced'] }); // Configure for a different method await protocol.connect(transport); @@ -791,11 +813,14 @@ describe('protocol tests', () => { it('should handle sequential batches of debounced notifications correctly', async () => { // ARRANGE - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ debouncedNotificationMethods: ['test/debounced'] }); await protocol.connect(transport); @@ -1003,18 +1028,21 @@ describe('mergeCapabilities', () => { }); describe('Task-based execution', () => { - let protocol: Protocol; + let protocol: Protocol>; let transport: MockTransport; let sendSpy: MockInstance; beforeEach(() => { transport = new MockTransport(); sendSpy = vi.spyOn(transport, 'send'); - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore: createMockTaskStore(), taskMessageQueue: new InMemoryTaskMessageQueue() }); }); @@ -1257,11 +1285,14 @@ describe('Task-based execution', () => { // rather than in _meta, and that task management is handled by tool implementors const mockTaskStore = createMockTaskStore(); - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore: mockTaskStore }); @@ -1326,11 +1357,14 @@ describe('Task-based execution', () => { } ); - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore: mockTaskStore }); @@ -1387,11 +1421,14 @@ describe('Task-based execution', () => { } ); - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore: mockTaskStore }); @@ -1433,11 +1470,14 @@ describe('Task-based execution', () => { onList: () => listedTasks.releaseLatch() }); - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore: mockTaskStore }); @@ -1466,11 +1506,14 @@ describe('Task-based execution', () => { const mockTaskStore = createMockTaskStore(); mockTaskStore.listTasks.mockRejectedValue(new Error('Invalid cursor: bad-cursor')); - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore: mockTaskStore }); @@ -1600,11 +1643,14 @@ describe('Task-based execution', () => { throw new Error('Task not found'); }); - const serverProtocol = new (class extends Protocol { + const serverProtocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore: mockTaskStore }); const serverTransport = new MockTransport(); @@ -1642,11 +1688,14 @@ describe('Task-based execution', () => { mockTaskStore.getTask.mockResolvedValue(null); - const serverProtocol = new (class extends Protocol { + const serverProtocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore: mockTaskStore }); const serverTransport = new MockTransport(); @@ -1690,11 +1739,14 @@ describe('Task-based execution', () => { mockTaskStore.updateTaskStatus.mockClear(); mockTaskStore.getTask.mockResolvedValue(completedTask); - const serverProtocol = new (class extends Protocol { + const serverProtocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore: mockTaskStore }); const serverTransport = new MockTransport(); @@ -1772,11 +1824,14 @@ describe('Task-based execution', () => { params: {} }); - const serverProtocol = new (class extends Protocol { + const serverProtocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore: mockTaskStore }); const serverTransport = new MockTransport(); @@ -1822,11 +1877,14 @@ describe('Task-based execution', () => { params: {} }); - const serverProtocol = new (class extends Protocol { + const serverProtocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore: mockTaskStore }); const serverTransport = new MockTransport(); @@ -1871,11 +1929,14 @@ describe('Task-based execution', () => { params: {} }); - const serverProtocol = new (class extends Protocol { + const serverProtocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore: mockTaskStore }); const serverTransport = new MockTransport(); @@ -1908,11 +1969,14 @@ describe('Task-based execution', () => { params: {} }); - const serverProtocol = new (class extends Protocol { + const serverProtocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore: mockTaskStore }); const serverTransport = new MockTransport(); @@ -1953,11 +2017,14 @@ describe('Task-based execution', () => { await mockTaskStore.storeTaskResult(task.taskId, 'completed', testResult); - const serverProtocol = new (class extends Protocol { + const serverProtocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore: mockTaskStore }); const serverTransport = new MockTransport(); @@ -1996,11 +2063,14 @@ describe('Task-based execution', () => { it('should propagate related-task metadata to handler sendRequest and sendNotification', async () => { const mockTaskStore = createMockTaskStore(); - const serverProtocol = new (class extends Protocol { + const serverProtocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore: mockTaskStore, taskMessageQueue: new InMemoryTaskMessageQueue() }); @@ -2010,9 +2080,9 @@ describe('Task-based execution', () => { await serverProtocol.connect(serverTransport); // Set up a handler that uses sendRequest and sendNotification - serverProtocol.setRequestHandler('tools/call', async (_request, extra) => { - // Send a notification using the extra.sendNotification - await extra.sendNotification({ + serverProtocol.setRequestHandler('tools/call', async (_request, ctx) => { + // Send a notification using the ctx.mcpReq.notify + await ctx.mcpReq.notify({ method: 'notifications/message', params: { level: 'info', data: 'test' } }); @@ -2072,18 +2142,21 @@ describe('Task-based execution', () => { }); describe('Request Cancellation vs Task Cancellation', () => { - let protocol: Protocol; + let protocol: Protocol>; let transport: MockTransport; let taskStore: TaskStore; beforeEach(() => { transport = new MockTransport(); taskStore = createMockTaskStore(); - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore }); }); @@ -2094,10 +2167,10 @@ describe('Request Cancellation vs Task Cancellation', () => { // Set up a request handler that checks if it was aborted let wasAborted = false; - protocol.setRequestHandler('ping', async (_request, extra) => { + protocol.setRequestHandler('ping', async (_request, ctx) => { // Simulate a long-running operation await new Promise(resolve => setTimeout(resolve, 100)); - wasAborted = extra.signal.aborted; + wasAborted = ctx.mcpReq.signal.aborted; return {}; }); @@ -2357,29 +2430,35 @@ describe('Request Cancellation vs Task Cancellation', () => { }); describe('Progress notification support for tasks', () => { - let protocol: Protocol; + let protocol: Protocol>; let transport: MockTransport; let sendSpy: MockInstance; beforeEach(() => { transport = new MockTransport(); sendSpy = vi.spyOn(transport, 'send'); - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })(); }); it('should maintain progress token association after CreateTaskResult is returned', async () => { const taskStore = createMockTaskStore(); - const protocol = new (class extends Protocol { + const protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore }); @@ -2468,11 +2547,14 @@ describe('Progress notification support for tasks', () => { it('should stop progress notifications when task reaches terminal status (completed)', async () => { const taskStore = createMockTaskStore(); - const protocol = new (class extends Protocol { + const protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore }); @@ -2481,13 +2563,14 @@ describe('Progress notification support for tasks', () => { await protocol.connect(transport); // Set up a request handler that will complete the task - protocol.setRequestHandler('tools/call', async (_request, extra) => { - if (extra.taskStore) { - const task = await extra.taskStore.createTask({ ttl: 60000 }); + protocol.setRequestHandler('tools/call', async (_request, ctx) => { + if (ctx.task?.store) { + const task = await ctx.task.store.createTask({ ttl: 60000 }); // Simulate async work then complete the task + const taskStore = ctx.task.store; setTimeout(async () => { - await extra.taskStore!.storeTaskResult(task.taskId, 'completed', { + await taskStore.storeTaskResult(task.taskId, 'completed', { content: [{ type: 'text', text: 'Done' }] }); }, 50); @@ -2603,11 +2686,14 @@ describe('Progress notification support for tasks', () => { it('should stop progress notifications when task reaches terminal status (failed)', async () => { const taskStore = createMockTaskStore(); - const protocol = new (class extends Protocol { + const protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore }); @@ -2701,11 +2787,14 @@ describe('Progress notification support for tasks', () => { it('should stop progress notifications when task is cancelled', async () => { const taskStore = createMockTaskStore(); - const protocol = new (class extends Protocol { + const protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore }); @@ -2796,11 +2885,14 @@ describe('Progress notification support for tasks', () => { it('should use the same progressToken throughout task lifetime', async () => { const taskStore = createMockTaskStore(); - const protocol = new (class extends Protocol { + const protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore }); @@ -3065,11 +3157,14 @@ describe('Message interception for task-related notifications', () => { it('should queue notifications with io.modelcontextprotocol/related-task metadata', async () => { const taskStore = createMockTaskStore(); const transport = new MockTransport(); - const server = new (class extends Protocol { + const server = new (class extends Protocol> { protected assertCapabilityForMethod(_method: string): void {} protected assertNotificationCapability(_method: string): void {} protected assertRequestHandlerCapability(_method: string): void {} protected assertTaskCapability(_method: string): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(_method: string): void {} })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue() }); @@ -3102,11 +3197,14 @@ describe('Message interception for task-related notifications', () => { it('should not queue notifications without related-task metadata', async () => { const taskStore = createMockTaskStore(); const transport = new MockTransport(); - const server = new (class extends Protocol { + const server = new (class extends Protocol> { protected assertCapabilityForMethod(_method: string): void {} protected assertNotificationCapability(_method: string): void {} protected assertRequestHandlerCapability(_method: string): void {} protected assertTaskCapability(_method: string): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(_method: string): void {} })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue() }); @@ -3129,11 +3227,14 @@ describe('Message interception for task-related notifications', () => { it('should propagate queue overflow errors without failing the task', async () => { const taskStore = createMockTaskStore(); const transport = new MockTransport(); - const server = new (class extends Protocol { + const server = new (class extends Protocol> { protected assertCapabilityForMethod(_method: string): void {} protected assertNotificationCapability(_method: string): void {} protected assertRequestHandlerCapability(_method: string): void {} protected assertTaskCapability(_method: string): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(_method: string): void {} })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue(), maxTaskQueueSize: 100 }); @@ -3176,11 +3277,14 @@ describe('Message interception for task-related notifications', () => { it('should extract task ID correctly from metadata', async () => { const taskStore = createMockTaskStore(); const transport = new MockTransport(); - const server = new (class extends Protocol { + const server = new (class extends Protocol> { protected assertCapabilityForMethod(_method: string): void {} protected assertNotificationCapability(_method: string): void {} protected assertRequestHandlerCapability(_method: string): void {} protected assertTaskCapability(_method: string): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(_method: string): void {} })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue() }); @@ -3209,11 +3313,14 @@ describe('Message interception for task-related notifications', () => { it('should preserve message order when queuing multiple notifications', async () => { const taskStore = createMockTaskStore(); const transport = new MockTransport(); - const server = new (class extends Protocol { + const server = new (class extends Protocol> { protected assertCapabilityForMethod(_method: string): void {} protected assertNotificationCapability(_method: string): void {} protected assertRequestHandlerCapability(_method: string): void {} protected assertTaskCapability(_method: string): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(_method: string): void {} })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue() }); @@ -3251,11 +3358,14 @@ describe('Message interception for task-related requests', () => { it('should queue requests with io.modelcontextprotocol/related-task metadata', async () => { const taskStore = createMockTaskStore(); const transport = new MockTransport(); - const server = new (class extends Protocol { + const server = new (class extends Protocol> { protected assertCapabilityForMethod(_method: string): void {} protected assertNotificationCapability(_method: string): void {} protected assertRequestHandlerCapability(_method: string): void {} protected assertTaskCapability(_method: string): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(_method: string): void {} })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue() }); @@ -3303,11 +3413,14 @@ describe('Message interception for task-related requests', () => { it('should not queue requests without related-task metadata', async () => { const taskStore = createMockTaskStore(); const transport = new MockTransport(); - const server = new (class extends Protocol { + const server = new (class extends Protocol> { protected assertCapabilityForMethod(_method: string): void {} protected assertNotificationCapability(_method: string): void {} protected assertRequestHandlerCapability(_method: string): void {} protected assertTaskCapability(_method: string): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(_method: string): void {} })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue() }); @@ -3342,11 +3455,14 @@ describe('Message interception for task-related requests', () => { it('should store request resolver for response routing', async () => { const taskStore = createMockTaskStore(); const transport = new MockTransport(); - const server = new (class extends Protocol { + const server = new (class extends Protocol> { protected assertCapabilityForMethod(_method: string): void {} protected assertNotificationCapability(_method: string): void {} protected assertRequestHandlerCapability(_method: string): void {} protected assertTaskCapability(_method: string): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(_method: string): void {} })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue() }); @@ -3395,11 +3511,14 @@ describe('Message interception for task-related requests', () => { const taskStore = createMockTaskStore(); const transport = new MockTransport(); const queue = new InMemoryTaskMessageQueue(); - const server = new (class extends Protocol { + const server = new (class extends Protocol> { protected assertCapabilityForMethod(_method: string): void {} protected assertNotificationCapability(_method: string): void {} protected assertRequestHandlerCapability(_method: string): void {} protected assertTaskCapability(_method: string): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(_method: string): void {} })({ taskStore, taskMessageQueue: queue }); @@ -3460,11 +3579,14 @@ describe('Message interception for task-related requests', () => { it('should log error when resolver is missing for side-channeled request', async () => { const taskStore = createMockTaskStore(); const transport = new MockTransport(); - const server = new (class extends Protocol { + const server = new (class extends Protocol> { protected assertCapabilityForMethod(_method: string): void {} protected assertNotificationCapability(_method: string): void {} protected assertRequestHandlerCapability(_method: string): void {} protected assertTaskCapability(_method: string): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(_method: string): void {} })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue() }); @@ -3536,11 +3658,14 @@ describe('Message interception for task-related requests', () => { it('should propagate queue overflow errors for requests without failing the task', async () => { const taskStore = createMockTaskStore(); const transport = new MockTransport(); - const server = new (class extends Protocol { + const server = new (class extends Protocol> { protected assertCapabilityForMethod(_method: string): void {} protected assertNotificationCapability(_method: string): void {} protected assertRequestHandlerCapability(_method: string): void {} protected assertTaskCapability(_method: string): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(_method: string): void {} })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue(), maxTaskQueueSize: 100 }); @@ -3590,18 +3715,21 @@ describe('Message interception for task-related requests', () => { }); describe('Message Interception', () => { - let protocol: Protocol; + let protocol: Protocol>; let transport: MockTransport; let mockTaskStore: TaskStore & { [K in keyof TaskStore]: MockInstance }; beforeEach(() => { transport = new MockTransport(); mockTaskStore = createMockTaskStore(); - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore: mockTaskStore, taskMessageQueue: new InMemoryTaskMessageQueue() }); }); @@ -4108,18 +4236,21 @@ describe('Message Interception', () => { }); describe('Queue lifecycle management', () => { - let protocol: Protocol; + let protocol: Protocol>; let transport: MockTransport; let mockTaskStore: TaskStore & { [K in keyof TaskStore]: MockInstance }; beforeEach(() => { transport = new MockTransport(); mockTaskStore = createMockTaskStore(); - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore: mockTaskStore, taskMessageQueue: new InMemoryTaskMessageQueue() }); }); @@ -4429,11 +4560,14 @@ describe('requestStream() method', () => { test('should yield result immediately for non-task requests', async () => { const transport = new MockTransport(); - const protocol = new (class extends Protocol { + const protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })(); await protocol.connect(transport); @@ -4472,11 +4606,14 @@ describe('requestStream() method', () => { test('should yield error message on request failure', async () => { const transport = new MockTransport(); - const protocol = new (class extends Protocol { + const protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })(); await protocol.connect(transport); @@ -4518,11 +4655,14 @@ describe('requestStream() method', () => { test('should handle cancellation via AbortSignal', async () => { const transport = new MockTransport(); - const protocol = new (class extends Protocol { + const protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })(); await protocol.connect(transport); @@ -4556,11 +4696,14 @@ describe('requestStream() method', () => { describe('Error responses', () => { test('should yield error as terminal message for server error response', async () => { const transport = new MockTransport(); - const protocol = new (class extends Protocol { + const protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })(); await protocol.connect(transport); @@ -4598,11 +4741,14 @@ describe('requestStream() method', () => { vi.useFakeTimers(); try { const transport = new MockTransport(); - const protocol = new (class extends Protocol { + const protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })(); await protocol.connect(transport); @@ -4637,11 +4783,14 @@ describe('requestStream() method', () => { test('should yield error as terminal message for cancellation', async () => { const transport = new MockTransport(); - const protocol = new (class extends Protocol { + const protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })(); await protocol.connect(transport); @@ -4670,11 +4819,14 @@ describe('requestStream() method', () => { test('should not yield any messages after error message', async () => { const transport = new MockTransport(); - const protocol = new (class extends Protocol { + const protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })(); await protocol.connect(transport); @@ -4722,11 +4874,14 @@ describe('requestStream() method', () => { test('should yield error as terminal message for task failure', async () => { const transport = new MockTransport(); const mockTaskStore = createMockTaskStore(); - const protocol = new (class extends Protocol { + const protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })({ taskStore: mockTaskStore }); await protocol.connect(transport); @@ -4782,11 +4937,14 @@ describe('requestStream() method', () => { test('should yield error as terminal message for network error', async () => { const transport = new MockTransport(); - const protocol = new (class extends Protocol { + const protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })(); await protocol.connect(transport); @@ -4810,11 +4968,14 @@ describe('requestStream() method', () => { test('should ensure error is always the final message', async () => { const transport = new MockTransport(); - const protocol = new (class extends Protocol { + const protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })(); await protocol.connect(transport); @@ -4855,7 +5016,7 @@ describe('requestStream() method', () => { }); describe('Error handling for missing resolvers', () => { - let protocol: Protocol; + let protocol: Protocol>; let transport: MockTransport; let taskStore: TaskStore & { [K in keyof TaskStore]: MockInstance }; let taskMessageQueue: TaskMessageQueue; @@ -4866,11 +5027,14 @@ describe('Error handling for missing resolvers', () => { taskMessageQueue = new InMemoryTaskMessageQueue(); errorHandler = vi.fn(); - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(_method: string): void {} protected assertNotificationCapability(_method: string): void {} protected assertRequestHandlerCapability(_method: string): void {} protected assertTaskCapability(_method: string): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(_method: string): void {} })({ taskStore, diff --git a/packages/core/test/shared/protocolTransportHandling.test.ts b/packages/core/test/shared/protocolTransportHandling.test.ts index b76c9a6d..889100e7 100644 --- a/packages/core/test/shared/protocolTransportHandling.test.ts +++ b/packages/core/test/shared/protocolTransportHandling.test.ts @@ -1,5 +1,6 @@ import { beforeEach, describe, expect, test } from 'vitest'; +import type { BaseContext } from '../../src/shared/protocol.js'; import { Protocol } from '../../src/shared/protocol.js'; import type { Transport } from '../../src/shared/transport.js'; import type { EmptyResult, JSONRPCMessage, Notification, Request, Result } from '../../src/types/types.js'; @@ -28,16 +29,19 @@ class MockTransport implements Transport { } describe('Protocol transport handling bug', () => { - let protocol: Protocol; + let protocol: Protocol>; let transportA: MockTransport; let transportB: MockTransport; beforeEach(() => { - protocol = new (class extends Protocol { + protocol = new (class extends Protocol> { protected assertCapabilityForMethod(): void {} protected assertNotificationCapability(): void {} protected assertRequestHandlerCapability(): void {} protected assertTaskCapability(): void {} + protected buildContext(ctx: any) { + return ctx; + } protected assertTaskHandlerCapability(): void {} })(); @@ -86,8 +90,8 @@ describe('Protocol transport handling bug', () => { const results: { transport: string; response: JSONRPCMessage[] }[] = []; // Set up handler with variable delay based on request id - protocol.setRequestHandler('ping', async (_request, extra) => { - const delay = extra.requestId === 1 ? 50 : 10; + protocol.setRequestHandler('ping', async (_request, ctx) => { + const delay = ctx.mcpReq.id === 1 ? 50 : 10; await new Promise(resolve => setTimeout(resolve, delay)); return {}; }); diff --git a/packages/middleware/node/test/streamableHttp.test.ts b/packages/middleware/node/test/streamableHttp.test.ts index 0ea3588b..ab3ba540 100644 --- a/packages/middleware/node/test/streamableHttp.test.ts +++ b/packages/middleware/node/test/streamableHttp.test.ts @@ -214,8 +214,10 @@ describe('Zod v4', () => { description: 'A user profile data tool', inputSchema: z.object({ active: z.boolean().describe('Profile status') }) }, - async ({ active }, { authInfo }): Promise => { - return { content: [{ type: 'text', text: `${active ? 'Active' : 'Inactive'} profile from token: ${authInfo?.token}!` }] }; + async ({ active }, ctx): Promise => { + return { + content: [{ type: 'text', text: `${active ? 'Active' : 'Inactive'} profile from token: ${ctx.http?.authInfo?.token}!` }] + }; } ); @@ -403,11 +405,11 @@ describe('Zod v4', () => { description: 'A simple test tool with request info', inputSchema: z.object({ name: z.string().describe('Name to greet') }) }, - async ({ name }, { requestInfo }): Promise => { + async ({ name }, ctx): Promise => { // Convert Headers object to plain object for JSON serialization // Headers is a Web API class that doesn't serialize with JSON.stringify const serializedRequestInfo = { - headers: Object.fromEntries(requestInfo?.headers ?? new Headers()) + headers: Object.fromEntries(ctx.http?.req?.headers ?? new Headers()) }; return { content: [ @@ -1831,7 +1833,7 @@ describe('Zod v4', () => { expect(text).not.toContain('retry:'); }); - it('should close POST SSE stream when extra.closeSSEStream is called', async () => { + it('should close POST SSE stream when ctx.http?.closeSSE is called', async () => { const result = await createTestServer({ sessionIdGenerator: () => randomUUID(), eventStore: createEventStore(), @@ -1849,10 +1851,10 @@ describe('Zod v4', () => { toolResolve = resolve; }); - // Register a tool that closes its own SSE stream via extra callback - mcpServer.registerTool('close-stream-tool', { description: 'Closes its own stream' }, async extra => { + // Register a tool that closes its own SSE stream via ctx callback + mcpServer.registerTool('close-stream-tool', { description: 'Closes its own stream' }, async ctx => { // Close the SSE stream for this request - extra.closeSSEStream?.(); + ctx.http?.closeSSE?.(); streamCloseCalled = true; // Wait before returning so we can observe the stream closure @@ -1903,7 +1905,7 @@ describe('Zod v4', () => { toolResolve!(); }); - it('should provide closeSSEStream callback in extra when eventStore is configured', async () => { + it('should provide closeSSEStream callback in ctx when eventStore is configured', async () => { const result = await createTestServer({ sessionIdGenerator: () => randomUUID(), eventStore: createEventStore(), @@ -1917,9 +1919,9 @@ describe('Zod v4', () => { // Track whether closeSSEStream callback was provided let receivedCloseSSEStream: (() => void) | undefined; - // Register a tool that captures the extra.closeSSEStream callback - mcpServer.registerTool('test-callback-tool', { description: 'Test tool' }, async extra => { - receivedCloseSSEStream = extra.closeSSEStream; + // Register a tool that captures the ctx.http?.closeSSE callback + mcpServer.registerTool('test-callback-tool', { description: 'Test tool' }, async ctx => { + receivedCloseSSEStream = ctx.http?.closeSSE; return { content: [{ type: 'text', text: 'Done' }] }; }); @@ -1976,10 +1978,10 @@ describe('Zod v4', () => { let receivedCloseSSEStream: (() => void) | undefined; let receivedCloseStandaloneSSEStream: (() => void) | undefined; - // Register a tool that captures the extra.closeSSEStream callback - mcpServer.registerTool('test-old-version-tool', { description: 'Test tool' }, async extra => { - receivedCloseSSEStream = extra.closeSSEStream; - receivedCloseStandaloneSSEStream = extra.closeStandaloneSSEStream; + // Register a tool that captures the ctx.http?.closeSSE callback + mcpServer.registerTool('test-old-version-tool', { description: 'Test tool' }, async ctx => { + receivedCloseSSEStream = ctx.http?.closeSSE; + receivedCloseStandaloneSSEStream = ctx.http?.closeStandaloneSSE; return { content: [{ type: 'text', text: 'Done' }] }; }); @@ -2035,9 +2037,9 @@ describe('Zod v4', () => { // Track whether closeSSEStream callback was provided let receivedCloseSSEStream: (() => void) | undefined; - // Register a tool that captures the extra.closeSSEStream callback - mcpServer.registerTool('test-no-callback-tool', { description: 'Test tool' }, async extra => { - receivedCloseSSEStream = extra.closeSSEStream; + // Register a tool that captures the ctx.http?.closeSSE callback + mcpServer.registerTool('test-no-callback-tool', { description: 'Test tool' }, async ctx => { + receivedCloseSSEStream = ctx.http?.closeSSE; return { content: [{ type: 'text', text: 'Done' }] }; }); @@ -2078,7 +2080,7 @@ describe('Zod v4', () => { expect(receivedCloseSSEStream).toBeUndefined(); }); - it('should provide closeStandaloneSSEStream callback in extra when eventStore is configured', async () => { + it('should provide closeStandaloneSSEStream callback in ctx when eventStore is configured', async () => { const result = await createTestServer({ sessionIdGenerator: () => randomUUID(), eventStore: createEventStore(), @@ -2092,9 +2094,9 @@ describe('Zod v4', () => { // Track whether closeStandaloneSSEStream callback was provided let receivedCloseStandaloneSSEStream: (() => void) | undefined; - // Register a tool that captures the extra.closeStandaloneSSEStream callback - mcpServer.registerTool('test-standalone-callback-tool', { description: 'Test tool' }, async extra => { - receivedCloseStandaloneSSEStream = extra.closeStandaloneSSEStream; + // Register a tool that captures the ctx.http?.closeStandaloneSSE callback + mcpServer.registerTool('test-standalone-callback-tool', { description: 'Test tool' }, async ctx => { + receivedCloseStandaloneSSEStream = ctx.http?.closeStandaloneSSE; return { content: [{ type: 'text', text: 'Done' }] }; }); @@ -2136,7 +2138,7 @@ describe('Zod v4', () => { expect(typeof receivedCloseStandaloneSSEStream).toBe('function'); }); - it('should close standalone GET SSE stream when extra.closeStandaloneSSEStream is called', async () => { + it('should close standalone GET SSE stream when ctx.http?.closeStandaloneSSE is called', async () => { const result = await createTestServer({ sessionIdGenerator: () => randomUUID(), eventStore: createEventStore(), @@ -2147,9 +2149,9 @@ describe('Zod v4', () => { baseUrl = result.baseUrl; mcpServer = result.mcpServer; - // Register a tool that closes the standalone SSE stream via extra callback - mcpServer.registerTool('close-standalone-stream-tool', { description: 'Closes standalone stream' }, async extra => { - extra.closeStandaloneSSEStream?.(); + // Register a tool that closes the standalone SSE stream via ctx callback + mcpServer.registerTool('close-standalone-stream-tool', { description: 'Closes standalone stream' }, async ctx => { + ctx.http?.closeStandaloneSSE?.(); return { content: [{ type: 'text', text: 'Stream closed' }] }; }); @@ -2217,7 +2219,7 @@ describe('Zod v4', () => { expect(done).toBe(true); }); - it('should allow client to reconnect after standalone SSE stream is closed via extra.closeStandaloneSSEStream', async () => { + it('should allow client to reconnect after standalone SSE stream is closed via ctx.http?.closeStandaloneSSE', async () => { const result = await createTestServer({ sessionIdGenerator: () => randomUUID(), eventStore: createEventStore(), @@ -2229,8 +2231,8 @@ describe('Zod v4', () => { mcpServer = result.mcpServer; // Register a tool that closes the standalone SSE stream - mcpServer.registerTool('close-standalone-for-reconnect', { description: 'Closes standalone stream' }, async extra => { - extra.closeStandaloneSSEStream?.(); + mcpServer.registerTool('close-standalone-for-reconnect', { description: 'Closes standalone stream' }, async ctx => { + ctx.http?.closeStandaloneSSE?.(); return { content: [{ type: 'text', text: 'Stream closed' }] }; }); diff --git a/packages/server/src/experimental/tasks/interfaces.ts b/packages/server/src/experimental/tasks/interfaces.ts index 56dc050e..5e45ad86 100644 --- a/packages/server/src/experimental/tasks/interfaces.ts +++ b/packages/server/src/experimental/tasks/interfaces.ts @@ -6,11 +6,11 @@ import type { AnySchema, CallToolResult, - CreateTaskRequestHandlerExtra, CreateTaskResult, + CreateTaskServerContext, GetTaskResult, Result, - TaskRequestHandlerExtra + TaskServerContext } from '@modelcontextprotocol/core'; import type { BaseToolCallback } from '../../server/mcp.js'; @@ -25,7 +25,7 @@ import type { BaseToolCallback } from '../../server/mcp.js'; */ export type CreateTaskRequestHandler = BaseToolCallback< SendResultT, - CreateTaskRequestHandlerExtra, + CreateTaskServerContext, Args >; @@ -35,7 +35,7 @@ export type CreateTaskRequestHandler = BaseToolCallback< SendResultT, - TaskRequestHandlerExtra, + TaskServerContext, Args >; diff --git a/packages/server/src/experimental/tasks/mcpServer.ts b/packages/server/src/experimental/tasks/mcpServer.ts index 2e16e82d..ccf3962f 100644 --- a/packages/server/src/experimental/tasks/mcpServer.ts +++ b/packages/server/src/experimental/tasks/mcpServer.ts @@ -55,16 +55,16 @@ export class ExperimentalMcpServerTasks { * inputSchema: z.object({ input: z.string() }), * execution: { taskSupport: 'required' } * }, { - * createTask: async (args, extra) => { - * const task = await extra.taskStore.createTask({ ttl: 300000 }); + * createTask: async (args, ctx) => { + * const task = await ctx.task.store.createTask({ ttl: 300000 }); * startBackgroundWork(task.taskId, args); * return { task }; * }, - * getTask: async (args, extra) => { - * return extra.taskStore.getTask(extra.taskId); + * getTask: async (args, ctx) => { + * return ctx.task.store.getTask(ctx.task.id); * }, - * getTaskResult: async (args, extra) => { - * return extra.taskStore.getTaskResult(extra.taskId); + * getTaskResult: async (args, ctx) => { + * return ctx.task.store.getTaskResult(ctx.task.id); * } * }); * ``` diff --git a/packages/server/src/server/mcp.ts b/packages/server/src/server/mcp.ts index d0f627ce..b7df78ae 100644 --- a/packages/server/src/server/mcp.ts +++ b/packages/server/src/server/mcp.ts @@ -6,8 +6,8 @@ import type { CompleteRequestPrompt, CompleteRequestResourceTemplate, CompleteResult, - CreateTaskRequestHandlerExtra, CreateTaskResult, + CreateTaskServerContext, GetPromptResult, Implementation, ListPromptsResult, @@ -18,13 +18,11 @@ import type { PromptArgument, PromptReference, ReadResourceResult, - RequestHandlerExtra, Resource, ResourceTemplateReference, Result, SchemaOutput, - ServerNotification, - ServerRequest, + ServerContext, Tool, ToolAnnotations, ToolExecution, @@ -152,7 +150,7 @@ export class McpServer { }) ); - this.server.setRequestHandler('tools/call', async (request, extra): Promise => { + this.server.setRequestHandler('tools/call', async (request, ctx): Promise => { try { const tool = this._registeredTools[request.params.name]; if (!tool) { @@ -184,12 +182,12 @@ export class McpServer { // Handle taskSupport 'optional' without task augmentation - automatic polling if (taskSupport === 'optional' && !isTaskRequest && isTaskHandler) { - return await this.handleAutomaticTaskPolling(tool, request, extra); + return await this.handleAutomaticTaskPolling(tool, request, ctx); } // Normal execution path const args = await this.validateToolInput(tool, request.params.arguments, request.params.name); - const result = await this.executeToolHandler(tool, args, extra); + const result = await this.executeToolHandler(tool, args, ctx); // Return CreateTaskResult immediately for task requests if (isTaskRequest) { @@ -293,13 +291,9 @@ export class McpServer { /** * Executes a tool handler (either regular or task-based). */ - private async executeToolHandler( - tool: RegisteredTool, - args: unknown, - extra: RequestHandlerExtra - ): Promise { + private async executeToolHandler(tool: RegisteredTool, args: unknown, ctx: ServerContext): Promise { // Executor encapsulates handler invocation with proper types - return tool.executor(args, extra); + return tool.executor(args, ctx); } /** @@ -308,15 +302,15 @@ export class McpServer { private async handleAutomaticTaskPolling( tool: RegisteredTool, request: RequestT, - extra: RequestHandlerExtra + ctx: ServerContext ): Promise { - if (!extra.taskStore) { + if (!ctx.task?.store) { throw new Error('No task store provided for task-capable tool.'); } // Validate input and create task using the executor const args = await this.validateToolInput(tool, request.params.arguments, request.params.name); - const createTaskResult = (await tool.executor(args, extra)) as CreateTaskResult; + const createTaskResult = (await tool.executor(args, ctx)) as CreateTaskResult; // Poll until completion const taskId = createTaskResult.task.taskId; @@ -325,7 +319,7 @@ export class McpServer { while (task.status !== 'completed' && task.status !== 'failed' && task.status !== 'cancelled') { await new Promise(resolve => setTimeout(resolve, pollInterval)); - const updatedTask = await extra.taskStore.getTask(taskId); + const updatedTask = await ctx.task.store.getTask(taskId); if (!updatedTask) { throw new ProtocolError(ProtocolErrorCode.InternalError, `Task ${taskId} not found during polling`); } @@ -333,7 +327,7 @@ export class McpServer { } // Return the final result - return (await extra.taskStore.getTaskResult(taskId)) as CallToolResult; + return (await ctx.task.store.getTaskResult(taskId)) as CallToolResult; } private _completionHandlerInitialized = false; @@ -439,7 +433,7 @@ export class McpServer { } }); - this.server.setRequestHandler('resources/list', async (_request, extra) => { + this.server.setRequestHandler('resources/list', async (_request, ctx) => { const resources = Object.entries(this._registeredResources) .filter(([_, resource]) => resource.enabled) .map(([uri, resource]) => ({ @@ -454,7 +448,7 @@ export class McpServer { continue; } - const result = await template.resourceTemplate.listCallback(extra); + const result = await template.resourceTemplate.listCallback(ctx); for (const resource of result.resources) { templateResources.push({ ...template.metadata, @@ -477,7 +471,7 @@ export class McpServer { return { resourceTemplates }; }); - this.server.setRequestHandler('resources/read', async (request, extra) => { + this.server.setRequestHandler('resources/read', async (request, ctx) => { const uri = new URL(request.params.uri); // First check for exact resource match @@ -486,14 +480,14 @@ export class McpServer { if (!resource.enabled) { throw new ProtocolError(ProtocolErrorCode.InvalidParams, `Resource ${uri} disabled`); } - return resource.readCallback(uri, extra); + return resource.readCallback(uri, ctx); } // Then check templates for (const template of Object.values(this._registeredResourceTemplates)) { const variables = template.resourceTemplate.uriTemplate.match(uri.toString()); if (variables) { - return template.readCallback(uri, variables, extra); + return template.readCallback(uri, variables, ctx); } } @@ -535,7 +529,7 @@ export class McpServer { }) ); - this.server.setRequestHandler('prompts/get', async (request, extra): Promise => { + this.server.setRequestHandler('prompts/get', async (request, ctx): Promise => { const prompt = this._registeredPrompts[request.params.name]; if (!prompt) { throw new ProtocolError(ProtocolErrorCode.InvalidParams, `Prompt ${request.params.name} not found`); @@ -546,7 +540,7 @@ export class McpServer { } // Handler encapsulates parsing and callback invocation with proper types - return prompt.handler(request.params.arguments, extra); + return prompt.handler(request.params.arguments, ctx); }); this._promptHandlersInitialized = true; @@ -988,20 +982,16 @@ export class ResourceTemplate { export type BaseToolCallback< SendResultT extends Result, - Extra extends RequestHandlerExtra, + Ctx extends ServerContext, Args extends AnySchema | undefined > = Args extends AnySchema - ? (args: SchemaOutput, extra: Extra) => SendResultT | Promise - : (extra: Extra) => SendResultT | Promise; + ? (args: SchemaOutput, ctx: Ctx) => SendResultT | Promise + : (ctx: Ctx) => SendResultT | Promise; /** * Callback for a tool handler registered with Server.tool(). */ -export type ToolCallback = BaseToolCallback< - CallToolResult, - RequestHandlerExtra, - Args ->; +export type ToolCallback = BaseToolCallback; /** * Supertype that can handle both regular tools (simple callback) and task-based tools (task handler object). @@ -1011,10 +1001,7 @@ export type AnyToolHandler = Too /** * Internal executor type that encapsulates handler invocation with proper types. */ -type ToolExecutor = ( - args: unknown, - extra: RequestHandlerExtra -) => Promise; +type ToolExecutor = (args: unknown, ctx: ServerContext) => Promise; export type RegisteredTool = { title?: string; @@ -1046,37 +1033,35 @@ export type RegisteredTool = { /** * Creates an executor that invokes the handler with the appropriate arguments. - * When inputSchema is defined, the handler is called with (args, extra). - * When inputSchema is undefined, the handler is called with just (extra). + * When inputSchema is defined, the handler is called with (args, ctx). + * When inputSchema is undefined, the handler is called with just (ctx). */ function createToolExecutor(inputSchema: AnySchema | undefined, handler: AnyToolHandler): ToolExecutor { const isTaskHandler = 'createTask' in handler; if (isTaskHandler) { const taskHandler = handler as TaskHandlerInternal; - return async (args, extra) => { - if (!extra.taskStore) { + return async (args, ctx) => { + if (!ctx.task?.store) { throw new Error('No task store provided.'); } - const taskExtra: CreateTaskRequestHandlerExtra = { ...extra, taskStore: extra.taskStore }; + const taskCtx: CreateTaskServerContext = { ...ctx, task: { store: ctx.task.store, requestedTtl: ctx.task?.requestedTtl } }; if (inputSchema) { - return taskHandler.createTask(args, taskExtra); + return taskHandler.createTask(args, taskCtx); } - // When no inputSchema, call with just extra (the handler expects (extra) signature) - return (taskHandler.createTask as (extra: CreateTaskRequestHandlerExtra) => CreateTaskResult | Promise)( - taskExtra - ); + // When no inputSchema, call with just ctx (the handler expects (ctx) signature) + return (taskHandler.createTask as (ctx: CreateTaskServerContext) => CreateTaskResult | Promise)(taskCtx); }; } if (inputSchema) { const callback = handler as ToolCallbackInternal; - return async (args, extra) => callback(args, extra); + return async (args, ctx) => callback(args, ctx); } - // When no inputSchema, call with just extra (the handler expects (extra) signature) - const callback = handler as (extra: RequestHandlerExtra) => CallToolResult | Promise; - return async (_args, extra) => callback(extra); + // When no inputSchema, call with just ctx (the handler expects (ctx) signature) + const callback = handler as (ctx: ServerContext) => CallToolResult | Promise; + return async (_args, ctx) => callback(ctx); } const EMPTY_OBJECT_JSON_SCHEMA = { @@ -1092,17 +1077,12 @@ export type ResourceMetadata = Omit; /** * Callback to list all resources matching a given template. */ -export type ListResourcesCallback = ( - extra: RequestHandlerExtra -) => ListResourcesResult | Promise; +export type ListResourcesCallback = (ctx: ServerContext) => ListResourcesResult | Promise; /** * Callback to read a resource at a given URI. */ -export type ReadResourceCallback = ( - uri: URL, - extra: RequestHandlerExtra -) => ReadResourceResult | Promise; +export type ReadResourceCallback = (uri: URL, ctx: ServerContext) => ReadResourceResult | Promise; export type RegisteredResource = { name: string; @@ -1129,7 +1109,7 @@ export type RegisteredResource = { export type ReadResourceTemplateCallback = ( uri: URL, variables: Variables, - extra: RequestHandlerExtra + ctx: ServerContext ) => ReadResourceResult | Promise; export type RegisteredResourceTemplate = { @@ -1152,28 +1132,19 @@ export type RegisteredResourceTemplate = { }; export type PromptCallback = Args extends AnySchema - ? ( - args: SchemaOutput, - extra: RequestHandlerExtra - ) => GetPromptResult | Promise - : (extra: RequestHandlerExtra) => GetPromptResult | Promise; + ? (args: SchemaOutput, ctx: ServerContext) => GetPromptResult | Promise + : (ctx: ServerContext) => GetPromptResult | Promise; /** * Internal handler type that encapsulates parsing and callback invocation. * This allows type-safe handling without runtime type assertions. */ -type PromptHandler = ( - args: Record | undefined, - extra: RequestHandlerExtra -) => Promise; +type PromptHandler = (args: Record | undefined, ctx: ServerContext) => Promise; -type ToolCallbackInternal = ( - args: unknown, - extra: RequestHandlerExtra -) => CallToolResult | Promise; +type ToolCallbackInternal = (args: unknown, ctx: ServerContext) => CallToolResult | Promise; type TaskHandlerInternal = { - createTask: (args: unknown, extra: CreateTaskRequestHandlerExtra) => CreateTaskResult | Promise; + createTask: (args: unknown, ctx: CreateTaskServerContext) => CreateTaskResult | Promise; }; export type RegisteredPrompt = { @@ -1206,26 +1177,21 @@ function createPromptHandler( callback: PromptCallback ): PromptHandler { if (argsSchema) { - const typedCallback = callback as ( - args: SchemaOutput, - extra: RequestHandlerExtra - ) => GetPromptResult | Promise; + const typedCallback = callback as (args: SchemaOutput, ctx: ServerContext) => GetPromptResult | Promise; - return async (args, extra) => { + return async (args, ctx) => { const parseResult = await parseSchemaAsync(argsSchema, args); if (!parseResult.success) { const errorMessage = parseResult.error.issues.map((i: { message: string }) => i.message).join(', '); throw new ProtocolError(ProtocolErrorCode.InvalidParams, `Invalid arguments for prompt ${name}: ${errorMessage}`); } - return typedCallback(parseResult.data as SchemaOutput, extra); + return typedCallback(parseResult.data as SchemaOutput, ctx); }; } else { - const typedCallback = callback as ( - extra: RequestHandlerExtra - ) => GetPromptResult | Promise; + const typedCallback = callback as (ctx: ServerContext) => GetPromptResult | Promise; - return async (_args, extra) => { - return typedCallback(extra); + return async (_args, ctx) => { + return typedCallback(ctx); }; } } diff --git a/packages/server/src/server/server.ts b/packages/server/src/server/server.ts index 8bfe8bfe..ba9612ec 100644 --- a/packages/server/src/server/server.ts +++ b/packages/server/src/server/server.ts @@ -1,4 +1,5 @@ import type { + BaseContext, ClientCapabilities, CreateMessageRequest, CreateMessageRequestParamsBase, @@ -16,17 +17,18 @@ import type { ListRootsRequest, LoggingLevel, LoggingMessageNotification, + MessageExtraInfo, Notification, NotificationOptions, ProtocolOptions, Request, - RequestHandlerExtra, RequestMethod, RequestOptions, RequestTypeMap, ResourceUpdatedNotification, Result, ServerCapabilities, + ServerContext, ServerNotification, ServerRequest, ServerResult, @@ -110,7 +112,12 @@ export class Server< RequestT extends Request = Request, NotificationT extends Notification = Notification, ResultT extends Result = Result -> extends Protocol { +> extends Protocol< + ServerRequest | RequestT, + ServerNotification | NotificationT, + ServerResult | ResultT, + ServerContext +> { private _clientCapabilities?: ClientCapabilities; private _clientVersion?: Implementation; private _capabilities: ServerCapabilities; @@ -139,9 +146,9 @@ export class Server< this.setNotificationHandler('notifications/initialized', () => this.oninitialized?.()); if (this._capabilities.logging) { - this.setRequestHandler('logging/setLevel', async (request, extra) => { + this.setRequestHandler('logging/setLevel', async (request, ctx) => { const transportSessionId: string | undefined = - extra.sessionId || (extra.requestInfo?.headers.get('mcp-session-id') as string) || undefined; + ctx.sessionId || (ctx.http?.req?.headers.get('mcp-session-id') as string) || undefined; const { level } = request.params; const parseResult = parseSchema(LoggingLevelSchema, level); if (parseResult.success) { @@ -152,6 +159,32 @@ export class Server< } } + protected override buildContext( + ctx: BaseContext, + transportInfo?: MessageExtraInfo + ): ServerContext { + // Only create http when there's actual HTTP transport info or auth info + const hasHttpInfo = + ctx.http || transportInfo?.requestInfo || transportInfo?.closeSSEStream || transportInfo?.closeStandaloneSSEStream; + return { + ...ctx, + mcpReq: { + ...ctx.mcpReq, + log: (level, data, logger) => this.sendLoggingMessage({ level, data, logger }), + elicitInput: (params, options) => this.elicitInput(params, options), + requestSampling: (params, options) => this.createMessage(params, options) + }, + http: hasHttpInfo + ? { + ...ctx.http, + req: transportInfo?.requestInfo, + closeSSE: transportInfo?.closeSSEStream, + closeStandaloneSSE: transportInfo?.closeStandaloneSSEStream + } + : undefined + }; + } + /** * Access experimental features. * @@ -199,13 +232,13 @@ export class Server< method: M, handler: ( request: RequestTypeMap[M], - extra: RequestHandlerExtra + ctx: ServerContext ) => ServerResult | ResultT | Promise ): void { if (method === 'tools/call') { const wrappedHandler = async ( request: RequestTypeMap[M], - extra: RequestHandlerExtra + ctx: ServerContext ): Promise => { const validatedRequest = parseSchema(CallToolRequestSchema, request); if (!validatedRequest.success) { @@ -216,7 +249,7 @@ export class Server< const { params } = validatedRequest.data; - const result = await Promise.resolve(handler(request, extra)); + const result = await Promise.resolve(handler(request, ctx)); // When task creation is requested, validate and return CreateTaskResult if (params.task) { @@ -243,7 +276,7 @@ export class Server< }; // Install the wrapped handler - return super.setRequestHandler(method, wrappedHandler as unknown as typeof handler); + return super.setRequestHandler(method, wrappedHandler); } // Other handlers use default behavior diff --git a/test/conformance/src/everythingServer.ts b/test/conformance/src/everythingServer.ts index c73b3582..eae4e27e 100644 --- a/test/conformance/src/everythingServer.ts +++ b/test/conformance/src/everythingServer.ts @@ -202,8 +202,8 @@ function createMcpServer() { description: 'Tests tool that emits log messages during execution', inputSchema: z.object({}) }, - async (_args, extra): Promise => { - await extra.sendNotification({ + async (_args, ctx): Promise => { + await ctx.mcpReq.notify({ method: 'notifications/message', params: { level: 'info', @@ -212,7 +212,7 @@ function createMcpServer() { }); await new Promise(resolve => setTimeout(resolve, 50)); - await extra.sendNotification({ + await ctx.mcpReq.notify({ method: 'notifications/message', params: { level: 'info', @@ -221,7 +221,7 @@ function createMcpServer() { }); await new Promise(resolve => setTimeout(resolve, 50)); - await extra.sendNotification({ + await ctx.mcpReq.notify({ method: 'notifications/message', params: { level: 'info', @@ -241,10 +241,10 @@ function createMcpServer() { description: 'Tests tool that reports progress notifications', inputSchema: z.object({}) }, - async (_args, extra): Promise => { - const progressToken = extra._meta?.progressToken ?? 0; + async (_args, ctx): Promise => { + const progressToken = ctx.mcpReq._meta?.progressToken ?? 0; console.log('Progress token:', progressToken); - await extra.sendNotification({ + await ctx.mcpReq.notify({ method: 'notifications/progress', params: { progressToken, @@ -255,7 +255,7 @@ function createMcpServer() { }); await new Promise(resolve => setTimeout(resolve, 50)); - await extra.sendNotification({ + await ctx.mcpReq.notify({ method: 'notifications/progress', params: { progressToken, @@ -266,7 +266,7 @@ function createMcpServer() { }); await new Promise(resolve => setTimeout(resolve, 50)); - await extra.sendNotification({ + await ctx.mcpReq.notify({ method: 'notifications/progress', params: { progressToken, @@ -301,23 +301,23 @@ function createMcpServer() { 'Tests SSE stream disconnection and client reconnection (SEP-1699). Server will close the stream mid-call and send the result after client reconnects.', inputSchema: z.object({}) }, - async (_args, extra): Promise => { + async (_args, ctx): Promise => { const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms)); - console.log(`[${extra.sessionId}] Starting test_reconnection tool...`); + console.log(`[${ctx.sessionId}] Starting test_reconnection tool...`); // Get the transport for this session - const transport = extra.sessionId ? transports[extra.sessionId] : undefined; - if (transport && extra.requestId) { + const transport = ctx.sessionId ? transports[ctx.sessionId] : undefined; + if (transport && ctx.mcpReq.id) { // Close the SSE stream to trigger client reconnection - console.log(`[${extra.sessionId}] Closing SSE stream to trigger client polling...`); - transport.closeSSEStream(extra.requestId); + console.log(`[${ctx.sessionId}] Closing SSE stream to trigger client polling...`); + transport.closeSSEStream(ctx.mcpReq.id); } // Wait for client to reconnect (should respect retry field) await sleep(100); - console.log(`[${extra.sessionId}] test_reconnection tool complete`); + console.log(`[${ctx.sessionId}] test_reconnection tool complete`); return { content: [ @@ -339,10 +339,10 @@ function createMcpServer() { prompt: z.string().describe('The prompt to send to the LLM') }) }, - async (args: { prompt: string }, extra): Promise => { + async (args: { prompt: string }, ctx): Promise => { try { // Request sampling from client - const result = (await extra.sendRequest( + const result = (await ctx.mcpReq.send( { method: 'sampling/createMessage', params: { @@ -393,10 +393,10 @@ function createMcpServer() { message: z.string().describe('The message to show the user') }) }, - async (args: { message: string }, extra): Promise => { + async (args: { message: string }, ctx): Promise => { try { // Request user input from client - const result = await extra.sendRequest( + const result = await ctx.mcpReq.send( { method: 'elicitation/create', params: { @@ -445,10 +445,10 @@ function createMcpServer() { description: 'Tests elicitation with default values per SEP-1034', inputSchema: z.object({}) }, - async (_args, extra): Promise => { + async (_args, ctx): Promise => { try { // Request user input with default values for all primitive types - const result = await extra.sendRequest( + const result = await ctx.mcpReq.send( { method: 'elicitation/create', params: { @@ -519,10 +519,10 @@ function createMcpServer() { description: 'Tests elicitation with enum schema improvements per SEP-1330', inputSchema: z.object({}) }, - async (_args, extra): Promise => { + async (_args, ctx): Promise => { try { // Request user input with all 5 enum schema variants - const result = await extra.sendRequest( + const result = await ctx.mcpReq.send( { method: 'elicitation/create', params: { diff --git a/test/integration/test/client/client.test.ts b/test/integration/test/client/client.test.ts index e0802ff1..3c57fc54 100644 --- a/test/integration/test/client/client.test.ts +++ b/test/integration/test/client/client.test.ts @@ -1096,10 +1096,10 @@ test('should handle request timeout', async () => { ); // Set up server with a delayed response - server.setRequestHandler('resources/list', async (_request, extra) => { + server.setRequestHandler('resources/list', async (_request, ctx) => { const timer = new Promise(resolve => { const timeout = setTimeout(resolve, 100); - extra.signal.addEventListener('abort', () => clearTimeout(timeout)); + ctx.mcpReq.signal.addEventListener('abort', () => clearTimeout(timeout)); }); await timer; @@ -2269,27 +2269,27 @@ describe('Task-based execution', () => { inputSchema: z.object({}) }, { - async createTask(_args, extra) { - const task = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + async createTask(_args, ctx) { + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); const result = { content: [{ type: 'text', text: 'Tool executed successfully!' }] }; - await extra.taskStore.storeTaskResult(task.taskId, 'completed', result); + await ctx.task.store.storeTaskResult(task.taskId, 'completed', result); return { task }; }, - async getTask(_args, extra) { - const task = await extra.taskStore.getTask(extra.taskId); + async getTask(_args, ctx) { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { - throw new Error(`Task ${extra.taskId} not found`); + throw new Error(`Task ${ctx.task.id} not found`); } return task; }, - async getTaskResult(_args, extra) { - const result = await extra.taskStore.getTaskResult(extra.taskId); + async getTaskResult(_args, ctx) { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as { content: Array<{ type: 'text'; text: string }> }; } } @@ -2345,27 +2345,27 @@ describe('Task-based execution', () => { inputSchema: z.object({}) }, { - async createTask(_args, extra) { - const task = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + async createTask(_args, ctx) { + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); const result = { content: [{ type: 'text', text: 'Success!' }] }; - await extra.taskStore.storeTaskResult(task.taskId, 'completed', result); + await ctx.task.store.storeTaskResult(task.taskId, 'completed', result); return { task }; }, - async getTask(_args, extra) { - const task = await extra.taskStore.getTask(extra.taskId); + async getTask(_args, ctx) { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { - throw new Error(`Task ${extra.taskId} not found`); + throw new Error(`Task ${ctx.task.id} not found`); } return task; }, - async getTaskResult(_args, extra) { - const result = await extra.taskStore.getTaskResult(extra.taskId); + async getTaskResult(_args, ctx) { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as { content: Array<{ type: 'text'; text: string }> }; } } @@ -2422,27 +2422,27 @@ describe('Task-based execution', () => { inputSchema: z.object({}) }, { - async createTask(_args, extra) { - const task = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + async createTask(_args, ctx) { + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); const result = { content: [{ type: 'text', text: 'Result data!' }] }; - await extra.taskStore.storeTaskResult(task.taskId, 'completed', result); + await ctx.task.store.storeTaskResult(task.taskId, 'completed', result); return { task }; }, - async getTask(_args, extra) { - const task = await extra.taskStore.getTask(extra.taskId); + async getTask(_args, ctx) { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { - throw new Error(`Task ${extra.taskId} not found`); + throw new Error(`Task ${ctx.task.id} not found`); } return task; }, - async getTaskResult(_args, extra) { - const result = await extra.taskStore.getTaskResult(extra.taskId); + async getTaskResult(_args, ctx) { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as { content: Array<{ type: 'text'; text: string }> }; } } @@ -2506,27 +2506,27 @@ describe('Task-based execution', () => { inputSchema: z.object({}) }, { - async createTask(_args, extra) { - const task = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + async createTask(_args, ctx) { + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); const result = { content: [{ type: 'text', text: 'Success!' }] }; - await extra.taskStore.storeTaskResult(task.taskId, 'completed', result); + await ctx.task.store.storeTaskResult(task.taskId, 'completed', result); return { task }; }, - async getTask(_args, extra) { - const task = await extra.taskStore.getTask(extra.taskId); + async getTask(_args, ctx) { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { - throw new Error(`Task ${extra.taskId} not found`); + throw new Error(`Task ${ctx.task.id} not found`); } return task; }, - async getTaskResult(_args, extra) { - const result = await extra.taskStore.getTaskResult(extra.taskId); + async getTaskResult(_args, ctx) { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as { content: Array<{ type: 'text'; text: string }> }; } } @@ -2603,18 +2603,18 @@ describe('Task-based execution', () => { } ); - client.setRequestHandler('elicitation/create', async (request, extra) => { + client.setRequestHandler('elicitation/create', async (request, ctx) => { const result = { action: 'accept', content: { username: 'list-user' } }; // Check if task creation is requested - if (request.params.task && extra.taskStore) { - const task = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + if (request.params.task && ctx.task?.store) { + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); - await extra.taskStore.storeTaskResult(task.taskId, 'completed', result); + await ctx.task.store.storeTaskResult(task.taskId, 'completed', result); // Return CreateTaskResult when task creation is requested return { task }; } @@ -2696,18 +2696,18 @@ describe('Task-based execution', () => { } ); - client.setRequestHandler('elicitation/create', async (request, extra) => { + client.setRequestHandler('elicitation/create', async (request, ctx) => { const result = { action: 'accept', content: { username: 'list-user' } }; // Check if task creation is requested - if (request.params.task && extra.taskStore) { - const task = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + if (request.params.task && ctx.task?.store) { + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); - await extra.taskStore.storeTaskResult(task.taskId, 'completed', result); + await ctx.task.store.storeTaskResult(task.taskId, 'completed', result); // Return CreateTaskResult when task creation is requested return { task }; } @@ -2788,18 +2788,18 @@ describe('Task-based execution', () => { } ); - client.setRequestHandler('elicitation/create', async (request, extra) => { + client.setRequestHandler('elicitation/create', async (request, ctx) => { const result = { action: 'accept', content: { username: 'result-user' } }; // Check if task creation is requested - if (request.params.task && extra.taskStore) { - const task = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + if (request.params.task && ctx.task?.store) { + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); - await extra.taskStore.storeTaskResult(task.taskId, 'completed', result); + await ctx.task.store.storeTaskResult(task.taskId, 'completed', result); // Return CreateTaskResult when task creation is requested return { task }; } @@ -2879,18 +2879,18 @@ describe('Task-based execution', () => { } ); - client.setRequestHandler('elicitation/create', async (request, extra) => { + client.setRequestHandler('elicitation/create', async (request, ctx) => { const result = { action: 'accept', content: { username: 'list-user' } }; // Check if task creation is requested - if (request.params.task && extra.taskStore) { - const task = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + if (request.params.task && ctx.task?.store) { + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); - await extra.taskStore.storeTaskResult(task.taskId, 'completed', result); + await ctx.task.store.storeTaskResult(task.taskId, 'completed', result); // Return CreateTaskResult when task creation is requested return { task }; } @@ -2991,27 +2991,27 @@ describe('Task-based execution', () => { }) }, { - async createTask({ id }, extra) { - const task = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + async createTask({ id }, ctx) { + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); const result = { content: [{ type: 'text', text: `Result for ${id || 'unknown'}` }] }; - await extra.taskStore.storeTaskResult(task.taskId, 'completed', result); + await ctx.task.store.storeTaskResult(task.taskId, 'completed', result); return { task }; }, - async getTask(_args, extra) { - const task = await extra.taskStore.getTask(extra.taskId); + async getTask(_args, ctx) { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { - throw new Error(`Task ${extra.taskId} not found`); + throw new Error(`Task ${ctx.task.id} not found`); } return task; }, - async getTaskResult(_args, extra) { - const result = await extra.taskStore.getTaskResult(extra.taskId); + async getTaskResult(_args, ctx) { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as { content: Array<{ type: 'text'; text: string }> }; } } @@ -3259,27 +3259,27 @@ test('should respect server task capabilities', async () => { inputSchema: z.object({}) }, { - async createTask(_args, extra) { - const task = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + async createTask(_args, ctx) { + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); const result = { content: [{ type: 'text', text: 'Success!' }] }; - await extra.taskStore.storeTaskResult(task.taskId, 'completed', result); + await ctx.task.store.storeTaskResult(task.taskId, 'completed', result); return { task }; }, - async getTask(_args, extra) { - const task = await extra.taskStore.getTask(extra.taskId); + async getTask(_args, ctx) { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { - throw new Error(`Task ${extra.taskId} not found`); + throw new Error(`Task ${ctx.task.id} not found`); } return task; }, - async getTaskResult(_args, extra) { - const result = await extra.taskStore.getTaskResult(extra.taskId); + async getTaskResult(_args, ctx) { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as { content: Array<{ type: 'text'; text: string }> }; } } diff --git a/test/integration/test/server.test.ts b/test/integration/test/server.test.ts index 7245265b..436b4427 100644 --- a/test/integration/test/server.test.ts +++ b/test/integration/test/server.test.ts @@ -1386,12 +1386,12 @@ test('should handle request timeout', async () => { } ); - client.setRequestHandler('sampling/createMessage', async (_request, extra) => { + client.setRequestHandler('sampling/createMessage', async (_request, ctx) => { await new Promise((resolve, reject) => { const timeout = setTimeout(resolve, 100); - extra.signal.addEventListener('abort', () => { + ctx.mcpReq.signal.addEventListener('abort', () => { clearTimeout(timeout); - reject(extra.signal.reason); + reject(ctx.mcpReq.signal.reason); }); }); @@ -2146,9 +2146,9 @@ describe('Task-based execution', () => { inputSchema: z.object({}) }, { - async createTask(_args, extra) { - const task = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + async createTask(_args, ctx) { + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); // Simulate some async work @@ -2157,20 +2157,20 @@ describe('Task-based execution', () => { const result = { content: [{ type: 'text', text: 'Tool executed successfully!' }] }; - await extra.taskStore.storeTaskResult(task.taskId, 'completed', result); + await ctx.task.store.storeTaskResult(task.taskId, 'completed', result); })(); return { task }; }, - async getTask(_args, extra) { - const task = await extra.taskStore.getTask(extra.taskId); + async getTask(_args, ctx) { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { - throw new Error(`Task ${extra.taskId} not found`); + throw new Error(`Task ${ctx.task.id} not found`); } return task; }, - async getTaskResult(_args, extra) { - const result = await extra.taskStore.getTaskResult(extra.taskId); + async getTaskResult(_args, ctx) { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as { content: Array<{ type: 'text'; text: string }> }; } } @@ -2341,13 +2341,13 @@ describe('Task-based execution', () => { let capturedElicitRequest: z.infer | null = null; // Set up client elicitation handler - client.setRequestHandler('elicitation/create', async (request, extra) => { + client.setRequestHandler('elicitation/create', async (request, ctx) => { let taskId: string | undefined; // Check if task creation is requested - if (request.params.task && extra.taskStore) { - const createdTask = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + if (request.params.task && ctx.task?.store) { + const createdTask = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); taskId = createdTask.taskId; } @@ -2371,15 +2371,15 @@ describe('Task-based execution', () => { inputSchema: z.object({}) }, { - async createTask(_args, extra) { - const task = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + async createTask(_args, ctx) { + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); // Perform async work that makes a nested request (async () => { - // During tool execution, make a nested request to the client using extra.sendRequest - const elicitResult = await extra.sendRequest( + // During tool execution, make a nested request to the client using ctx.mcpReq.send + const elicitResult = await ctx.mcpReq.send( { method: 'elicitation/create', params: { @@ -2405,20 +2405,20 @@ describe('Task-based execution', () => { } ] }; - await extra.taskStore.storeTaskResult(task.taskId, 'completed', result); + await ctx.task.store.storeTaskResult(task.taskId, 'completed', result); })(); return { task }; }, - async getTask(_args, extra) { - const task = await extra.taskStore.getTask(extra.taskId); + async getTask(_args, ctx) { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { - throw new Error(`Task ${extra.taskId} not found`); + throw new Error(`Task ${ctx.task.id} not found`); } return task; }, - async getTaskResult(_args, extra) { - const result = await extra.taskStore.getTaskResult(extra.taskId); + async getTaskResult(_args, ctx) { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as { content: Array<{ type: 'text'; text: string }> }; } } @@ -2498,18 +2498,18 @@ describe('Task-based execution', () => { } ); - client.setRequestHandler('elicitation/create', async (request, extra) => { + client.setRequestHandler('elicitation/create', async (request, ctx) => { const result = { action: 'accept', content: { username: 'server-test-user', confirmed: true } }; // Check if task creation is requested - if (request.params.task && extra.taskStore) { - const task = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + if (request.params.task && ctx.task?.store) { + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); - await extra.taskStore.storeTaskResult(task.taskId, 'completed', result); + await ctx.task.store.storeTaskResult(task.taskId, 'completed', result); // Return CreateTaskResult when task creation is requested return { task }; } @@ -2579,18 +2579,18 @@ describe('Task-based execution', () => { } ); - client.setRequestHandler('elicitation/create', async (request, extra) => { + client.setRequestHandler('elicitation/create', async (request, ctx) => { const result = { action: 'accept', content: { username: 'list-user' } }; // Check if task creation is requested - if (request.params.task && extra.taskStore) { - const task = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + if (request.params.task && ctx.task?.store) { + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); - await extra.taskStore.storeTaskResult(task.taskId, 'completed', result); + await ctx.task.store.storeTaskResult(task.taskId, 'completed', result); // Return CreateTaskResult when task creation is requested return { task }; } @@ -2658,18 +2658,18 @@ describe('Task-based execution', () => { } ); - client.setRequestHandler('elicitation/create', async (request, extra) => { + client.setRequestHandler('elicitation/create', async (request, ctx) => { const result = { action: 'accept', content: { username: 'result-user', confirmed: true } }; // Check if task creation is requested - if (request.params.task && extra.taskStore) { - const task = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + if (request.params.task && ctx.task?.store) { + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); - await extra.taskStore.storeTaskResult(task.taskId, 'completed', result); + await ctx.task.store.storeTaskResult(task.taskId, 'completed', result); // Return CreateTaskResult when task creation is requested return { task }; } @@ -2739,18 +2739,18 @@ describe('Task-based execution', () => { } ); - client.setRequestHandler('elicitation/create', async (request, extra) => { + client.setRequestHandler('elicitation/create', async (request, ctx) => { const result = { action: 'accept', content: { username: 'list-user' } }; // Check if task creation is requested - if (request.params.task && extra.taskStore) { - const task = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + if (request.params.task && ctx.task?.store) { + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); - await extra.taskStore.storeTaskResult(task.taskId, 'completed', result); + await ctx.task.store.storeTaskResult(task.taskId, 'completed', result); // Return CreateTaskResult when task creation is requested return { task }; } @@ -2853,9 +2853,9 @@ describe('Task-based execution', () => { }) }, { - async createTask({ delay, taskNum }, extra) { - const task = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + async createTask({ delay, taskNum }, ctx) { + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); // Simulate async work @@ -2864,20 +2864,20 @@ describe('Task-based execution', () => { const result = { content: [{ type: 'text', text: `Completed task ${taskNum || 'unknown'}` }] }; - await extra.taskStore.storeTaskResult(task.taskId, 'completed', result); + await ctx.task.store.storeTaskResult(task.taskId, 'completed', result); })(); return { task }; }, - async getTask(_args, extra) { - const task = await extra.taskStore.getTask(extra.taskId); + async getTask(_args, ctx) { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { - throw new Error(`Task ${extra.taskId} not found`); + throw new Error(`Task ${ctx.task.id} not found`); } return task; }, - async getTaskResult(_args, extra) { - const result = await extra.taskStore.getTaskResult(extra.taskId); + async getTaskResult(_args, ctx) { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as { content: Array<{ type: 'text'; text: string }> }; } } @@ -3082,18 +3082,18 @@ test('should respect client task capabilities', async () => { } ); - client.setRequestHandler('elicitation/create', async (request, extra) => { + client.setRequestHandler('elicitation/create', async (request, ctx) => { const result = { action: 'accept', content: { username: 'test-user' } }; // Check if task creation is requested - if (request.params.task && extra.taskStore) { - const task = await extra.taskStore.createTask({ - ttl: extra.taskRequestedTtl + if (request.params.task && ctx.task?.store) { + const task = await ctx.task.store.createTask({ + ttl: ctx.task.requestedTtl }); - await extra.taskStore.storeTaskResult(task.taskId, 'completed', result); + await ctx.task.store.storeTaskResult(task.taskId, 'completed', result); // Return CreateTaskResult when task creation is requested return { task }; } diff --git a/test/integration/test/server/mcp.test.ts b/test/integration/test/server/mcp.test.ts index 2b4feb0e..58e67792 100644 --- a/test/integration/test/server/mcp.test.ts +++ b/test/integration/test/server/mcp.test.ts @@ -94,6 +94,159 @@ describe('Zod v4', () => { ]); }); + /*** + * Test: ctx.mcpReq.log convenience method + */ + test('should send logging messages via ctx.mcpReq.log() convenience method', async () => { + const mcpServer = new McpServer( + { + name: 'test server', + version: '1.0' + }, + { capabilities: { logging: {} } } + ); + + const notifications: Notification[] = []; + const client = new Client({ + name: 'test client', + version: '1.0' + }); + client.fallbackNotificationHandler = async notification => { + notifications.push(notification); + }; + + mcpServer.registerTool( + 'log-test', + { + description: 'A tool that logs via ctx.mcpReq.log()' + }, + async ctx => { + await ctx.mcpReq.log('info', 'Log from convenience method', 'test-logger'); + return { + content: [{ type: 'text' as const, text: 'done' }] + }; + } + ); + + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + await Promise.all([client.connect(clientTransport), mcpServer.server.connect(serverTransport)]); + + await client.callTool({ name: 'log-test' }); + + expect(notifications).toMatchObject([ + { + method: 'notifications/message', + params: { + level: 'info', + data: 'Log from convenience method', + logger: 'test-logger' + } + } + ]); + }); + + /*** + * Test: ctx.mcpReq.elicitInput convenience method + */ + test('should elicit input via ctx.mcpReq.elicitInput() convenience method', async () => { + const mcpServer = new McpServer({ + name: 'test server', + version: '1.0' + }); + + let elicitResult: unknown = null; + + mcpServer.registerTool( + 'elicit-test', + { + description: 'A tool that elicits input via ctx.mcpReq.elicitInput()' + }, + async ctx => { + elicitResult = await ctx.mcpReq.elicitInput({ + message: 'Please confirm', + requestedSchema: { + type: 'object', + properties: { + confirmed: { type: 'boolean' } + } + } + }); + return { + content: [{ type: 'text' as const, text: 'done' }] + }; + } + ); + + const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { elicitation: {} } }); + + client.setRequestHandler('elicitation/create', async () => ({ + action: 'accept', + content: { confirmed: true } + })); + + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + await Promise.all([client.connect(clientTransport), mcpServer.server.connect(serverTransport)]); + + await client.callTool({ name: 'elicit-test' }); + + expect(elicitResult).toMatchObject({ + action: 'accept', + content: { confirmed: true } + }); + }); + + /*** + * Test: ctx.mcpReq.requestSampling convenience method + */ + test('should request sampling via ctx.mcpReq.requestSampling() convenience method', async () => { + const mcpServer = new McpServer({ + name: 'test server', + version: '1.0' + }); + + let samplingResult: unknown = null; + + mcpServer.registerTool( + 'sampling-test', + { + description: 'A tool that requests sampling via ctx.mcpReq.requestSampling()' + }, + async ctx => { + samplingResult = await ctx.mcpReq.requestSampling({ + messages: [ + { + role: 'user', + content: { type: 'text', text: 'Hello' } + } + ], + maxTokens: 100 + }); + return { + content: [{ type: 'text' as const, text: 'done' }] + }; + } + ); + + const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: {} } }); + + client.setRequestHandler('sampling/createMessage', async () => ({ + model: 'test-model', + role: 'assistant' as const, + content: { type: 'text' as const, text: 'Hello back' } + })); + + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + await Promise.all([client.connect(clientTransport), mcpServer.server.connect(serverTransport)]); + + await client.callTool({ name: 'sampling-test' }); + + expect(samplingResult).toMatchObject({ + model: 'test-model', + role: 'assistant', + content: { type: 'text', text: 'Hello back' } + }); + }); + /*** * Test: Progress Notification with Message Field */ @@ -112,13 +265,13 @@ describe('Zod v4', () => { steps: z.number().min(1).describe('Number of steps to perform') }) }, - async ({ steps }, { sendNotification, _meta }) => { - const progressToken = _meta?.progressToken; + async ({ steps }, ctx) => { + const progressToken = ctx.mcpReq._meta?.progressToken; if (progressToken) { // Send progress notification for each step for (let i = 1; i <= steps; i++) { - await sendNotification({ + await ctx.mcpReq.notify({ method: 'notifications/progress', params: { progressToken, @@ -1320,7 +1473,7 @@ describe('Zod v4', () => { /*** * Test: Pass Session ID to Tool Callback */ - test('should pass sessionId to tool callback via RequestHandlerExtra', async () => { + test('should pass sessionId to tool callback via ServerContext', async () => { const mcpServer = new McpServer({ name: 'test server', version: '1.0' @@ -1332,8 +1485,8 @@ describe('Zod v4', () => { }); let receivedSessionId: string | undefined; - mcpServer.registerTool('test-tool', {}, async extra => { - receivedSessionId = extra.sessionId; + mcpServer.registerTool('test-tool', {}, async ctx => { + receivedSessionId = ctx.sessionId; return { content: [ { @@ -1366,7 +1519,7 @@ describe('Zod v4', () => { /*** * Test: Pass Request ID to Tool Callback */ - test('should pass requestId to tool callback via RequestHandlerExtra', async () => { + test('should pass requestId to tool callback via ServerContext', async () => { const mcpServer = new McpServer({ name: 'test server', version: '1.0' @@ -1378,13 +1531,13 @@ describe('Zod v4', () => { }); let receivedRequestId: string | number | undefined; - mcpServer.registerTool('request-id-test', {}, async extra => { - receivedRequestId = extra.requestId; + mcpServer.registerTool('request-id-test', {}, async ctx => { + receivedRequestId = ctx.mcpReq.id; return { content: [ { type: 'text', - text: `Received request ID: ${extra.requestId}` + text: `Received request ID: ${ctx.mcpReq.id}` } ] }; @@ -1440,8 +1593,8 @@ describe('Zod v4', () => { receivedLogMessage = notification.params.data as string; }); - mcpServer.registerTool('test-tool', {}, async extra => { - await extra.sendNotification({ + mcpServer.registerTool('test-tool', {}, async ctx => { + await ctx.mcpReq.notify({ method: 'notifications/message', params: { level: 'debug', data: loggingMessage } }); @@ -1787,17 +1940,17 @@ describe('Zod v4', () => { } }, { - createTask: async (_args, extra) => { - const task = await extra.taskStore.createTask({ ttl: 60_000 }); + createTask: async (_args, ctx) => { + const task = await ctx.task.store.createTask({ ttl: 60_000 }); return { task }; }, - getTask: async (_args, extra) => { - const task = await extra.taskStore.getTask(extra.taskId); + getTask: async (_args, ctx) => { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) throw new Error('Task not found'); return task; }, - getTaskResult: async (_args, extra) => { - return (await extra.taskStore.getTaskResult(extra.taskId)) as CallToolResult; + getTaskResult: async (_args, ctx) => { + return (await ctx.task.store.getTaskResult(ctx.task.id)) as CallToolResult; } } ); @@ -1856,17 +2009,17 @@ describe('Zod v4', () => { } }, { - createTask: async (_args, extra) => { - const task = await extra.taskStore.createTask({ ttl: 60_000 }); + createTask: async (_args, ctx) => { + const task = await ctx.task.store.createTask({ ttl: 60_000 }); return { task }; }, - getTask: async (_args, extra) => { - const task = await extra.taskStore.getTask(extra.taskId); + getTask: async (_args, ctx) => { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) throw new Error('Task not found'); return task; }, - getTaskResult: async (_args, extra) => { - return (await extra.taskStore.getTaskResult(extra.taskId)) as CallToolResult; + getTaskResult: async (_args, ctx) => { + return (await ctx.task.store.getTaskResult(ctx.task.id)) as CallToolResult; } } ); @@ -2841,7 +2994,7 @@ describe('Zod v4', () => { /*** * Test: Pass Request ID to Resource Callback */ - test('should pass requestId to resource callback via RequestHandlerExtra', async () => { + test('should pass requestId to resource callback via ServerContext', async () => { const mcpServer = new McpServer({ name: 'test server', version: '1.0' @@ -2853,13 +3006,13 @@ describe('Zod v4', () => { }); let receivedRequestId: string | number | undefined; - mcpServer.registerResource('request-id-test', 'test://resource', {}, async (_uri, extra) => { - receivedRequestId = extra.requestId; + mcpServer.registerResource('request-id-test', 'test://resource', {}, async (_uri, ctx) => { + receivedRequestId = ctx.mcpReq.id; return { contents: [ { uri: 'test://resource', - text: `Received request ID: ${extra.requestId}` + text: `Received request ID: ${ctx.mcpReq.id}` } ] }; @@ -3779,7 +3932,7 @@ describe('Zod v4', () => { /*** * Test: Pass Request ID to Prompt Callback */ - test('should pass requestId to prompt callback via RequestHandlerExtra', async () => { + test('should pass requestId to prompt callback via ServerContext', async () => { const mcpServer = new McpServer({ name: 'test server', version: '1.0' @@ -3791,15 +3944,15 @@ describe('Zod v4', () => { }); let receivedRequestId: string | number | undefined; - mcpServer.registerPrompt('request-id-test', {}, async extra => { - receivedRequestId = extra.requestId; + mcpServer.registerPrompt('request-id-test', {}, async ctx => { + receivedRequestId = ctx.mcpReq.id; return { messages: [ { role: 'assistant', content: { type: 'text', - text: `Received request ID: ${extra.requestId}` + text: `Received request ID: ${ctx.mcpReq.id}` } } ] @@ -6219,11 +6372,11 @@ describe('Zod v4', () => { } }, { - createTask: async ({ input }, extra) => { - const task = await extra.taskStore.createTask({ ttl: 60_000, pollInterval: 100 }); + createTask: async ({ input }, ctx) => { + const task = await ctx.task.store.createTask({ ttl: 60_000, pollInterval: 100 }); // Capture taskStore for use in setTimeout - const store = extra.taskStore; + const store = ctx.task.store; // Simulate async work setTimeout(async () => { @@ -6234,15 +6387,15 @@ describe('Zod v4', () => { return { task }; }, - getTask: async (_args, extra) => { - const task = await extra.taskStore.getTask(extra.taskId); + getTask: async (_args, ctx) => { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { throw new Error('Task not found'); } return task; }, - getTaskResult: async (_input, extra) => { - const result = await extra.taskStore.getTaskResult(extra.taskId); + getTaskResult: async (_input, ctx) => { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as CallToolResult; } } @@ -6324,11 +6477,11 @@ describe('Zod v4', () => { } }, { - createTask: async ({ value }, extra) => { - const task = await extra.taskStore.createTask({ ttl: 60_000, pollInterval: 100 }); + createTask: async ({ value }, ctx) => { + const task = await ctx.task.store.createTask({ ttl: 60_000, pollInterval: 100 }); // Capture taskStore for use in setTimeout - const store = extra.taskStore; + const store = ctx.task.store; // Simulate async work setTimeout(async () => { @@ -6340,15 +6493,15 @@ describe('Zod v4', () => { return { task }; }, - getTask: async (_args, extra) => { - const task = await extra.taskStore.getTask(extra.taskId); + getTask: async (_args, ctx) => { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { throw new Error('Task not found'); } return task; }, - getTaskResult: async (_value, extra) => { - const result = await extra.taskStore.getTaskResult(extra.taskId); + getTaskResult: async (_value, ctx) => { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as CallToolResult; } } @@ -6432,11 +6585,11 @@ describe('Zod v4', () => { } }, { - createTask: async ({ data }, extra) => { - const task = await extra.taskStore.createTask({ ttl: 60_000, pollInterval: 100 }); + createTask: async ({ data }, ctx) => { + const task = await ctx.task.store.createTask({ ttl: 60_000, pollInterval: 100 }); // Capture taskStore for use in setTimeout - const store = extra.taskStore; + const store = ctx.task.store; // Simulate async work setTimeout(async () => { @@ -6448,15 +6601,15 @@ describe('Zod v4', () => { return { task }; }, - getTask: async (_args, extra) => { - const task = await extra.taskStore.getTask(extra.taskId); + getTask: async (_args, ctx) => { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { throw new Error('Task not found'); } return task; }, - getTaskResult: async (_data, extra) => { - const result = await extra.taskStore.getTaskResult(extra.taskId); + getTaskResult: async (_data, ctx) => { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as CallToolResult; } } @@ -6549,11 +6702,11 @@ describe('Zod v4', () => { } }, { - createTask: async extra => { - const task = await extra.taskStore.createTask({ ttl: 60_000, pollInterval: 100 }); + createTask: async ctx => { + const task = await ctx.task.store.createTask({ ttl: 60_000, pollInterval: 100 }); // Capture taskStore for use in setTimeout - const store = extra.taskStore; + const store = ctx.task.store; // Simulate async failure setTimeout(async () => { @@ -6566,15 +6719,15 @@ describe('Zod v4', () => { return { task }; }, - getTask: async extra => { - const task = await extra.taskStore.getTask(extra.taskId); + getTask: async ctx => { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { throw new Error('Task not found'); } return task; }, - getTaskResult: async extra => { - const result = await extra.taskStore.getTaskResult(extra.taskId); + getTaskResult: async ctx => { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as CallToolResult; } } @@ -6655,11 +6808,11 @@ describe('Zod v4', () => { } }, { - createTask: async extra => { - const task = await extra.taskStore.createTask({ ttl: 60_000, pollInterval: 100 }); + createTask: async ctx => { + const task = await ctx.task.store.createTask({ ttl: 60_000, pollInterval: 100 }); // Capture taskStore for use in setTimeout - const store = extra.taskStore; + const store = ctx.task.store; // Simulate async cancellation setTimeout(async () => { @@ -6669,15 +6822,15 @@ describe('Zod v4', () => { return { task }; }, - getTask: async extra => { - const task = await extra.taskStore.getTask(extra.taskId); + getTask: async ctx => { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { throw new Error('Task not found'); } return task; }, - getTaskResult: async extra => { - const result = await extra.taskStore.getTaskResult(extra.taskId); + getTaskResult: async ctx => { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as CallToolResult; } } @@ -6742,19 +6895,19 @@ describe('Zod v4', () => { } }, { - createTask: async (_args, extra) => { - const task = await extra.taskStore.createTask({ ttl: 60_000, pollInterval: 100 }); + createTask: async (_args, ctx) => { + const task = await ctx.task.store.createTask({ ttl: 60_000, pollInterval: 100 }); return { task }; }, - getTask: async (_args, extra) => { - const task = await extra.taskStore.getTask(extra.taskId); + getTask: async (_args, ctx) => { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { throw new Error('Task not found'); } return task; }, - getTaskResult: async (_args, extra) => { - const result = await extra.taskStore.getTaskResult(extra.taskId); + getTaskResult: async (_args, ctx) => { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as CallToolResult; } } diff --git a/test/integration/test/taskLifecycle.test.ts b/test/integration/test/taskLifecycle.test.ts index b4d3ce6d..622cf748 100644 --- a/test/integration/test/taskLifecycle.test.ts +++ b/test/integration/test/taskLifecycle.test.ts @@ -63,8 +63,8 @@ describe('Task Lifecycle Integration Tests', () => { }) }, { - async createTask({ duration, shouldFail }, extra) { - const task = await extra.taskStore.createTask({ + async createTask({ duration, shouldFail }, ctx) { + const task = await ctx.task.store.createTask({ ttl: 60_000, pollInterval: 100 }); @@ -75,11 +75,11 @@ describe('Task Lifecycle Integration Tests', () => { try { await (shouldFail - ? extra.taskStore.storeTaskResult(task.taskId, 'failed', { + ? ctx.task.store.storeTaskResult(task.taskId, 'failed', { content: [{ type: 'text', text: 'Task failed as requested' }], isError: true }) - : extra.taskStore.storeTaskResult(task.taskId, 'completed', { + : ctx.task.store.storeTaskResult(task.taskId, 'completed', { content: [{ type: 'text', text: `Completed after ${duration}ms` }] })); } catch { @@ -89,15 +89,15 @@ describe('Task Lifecycle Integration Tests', () => { return { task }; }, - async getTask(_args, extra) { - const task = await extra.taskStore.getTask(extra.taskId); + async getTask(_args, ctx) { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { - throw new Error(`Task ${extra.taskId} not found`); + throw new Error(`Task ${ctx.task.id} not found`); } return task; }, - async getTaskResult(_args, extra) { - const result = await extra.taskStore.getTaskResult(extra.taskId); + async getTaskResult(_args, ctx) { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as { content: Array<{ type: 'text'; text: string }> }; } } @@ -114,8 +114,8 @@ describe('Task Lifecycle Integration Tests', () => { }) }, { - async createTask({ userName }, extra) { - const task = await extra.taskStore.createTask({ + async createTask({ userName }, ctx) { + const task = await ctx.task.store.createTask({ ttl: 60_000, pollInterval: 100 }); @@ -128,14 +128,14 @@ describe('Task Lifecycle Integration Tests', () => { if (userName) { // Complete immediately if userName was provided try { - await extra.taskStore.storeTaskResult(task.taskId, 'completed', { + await ctx.task.store.storeTaskResult(task.taskId, 'completed', { content: [{ type: 'text', text: `Hello, ${userName}!` }] }); } catch { // Task may have been cleaned up if test ended } } else { - const elicitationResult = await extra.sendRequest( + const elicitationResult = await ctx.mcpReq.send( { method: 'elicitation/create', params: { @@ -160,7 +160,7 @@ describe('Task Lifecycle Integration Tests', () => { ? elicitationResult.content.userName : 'Unknown'; try { - await extra.taskStore.storeTaskResult(task.taskId, 'completed', { + await ctx.task.store.storeTaskResult(task.taskId, 'completed', { content: [{ type: 'text', text: `Hello, ${name}!` }] }); } catch { @@ -171,15 +171,15 @@ describe('Task Lifecycle Integration Tests', () => { return { task }; }, - async getTask(_args, extra) { - const task = await extra.taskStore.getTask(extra.taskId); + async getTask(_args, ctx) { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { - throw new Error(`Task ${extra.taskId} not found`); + throw new Error(`Task ${ctx.task.id} not found`); } return task; }, - async getTaskResult(_args, extra) { - const result = await extra.taskStore.getTaskResult(extra.taskId); + async getTaskResult(_args, ctx) { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as { content: Array<{ type: 'text'; text: string }> }; } } @@ -418,8 +418,8 @@ describe('Task Lifecycle Integration Tests', () => { }) }, { - async createTask({ requestCount }, extra) { - const task = await extra.taskStore.createTask({ + async createTask({ requestCount }, ctx) { + const task = await ctx.task.store.createTask({ ttl: 60_000, pollInterval: 100 }); @@ -432,7 +432,7 @@ describe('Task Lifecycle Integration Tests', () => { // Send multiple elicitation requests for (let i = 0; i < requestCount; i++) { - const elicitationResult = await extra.sendRequest( + const elicitationResult = await ctx.mcpReq.send( { method: 'elicitation/create', params: { @@ -458,7 +458,7 @@ describe('Task Lifecycle Integration Tests', () => { // Complete with all responses try { - await extra.taskStore.storeTaskResult(task.taskId, 'completed', { + await ctx.task.store.storeTaskResult(task.taskId, 'completed', { content: [{ type: 'text', text: `Received responses: ${responses.join(', ')}` }] }); } catch { @@ -468,15 +468,15 @@ describe('Task Lifecycle Integration Tests', () => { return { task }; }, - async getTask(_args, extra) { - const task = await extra.taskStore.getTask(extra.taskId); + async getTask(_args, ctx) { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { - throw new Error(`Task ${extra.taskId} not found`); + throw new Error(`Task ${ctx.task.id} not found`); } return task; }, - async getTaskResult(_args, extra) { - const result = await extra.taskStore.getTaskResult(extra.taskId); + async getTaskResult(_args, ctx) { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as { content: Array<{ type: 'text'; text: string }> }; } } @@ -907,8 +907,8 @@ describe('Task Lifecycle Integration Tests', () => { }) }, { - async createTask({ messageCount }, extra) { - const task = await extra.taskStore.createTask({ + async createTask({ messageCount }, ctx) { + const task = await ctx.task.store.createTask({ ttl: 60_000, pollInterval: 100 }); @@ -921,8 +921,8 @@ describe('Task Lifecycle Integration Tests', () => { // Queue multiple elicitation requests for (let i = 0; i < messageCount; i++) { // Send request but don't await - let it queue - extra - .sendRequest( + ctx.mcpReq + .send( { method: 'elicitation/create', params: { @@ -957,15 +957,15 @@ describe('Task Lifecycle Integration Tests', () => { return { task }; }, - async getTask(_args, extra) { - const task = await extra.taskStore.getTask(extra.taskId); + async getTask(_args, ctx) { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { - throw new Error(`Task ${extra.taskId} not found`); + throw new Error(`Task ${ctx.task.id} not found`); } return task; }, - async getTaskResult(_args, extra) { - const result = await extra.taskStore.getTaskResult(extra.taskId); + async getTaskResult(_args, ctx) { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as { content: Array<{ type: 'text'; text: string }> }; } } @@ -1105,8 +1105,8 @@ describe('Task Lifecycle Integration Tests', () => { }) }, { - async createTask({ messageCount, delayBetweenMessages }, extra) { - const task = await extra.taskStore.createTask({ + async createTask({ messageCount, delayBetweenMessages }, ctx) { + const task = await ctx.task.store.createTask({ ttl: 60_000, pollInterval: 100 }); @@ -1121,7 +1121,7 @@ describe('Task Lifecycle Integration Tests', () => { // Send messages with delays between them for (let i = 0; i < messageCount; i++) { - const elicitationResult = await extra.sendRequest( + const elicitationResult = await ctx.mcpReq.send( { method: 'elicitation/create', params: { @@ -1152,7 +1152,7 @@ describe('Task Lifecycle Integration Tests', () => { // Complete with all responses try { - await extra.taskStore.storeTaskResult(task.taskId, 'completed', { + await ctx.task.store.storeTaskResult(task.taskId, 'completed', { content: [{ type: 'text', text: `Received all responses: ${responses.join(', ')}` }] }); } catch { @@ -1161,7 +1161,7 @@ describe('Task Lifecycle Integration Tests', () => { } catch (error) { // Handle errors try { - await extra.taskStore.storeTaskResult(task.taskId, 'failed', { + await ctx.task.store.storeTaskResult(task.taskId, 'failed', { content: [{ type: 'text', text: `Error: ${error}` }], isError: true }); @@ -1173,15 +1173,15 @@ describe('Task Lifecycle Integration Tests', () => { return { task }; }, - async getTask(_args, extra) { - const task = await extra.taskStore.getTask(extra.taskId); + async getTask(_args, ctx) { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { - throw new Error(`Task ${extra.taskId} not found`); + throw new Error(`Task ${ctx.task.id} not found`); } return task; }, - async getTaskResult(_args, extra) { - const result = await extra.taskStore.getTaskResult(extra.taskId); + async getTaskResult(_args, ctx) { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as { content: Array<{ type: 'text'; text: string }> }; } } @@ -1321,8 +1321,8 @@ describe('Task Lifecycle Integration Tests', () => { }) }, { - async createTask({ messageCount }, extra) { - const task = await extra.taskStore.createTask({ + async createTask({ messageCount }, ctx) { + const task = await ctx.task.store.createTask({ ttl: 60_000, pollInterval: 100 }); @@ -1335,8 +1335,8 @@ describe('Task Lifecycle Integration Tests', () => { for (let i = 0; i < messageCount; i++) { // Start the request but don't wait for response // The request gets queued when sendRequest is called - extra - .sendRequest( + ctx.mcpReq + .send( { method: 'elicitation/create', params: { @@ -1361,7 +1361,7 @@ describe('Task Lifecycle Integration Tests', () => { // Complete the task after all messages are queued try { - await extra.taskStore.storeTaskResult(task.taskId, 'completed', { + await ctx.task.store.storeTaskResult(task.taskId, 'completed', { content: [{ type: 'text', text: 'Task completed quickly' }] }); } catch { @@ -1370,7 +1370,7 @@ describe('Task Lifecycle Integration Tests', () => { } catch (error) { // Handle errors try { - await extra.taskStore.storeTaskResult(task.taskId, 'failed', { + await ctx.task.store.storeTaskResult(task.taskId, 'failed', { content: [{ type: 'text', text: `Error: ${error}` }], isError: true }); @@ -1382,15 +1382,15 @@ describe('Task Lifecycle Integration Tests', () => { return { task }; }, - async getTask(_args, extra) { - const task = await extra.taskStore.getTask(extra.taskId); + async getTask(_args, ctx) { + const task = await ctx.task.store.getTask(ctx.task.id); if (!task) { - throw new Error(`Task ${extra.taskId} not found`); + throw new Error(`Task ${ctx.task.id} not found`); } return task; }, - async getTaskResult(_args, extra) { - const result = await extra.taskStore.getTaskResult(extra.taskId); + async getTaskResult(_args, ctx) { + const result = await ctx.task.store.getTaskResult(ctx.task.id); return result as { content: Array<{ type: 'text'; text: string }> }; } } diff --git a/test/integration/test/taskResumability.test.ts b/test/integration/test/taskResumability.test.ts index f54fa4a9..45db390e 100644 --- a/test/integration/test/taskResumability.test.ts +++ b/test/integration/test/taskResumability.test.ts @@ -67,9 +67,9 @@ describe('Zod v4', () => { message: z.string().describe('Message to send').default('Test notification') }) }, - async ({ message }, { sendNotification }) => { + async ({ message }, ctx) => { // Send notification immediately - await sendNotification({ + await ctx.mcpReq.notify({ method: 'notifications/message', params: { level: 'info', @@ -93,10 +93,10 @@ describe('Zod v4', () => { interval: z.number().describe('Interval between notifications in ms').default(50) }) }, - async ({ count, interval }, { sendNotification }) => { + async ({ count, interval }, ctx) => { // Send notifications at specified intervals for (let i = 0; i < count; i++) { - await sendNotification({ + await ctx.mcpReq.notify({ method: 'notifications/message', params: { level: 'info',