[RFC] refactor: extract task orchestration from Protocol into TaskManager #1449
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
🦋 Changeset detectedLatest commit: 9882b72 The changes in this PR will be included in the next version bump. This PR includes changesets to release 6 packages
Not sure what this means? Click here to learn what changesets are. Click here if you're a maintainer who wants to add another changeset to this PR |
@modelcontextprotocol/client
@modelcontextprotocol/server
@modelcontextprotocol/express
@modelcontextprotocol/hono
@modelcontextprotocol/node
commit: |
Extract task state management, message queuing, polling, and handler bodies from Protocol into a dedicated TaskManager class. This is the first step in isolating task concerns from core JSON-RPC framing. TaskManager provides: - Public API: requestStream(), getTask(), getTaskResult(), listTasks(), cancelTask() - Handler bodies: handleGetTask(), handleGetTaskPayload(), handleListTasks(), handleCancelTask() - Delegation methods: prepareOutboundRequest(), extractInboundTaskContext(), handleResponse(), shouldPreserveProgressHandler(), routeNotification(), routeResponse() - Request-scoped store: createRequestTaskStore() - Types: TaskManagerOptions, RequestTaskStore, TaskContext, TaskManagerHost Protocol.ts is not yet modified - this is an additive change only.
Protocol.ts now delegates task orchestration to TaskManager at 5 points: - request(): task param augmentation and queue routing - _onrequest(): task context extraction and response routing - _onresponse(): side-channel resolver handling and progress preservation - notification(): task-related notification queueing - _onclose(): task state cleanup Key changes: - Removed ~700 lines of task code from Protocol - Protocol constructor accepts optional TaskManager instance - RequestHandlerExtra.task? groups task fields (breaking: was flat taskId/taskStore/taskRequestedTtl) - Removed abstract assertTaskCapability/assertTaskHandlerCapability methods - ProtocolOptions no longer has task fields - Re-exports TaskManager types for import compatibility Protocol.ts: 1,667 -> 964 lines (-42%) Note: Client/Server not yet updated - they will fail to compile until the abstract method removal is reflected in their implementations.
felixweinberger
changed the title
- Client/Server create TaskManager from options and pass to Protocol - Wire task capability assertions on TaskManager instance - Remove assertTaskCapability/assertTaskHandlerCapability overrides - Add Partial<TaskManagerOptions> to ClientOptions and ServerOptions - All source code compiles cleanly (39 errors remain in protocol.test.ts)
- Import RequestTaskStore from taskManager.ts instead of protocol.ts - CreateTaskRequestHandlerExtra.task is now required TaskContext - TaskRequestHandlerExtra.task is now required TaskContext with taskId
- Update protocol.test.ts to construct TaskManager and pass to Protocol - Remove assertTaskCapability/assertTaskHandlerCapability from test subclasses - Update McpServer to use extra.task?.taskStore instead of extra.taskStore - Update example server to use extra.task.taskStore/taskId/requestedTtl - Update core task interfaces for TaskContext grouping All packages now typecheck cleanly with zero errors.
…emain) Update TestProtocol interface and all test references to access task state through protocol.tasks instead of directly on Protocol. Remaining 24 failures fall into 3 categories: 1. Progress + tasks (2): CreateTaskResult detection needs verification 2. requestStream (8): Tests create Protocol without TaskManager 3. Error handling (14): Tests need TaskManager for resolver access All are test-only issues - source code compiles and typechecks cleanly.
- Add TaskManager to requestStream test protocols (they need it for tasks API) - Fix IIFE closures wrongly matched by sed (restore to no-arg) - Add removeProgressHandler to TaskManagerHost interface - Fix prepareOutboundRequest to propagate queue overflow errors - Update requestTaskStore references to createRequestTaskStore - Fix optional chaining to non-null assertions in test casts 430/430 core tests pass. Zero type errors.
- Rewrite ExperimentalClientTasks to delegate to client.tasks (no more unsafe casts) - Rewrite ExperimentalServerTasks to delegate to server.tasks (no more unsafe casts) - Make taskStore optional in TaskManagerOptions (clients need TaskManager for outbound API) - Always create TaskManager in Client/Server (outbound API always available) - Fix extractInboundTaskContext to create context for task creation requests too - Update integration test references (extra.task!.taskStore etc.) Core: 430/430 pass. Integration: 475/487 pass (12 failures in message queuing/auto-poll).
- Fix handleAutomaticTaskPolling to create task store from server.tasks when extra.task is not set (auto-poll creates tasks internally) - Fix sendRequest wrapping to activate on taskContext (not just relatedTaskId) so task creation handlers can properly set input_required status Integration: 482/487 pass (5 failures in HTTP message queuing delivery).
Root cause: extractInboundTaskContext tried to create a RequestTaskStore even when the Protocol had no taskStore configured (client-side). When a client received a task-related elicitation request, the store creation threw, silently breaking the handler response flow. Fix: Only create taskContext (with RequestTaskStore) when _taskStore is configured. Clients receiving task-related requests (elicitation, sampling) still get relatedTaskId for metadata propagation but don't get a taskStore. Also: - Fix McpServer handleAutomaticTaskPolling to construct JSONRPCRequest for createRequestTaskStore - Remove debug logging Results: - Core: 430/430 pass - Integration: 487/487 pass - Type errors: 0
- Remove duplicate exports (taskManager.ts exported from both protocol.ts and index.ts). Keep only index.ts export, use type-only import in protocol.ts - Remove unused imports (RelatedTaskMetadata in taskManager.ts, RequestTaskStore in interfaces.ts) - Fix import sorting and formatting (auto-fixed by eslint/prettier) - Remove isTaskAugmentedRequestParams unused import from protocol.ts
- Add await to _clearTaskQueue calls in handleGetTaskPayload and handleCancelTask (was fire-and-forget, errors became unhandled rejections) - Remove unused relatedRequestId parameter from sendOnResponseStream callback (Protocol already captures extra.requestId in closure) - Remove dead assertTaskCapability/assertTaskHandlerCapability overrides from protocolTransportHandling.test.ts
- Move assertTaskCapability/assertTaskHandlerCapability from mutable properties on TaskManager to TaskManagerOptions constructor params, eliminating fragile post-construction wiring - Add _requireHost getter with clear error message instead of _host! non-null assertions throughout TaskManager - Extract duplicated relatedTask metadata injection in Protocol.notification() into _buildJsonRpcNotification() helper - Add @internal JSDoc to TaskManagerHost (will be removed from public exports in PR #1447) - Add changeset for the breaking changes - Revert unrelated pnpm-workspace.yaml change
felixweinberger
force-pushed
the
fweinberger/extract-tasks-from-protocol
branch
from
c3104d6 to
b783fcd
Compare
The string-based setRequestHandler API (from #1446) requires methods to be in the RequestTypeMap. Replace custom test method names (test/longRunning, test/method, etc.) with 'tools/call' which is a valid protocol method.
Tests that register handlers with empty params need to use 'ping' instead of 'tools/call' since the string-based setRequestHandler API validates against CallToolRequestSchema which requires params.name.
…le methods Replace 12 fine-grained conditional delegation calls in Protocol with 4 coarse-grained lifecycle methods on TaskManager, and introduce NullTaskManager (Null Object pattern) to eliminate all conditional checks. Protocol integration points: - _onrequest: 6 conditionals → 1 processInboundRequest() call - _onresponse: 2 conditionals → 1 processInboundResponse() call - request(): 2 conditionals → 1 processOutboundRequest() call - notification(): 1 conditional → 1 processOutboundNotification() call - _onclose: 1 conditional → 1 unconditional onClose() call NullTaskManager returns neutral/pass-through values, so Protocol always calls TaskManager methods without checking if tasks are enabled. No changes to external API surface (client.tasks / server.tasks still return TaskManager | undefined).
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Extracts all task-related logic from
protocol.tsinto a dedicatedTaskManagerclass, reducing Protocol from 1,667 to 913 lines (-45%). Uses a Null Object pattern (NullTaskManager) so Protocol has zero conditional task checks.Motivation and Context
protocol.tshad grown to 1,667 lines, with ~70% being task-related code spread across every major method (request(),_onrequest(),_onresponse(),notification(), constructor). This created:RequestHandlerExtrahad 3 task-specific optional fields (taskId,taskStore,taskRequestedTtl) that wereundefinedfor all non-task handlersExperimentalClientTasks/ExperimentalServerTasksusedas unknown ascasts to access Protocol's protected methodsprotocol.ts↔experimental/tasks/interfaces.tsHow Has This Been Tested?
Breaking Changes
extra.taskId→extra.task?.taskIdextra.taskStore→extra.task?.taskStoreextra.taskRequestedTtl→extra.task?.requestedTtlProtocolOptionsno longer acceptstaskStore/taskMessageQueue— pass viaTaskManagerOptionsinClientOptions/ServerOptionsassertTaskCapability/assertTaskHandlerCapabilityremoved from ProtocolTypes of changes
Key design decisions
TaskManageris owned by Protocol, not a subclassNullTaskManagerprovides no-op implementations of all lifecycle methods, eliminating everyif (this._taskManager)conditional in ProtocolprocessInboundRequest(),processOutboundRequest(),processInboundResponse(),processOutboundNotification(), andonClose()— each consolidating multiple fine-grained operationsTaskManagerOptions.taskStoreis optional: clients need TaskManager for outbound API (getTask,requestStream) but not for storing tasksRequestHandlerExtra.task?groups task fields cleanly — non-task handlers seeundefinedclient.tasks/server.tasksprovides clean public API, eliminating unsafe casts in experimental wrappersChecklist
Additional context
if (taskManager)checksas unknown ascastsRequestHandlerExtratask fieldstask?