Context API in callbacks

src/server/context.ts

@@ -0,0 +1,257 @@
+import {
+    CreateMessageRequest,
+    CreateMessageResult,
+    ElicitRequest,
+    ElicitResult,
+    ElicitResultSchema,
+    JSONRPCRequest,
+    LoggingMessageNotification,
+    Notification,
+    Request,
+    RequestId,
+    RequestInfo,
+    RequestMeta,
+    Result,
+    ServerNotification,
+    ServerRequest
+} from '../types.js';
+import { RequestHandlerExtra, RequestOptions, RequestTaskStore } from '../shared/protocol.js';
+import { Server } from './index.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> {
+    elicit(params: ElicitRequest['params'], options?: RequestOptions): Promise<ElicitResult>;
+    requestSampling: (params: CreateMessageRequest['params'], options?: RequestOptions) => Promise<CreateMessageResult>;
+    log(params: LoggingMessageNotification['params'], sessionId?: string): Promise<void>;
+    debug(message: string, extraLogData?: Record<string, unknown>, sessionId?: string): Promise<void>;
+    info(message: string, extraLogData?: Record<string, unknown>, sessionId?: string): Promise<void>;
+    warning(message: string, extraLogData?: Record<string, unknown>, sessionId?: string): Promise<void>;
+    error(message: string, extraLogData?: Record<string, unknown>, sessionId?: string): Promise<void>;
+}
+/**
+ * A context object that is passed to request handlers.
+ *
+ * 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>
+{
+    private readonly server: Server<RequestT, NotificationT, ResultT>;
+
+    /**
+     * The request context.
+     * A type-safe context that is passed to request handlers.
+     */
+    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>;
+        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;
+    }
+
+    public get signal(): AbortSignal {
+        return this.requestCtx.signal;
+    }
+
+    public get authInfo(): AuthInfo | undefined {
+        return this.requestCtx.authInfo;
+    }
+
+    public get requestInfo(): RequestInfo | undefined {
+        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.mcpContext.sessionId;
+    }
+
+    public get taskId(): string | undefined {
+        return this.requestCtx.taskId;
+    }
+
+    public get taskStore(): RequestTaskStore | undefined {
+        return this.requestCtx.taskStore;
+    }
+
+    public get taskRequestedTtl(): number | undefined {
+        return this.requestCtx.taskRequestedTtl ?? undefined;
+    }
+
+    public get closeSSEStream(): (() => void) | undefined {
+        return this.requestCtx.closeSSEStream;
+    }
+
+    public get closeStandaloneSSEStream(): (() => void) | undefined {
+        return this.requestCtx.closeStandaloneSSEStream;
+    }
+
+    /**
+     * Sends a notification that relates to the current request being handled.
+     *
+     * This is used by certain transports to correctly associate related messages.
+     */
+    public sendNotification = (notification: NotificationT | ServerNotification): Promise<void> => {
+        return this.requestCtx.sendNotification(notification);
+    };
+
+    /**
+     * Sends a request that relates to the current request being handled.
+     *
+     * This is used by certain transports to correctly associate related messages.
+     */
+    public sendRequest = <U extends AnySchema>(
+        request: RequestT | ServerRequest,
+        resultSchema: U,
+        options?: RequestOptions
+    ): Promise<SchemaOutput<U>> => {
+        return this.requestCtx.sendRequest(request, resultSchema, { ...options, relatedRequestId: this.requestId });
+    };
+
+    /**
+     * Sends a request to sample an LLM via the client.
+     */
+    public requestSampling(params: CreateMessageRequest['params'], options?: RequestOptions) {
+        return this.server.createMessage(params, options);
+    }
+
+    /**
+     * Sends an elicitation request to the client.
+     */
+    public async elicit(params: ElicitRequest['params'], options?: RequestOptions): Promise<ElicitResult> {
+        const request: ElicitRequest = {
+            method: 'elicitation/create',
+            params
+        };
+        return await this.server.request(request, ElicitResultSchema, { ...options, relatedRequestId: this.requestId });
+    }
+
+    /**
+     * Sends a logging message.
+     */
+    public async log(params: LoggingMessageNotification['params'], sessionId?: string) {
+        await this.server.sendLoggingMessage(params, sessionId);
+    }
+
+    /**
+     * Sends a debug log message.
+     */
+    public async debug(message: string, extraLogData?: Record<string, unknown>, sessionId?: string) {
+        await this.log(
+            {
+                level: 'debug',
+                data: {
+                    ...extraLogData,
+                    message
+                },
+                logger: 'server'
+            },
+            sessionId
+        );
+    }
+
+    /**
+     * Sends an info log message.
+     */
+    public async info(message: string, extraLogData?: Record<string, unknown>, sessionId?: string) {
+        await this.log(
+            {
+                level: 'info',
+                data: {
+                    ...extraLogData,
+                    message
+                },
+                logger: 'server'
+            },
+            sessionId
+        );
+    }
+
+    /**
+     * Sends a warning log message.
+     */
+    public async warning(message: string, extraLogData?: Record<string, unknown>, sessionId?: string) {
+        await this.log(
+            {
+                level: 'warning',
+                data: {
+                    ...extraLogData,
+                    message
+                },
+                logger: 'server'
+            },
+            sessionId
+        );
+    }
+
+    /**
+     * Sends an error log message.
+     */
+    public async error(message: string, extraLogData?: Record<string, unknown>, sessionId?: string) {
+        await this.log(
+            {
+                level: 'error',
+                data: {
+                    ...extraLogData,
+                    message
+                },
+                logger: 'server'
+            },
+            sessionId
+        );
+    }
+}

