draft: request context utility - utility methods, lifespan context as per python sdk

Author: KKonstantinov

src/server/context.ts

@@ -0,0 +1,204 @@
+import {
+    CreateMessageRequest,
+    CreateMessageResultSchema,
+    ElicitRequest,
+    ElicitResultSchema,
+    LoggingMessageNotification,
+    Notification,
+    Request,
+    RequestId,
+    RequestInfo,
+    RequestMeta,
+    Result,
+    ServerNotification,
+    ServerRequest,
+    ServerResult
+} from '../types.js';
+import { RequestHandlerExtra, RequestOptions } from '../shared/protocol.js';
+import { ZodType } from 'zod';
+import type { z } from 'zod';
+import { Server } from './index.js';
+import { LifespanContext, RequestContext } from '../shared/requestContext.js';
+import { AuthInfo } from './auth/types.js';
+
+/**
+ * A context object that is passed to request handlers.
+ *
+ * Implements the RequestHandlerExtra interface for backwards compatibility.
+ * Notes:
+ * Keeps this backwards compatible with the old RequestHandlerExtra interface and provides getter methods for backwards compatibility.
+ * In a breaking change, this can be removed and the RequestContext can be used directly via ctx.requestContext.
+ *
+ * TODO(Konstantin): Could be restructured better when breaking changes are allowed. More
+ */
+export class Context<
+    LifespanContextT extends LifespanContext | undefined = undefined,
+    RequestT extends Request = Request,
+    NotificationT extends Notification = Notification,
+    ResultT extends Result = Result
+> implements RequestHandlerExtra<ServerRequest | RequestT, ServerNotification | NotificationT>
+{
+    private readonly server: Server<RequestT, NotificationT, ResultT, LifespanContextT>;
+
+    /**
+     * The request context.
+     * A type-safe context that is passed to request handlers.
+     */
+    public readonly requestCtx: RequestContext<
+        LifespanContextT,
+        RequestT | ServerRequest,
+        NotificationT | ServerNotification,
+        ResultT | ServerResult
+    >;
+
+    constructor(args: {
+        server: Server<RequestT, NotificationT, ResultT, LifespanContextT>;
+        requestCtx: RequestContext<LifespanContextT, RequestT | ServerRequest, NotificationT | ServerNotification, ResultT | ServerResult>;
+    }) {
+        this.server = args.server;
+        this.requestCtx = args.requestCtx;
+    }
+
+    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;
+    }
+
+    public get _meta(): RequestMeta | undefined {
+        return this.requestCtx._meta;
+    }
+
+    public get sessionId(): string | undefined {
+        return this.requestCtx.sessionId;
+    }
+
+    /**
+     * 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 ZodType<object>>(
+        request: RequestT | ServerRequest,
+        resultSchema: U,
+        options?: RequestOptions
+    ): Promise<z.infer<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) {
+        const request: CreateMessageRequest = {
+            method: 'sampling/createMessage',
+            params
+        };
+        return this.server.request(request, CreateMessageResultSchema, { ...options, relatedRequestId: this.requestId });
+    }
+
+    /**
+     * Sends an elicitation request to the client.
+     */
+    public elicit(params: ElicitRequest['params'], options?: RequestOptions) {
+        const request: ElicitRequest = {
+            method: 'elicitation/create',
+            params
+        };
+        return 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

@@ -1,4 +1,4 @@
-import { mergeCapabilities, Protocol, ProtocolOptions, RequestOptions } from '../shared/protocol.js';
+import { mergeCapabilities, Protocol, ProtocolOptions, RequestHandlerExtra, RequestOptions } from '../shared/protocol.js';
 import {
     ClientCapabilities,
     CreateMessageRequest,
@@ -29,11 +29,18 @@ import {
     SUPPORTED_PROTOCOL_VERSIONS,
     LoggingLevel,
     SetLevelRequestSchema,
-    LoggingLevelSchema
+    LoggingLevelSchema,
+    JSONRPCRequest,
+    MessageExtraInfo
 } from '../types.js';
 import Ajv from 'ajv';
+import { Context } from './context.js';
+import { LifespanContext, RequestContext } from '../shared/requestContext.js';
+import { ZodLiteral, ZodObject } from 'zod';
+import { z } from 'zod';
+import { Transport } from 'src/shared/transport.js';
 
-export type ServerOptions = ProtocolOptions & {
+export type ServerOptions<LifespanContextT extends LifespanContext | undefined = undefined> = ProtocolOptions & {
     /**
      * Capabilities to advertise as being supported by this server.
      */
@@ -43,6 +50,11 @@ export type ServerOptions = ProtocolOptions & {
      * Optional instructions describing how to use the server and its features.
      */
     instructions?: string;
+
+    /**
+     * Optional lifespan context type.
+     */
+    lifespan?: LifespanContextT;
 };
 
 /**
@@ -73,12 +85,14 @@ export type ServerOptions = ProtocolOptions & {
 export class Server<
     RequestT extends Request = Request,
     NotificationT extends Notification = Notification,
-    ResultT extends Result = Result
+    ResultT extends Result = Result,
+    LifespanContextT extends LifespanContext | undefined = undefined
 > extends Protocol<ServerRequest | RequestT, ServerNotification | NotificationT, ServerResult | ResultT> {
     private _clientCapabilities?: ClientCapabilities;
     private _clientVersion?: Implementation;
     private _capabilities: ServerCapabilities;
     private _instructions?: string;
+    private _lifespan?: LifespanContextT;
 
     /**
      * Callback for when initialization has fully completed (i.e., the client has sent an `initialized` notification).
@@ -90,12 +104,12 @@ export class Server<
      */
     constructor(
         private _serverInfo: Implementation,
-        options?: ServerOptions
+        options?: ServerOptions<LifespanContextT>
     ) {
         super(options);
         this._capabilities = options?.capabilities ?? {};
         this._instructions = options?.instructions;
-
+        this._lifespan = options?.lifespan;
         this.setRequestHandler(InitializeRequestSchema, request => this._oninitialize(request));
         this.setNotificationHandler(InitializedNotificationSchema, () => this.oninitialized?.());
 
@@ -125,6 +139,13 @@ export class Server<
         return currentLevel ? this.LOG_LEVEL_SEVERITY.get(level)! < this.LOG_LEVEL_SEVERITY.get(currentLevel)! : false;
     };
 
+    // Runtime type guard: ensure extra is our Context
+    private isContextExtra(
+        extra: RequestHandlerExtra<ServerRequest | RequestT, ServerNotification | NotificationT>
+    ): extra is Context<LifespanContextT, RequestT, NotificationT, ResultT> {
+        return extra instanceof Context;
+    }
+
     /**
      * Registers new capabilities. This can only be called before connecting to a transport.
      *
@@ -277,6 +298,65 @@ export class Server<
         return this._capabilities;
     }
 
+    /**
+     * Registers a handler where `extra` is typed as `Context` for ergonomic server-side usage.
+     * Internally, this wraps the handler and forwards to the base implementation.
+     */
+    public override setRequestHandler<
+        T extends ZodObject<{
+            method: ZodLiteral<string>;
+        }>
+    >(
+        requestSchema: T,
+        handler: (
+            request: z.infer<T>,
+            extra: Context<LifespanContextT, RequestT, NotificationT, ResultT>
+        ) => ServerResult | ResultT | Promise<ServerResult | ResultT>
+    ): void {
+        super.setRequestHandler(requestSchema, (request, extra) => {
+            if (!this.isContextExtra(extra)) {
+                throw new Error('Internal error: Expected Context for request handler extra');
+            }
+            return handler(request, extra);
+        });
+    }
+
+    protected override createRequestExtra(
+        request: JSONRPCRequest,
+        abortController: AbortController,
+        capturedTransport: Transport | undefined,
+        extra?: MessageExtraInfo
+    ): RequestHandlerExtra<ServerRequest | RequestT, ServerNotification | NotificationT> {
+        const base = super.createRequestExtra(request, abortController, capturedTransport, extra) as RequestHandlerExtra<
+            ServerRequest | RequestT,
+            ServerNotification | NotificationT
+        >;
+
+        // Wrap base in Context to add server utilities while preserving shape
+        const requestCtx = new RequestContext<
+            LifespanContextT,
+            ServerRequest | RequestT,
+            ServerNotification | NotificationT,
+            ServerResult | ResultT
+        >({
+            signal: base.signal,
+            authInfo: base.authInfo,
+            requestInfo: base.requestInfo,
+            requestId: base.requestId,
+            _meta: base._meta,
+            sessionId: base.sessionId,
+            lifespanContext: this._lifespan as LifespanContextT,
+            protocol: this
+        });
+
+        const ctx = new Context<LifespanContextT, RequestT, NotificationT, ResultT>({
+            server: this,
+            requestCtx
+        });
+
+        return ctx;
+    }
+
     async ping() {
         return this.request({ method: 'ping' }, EmptyResultSchema);
     }