Skip to content

Commit 3571fc3

Browse files
committed
feat: support extending McpServer with subclass
1 parent 592c91f commit 3571fc3

File tree

4 files changed

+27
-19
lines changed

4 files changed

+27
-19
lines changed

package-lock.json

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

src/server/mcp.ts

Lines changed: 13 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -54,12 +54,17 @@ export class McpServer {
5454
*/
5555
public readonly server: Server;
5656

57-
private _registeredResources: { [uri: string]: RegisteredResource } = {};
58-
private _registeredResourceTemplates: {
57+
protected _registeredResources: { [uri: string]: RegisteredResource } = {};
58+
protected _registeredResourceTemplates: {
5959
[name: string]: RegisteredResourceTemplate;
6060
} = {};
61-
private _registeredTools: { [name: string]: RegisteredTool } = {};
62-
private _registeredPrompts: { [name: string]: RegisteredPrompt } = {};
61+
protected _registeredTools: { [name: string]: RegisteredTool } = {};
62+
protected _registeredPrompts: { [name: string]: RegisteredPrompt } = {};
63+
64+
protected _toolHandlersInitialized = false;
65+
protected _completionHandlerInitialized = false;
66+
protected _resourceHandlersInitialized = false;
67+
protected _promptHandlersInitialized = false;
6368

6469
constructor(serverInfo: Implementation, options?: ServerOptions) {
6570
this.server = new Server(serverInfo, options);
@@ -81,13 +86,11 @@ export class McpServer {
8186
await this.server.close();
8287
}
8388

84-
private _toolHandlersInitialized = false;
85-
8689
private setToolRequestHandlers() {
8790
if (this._toolHandlersInitialized) {
8891
return;
8992
}
90-
93+
9194
this.server.assertCanSetRequestHandler(
9295
ListToolsRequestSchema.shape.method.value,
9396
);
@@ -177,8 +180,6 @@ export class McpServer {
177180
this._toolHandlersInitialized = true;
178181
}
179182

180-
private _completionHandlerInitialized = false;
181-
182183
private setCompletionRequestHandler() {
183184
if (this._completionHandlerInitialized) {
184185
return;
@@ -267,8 +268,6 @@ export class McpServer {
267268
return createCompletionResult(suggestions);
268269
}
269270

270-
private _resourceHandlersInitialized = false;
271-
272271
private setResourceRequestHandlers() {
273272
if (this._resourceHandlersInitialized) {
274273
return;
@@ -366,12 +365,10 @@ export class McpServer {
366365
);
367366

368367
this.setCompletionRequestHandler();
369-
368+
370369
this._resourceHandlersInitialized = true;
371370
}
372371

373-
private _promptHandlersInitialized = false;
374-
375372
private setPromptRequestHandlers() {
376373
if (this._promptHandlersInitialized) {
377374
return;
@@ -438,7 +435,7 @@ export class McpServer {
438435
);
439436

440437
this.setCompletionRequestHandler();
441-
438+
442439
this._promptHandlersInitialized = true;
443440
}
444441

@@ -770,7 +767,7 @@ type RegisteredPrompt = {
770767
callback: PromptCallback<undefined | PromptArgsRawShape>;
771768
};
772769

773-
function promptArgumentsFromSchema(
770+
export function promptArgumentsFromSchema(
774771
schema: ZodObject<PromptArgsRawShape>,
775772
): PromptArgument[] {
776773
return Object.entries(schema.shape).map(

src/shared/protocol.ts

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,11 @@ export type RequestHandlerExtra = {
9393
* The session ID from the transport, if available.
9494
*/
9595
sessionId?: string;
96+
97+
/**
98+
* The authenticated user, if available.
99+
*/
100+
user?: unknown;
96101
};
97102

98103
/**
@@ -316,6 +321,7 @@ export abstract class Protocol<
316321
const extra: RequestHandlerExtra = {
317322
signal: abortController.signal,
318323
sessionId: this._transport?.sessionId,
324+
user: this._transport?.user,
319325
};
320326

321327
// Starting with Promise.resolve() puts any synchronous errors into the monad as well.
@@ -361,7 +367,7 @@ export abstract class Protocol<
361367
private _onprogress(notification: ProgressNotification): void {
362368
const { progressToken, ...params } = notification.params;
363369
const messageId = Number(progressToken);
364-
370+
365371
const handler = this._progressHandlers.get(messageId);
366372
if (!handler) {
367373
this._onerror(new Error(`Received a progress notification for an unknown token: ${JSON.stringify(notification)}`));

src/shared/transport.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,4 +46,9 @@ export interface Transport {
4646
* The session ID generated for this connection.
4747
*/
4848
sessionId?: string;
49+
50+
/**
51+
* The authenticated user for this transport session.
52+
*/
53+
user?: unknown;
4954
}

0 commit comments

Comments
 (0)