src/server/index.ts

@@ -40,7 +40,10 @@ import {
     type ToolUseContent,
     CallToolRequestSchema,
     CallToolResultSchema,
-    CreateTaskResultSchema
+    CreateTaskResultSchema,
+    JSONRPCRequest,
+    TaskCreationParams,
+    MessageExtraInfo
 } from '../types.js';
 import { AjvJsonSchemaValidator } from '../validation/ajv-provider.js';
 import type { JsonSchemaType, jsonSchemaValidator } from '../validation/types.js';
@@ -56,6 +59,9 @@ import {
 import { RequestHandlerExtra } from '../shared/protocol.js';
 import { ExperimentalServerTasks } from '../experimental/tasks/server.js';
 import { assertToolsCallTaskCapability, assertClientRequestTaskCapability } from '../experimental/tasks/helpers.js';
+import { Context } from './context.js';
+import { TaskStore } from '../experimental/index.js';
+import { Transport } from '../shared/transport.js';
 
 export type ServerOptions = ProtocolOptions & {
     /**
@@ -219,9 +225,31 @@ export class Server<
         requestSchema: T,
         handler: (
             request: SchemaOutput<T>,
-            extra: RequestHandlerExtra<ServerRequest | RequestT, ServerNotification | NotificationT>
+            extra: Context<RequestT, NotificationT, ResultT>
         ) => ServerResult | ResultT | Promise<ServerResult | ResultT>
     ): void {
+        // 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>
+            ) => {
+                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;
         if (!methodSchema) {
@@ -259,7 +287,7 @@ export class Server<
 
                 const { params } = validatedRequest.data;
 
-                const result = await Promise.resolve(handler(request, extra));
+                const result = await Promise.resolve(handlerDecoratorFactory(handler)(request, extra));
 
                 // When task creation is requested, validate and return CreateTaskResult
                 if (params.task) {
@@ -286,11 +314,18 @@ export class Server<
             };
 
             // Install the wrapped handler
-            return super.setRequestHandler(requestSchema, wrappedHandler as unknown as typeof handler);
+            return super.setRequestHandler(requestSchema, handlerDecoratorFactory(wrappedHandler));
         }
 
         // Other handlers use default behavior
-        return super.setRequestHandler(requestSchema, handler);
+        return super.setRequestHandler(requestSchema, handlerDecoratorFactory(handler));
+    }
+
+    // Runtime type guard: ensure extra is our Context
+    private isContextExtra(
+        extra: RequestHandlerExtra<ServerRequest | RequestT, ServerNotification | NotificationT>
+    ): extra is Context<RequestT, NotificationT, ResultT> {
+        return extra instanceof Context;
     }
 
     protected assertCapabilityForMethod(method: RequestT['method']): void {
@@ -468,6 +503,25 @@ 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>;
+
+        // Expose a Context instance to handlers, which implements RequestHandlerExtra
+        return new Context<RequestT, NotificationT, ResultT>({
+            server: this,
+            request: args.request,
+            requestCtx: base
+        });
+    }
+
     async ping() {
         return this.request({ method: 'ping' }, EmptyResultSchema);
     }