context API - backwards compatible introduction

Author: KKonstantinov

src/server/context.ts

@@ -4,6 +4,7 @@ import {
     ElicitRequest,
     ElicitResult,
     ElicitResultSchema,
+    JSONRPCRequest,
     LoggingMessageNotification,
     Notification,
     Request,
@@ -12,16 +13,15 @@ import {
     RequestMeta,
     Result,
     ServerNotification,
-    ServerRequest,
-    ServerResult
+    ServerRequest
 } from '../types.js';
 import { RequestHandlerExtra, RequestOptions, RequestTaskStore } from '../shared/protocol.js';
 import { Server } from './index.js';
-import { RequestContext } from './requestContext.js';
 import { AuthInfo } from './auth/types.js';
 import { AnySchema, SchemaOutput } from './zod-compat.js';
 
-export interface ContextInterface<RequestT extends Request = Request, NotificationT extends Notification = Notification> extends RequestHandlerExtra<ServerRequest | RequestT, NotificationT | ServerNotification> {
+export interface ContextInterface<RequestT extends Request = Request, NotificationT extends Notification = Notification>
+    extends RequestHandlerExtra<ServerRequest | RequestT, NotificationT | ServerNotification> {
     elicit(params: ElicitRequest['params'], options?: RequestOptions): Promise<ElicitResult>;
     requestSampling: (params: CreateMessageRequest['params'], options?: RequestOptions) => Promise<CreateMessageResult>;
     log(params: LoggingMessageNotification['params'], sessionId?: string): Promise<void>;
@@ -35,32 +35,61 @@ export interface ContextInterface<RequestT extends Request = Request, Notificati
  *
  * Implements the RequestHandlerExtra interface for backwards compatibility.
  */
-export class Context<
-    RequestT extends Request = Request,
-    NotificationT extends Notification = Notification,
-    ResultT extends Result = Result
-> implements ContextInterface<RequestT, NotificationT>
+export class Context<RequestT extends Request = Request, NotificationT extends Notification = Notification, ResultT extends Result = Result>
+    implements ContextInterface<RequestT, NotificationT>
 {
     private readonly server: Server<RequestT, NotificationT, ResultT>;
 
     /**
      * The request context.
      * A type-safe context that is passed to request handlers.
      */
-    public readonly requestCtx: RequestContext< 
-        RequestT | ServerRequest,
-        NotificationT | ServerNotification,
-        ResultT | ServerResult
-    >;
+    private readonly requestCtx: RequestHandlerExtra<ServerRequest | RequestT, ServerNotification | NotificationT>;
+
+    /**
+     * The MCP context - Contains information about the current MCP request and session.
+     */
+    public readonly mcpContext: {
+        /**
+         * The JSON-RPC ID of the request being handled.
+         * This can be useful for tracking or logging purposes.
+         */
+        requestId: RequestId;
+        /**
+         * The method of the request.
+         */
+        method: string;
+        /**
+         * The metadata of the request.
+         */
+        _meta?: RequestMeta;
+        /**
+         * The session ID of the request.
+         */
+        sessionId?: string;
+    };
 
     constructor(args: {
         server: Server<RequestT, NotificationT, ResultT>;
-        requestCtx: RequestContext<RequestT | ServerRequest, NotificationT | ServerNotification, ResultT | ServerResult>;
+        request: JSONRPCRequest;
+        requestCtx: RequestHandlerExtra<ServerRequest | RequestT, ServerNotification | NotificationT>;
     }) {
         this.server = args.server;
         this.requestCtx = args.requestCtx;
+        this.mcpContext = {
+            requestId: args.requestCtx.requestId,
+            method: args.request.method,
+            _meta: args.requestCtx._meta,
+            sessionId: args.requestCtx.sessionId
+        };
     }
 
+    /**
+     * The JSON-RPC ID of the request being handled.
+     * This can be useful for tracking or logging purposes.
+     *
+     * @deprecated Use {@link mcpContext.requestId} instead.
+     */
     public get requestId(): RequestId {
         return this.requestCtx.requestId;
     }
@@ -77,12 +106,18 @@ export class Context<
         return this.requestCtx.requestInfo;
     }
 
+    /**
+     * @deprecated Use {@link mcpContext._meta} instead.
+     */
     public get _meta(): RequestMeta | undefined {
         return this.requestCtx._meta;
     }
 
+    /**
+     * @deprecated Use {@link mcpContext.sessionId} instead.
+     */
     public get sessionId(): string | undefined {
-        return this.requestCtx.sessionId;
+        return this.mcpContext.sessionId;
     }
 
     public get taskId(): string | undefined {
@@ -97,12 +132,12 @@ export class Context<
         return this.requestCtx.taskRequestedTtl ?? undefined;
     }
 
-    public closeSSEStream = (): void => {
-        return this.requestCtx?.closeSSEStream();
+    public get closeSSEStream(): (() => void) | undefined {
+        return this.requestCtx.closeSSEStream;
     }
 
-    public closeStandaloneSSEStream = (): void => {
-        return this.requestCtx?.closeStandaloneSSEStream();
+    public get closeStandaloneSSEStream(): (() => void) | undefined {
+        return this.requestCtx.closeStandaloneSSEStream;
     }
 
     /**
@@ -111,7 +146,7 @@ export class Context<
      * This is used by certain transports to correctly associate related messages.
      */
     public sendNotification = (notification: NotificationT | ServerNotification): Promise<void> => {
-        return this.requestCtx.sendNotification(notification);
+        return this.server.notification(notification);
     };
 
     /**
@@ -124,7 +159,7 @@ export class Context<
         resultSchema: U,
         options?: RequestOptions
     ): Promise<SchemaOutput<U>> => {
-        return this.requestCtx.sendRequest(request, resultSchema, { ...options, relatedRequestId: this.requestId });
+        return this.server.request(request, resultSchema, { ...options, relatedRequestId: this.requestId });
     };
 
     /**
@@ -219,4 +254,4 @@ export class Context<
             sessionId
         );
     }
-}
+}

src/server/index.ts

@@ -62,7 +62,6 @@ import { assertToolsCallTaskCapability, assertClientRequestTaskCapability } from
 import { Context } from './context.js';
 import { TaskStore } from '../experimental/index.js';
 import { Transport } from '../shared/transport.js';
-import { RequestContext } from './requestContext.js';
 
 export type ServerOptions = ProtocolOptions & {
     /**
@@ -232,16 +231,24 @@ export class Server<
         // Wrap the handler to ensure the extra is a Context and return a decorated handler that can be passed to the base implementation
 
         // Factory function to create a handler decorator that ensures the extra is a Context and returns a decorated handler that can be passed to the base implementation
-        const handlerDecoratorFactory = (innerHandler: (request: SchemaOutput<T>, extra: Context<RequestT, NotificationT, ResultT>) => ServerResult | ResultT | Promise<ServerResult | ResultT>) => {
-            const decoratedHandler = (request: SchemaOutput<T>, extra: RequestHandlerExtra<ServerRequest | RequestT, ServerNotification | NotificationT>) => {
+        const handlerDecoratorFactory = (
+            innerHandler: (
+                request: SchemaOutput<T>,
+                extra: Context<RequestT, NotificationT, ResultT>
+            ) => ServerResult | ResultT | Promise<ServerResult | ResultT>
+        ) => {
+            const decoratedHandler = (
+                request: SchemaOutput<T>,
+                extra: RequestHandlerExtra<ServerRequest | RequestT, ServerNotification | NotificationT>
+            ) => {
                 if (!this.isContextExtra(extra)) {
                     throw new Error('Internal error: Expected Context for request handler extra');
                 }
                 return innerHandler(request, extra);
-            }
+            };
 
             return decoratedHandler;
-        }
+        };
 
         const shape = getObjectShape(requestSchema);
         const methodSchema = shape?.method;
@@ -496,45 +503,23 @@ export class Server<
         return this._capabilities;
     }
 
-    protected createRequestExtra(
-        args: {
-            request: JSONRPCRequest,
-            taskStore: TaskStore | undefined,
-            relatedTaskId: string | undefined,
-            taskCreationParams: TaskCreationParams | undefined,
-            abortController: AbortController,
-            capturedTransport: Transport | undefined,
-            extra?: MessageExtraInfo
-        }
-    ): RequestHandlerExtra<ServerRequest | RequestT, ServerNotification | NotificationT> {
-        const base = super.createRequestExtra(args) as RequestHandlerExtra<
-            ServerRequest | RequestT,
-            ServerNotification | NotificationT
-        >;
-
-        // Wrap base in Context to add server utilities while preserving shape
-        const requestCtx = new RequestContext<
-            ServerRequest | RequestT,
-            ServerNotification | NotificationT,
-            ServerResult | ResultT
-        >({
-            signal: base.signal,
-            authInfo: base.authInfo,
-            requestInfo: base.requestInfo,
-            requestId: base.requestId,
-            _meta: base._meta,
-            sessionId: base.sessionId,
-            protocol: this,
-            closeSSEStream: base.closeSSEStream ?? undefined,
-            closeStandaloneSSEStream: base.closeStandaloneSSEStream ?? undefined
-        });
-
-        const ctx = new Context<RequestT, NotificationT, ResultT>({
+    protected createRequestExtra(args: {
+        request: JSONRPCRequest;
+        taskStore: TaskStore | undefined;
+        relatedTaskId: string | undefined;
+        taskCreationParams: TaskCreationParams | undefined;
+        abortController: AbortController;
+        capturedTransport: Transport | undefined;
+        extra?: MessageExtraInfo;
+    }): RequestHandlerExtra<ServerRequest | RequestT, ServerNotification | NotificationT> {
+        const base = super.createRequestExtra(args) as RequestHandlerExtra<ServerRequest | RequestT, ServerNotification | NotificationT>;
+
+        // Expose a Context instance to handlers, which implements RequestHandlerExtra
+        return new Context<RequestT, NotificationT, ResultT>({
             server: this,
-            requestCtx
+            request: args.request,
+            requestCtx: base
         });
-
-        return ctx;
     }
 
     async ping() {

src/server/mcp.ts

@@ -63,6 +63,7 @@ import { validateAndWarnToolName } from '../shared/toolNameValidation.js';
 import { ExperimentalMcpServerTasks } from '../experimental/tasks/mcp-server.js';
 import type { ToolTaskHandler } from '../experimental/tasks/interfaces.js';
 import { ZodOptional } from 'zod';
+import { ContextInterface } from './context.js';
 
 /**
  * High-level MCP server that provides a simpler API for working with resources, tools, and prompts.
@@ -324,7 +325,7 @@ export class McpServer {
     private async executeToolHandler(
         tool: RegisteredTool,
         args: unknown,
-        extra: RequestHandlerExtra<ServerRequest, ServerNotification>
+        extra: ContextInterface<ServerRequest, ServerNotification>
     ): Promise<CallToolResult | CreateTaskResult> {
         const handler = tool.handler as AnyToolHandler<ZodRawShapeCompat | undefined>;
         const isTaskHandler = 'createTask' in handler;
@@ -1270,7 +1271,7 @@ export class ResourceTemplate {
 
 export type BaseToolCallback<
     SendResultT extends Result,
-    Extra extends RequestHandlerExtra<ServerRequest, ServerNotification>,
+    Extra extends ContextInterface<ServerRequest, ServerNotification>,
     Args extends undefined | ZodRawShapeCompat | AnySchema
 > = Args extends ZodRawShapeCompat
     ? (args: ShapeOutput<Args>, extra: Extra) => SendResultT | Promise<SendResultT>
@@ -1290,7 +1291,7 @@ export type BaseToolCallback<
  */
 export type ToolCallback<Args extends undefined | ZodRawShapeCompat | AnySchema = undefined> = BaseToolCallback<
     CallToolResult,
-    RequestHandlerExtra<ServerRequest, ServerNotification>,
+    ContextInterface<ServerRequest, ServerNotification>,
     Args
 >;
 
@@ -1409,15 +1410,15 @@ export type ResourceMetadata = Omit<Resource, 'uri' | 'name'>;
  * Callback to list all resources matching a given template.
  */
 export type ListResourcesCallback = (
-    extra: RequestHandlerExtra<ServerRequest, ServerNotification>
+    extra: ContextInterface<ServerRequest, ServerNotification>
 ) => ListResourcesResult | Promise<ListResourcesResult>;
 
 /**
  * Callback to read a resource at a given URI.
  */
 export type ReadResourceCallback = (
     uri: URL,
-    extra: RequestHandlerExtra<ServerRequest, ServerNotification>
+    extra: ContextInterface<ServerRequest, ServerNotification>
 ) => ReadResourceResult | Promise<ReadResourceResult>;
 
 export type RegisteredResource = {
@@ -1445,7 +1446,7 @@ export type RegisteredResource = {
 export type ReadResourceTemplateCallback = (
     uri: URL,
     variables: Variables,
-    extra: RequestHandlerExtra<ServerRequest, ServerNotification>
+    extra: ContextInterface<ServerRequest, ServerNotification>
 ) => ReadResourceResult | Promise<ReadResourceResult>;
 
 export type RegisteredResourceTemplate = {
@@ -1470,8 +1471,8 @@ export type RegisteredResourceTemplate = {
 type PromptArgsRawShape = ZodRawShapeCompat;
 
 export type PromptCallback<Args extends undefined | PromptArgsRawShape = undefined> = Args extends PromptArgsRawShape
-    ? (args: ShapeOutput<Args>, extra: RequestHandlerExtra<ServerRequest, ServerNotification>) => GetPromptResult | Promise<GetPromptResult>
-    : (extra: RequestHandlerExtra<ServerRequest, ServerNotification>) => GetPromptResult | Promise<GetPromptResult>;
+    ? (args: ShapeOutput<Args>, extra: ContextInterface<ServerRequest, ServerNotification>) => GetPromptResult | Promise<GetPromptResult>
+    : (extra: ContextInterface<ServerRequest, ServerNotification>) => GetPromptResult | Promise<GetPromptResult>;
 
 export type RegisteredPrompt = {
     title?: string;