[v2] Context API in callbacks

.gitignore

@@ -133,3 +133,4 @@ dist/
 
 # IDE
 .idea/
+.cursor/

examples/server/src/simpleStreamableHttp.ts

@@ -480,6 +480,7 @@ const getServer = () => {
         {
             async createTask({ duration }, { taskStore, taskRequestedTtl }) {
                 // Create the task
+                if (!taskStore) throw new Error('Task store not found');
                 const task = await taskStore.createTask({
                     ttl: taskRequestedTtl
                 });
@@ -503,10 +504,12 @@ const getServer = () => {
                 };
             },
             async getTask(_args, { taskId, taskStore }) {
-                return await taskStore.getTask(taskId);
+                if (!taskStore) throw new Error('Task store not found');
+                return await taskStore.getTask(taskId!);
             },
             async getTaskResult(_args, { taskId, taskStore }) {
-                const result = await taskStore.getTaskResult(taskId);
+                if (!taskStore) throw new Error('Task store not found');
+                const result = await taskStore.getTaskResult(taskId!);
                 return result as CallToolResult;
             }
         }

packages/core/src/experimental/tasks/interfaces.ts

@@ -3,7 +3,6 @@
  * WARNING: These APIs are experimental and may change without notice.
  */
 
-import type { RequestHandlerExtra, RequestTaskStore } from '../../shared/protocol.js';
 import type {
     JSONRPCErrorResponse,
     JSONRPCNotification,
@@ -12,8 +11,6 @@ import type {
     Request,
     RequestId,
     Result,
-    ServerNotification,
-    ServerRequest,
     Task,
     ToolExecution
 } from '../../types/types.js';
@@ -22,23 +19,6 @@ import type {
 // Task Handler Types (for registerToolTask)
 // ============================================================================
 
-/**
- * Extended handler extra with task store for task creation.
- * @experimental
- */
-export interface CreateTaskRequestHandlerExtra extends RequestHandlerExtra<ServerRequest, ServerNotification> {
-    taskStore: RequestTaskStore;
-}
-
-/**
- * Extended handler extra with task ID and store for task operations.
- * @experimental
- */
-export interface TaskRequestHandlerExtra extends RequestHandlerExtra<ServerRequest, ServerNotification> {
-    taskId: string;
-    taskStore: RequestTaskStore;
-}
-
 /**
  * Task-specific execution configuration.
  * taskSupport cannot be 'forbidden' for task-based tools.

packages/core/src/shared/protocol.ts

@@ -237,6 +237,8 @@ export interface RequestTaskStore {
 
 /**
  * Extra data given to request handlers.
+ *
+ * @deprecated Use {@link ContextInterface} from {@link Context} instead. Future major versions will remove this type.
  */
 export type RequestHandlerExtra<SendRequestT extends Request, SendNotificationT extends Notification> = {
     /**
@@ -718,43 +720,15 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
         const taskCreationParams = isTaskAugmentedRequestParams(request.params) ? request.params.task : undefined;
         const taskStore = this._taskStore ? this.requestTaskStore(request, capturedTransport?.sessionId) : undefined;
 
-        const fullExtra: RequestHandlerExtra<SendRequestT, SendNotificationT> = {
-            signal: abortController.signal,
-            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 };
-                }
-                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
-        };
+        const fullExtra: RequestHandlerExtra<SendRequestT, SendNotificationT> = this.createRequestExtra({
+            request,
+            taskStore,
+            relatedTaskId,
+            taskCreationParams,
+            abortController,
+            capturedTransport,
+            extra
+        });
 
         // Starting with Promise.resolve() puts any synchronous errors into the monad as well.
         Promise.resolve()
@@ -832,6 +806,60 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
             });
     }
 
+    /**
+     * Creates the RequestHandlerExtra passed to handlers. Subclasses may override to
+     * provide a richer context object as long as it satisfies the same structural contract.
+     */
+    protected createRequestExtra(args: {
+        request: JSONRPCRequest;
+        taskStore: TaskStore | undefined;
+        relatedTaskId: string | undefined;
+        taskCreationParams: TaskCreationParams | undefined;
+        abortController: AbortController;
+        capturedTransport: Transport | undefined;
+        extra?: MessageExtraInfo;
+    }): RequestHandlerExtra<SendRequestT, SendNotificationT> {
+        const { request, taskStore, relatedTaskId, taskCreationParams, abortController, capturedTransport, extra } = args;
+
+        return {
+            signal: abortController.signal,
+            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 };
+                }
+                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
+        } as RequestHandlerExtra<SendRequestT, SendNotificationT>;
+    }
+
     private _onprogress(notification: ProgressNotification): void {
         const { progressToken, ...params } = notification.params;
         const messageId = Number(progressToken);

packages/core/test/shared/protocol.test.ts

@@ -2095,7 +2095,7 @@ describe('Request Cancellation vs Task Cancellation', () => {
             let wasAborted = false;
             const TestRequestSchema = z.object({
                 method: z.literal('test/longRunning'),
-                params: z.optional(z.record(z.unknown()))
+                params: z.optional(z.record(z.string(), z.unknown()))
             });
             protocol.setRequestHandler(TestRequestSchema, async (_request, extra) => {
                 // Simulate a long-running operation
@@ -2310,7 +2310,7 @@ describe('Request Cancellation vs Task Cancellation', () => {
             let requestCompleted = false;
             const TestMethodSchema = z.object({
                 method: z.literal('test/method'),
-                params: z.optional(z.record(z.unknown()))
+                params: z.optional(z.record(z.string(), z.unknown()))
             });
             protocol.setRequestHandler(TestMethodSchema, async () => {
                 await new Promise(resolve => setTimeout(resolve, 50));
@@ -3690,7 +3690,7 @@ describe('Message Interception', () => {
                 method: z.literal('test/taskRequest'),
                 params: z
                     .object({
-                        _meta: z.optional(z.record(z.unknown()))
+                        _meta: z.optional(z.record(z.string(), z.unknown()))
                     })
                     .passthrough()
             });
@@ -3737,7 +3737,7 @@ describe('Message Interception', () => {
                 method: z.literal('test/taskRequestError'),
                 params: z
                     .object({
-                        _meta: z.optional(z.record(z.unknown()))
+                        _meta: z.optional(z.record(z.string(), z.unknown()))
                     })
                     .passthrough()
             });
@@ -3817,7 +3817,7 @@ describe('Message Interception', () => {
             // Set up a request handler
             const TestRequestSchema = z.object({
                 method: z.literal('test/normalRequest'),
-                params: z.optional(z.record(z.unknown()))
+                params: z.optional(z.record(z.string(), z.unknown()))
             });
 
             protocol.setRequestHandler(TestRequestSchema, async () => {

packages/server/src/experimental/tasks/interfaces.ts

@@ -6,14 +6,15 @@
 import type {
     AnySchema,
     CallToolResult,
-    CreateTaskRequestHandlerExtra,
     CreateTaskResult,
     GetTaskResult,
     Result,
-    TaskRequestHandlerExtra,
+    ServerNotification,
+    ServerRequest,
     ZodRawShapeCompat
 } from '@modelcontextprotocol/core';
 
+import type { ContextInterface } from '../../server/context.js';
 import type { BaseToolCallback } from '../../server/mcp.js';
 
 // ============================================================================
@@ -27,7 +28,7 @@ import type { BaseToolCallback } from '../../server/mcp.js';
 export type CreateTaskRequestHandler<
     SendResultT extends Result,
     Args extends undefined | ZodRawShapeCompat | AnySchema = undefined
-> = BaseToolCallback<SendResultT, CreateTaskRequestHandlerExtra, Args>;
+> = BaseToolCallback<SendResultT, ContextInterface<ServerRequest, ServerNotification>, Args>;
 
 /**
  * Handler for task operations (get, getResult).
@@ -36,7 +37,7 @@ export type CreateTaskRequestHandler<
 export type TaskRequestHandler<
     SendResultT extends Result,
     Args extends undefined | ZodRawShapeCompat | AnySchema = undefined
-> = BaseToolCallback<SendResultT, TaskRequestHandlerExtra, Args>;
+> = BaseToolCallback<SendResultT, ContextInterface<ServerRequest, ServerNotification>, Args>;
 
 /**
  * Interface for task-based tool handlers.

packages/server/src/index.ts

@@ -1,4 +1,5 @@
 export * from './server/completable.js';
+export * from './server/context.js';
 export * from './server/express.js';
 export * from './server/mcp.js';
 export * from './server/server.js';