Skip to content

Commit 1cf942d

Browse files
refactor: extract task orchestration from Protocol into TaskManager
Extract all task-related logic from Protocol into a dedicated TaskManager class with consolidated lifecycle API and NullTaskManager (Null Object pattern). Protocol delegates to TaskManager via 4 lifecycle methods: - processInboundRequest: context extraction + function wrapping - processOutboundRequest: metadata augmentation + queue routing - processInboundResponse: side-channel handling + progress preservation - processOutboundNotification: task-related notification routing NullTaskManager provides no-op defaults, eliminating all conditional checks in Protocol.
1 parent faeb03e commit 1cf942d

File tree

14 files changed

+1218
-1131
lines changed

14 files changed

+1218
-1131
lines changed

.changeset/extract-task-manager.md

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
---
2+
"@modelcontextprotocol/core": minor
3+
"@modelcontextprotocol/client": minor
4+
"@modelcontextprotocol/server": minor
5+
---
6+
7+
refactor: extract task orchestration from Protocol into TaskManager
8+
9+
**Breaking changes:**
10+
- `extra.taskId``extra.task?.taskId`
11+
- `extra.taskStore``extra.task?.taskStore`
12+
- `extra.taskRequestedTtl``extra.task?.requestedTtl`
13+
- `ProtocolOptions` no longer accepts `taskStore`/`taskMessageQueue` — pass via `TaskManagerOptions` in `ClientOptions`/`ServerOptions`
14+
- Abstract methods `assertTaskCapability`/`assertTaskHandlerCapability` removed from Protocol

packages/client/src/experimental/tasks/client.ts

Lines changed: 6 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -25,11 +25,6 @@ import type { Client } from '../../client/client.js';
2525
* @internal
2626
*/
2727
interface ClientInternal {
28-
requestStream<T extends AnyObjectSchema>(
29-
request: Request,
30-
resultSchema: T,
31-
options?: RequestOptions
32-
): AsyncGenerator<ResponseMessage<SchemaOutput<T>>, void, void>;
3328
isToolTask(toolName: string): boolean;
3429
getToolOutputValidator(toolName: string): ((data: unknown) => { valid: boolean; errorMessage?: string }) | undefined;
3530
}
@@ -90,7 +85,6 @@ export class ExperimentalClientTasks {
9085
): AsyncGenerator<ResponseMessage<SchemaOutput<typeof CallToolResultSchema>>, void, void> {
9186
// Access Client's internal methods
9287
const clientInternal = this._client as unknown as ClientInternal;
93-
9488
// Add task creation parameters if server supports it and not explicitly provided
9589
const optionsWithTask = {
9690
...options,
@@ -99,7 +93,7 @@ export class ExperimentalClientTasks {
9993
task: options?.task ?? (clientInternal.isToolTask(params.name) ? {} : undefined)
10094
};
10195

102-
const stream = clientInternal.requestStream({ method: 'tools/call', params }, CallToolResultSchema, optionsWithTask);
96+
const stream = this._client.tasks.requestStream({ method: 'tools/call', params }, CallToolResultSchema, optionsWithTask);
10397

10498
// Get the validator for this tool (if it has an output schema)
10599
const validator = clientInternal.getToolOutputValidator(params.name);
@@ -170,9 +164,7 @@ export class ExperimentalClientTasks {
170164
* @experimental
171165
*/
172166
async getTask(taskId: string, options?: RequestOptions): Promise<GetTaskResult> {
173-
// Delegate to the client's underlying Protocol method
174-
type ClientWithGetTask = { getTask(params: { taskId: string }, options?: RequestOptions): Promise<GetTaskResult> };
175-
return (this._client as unknown as ClientWithGetTask).getTask({ taskId }, options);
167+
return this._client.tasks.getTask({ taskId }, options);
176168
}
177169

178170
/**
@@ -186,16 +178,7 @@ export class ExperimentalClientTasks {
186178
* @experimental
187179
*/
188180
async getTaskResult<T extends AnyObjectSchema>(taskId: string, resultSchema?: T, options?: RequestOptions): Promise<SchemaOutput<T>> {
189-
// Delegate to the client's underlying Protocol method
190-
return (
191-
this._client as unknown as {
192-
getTaskResult: <U extends AnyObjectSchema>(
193-
params: { taskId: string },
194-
resultSchema?: U,
195-
options?: RequestOptions
196-
) => Promise<SchemaOutput<U>>;
197-
}
198-
).getTaskResult({ taskId }, resultSchema, options);
181+
return this._client.tasks.getTaskResult({ taskId }, resultSchema!, options);
199182
}
200183

201184
/**
@@ -208,12 +191,7 @@ export class ExperimentalClientTasks {
208191
* @experimental
209192
*/
210193
async listTasks(cursor?: string, options?: RequestOptions): Promise<ListTasksResult> {
211-
// Delegate to the client's underlying Protocol method
212-
return (
213-
this._client as unknown as {
214-
listTasks: (params?: { cursor?: string }, options?: RequestOptions) => Promise<ListTasksResult>;
215-
}
216-
).listTasks(cursor ? { cursor } : undefined, options);
194+
return this._client.tasks.listTasks(cursor ? { cursor } : undefined, options);
217195
}
218196

219197
/**
@@ -225,12 +203,7 @@ export class ExperimentalClientTasks {
225203
* @experimental
226204
*/
227205
async cancelTask(taskId: string, options?: RequestOptions): Promise<CancelTaskResult> {
228-
// Delegate to the client's underlying Protocol method
229-
return (
230-
this._client as unknown as {
231-
cancelTask: (params: { taskId: string }, options?: RequestOptions) => Promise<CancelTaskResult>;
232-
}
233-
).cancelTask({ taskId }, options);
206+
return this._client.tasks.cancelTask({ taskId }, options);
234207
}
235208

236209
/**
@@ -252,14 +225,6 @@ export class ExperimentalClientTasks {
252225
resultSchema: T,
253226
options?: RequestOptions
254227
): AsyncGenerator<ResponseMessage<SchemaOutput<T>>, void, void> {
255-
// Delegate to the client's underlying Protocol method
256-
type ClientWithRequestStream = {
257-
requestStream<U extends AnyObjectSchema>(
258-
request: Request,
259-
resultSchema: U,
260-
options?: RequestOptions
261-
): AsyncGenerator<ResponseMessage<SchemaOutput<U>>, void, void>;
262-
};
263-
return (this._client as unknown as ClientWithRequestStream).requestStream(request, resultSchema, options);
228+
return this._client.tasks.requestStream(request, resultSchema, options);
264229
}
265230
}

packages/core/src/experimental/tasks/interfaces.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,8 @@
33
* WARNING: These APIs are experimental and may change without notice.
44
*/
55

6-
import type { RequestTaskStore, ServerContext } from '../../shared/protocol.js';
6+
import type { ServerContext } from '../../shared/protocol.js';
7+
import type { RequestTaskStore } from '../../shared/taskManager.js';
78
import type {
89
JSONRPCErrorResponse,
910
JSONRPCNotification,

packages/core/src/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ export * from './shared/metadataUtils.js';
66
export * from './shared/protocol.js';
77
export * from './shared/responseMessage.js';
88
export * from './shared/stdio.js';
9+
export * from './shared/taskManager.js';
910
export * from './shared/toolNameValidation.js';
1011
export * from './shared/transport.js';
1112
export * from './shared/uriTemplate.js';

0 commit comments

Comments
 (0)