Skip to content

Commit d27775a

Browse files
committed
feat: group task-related fields into single object
1 parent f27cb52 commit d27775a

File tree

10 files changed

+206
-191
lines changed

10 files changed

+206
-191
lines changed

src/examples/server/simpleStreamableHttp.ts

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -472,16 +472,16 @@ const getServer = () => {
472472
}
473473
},
474474
{
475-
async createTask({ duration }, { taskStore, taskRequestedTtl }) {
475+
async createTask({ duration }, { task }) {
476476
// Create the task
477-
const task = await taskStore.createTask({
478-
ttl: taskRequestedTtl
477+
const newTask = await task.store.createTask({
478+
ttl: task.requestedTtl
479479
});
480480

481481
// Simulate out-of-band work
482482
(async () => {
483483
await new Promise(resolve => setTimeout(resolve, duration));
484-
await taskStore.storeTaskResult(task.taskId, 'completed', {
484+
await task.store.storeTaskResult(newTask.taskId, 'completed', {
485485
content: [
486486
{
487487
type: 'text',
@@ -493,14 +493,14 @@ const getServer = () => {
493493

494494
// Return CreateTaskResult with the created task
495495
return {
496-
task
496+
task: newTask
497497
};
498498
},
499-
async getTask(_args, { taskId, taskStore }) {
500-
return await taskStore.getTask(taskId);
499+
async getTask(_args, { task }) {
500+
return await task.store.getTask(task.id);
501501
},
502-
async getTaskResult(_args, { taskId, taskStore }) {
503-
const result = await taskStore.getTaskResult(taskId);
502+
async getTaskResult(_args, { task }) {
503+
const result = await task.store.getTaskResult(task.id);
504504
return result as CallToolResult;
505505
}
506506
}

src/experimental/tasks/interfaces.ts

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,28 +19,27 @@ import {
1919
Request
2020
} from '../../types.js';
2121
import { CreateTaskResult } from './types.js';
22-
import type { RequestHandlerExtra, RequestTaskStore } from '../../shared/protocol.js';
22+
import type { RequestHandlerExtra, TaskContext } from '../../shared/protocol.js';
2323
import type { ZodRawShapeCompat, AnySchema, ShapeOutput } from '../../server/zod-compat.js';
2424

2525
// ============================================================================
2626
// Task Handler Types (for registerToolTask)
2727
// ============================================================================
2828

2929
/**
30-
* Extended handler extra with task store for task creation.
30+
* Extended handler extra with task context for task creation.
3131
* @experimental
3232
*/
3333
export interface CreateTaskRequestHandlerExtra extends RequestHandlerExtra<ServerRequest, ServerNotification> {
34-
taskStore: RequestTaskStore;
34+
task: TaskContext;
3535
}
3636

3737
/**
3838
* Extended handler extra with task ID and store for task operations.
3939
* @experimental
4040
*/
4141
export interface TaskRequestHandlerExtra extends RequestHandlerExtra<ServerRequest, ServerNotification> {
42-
taskId: string;
43-
taskStore: RequestTaskStore;
42+
task: TaskContext & { id: string };
4443
}
4544

4645
/**

src/experimental/tasks/mcp-server.ts

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -56,15 +56,15 @@ export class ExperimentalMcpServerTasks {
5656
* execution: { taskSupport: 'required' }
5757
* }, {
5858
* createTask: async (args, extra) => {
59-
* const task = await extra.taskStore.createTask({ ttl: 300000 });
60-
* startBackgroundWork(task.taskId, args);
61-
* return { task };
59+
* const newTask = await extra.task.store.createTask({ ttl: 300000 });
60+
* startBackgroundWork(newTask.taskId, args);
61+
* return { task: newTask };
6262
* },
6363
* getTask: async (args, extra) => {
64-
* return extra.taskStore.getTask(extra.taskId);
64+
* return extra.task.store.getTask(extra.task.id);
6565
* },
6666
* getTaskResult: async (args, extra) => {
67-
* return extra.taskStore.getTaskResult(extra.taskId);
67+
* return extra.task.store.getTaskResult(extra.task.id);
6868
* }
6969
* });
7070
* ```

src/server/mcp.ts

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -330,10 +330,10 @@ export class McpServer {
330330
const isTaskHandler = 'createTask' in handler;
331331

332332
if (isTaskHandler) {
333-
if (!extra.taskStore) {
333+
if (!extra.task?.store) {
334334
throw new Error('No task store provided.');
335335
}
336-
const taskExtra = { ...extra, taskStore: extra.taskStore };
336+
const taskExtra = { ...extra, task: extra.task };
337337

338338
if (tool.inputSchema) {
339339
const typedHandler = handler as ToolTaskHandler<ZodRawShapeCompat>;
@@ -365,14 +365,14 @@ export class McpServer {
365365
request: RequestT,
366366
extra: RequestHandlerExtra<ServerRequest, ServerNotification>
367367
): Promise<CallToolResult> {
368-
if (!extra.taskStore) {
368+
if (!extra.task?.store) {
369369
throw new Error('No task store provided for task-capable tool.');
370370
}
371371

372372
// Validate input and create task
373373
const args = await this.validateToolInput(tool, request.params.arguments, request.params.name);
374374
const handler = tool.handler as ToolTaskHandler<ZodRawShapeCompat | undefined>;
375-
const taskExtra = { ...extra, taskStore: extra.taskStore };
375+
const taskExtra = { ...extra, task: extra.task };
376376

377377
const createTaskResult: CreateTaskResult = args // undefined only if tool.inputSchema is undefined
378378
? await Promise.resolve((handler as ToolTaskHandler<ZodRawShapeCompat>).createTask(args, taskExtra))
@@ -386,15 +386,15 @@ export class McpServer {
386386

387387
while (task.status !== 'completed' && task.status !== 'failed' && task.status !== 'cancelled') {
388388
await new Promise(resolve => setTimeout(resolve, pollInterval));
389-
const updatedTask = await extra.taskStore.getTask(taskId);
389+
const updatedTask = await extra.task.store.getTask(taskId);
390390
if (!updatedTask) {
391391
throw new McpError(ErrorCode.InternalError, `Task ${taskId} not found during polling`);
392392
}
393393
task = updatedTask;
394394
}
395395

396396
// Return the final result
397-
return (await extra.taskStore.getTaskResult(taskId)) as CallToolResult;
397+
return (await extra.task.store.getTaskResult(taskId)) as CallToolResult;
398398
}
399399

400400
private _completionHandlerInitialized = false;

src/shared/protocol.ts

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -231,6 +231,18 @@ export interface RequestTaskStore {
231231
listTasks(cursor?: string): Promise<{ tasks: Task[]; nextCursor?: string }>;
232232
}
233233

234+
/**
235+
* Context for task-related operations in request handlers.
236+
*/
237+
export interface TaskContext {
238+
/** The related task identifier (present when operating on existing task) */
239+
id?: string;
240+
/** Task store for managing task state */
241+
store: RequestTaskStore;
242+
/** Requested TTL in milliseconds (from client's task creation params) */
243+
requestedTtl?: number;
244+
}
245+
234246
/**
235247
* Extra data given to request handlers.
236248
*/
@@ -261,11 +273,11 @@ export type RequestHandlerExtra<SendRequestT extends Request, SendNotificationT
261273
*/
262274
requestId: RequestId;
263275

264-
taskId?: string;
265-
266-
taskStore?: RequestTaskStore;
267-
268-
taskRequestedTtl?: number;
276+
/**
277+
* Task context for task-related operations.
278+
* Present when the server has task storage enabled.
279+
*/
280+
task?: TaskContext;
269281

270282
/**
271283
* The original HTTP request.
@@ -745,9 +757,13 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
745757
authInfo: extra?.authInfo,
746758
requestId: request.id,
747759
requestInfo: extra?.requestInfo,
748-
taskId: relatedTaskId,
749-
taskStore: taskStore,
750-
taskRequestedTtl: taskCreationParams?.ttl,
760+
task: taskStore
761+
? {
762+
id: relatedTaskId,
763+
store: taskStore,
764+
requestedTtl: taskCreationParams?.ttl
765+
}
766+
: undefined,
751767
closeSSEStream: extra?.closeSSEStream,
752768
closeStandaloneSSEStream: extra?.closeStandaloneSSEStream
753769
};

0 commit comments

Comments
 (0)