v2 - RFC Extract Tasks out of protocol.ts into TaskManager (#1673)
Co-authored-by: Felix Weinberger <fweinberger@anthropic.com>
This commit is contained in:
parent
7d9c72eee5
commit
462c3fc47d
2 changed files with 2122 additions and 1942 deletions
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
>;
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue