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

@ -30,6 +30,7 @@ import type {
ResultTypeMap,
ServerCapabilities,
SubscribeRequest,
TaskManagerOptions,
Tool,
Transport,
UnsubscribeRequest
@ -46,6 +47,7 @@ import {
ElicitRequestSchema,
ElicitResultSchema,
EmptyResultSchema,
extractTaskManagerOptions,
GetPromptResultSchema,
InitializeResultSchema,
LATEST_PROTOCOL_VERSION,
@ -140,11 +142,19 @@ export function getSupportedElicitationModes(capabilities: ClientCapabilities['e
return { supportsFormMode, supportsUrlMode };
}
/**
* Extended tasks capability that includes runtime configuration (store, messageQueue).
* The runtime-only fields are stripped before advertising capabilities to servers.
*/
export type ClientTasksCapabilityWithRuntime = NonNullable<ClientCapabilities['tasks']> & TaskManagerOptions;
export type ClientOptions = ProtocolOptions & {
/**
* Capabilities to advertise as being supported by this client.
*/
capabilities?: ClientCapabilities;
capabilities?: Omit<ClientCapabilities, 'tasks'> & {
tasks?: ClientTasksCapabilityWithRuntime;
};
/**
* JSON Schema validator for tool output validation.
@ -213,11 +223,22 @@ export class Client extends Protocol<ClientContext> {
private _clientInfo: Implementation,
options?: ClientOptions
) {
super(options);
this._capabilities = options?.capabilities ?? {};
super({
...options,
tasks: extractTaskManagerOptions(options?.capabilities?.tasks)
});
this._capabilities = options?.capabilities ? { ...options.capabilities } : {};
this._jsonSchemaValidator = options?.jsonSchemaValidator ?? new DefaultJsonSchemaValidator();
this._enforceStrictCapabilities = options?.enforceStrictCapabilities ?? false;
// 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;
}
// Store list changed config for setup after connection (when we know server capabilities)
if (options?.listChanged) {
this._pendingListChangedConfig = options.listChanged;
@ -650,12 +671,6 @@ export class Client extends Protocol<ClientContext> {
}
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 'sampling/createMessage': {
if (!this._capabilities.sampling) {
@ -687,19 +702,6 @@ export class Client extends Protocol<ClientContext> {
break;
}
case 'tasks/get':
case 'tasks/list':
case 'tasks/result':
case 'tasks/cancel': {
if (!this._capabilities.tasks) {
throw new SdkError(
SdkErrorCode.CapabilityNotSupported,
`Client does not support tasks capability (required for ${method})`
);
}
break;
}
case 'ping': {
// No specific capability required for ping
break;
@ -712,16 +714,9 @@ export class Client extends Protocol<ClientContext> {
}
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;
}
assertClientRequestTaskCapability(this._capabilities.tasks?.requests, method, 'Client');
assertClientRequestTaskCapability(this._capabilities?.tasks?.requests, method, 'Client');
}
/** Sends a ping to the server to check connectivity. */
async ping(options?: RequestOptions) {
return this._requestWithSchema({ method: 'ping' }, EmptyResultSchema, options);
}

View file

@ -6,6 +6,7 @@
*/
import type {
AnyObjectSchema,
CallToolRequest,
CallToolResult,
CancelTaskResult,
@ -13,12 +14,19 @@ import type {
GetTaskPayloadResult,
GetTaskResult,
ListTasksResult,
Request,
RequestMethod,
RequestOptions,
ResponseMessage,
ResultTypeMap
} from '@modelcontextprotocol/core';
import { GetTaskPayloadResultSchema, ProtocolError, ProtocolErrorCode } from '@modelcontextprotocol/core';
import {
CallToolResultSchema,
getResultSchema,
GetTaskPayloadResultSchema,
ProtocolError,
ProtocolErrorCode
} from '@modelcontextprotocol/core';
import type { Client } from '../../client/client.js';
@ -27,10 +35,6 @@ import type { Client } from '../../client/client.js';
* @internal
*/
interface ClientInternal {
requestStream<M extends RequestMethod>(
request: { method: M; params?: Record<string, unknown> },
options?: RequestOptions
): AsyncGenerator<ResponseMessage<ResultTypeMap[M]>, void, void>;
isToolTask(toolName: string): boolean;
getToolOutputValidator(toolName: string): ((data: unknown) => { valid: boolean; errorMessage?: string }) | undefined;
}
@ -49,6 +53,10 @@ interface ClientInternal {
export class ExperimentalClientTasks {
constructor(private readonly _client: Client) {}
private get _module() {
return this._client.taskManager;
}
/**
* Calls a tool and returns an AsyncGenerator that yields response messages.
* The generator is guaranteed to end with either a `'result'` or `'error'` message.
@ -103,7 +111,7 @@ export class ExperimentalClientTasks {
task: options?.task ?? (clientInternal.isToolTask(params.name) ? {} : undefined)
};
const stream = clientInternal.requestStream({ method: 'tools/call', params }, optionsWithTask);
const stream = this._module.requestStream({ method: 'tools/call', params }, CallToolResultSchema, optionsWithTask);
// Get the validator for this tool (if it has an output schema)
const validator = clientInternal.getToolOutputValidator(params.name);
@ -175,9 +183,7 @@ export class ExperimentalClientTasks {
* @experimental
*/
async getTask(taskId: string, options?: RequestOptions): Promise<GetTaskResult> {
// Delegate to the client's underlying Protocol method
type ClientWithGetTask = { getTask(params: { taskId: string }, options?: RequestOptions): Promise<GetTaskResult> };
return (this._client as unknown as ClientWithGetTask).getTask({ taskId }, options);
return this._module.getTask({ taskId }, options);
}
/**
@ -191,15 +197,7 @@ export class ExperimentalClientTasks {
* @experimental
*/
async getTaskResult(taskId: string, options?: RequestOptions): Promise<GetTaskPayloadResult> {
return (
this._client 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);
}
/**
@ -212,12 +210,7 @@ export class ExperimentalClientTasks {
* @experimental
*/
async listTasks(cursor?: string, options?: RequestOptions): Promise<ListTasksResult> {
// Delegate to the client's underlying Protocol method
return (
this._client as unknown as {
listTasks: (params?: { cursor?: string }, options?: RequestOptions) => Promise<ListTasksResult>;
}
).listTasks(cursor ? { cursor } : undefined, options);
return this._module.listTasks(cursor ? { cursor } : undefined, options);
}
/**
@ -229,12 +222,7 @@ export class ExperimentalClientTasks {
* @experimental
*/
async cancelTask(taskId: string, options?: RequestOptions): Promise<CancelTaskResult> {
// Delegate to the client's underlying Protocol method
return (
this._client as unknown as {
cancelTask: (params: { taskId: string }, options?: RequestOptions) => Promise<CancelTaskResult>;
}
).cancelTask({ taskId }, options);
return this._module.cancelTask({ taskId }, options);
}
/**
@ -279,7 +267,11 @@ export class ExperimentalClientTasks {
request: { method: M; params?: Record<string, unknown> },
options?: RequestOptions
): AsyncGenerator<ResponseMessage<ResultTypeMap[M]>, void, void> {
// Delegate to the client's underlying Protocol method
return (this._client as unknown as ClientInternal).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
>;
}
}