v2 - RFC Extract Tasks out of protocol.ts into TaskManager (#1673)

Co-authored-by: Felix Weinberger <fweinberger@anthropic.com>
This commit is contained in:
Konstantin Konstantinov 2026-03-25 21:58:01 +02:00 • committed by GitHub
commit 462c3fc47d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 2122 additions and 1942 deletions

View file

@ -6,6 +6,7 @@
*/
import type {
AnyObjectSchema,
CancelTaskResult,
CreateMessageRequestParams,
CreateMessageResult,
@ -15,12 +16,13 @@ import type {
GetTaskPayloadResult,
GetTaskResult,
ListTasksResult,
Request,
RequestMethod,
RequestOptions,
ResponseMessage,
ResultTypeMap
} from '@modelcontextprotocol/core';
import { GetTaskPayloadResultSchema } from '@modelcontextprotocol/core';
import { getResultSchema, GetTaskPayloadResultSchema } from '@modelcontextprotocol/core';
import type { Server } from '../../server/server.js';
@ -39,6 +41,10 @@ import type { Server } from '../../server/server.js';
export class ExperimentalServerTasks {
constructor(private readonly _server: Server) {}
private get _module() {
return this._server.taskManager;
}
/**
* Sends a request and returns an AsyncGenerator that yields response messages.
* The generator is guaranteed to end with either a `'result'` or `'error'` message.
@ -56,14 +62,12 @@ export class ExperimentalServerTasks {
request: { method: M; params?: Record<string, unknown> },
options?: RequestOptions
): AsyncGenerator<ResponseMessage<ResultTypeMap[M]>, void, void> {
// Delegate to the server's underlying Protocol method
type ServerWithRequestStream = {
requestStream<N extends RequestMethod>(
request: { method: N; params?: Record<string, unknown> },
options?: RequestOptions
): AsyncGenerator<ResponseMessage<ResultTypeMap[N]>, void, void>;
};
return (this._server as unknown as ServerWithRequestStream).requestStream(request, options);
const resultSchema = getResultSchema(request.method) as unknown as AnyObjectSchema;
return this._module.requestStream(request as Request, resultSchema, options) as AsyncGenerator<
ResponseMessage<ResultTypeMap[M]>,
void,
void
>;
}
/**
@ -250,8 +254,7 @@ export class ExperimentalServerTasks {
* @experimental
*/
async getTask(taskId: string, options?: RequestOptions): Promise<GetTaskResult> {
type ServerWithGetTask = { getTask(params: { taskId: string }, options?: RequestOptions): Promise<GetTaskResult> };
return (this._server as unknown as ServerWithGetTask).getTask({ taskId }, options);
return this._module.getTask({ taskId }, options);
}
/**
@ -265,15 +268,7 @@ export class ExperimentalServerTasks {
* @experimental
*/
async getTaskResult(taskId: string, options?: RequestOptions): Promise<GetTaskPayloadResult> {
return (
this._server as unknown as {
getTaskResult: (
params: { taskId: string },
resultSchema: typeof GetTaskPayloadResultSchema,
options?: RequestOptions
) => Promise<GetTaskPayloadResult>;
}
).getTaskResult({ taskId }, GetTaskPayloadResultSchema, options);
return this._module.getTaskResult({ taskId }, GetTaskPayloadResultSchema, options);
}
/**
@ -286,11 +281,7 @@ export class ExperimentalServerTasks {
* @experimental
*/
async listTasks(cursor?: string, options?: RequestOptions): Promise<ListTasksResult> {
return (
this._server as unknown as {
listTasks: (params?: { cursor?: string }, options?: RequestOptions) => Promise<ListTasksResult>;
}
).listTasks(cursor ? { cursor } : undefined, options);
return this._module.listTasks(cursor ? { cursor } : undefined, options);
}
/**
@ -302,10 +293,6 @@ export class ExperimentalServerTasks {
* @experimental
*/
async cancelTask(taskId: string, options?: RequestOptions): Promise<CancelTaskResult> {
return (
this._server as unknown as {
cancelTask: (params: { taskId: string }, options?: RequestOptions) => Promise<CancelTaskResult>;
}
).cancelTask({ taskId }, options);
return this._module.cancelTask({ taskId }, options);
}
}

View file

@ -29,6 +29,7 @@ import type {
ServerCapabilities,
ServerContext,
ServerResult,
TaskManagerOptions,
ToolResultContent,
ToolUseContent
} from '@modelcontextprotocol/core';
@ -42,6 +43,7 @@ import {
CreateTaskResultSchema,
ElicitResultSchema,
EmptyResultSchema,
extractTaskManagerOptions,
LATEST_PROTOCOL_VERSION,
ListRootsResultSchema,
LoggingLevelSchema,
@ -57,11 +59,19 @@ import { DefaultJsonSchemaValidator } from '@modelcontextprotocol/server/_shims'
import { ExperimentalServerTasks } from '../experimental/tasks/server.js';
/**
* Extended tasks capability that includes runtime configuration (store, messageQueue).
* The runtime-only fields are stripped before advertising capabilities to clients.
*/
export type ServerTasksCapabilityWithRuntime = NonNullable<ServerCapabilities['tasks']> & TaskManagerOptions;
export type ServerOptions = ProtocolOptions & {
/**
* Capabilities to advertise as being supported by this server.
*/
capabilities?: ServerCapabilities;
capabilities?: Omit<ServerCapabilities, 'tasks'> & {
tasks?: ServerTasksCapabilityWithRuntime;
};
/**
* Optional instructions describing how to use the server and its features.
@ -106,11 +116,22 @@ export class Server extends Protocol<ServerContext> {
private _serverInfo: Implementation,
options?: ServerOptions
) {
super(options);
this._capabilities = options?.capabilities ?? {};
super({
...options,
tasks: extractTaskManagerOptions(options?.capabilities?.tasks)
});
this._capabilities = options?.capabilities ? { ...options.capabilities } : {};
this._instructions = options?.instructions;
this._jsonSchemaValidator = options?.jsonSchemaValidator ?? new DefaultJsonSchemaValidator();
// Strip runtime-only fields from advertised capabilities
if (options?.capabilities?.tasks) {
// eslint-disable-next-line @typescript-eslint/no-unused-vars
const { taskStore, taskMessageQueue, defaultTaskPollInterval, maxTaskQueueSize, ...wireCapabilities } =
options.capabilities.tasks;
this._capabilities.tasks = wireCapabilities;
}
this.setRequestHandler('initialize', request => this._oninitialize(request));
this.setNotificationHandler('notifications/initialized', () => this.oninitialized?.());
@ -347,12 +368,6 @@ export class Server extends Protocol<ServerContext> {
}
protected assertRequestHandlerCapability(method: string): void {
// Task handlers are registered in Protocol constructor before _capabilities is initialized
// Skip capability check for task methods during initialization
if (!this._capabilities) {
return;
}
switch (method) {
case 'completion/complete': {
if (!this._capabilities.completions) {
@ -393,19 +408,6 @@ export class Server extends Protocol<ServerContext> {
break;
}
case 'tasks/get':
case 'tasks/list':
case 'tasks/result':
case 'tasks/cancel': {
if (!this._capabilities.tasks) {
throw new SdkError(
SdkErrorCode.CapabilityNotSupported,
`Server does not support tasks capability (required for ${method})`
);
}
break;
}
case 'ping':
case 'initialize': {
// No specific capability required for these methods
@ -419,13 +421,7 @@ export class Server extends Protocol<ServerContext> {
}
protected assertTaskHandlerCapability(method: string): void {
// Task handlers are registered in Protocol constructor before _capabilities is initialized
// Skip capability check for task methods during initialization
if (!this._capabilities) {
return;
}
assertToolsCallTaskCapability(this._capabilities.tasks?.requests, method, 'Server');
assertToolsCallTaskCapability(this._capabilities?.tasks?.requests, method, 'Server');
}
private async _oninitialize(request: InitializeRequest): Promise<InitializeResult> {