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
|
|
@ -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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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> {
|
||||
|
|
|
|||
Loading…
Reference in a new issue