feat: support extending McpServer with subclass

This commit is contained in:
wangshijun 2025-04-01 12:21:08 +08:00
commit 3571fc353f
3 changed files with 27 additions and 19 deletions

View file

@ -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(

View file

@ -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)}`));

View file

@ -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;
}