fixup! simplify: remove OutboundMiddleware/useOutbound; StreamDriver calls TaskManager directly

felixweinberger · felixweinberger · commit 32669c3916c4 · 2026-04-21T20:31:00.000Z
Invert the dependency per review: StreamDriver is now fully task-agnostic.
- StreamDriver exposes generic per-request `intercept` hook (RequestOptions) and
  per-connection `onresponse` tap returning {consumed, preserveProgress}.
- TaskManager.sendRequest/sendNotification helpers wrap Outbound.request/notification,
  threading processOutboundRequest through `intercept`.
- Protocol/McpServer/Client wire TaskManager via these helpers + the onresponse tap.
- McpServer request-shaped Outbound also honours `intercept`.
diff --git a/packages/client/src/client/client.ts b/packages/client/src/client/client.ts
@@ -282,11 +282,14 @@ export class Client extends Dispatcher<ClientContext> {
         if (isChannelTransport(transport)) {
             const driverOpts: StreamDriverOptions = {
                 supportedProtocolVersions: this._supportedProtocolVersions,
-                debouncedNotificationMethods: this._options?.debouncedNotificationMethods,
-                taskManager: this._taskManager
+                debouncedNotificationMethods: this._options?.debouncedNotificationMethods
             };
             this._ct = channelAsClientTransport(transport, this, driverOpts);
-            this._ct.driver!.onclose = () => this.onclose?.();
+            this._ct.driver!.onresponse = (r, id) => this._taskManager.processInboundResponse(r, id);
+            this._ct.driver!.onclose = () => {
+                this._taskManager.onClose();
+                this.onclose?.();
+            };
             this._ct.driver!.onerror = e => this.onerror?.(e);
             const skipInit = transport.sessionId !== undefined;
             if (skipInit) {
diff --git a/packages/core/src/shared/context.ts b/packages/core/src/shared/context.ts
@@ -8,6 +8,8 @@ import type {
     ElicitRequestURLParams,
     ElicitResult,
     JSONRPCMessage,
+    JSONRPCRequest,
+    JSONRPCResultResponse,
     LoggingLevel,
     Notification,
     Progress,
@@ -148,6 +150,20 @@ export type RequestOptions = {
      * If provided, associates this request with a related task.
      */
     relatedTask?: RelatedTaskMetadata;
+
+    /**
+     * @internal Called by the channel adapter after the wire request is built
+     * (id assigned, `_meta.progressToken` set, response/progress handlers registered)
+     * but before send. Return `true` to skip the send; the caller takes ownership of
+     * delivering `wire` later. The registered handlers stay live. `settle` injects
+     * a result/error into the registered response handler out-of-band.
+     */
+    intercept?: (
+        wire: JSONRPCRequest,
+        messageId: number,
+        settle: (response: JSONRPCResultResponse | Error) => void,
+        onError: (error: unknown) => void
+    ) => boolean;
 } & TransportSendOptions;
 
 /**
diff --git a/packages/core/src/shared/protocol.ts b/packages/core/src/shared/protocol.ts
@@ -162,12 +162,13 @@ export abstract class Protocol<ContextT extends BaseContext> {
         const driver = new StreamDriver(this._dispatcher, transport, {
             supportedProtocolVersions: this._supportedProtocolVersions,
             debouncedNotificationMethods: this._options?.debouncedNotificationMethods,
-            buildEnv: (extra, base) => ({ ...base, _transportExtra: extra }),
-            taskManager: this._ownTaskManager
+            buildEnv: (extra, base) => ({ ...base, _transportExtra: extra })
         });
         this._outbound = driver;
+        driver.onresponse = (r, id) => this._ownTaskManager.processInboundResponse(r, id);
         driver.onclose = () => {
             if (this._outbound === driver) this._outbound = undefined;
+            this._ownTaskManager.onClose();
             this.onclose?.();
         };
         driver.onerror = error => this.onerror?.(error);
@@ -215,7 +216,7 @@ export abstract class Protocol<ContextT extends BaseContext> {
         if (this._options?.enforceStrictCapabilities === true) {
             this.assertCapabilityForMethod(request.method as RequestMethod);
         }
-        return this._outbound.request(request, resultSchema, options);
+        return this._ownTaskManager.sendRequest(request, resultSchema, options, this._outbound);
     }
 
     /**
@@ -226,7 +227,7 @@ export abstract class Protocol<ContextT extends BaseContext> {
             throw new SdkError(SdkErrorCode.NotConnected, 'Not connected');
         }
         this.assertNotificationCapability(notification.method as NotificationMethod);
-        return this._outbound.notification(notification, options);
+        return this._ownTaskManager.sendNotification(notification, options, this._outbound);
     }
 
     // ───────────────────────────────────────────────────────────────────────
diff --git a/packages/core/src/shared/streamDriver.ts b/packages/core/src/shared/streamDriver.ts
@@ -28,7 +28,6 @@ import { parseSchema } from '../util/schema.js';
 import type { NotificationOptions, Outbound, ProgressCallback, RequestEnv, RequestOptions } from './context.js';
 import { DEFAULT_REQUEST_TIMEOUT_MSEC } from './context.js';
 import type { Dispatcher } from './dispatcher.js';
-import type { TaskManager } from './taskManager.js';
 import type { AttachOptions, Transport } from './transport.js';
 
 type TimeoutInfo = {
@@ -48,13 +47,6 @@ export type StreamDriverOptions = {
      * {@linkcode MessageExtraInfo} (e.g. auth, http req).
      */
     buildEnv?: (extra: MessageExtraInfo | undefined, base: RequestEnv) => RequestEnv;
-    /**
-     * Optional {@linkcode TaskManager}. When provided, the driver calls its
-     * `processOutbound*` / `processInboundResponse` / `onClose` hooks at the
-     * request-correlation seam. Inbound task processing happens via the
-     * TaskManager's dispatch middleware (registered by the caller), not here.
-     */
-    taskManager?: TaskManager;
 };
 
 /**
@@ -73,10 +65,18 @@ export class StreamDriver implements Outbound {
     private _pendingDebouncedNotifications = new Set<string>();
     private _closed = false;
     private _supportedProtocolVersions: string[];
-    private _taskManager?: TaskManager;
 
     onclose?: () => void;
     onerror?: (error: Error) => void;
+    /**
+     * Tap for every inbound response. Return `consumed: true` to claim it (suppresses the
+     * matched-handler dispatch / unknown-id error). Return `preserveProgress: true` to keep
+     * the progress handler registered after the matched handler runs. Set by the owner.
+     */
+    onresponse?: (
+        response: JSONRPCResponse | JSONRPCErrorResponse,
+        messageId: number
+    ) => { consumed: boolean; preserveProgress?: boolean };
 
     constructor(
         // eslint-disable-next-line @typescript-eslint/no-explicit-any -- driver is context-agnostic; subclass owns ContextT
@@ -85,7 +85,6 @@ export class StreamDriver implements Outbound {
         private _options: StreamDriverOptions = {}
     ) {
         this._supportedProtocolVersions = _options.supportedProtocolVersions ?? SUPPORTED_PROTOCOL_VERSIONS;
-        this._taskManager = _options.taskManager;
     }
 
     /** {@linkcode Outbound.removeProgressHandler}. */
@@ -208,31 +207,19 @@ export class StreamDriver implements Outbound {
                 options?.resetTimeoutOnProgress ?? false
             );
 
-            let queued = false;
-            if (this._taskManager) {
-                const sideChannelResponse = (resp: JSONRPCResultResponse | Error) => {
-                    const h = this._responseHandlers.get(messageId);
-                    if (h) h(resp);
-                    else this._onerror(new Error(`Response handler missing for side-channeled request ${messageId}`));
-                };
-                try {
-                    queued = this._taskManager.processOutboundRequest(jsonrpcRequest, options, messageId, sideChannelResponse, error => {
-                        this._progressHandlers.delete(messageId);
-                        reject(error);
-                    }).queued;
-                } catch (error) {
+            if (options?.intercept) {
+                const settle = (r: JSONRPCResultResponse | Error) => this._responseHandlers.get(messageId)?.(r);
+                const onError = (e: unknown) => {
                     this._progressHandlers.delete(messageId);
-                    reject(error);
-                    return;
-                }
+                    reject(e);
+                };
+                if (options.intercept(jsonrpcRequest, messageId, settle, onError)) return;
             }
 
-            if (!queued) {
-                this.pipe.send(jsonrpcRequest, { relatedRequestId, resumptionToken, onresumptiontoken }).catch(error => {
-                    this._progressHandlers.delete(messageId);
-                    reject(error);
-                });
-            }
+            this.pipe.send(jsonrpcRequest, { relatedRequestId, resumptionToken, onresumptiontoken }).catch(error => {
+                this._progressHandlers.delete(messageId);
+                reject(error);
+            });
         }).finally(() => {
             if (onAbort) options?.signal?.removeEventListener('abort', onAbort);
             if (cleanupId !== undefined) {
@@ -246,9 +233,8 @@ export class StreamDriver implements Outbound {
      * Sends a notification over the pipe. Supports debouncing per the constructor option.
      */
     async notification(notification: Notification, options?: NotificationOptions): Promise<void> {
-        const taskResult = await this._taskManager?.processOutboundNotification(notification, options);
-        if (taskResult?.queued || this._closed) return;
-        const jsonrpc: JSONRPCNotification = taskResult?.jsonrpcNotification ?? {
+        if (this._closed) return;
+        const jsonrpc: JSONRPCNotification = {
             jsonrpc: '2.0',
             method: notification.method,
             params: notification.params
@@ -345,19 +331,16 @@ export class StreamDriver implements Outbound {
 
     private _onresponse(response: JSONRPCResponse | JSONRPCErrorResponse): void {
         const messageId = Number(response.id);
-        const taskResult = this._taskManager?.processInboundResponse(response, messageId);
-        if (taskResult?.consumed) return;
-
+        const tap = this.onresponse?.(response, messageId);
+        if (tap?.consumed) return;
         const handler = this._responseHandlers.get(messageId);
         if (handler === undefined) {
             this._onerror(new Error(`Received a response for an unknown message ID: ${JSON.stringify(response)}`));
             return;
         }
         this._responseHandlers.delete(messageId);
         this._cleanupTimeout(messageId);
-        if (!taskResult?.preserveProgress) {
-            this._progressHandlers.delete(messageId);
-        }
+        if (!tap?.preserveProgress) this._progressHandlers.delete(messageId);
         if (isJSONRPCResultResponse(response)) {
             handler(response);
         } else {
@@ -370,7 +353,6 @@ export class StreamDriver implements Outbound {
         const responseHandlers = this._responseHandlers;
         this._responseHandlers = new Map();
         this._progressHandlers.clear();
-        this._taskManager?.onClose();
         this._pendingDebouncedNotifications.clear();
         for (const info of this._timeoutInfo.values()) clearTimeout(info.timeoutId);
         this._timeoutInfo.clear();
@@ -439,12 +421,12 @@ export async function attachChannelTransport(
     const driver = new StreamDriver(dispatcher, pipe, {
         supportedProtocolVersions: options?.supportedProtocolVersions,
         debouncedNotificationMethods: options?.debouncedNotificationMethods,
-        taskManager: options?.taskManager,
         buildEnv: options?.buildEnv
     });
-    if (options?.onclose || options?.onerror) {
+    if (options?.onclose || options?.onerror || options?.onresponse) {
         driver.onclose = options.onclose;
         driver.onerror = options.onerror;
+        driver.onresponse = options.onresponse;
     }
     await driver.start();
     return driver;
diff --git a/packages/core/src/shared/taskManager.ts b/packages/core/src/shared/taskManager.ts
@@ -384,7 +384,7 @@ export class TaskManager {
     private _outboundRequest<T extends AnySchema>(req: Request, schema: T, opts?: RequestOptions): Promise<SchemaOutput<T>> {
         const ch = this._requireHooks.channel();
         if (!ch) throw new ProtocolError(ProtocolErrorCode.InternalError, 'Not connected');
-        return ch.request(req, schema, opts);
+        return this.sendRequest(req, schema, opts, ch);
     }
 
     // -- Public API (client-facing) --
@@ -804,7 +804,44 @@ export class TaskManager {
         };
     }
 
-    // -- Outbound-seam lifecycle methods (called directly by StreamDriver) --
+    // -- Outbound helpers (called by McpServer/Client/Protocol before delegating to Outbound) --
+
+    /**
+     * Task-aware request send: routes through {@linkcode RequestOptions.intercept} so the
+     * channel adapter builds the wire (id/progressToken/handlers) and TaskManager decides
+     * whether to queue it. Use this where instance-level outbound requests are made
+     * (Protocol/McpServer/Client), so the channel adapter stays task-agnostic.
+     */
+    sendRequest<T extends AnySchema>(
+        request: Request,
+        resultSchema: T,
+        options: RequestOptions | undefined,
+        outbound: Outbound
+    ): Promise<SchemaOutput<T>> {
+        if (!options?.relatedTask && !options?.task) {
+            return outbound.request(request, resultSchema, options);
+        }
+        return outbound.request(request, resultSchema, {
+            ...options,
+            intercept: (wire, messageId, settle, onError) =>
+                this.processOutboundRequest(wire, options, messageId, settle, onError).queued
+        });
+    }
+
+    /**
+     * Task-aware notification send: queues when `options.relatedTask` is set, otherwise
+     * delegates to `outbound.notification()` with related-task metadata attached.
+     */
+    async sendNotification(notification: Notification, options: NotificationOptions | undefined, outbound: Outbound): Promise<void> {
+        const result = await this.processOutboundNotification(notification, options);
+        if (result.queued) return;
+        await outbound.notification(
+            result.jsonrpcNotification
+                ? { method: result.jsonrpcNotification.method, params: result.jsonrpcNotification.params }
+                : notification,
+            options
+        );
+    }
 
     processOutboundRequest(
         jsonrpcRequest: JSONRPCRequest,
diff --git a/packages/core/src/shared/transport.ts b/packages/core/src/shared/transport.ts
@@ -8,7 +8,6 @@ import type {
     RequestId
 } from '../types/index.js';
 import type { RequestEnv } from './context.js';
-import type { TaskManager } from './taskManager.js';
 
 export type FetchLike = (url: string | URL, init?: RequestInit) => Promise<Response>;
 
@@ -163,10 +162,14 @@ export type Transport = ChannelTransport;
 export type AttachOptions = {
     supportedProtocolVersions?: string[];
     debouncedNotificationMethods?: string[];
-    taskManager?: TaskManager;
     buildEnv?: (extra: MessageExtraInfo | undefined, base: RequestEnv) => RequestEnv;
     onclose?: () => void;
     onerror?: (error: Error) => void;
+    /** Tap for every inbound response. See {@linkcode StreamDriver.onresponse}. */
+    onresponse?: (
+        response: JSONRPCResultResponse | JSONRPCErrorResponse,
+        messageId: number
+    ) => { consumed: boolean; preserveProgress?: boolean };
 };
 
 /**
diff --git a/packages/server/src/server/mcpServer.ts b/packages/server/src/server/mcpServer.ts
@@ -20,6 +20,7 @@ import type {
     JSONRPCErrorResponse,
     JSONRPCMessage,
     JSONRPCNotification,
+    JSONRPCResultResponse,
     JSONRPCRequest,
     JSONRPCResponse,
     JsonSchemaType,
@@ -311,36 +312,40 @@ export class McpServer extends Dispatcher<ServerContext> implements RegistriesHo
             outbound = {
                 close: () => transport.close(),
                 notification: transport.notify
-                    ? async (n, opts) => {
-                          const out = await this._taskManager.processOutboundNotification(n, opts);
-                          if (!out.queued)
-                              await transport.notify!(out.jsonrpcNotification ?? ({ jsonrpc: '2.0', ...n } as JSONRPCNotification));
-                      }
+                    ? async (n, _opts) => transport.notify!({ jsonrpc: '2.0', ...n } as JSONRPCNotification)
                     : noOutbound('notifications'),
                 request: transport.request
-                    ? async (r, schema, _opts) => {
-                          const id = this._nextOutboundId++;
-                          const resp = await transport.request!({
-                              jsonrpc: '2.0',
-                              id,
-                              method: r.method,
-                              params: r.params
-                          } as JSONRPCRequest);
-                          if ('error' in resp) throw ProtocolError.fromError(resp.error.code, resp.error.message, resp.error.data);
-                          const parsed = parseSchema(schema, resp.result);
-                          if (!parsed.success) throw parsed.error;
-                          return parsed.data as SchemaOutput<typeof schema>;
-                      }
+                    ? (r, schema, opts) =>
+                          new Promise((resolve, reject) => {
+                              const id = this._nextOutboundId++;
+                              const wire = { jsonrpc: '2.0', id, method: r.method, params: r.params } as JSONRPCRequest;
+                              const finish = (resp: JSONRPCResultResponse | Error) => {
+                                  if (resp instanceof Error) return reject(resp);
+                                  const parsed = parseSchema(schema, resp.result);
+                                  if (!parsed.success) return reject(parsed.error);
+                                  resolve(parsed.data as SchemaOutput<typeof schema>);
+                              };
+                              if (opts?.intercept?.(wire, id, finish, reject)) return;
+                              transport
+                                  .request!(wire)
+                                  .then(resp =>
+                                      'error' in resp
+                                          ? reject(ProtocolError.fromError(resp.error.code, resp.error.message, resp.error.data))
+                                          : finish(resp)
+                                  )
+                                  .catch(reject);
+                          })
                     : noOutbound('requests')
             };
         } else {
             outbound = await attachChannelTransport(transport, this, {
                 supportedProtocolVersions: this._supportedProtocolVersions,
                 debouncedNotificationMethods: this._options?.debouncedNotificationMethods,
                 buildEnv: (extra, base) => ({ ...base, _transportExtra: extra }),
-                taskManager: this._taskManager,
+                onresponse: (r, id) => this._taskManager.processInboundResponse(r, id),
                 onclose: () => {
                     if (this._outbound === outbound) this._outbound = undefined;
+                    this._taskManager.onClose();
                     this.onclose?.();
                 },
                 onerror: e => this.onerror?.(e)
@@ -594,7 +599,7 @@ export class McpServer extends Dispatcher<ServerContext> implements RegistriesHo
         if (this._options?.enforceStrictCapabilities === true) {
             assertCapabilityForMethod(req.method as RequestMethod, this._clientCapabilities);
         }
-        return this._requireOutbound().request(req, schema as never, options) as Promise<T>;
+        return this._taskManager.sendRequest(req, schema as never, options, this._requireOutbound()) as Promise<T>;
     }
 
     /**
@@ -748,7 +753,7 @@ export class McpServer extends Dispatcher<ServerContext> implements RegistriesHo
      */
     async notification(notification: Notification, options?: NotificationOptions): Promise<void> {
         assertNotificationCapability(notification.method as NotificationMethod, this._capabilities, this._clientCapabilities);
-        await this._outbound?.notification(notification, options);
+        if (this._outbound) await this._taskManager.sendNotification(notification, options, this._outbound);
     }
 
     async sendLoggingMessage(params: LoggingMessageNotification['params'], sessionId?: string): Promise<void> {
