feat: support extending McpServer with subclass

Author: wangshijun

package-lock.json

@@ -54,12 +54,17 @@ export class McpServer {
    */
   public readonly server: Server;
 
-  private _registeredResources: { [uri: string]: RegisteredResource } = {};
-  private _registeredResourceTemplates: {
+  protected _registeredResources: { [uri: string]: RegisteredResource } = {};
+  protected _registeredResourceTemplates: {
     [name: string]: RegisteredResourceTemplate;
   } = {};
-  private _registeredTools: { [name: string]: RegisteredTool } = {};
-  private _registeredPrompts: { [name: string]: RegisteredPrompt } = {};
+  protected _registeredTools: { [name: string]: RegisteredTool } = {};
+  protected _registeredPrompts: { [name: string]: RegisteredPrompt } = {};
+
+  protected _toolHandlersInitialized = false;
+  protected _completionHandlerInitialized = false;
+  protected _resourceHandlersInitialized = false;
+  protected _promptHandlersInitialized = false;
 
   constructor(serverInfo: Implementation, options?: ServerOptions) {
     this.server = new Server(serverInfo, options);
@@ -81,13 +86,11 @@ export class McpServer {
     await this.server.close();
   }
 
-  private _toolHandlersInitialized = false;
-
   private setToolRequestHandlers() {
     if (this._toolHandlersInitialized) {
       return;
     }
-    
+
     this.server.assertCanSetRequestHandler(
       ListToolsRequestSchema.shape.method.value,
     );
@@ -177,8 +180,6 @@ export class McpServer {
     this._toolHandlersInitialized = true;
   }
 
-  private _completionHandlerInitialized = false;
-
   private setCompletionRequestHandler() {
     if (this._completionHandlerInitialized) {
       return;
@@ -267,8 +268,6 @@ export class McpServer {
     return createCompletionResult(suggestions);
   }
 
-  private _resourceHandlersInitialized = false;
-
   private setResourceRequestHandlers() {
     if (this._resourceHandlersInitialized) {
       return;
@@ -366,12 +365,10 @@ export class McpServer {
     );
 
     this.setCompletionRequestHandler();
-    
+
     this._resourceHandlersInitialized = true;
   }
 
-  private _promptHandlersInitialized = false;
-
   private setPromptRequestHandlers() {
     if (this._promptHandlersInitialized) {
       return;
@@ -438,7 +435,7 @@ export class McpServer {
     );
 
     this.setCompletionRequestHandler();
-    
+
     this._promptHandlersInitialized = true;
   }
 
@@ -770,7 +767,7 @@ type RegisteredPrompt = {
   callback: PromptCallback<undefined | PromptArgsRawShape>;
 };
 
-function promptArgumentsFromSchema(
+export function promptArgumentsFromSchema(
   schema: ZodObject<PromptArgsRawShape>,
 ): PromptArgument[] {
   return Object.entries(schema.shape).map(

src/server/mcp.ts

@@ -93,6 +93,11 @@ export type RequestHandlerExtra = {
    * The session ID from the transport, if available.
    */
   sessionId?: string;
+
+  /**
+   * The authenticated user, if available.
+   */
+  user?: unknown;
 };
 
 /**
@@ -316,6 +321,7 @@ export abstract class Protocol<
     const extra: RequestHandlerExtra = {
       signal: abortController.signal,
       sessionId: this._transport?.sessionId,
+      user: this._transport?.user,
     };
 
     // Starting with Promise.resolve() puts any synchronous errors into the monad as well.
@@ -361,7 +367,7 @@ export abstract class Protocol<
   private _onprogress(notification: ProgressNotification): void {
     const { progressToken, ...params } = notification.params;
     const messageId = Number(progressToken);
-    
+
     const handler = this._progressHandlers.get(messageId);
     if (!handler) {
       this._onerror(new Error(`Received a progress notification for an unknown token: ${JSON.stringify(notification)}`));

src/shared/protocol.ts

@@ -46,4 +46,9 @@ export interface Transport {
    * The session ID generated for this connection.
    */
   sessionId?: string;
+
+  /**
+   * The authenticated user for this transport session.
+   */
+  user?: unknown;
 }