From 234e620ad2de981002be31067da68e1d95ef2d2f Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Sun, 6 Sep 2026 16:49:20 +0800 Subject: [PATCH] fix(mcp): pause active timeout during elicitation --- packages/plugins/mcp/src/sdk/invoke.test.ts | 144 +++++++++++++++- packages/plugins/mcp/src/sdk/invoke.ts | 180 ++++++++++++++++---- 2 files changed, 290 insertions(+), 34 deletions(-) diff --git a/packages/plugins/mcp/src/sdk/invoke.test.ts b/packages/plugins/mcp/src/sdk/invoke.test.ts index ee9100fb6b..2585776d2e 100644 --- a/packages/plugins/mcp/src/sdk/invoke.test.ts +++ b/packages/plugins/mcp/src/sdk/invoke.test.ts @@ -1,12 +1,15 @@ import { beforeAll, describe, expect, it } from "@effect/vitest"; import { Effect, Predicate } from "effect"; import { HttpServerResponse } from "effect/unstable/http"; +// oxlint-disable-next-line executor/no-vitest-import -- boundary: fake-clock coverage for the active-work deadline +import { afterEach, vi } from "vitest"; import { ProtocolError, SdkErrorCode, SdkHttpError, type OAuthClientProvider, + type ClientContext, } from "@modelcontextprotocol/client"; import { ElicitationResponse } from "@executor-js/sdk"; import { serveTestHttpApp } from "@executor-js/sdk/testing"; @@ -19,7 +22,7 @@ import { createMcpConnector, type McpConnection, type McpConnector } from "./con // that precondition here — these tests construct SDK errors directly. beforeAll(() => loadMcpClientSdk()); import { McpInvocationError, McpOAuthReauthorizationRequired } from "./errors"; -import { invokeMcpTool } from "./invoke"; +import { invokeMcpTool, makeActiveWorkDeadline, MCP_ACTIVE_WORK_TIMEOUT_MS } from "./invoke"; const acceptAll = () => Effect.succeed(ElicitationResponse.make({ action: "accept" })); @@ -148,6 +151,145 @@ const invocationRejectionCases = [ ]; describe("invokeMcpTool", () => { + afterEach(() => vi.useRealTimers()); + + it("pauses the active-work deadline across overlapping elicitations", () => { + vi.useFakeTimers({ toFake: ["Date", "setTimeout", "clearTimeout"] }); + const deadline = makeActiveWorkDeadline(100); + + vi.advanceTimersByTime(40); + deadline.pause(); + deadline.pause(); + vi.advanceTimersByTime(1_000); + expect(deadline.signal.aborted).toBe(false); + + deadline.resume(); + vi.advanceTimersByTime(100); + expect(deadline.signal.aborted).toBe(false); + + deadline.resume(); + vi.advanceTimersByTime(59); + expect(deadline.signal.aborted).toBe(false); + vi.advanceTimersByTime(1); + expect(deadline.signal.aborted).toBe(true); + deadline.dispose(); + }); + + it("uses the active signal for a tool call and excludes elicitation from its deadline", async () => { + vi.useFakeTimers({ toFake: ["Date", "setTimeout", "clearTimeout"] }); + + let requestHandler: + | ((request: { params: unknown }, context: ClientContext) => Promise) + | undefined; + let callOptions: { signal: AbortSignal; timeout: number } | undefined; + let finishElicitation: (() => void) | undefined; + let resolveElicitationStarted: (() => void) | undefined; + const elicitationStarted = new Promise((resolve) => { + resolveElicitationStarted = resolve; + }); + const connectionAbort = new AbortController(); + + const client = { + setRequestHandler: (_method: string, handler: unknown) => { + requestHandler = handler as typeof requestHandler; + }, + callTool: async (_request: unknown, options: { signal: AbortSignal; timeout: number }) => { + callOptions = options; + await requestHandler!( + { + params: { mode: "form", message: "Approve?", requestedSchema: {} }, + }, + { mcpReq: { signal: connectionAbort.signal } } as ClientContext, + ); + // oxlint-disable-next-line executor/no-promise-reject -- boundary: fake MCP client models SDK abort rejection + return await new Promise((_resolve, reject) => { + // oxlint-disable-next-line executor/no-promise-reject -- boundary: fake MCP client models SDK abort rejection + options.signal.addEventListener("abort", () => reject(options.signal.reason), { + once: true, + }); + }); + }, + }; + + const invocation = Effect.runPromise( + invokeMcpTool({ + toolId: "slow", + toolName: "slow", + args: {}, + transport: "streamable-http", + connector: Effect.succeed({ + // oxlint-disable-next-line executor/no-double-cast -- boundary: minimal fake MCP client implements only invokeMcpTool's surface + client: client as unknown as McpConnection["client"], + close: () => Promise.resolve(), + }), + elicit: () => + Effect.callback((resume) => { + resolveElicitationStarted!(); + finishElicitation = () => + resume(Effect.succeed(ElicitationResponse.make({ action: "accept" }))); + }), + }), + ).then( + () => "completed" as const, + () => "failed" as const, + ); + + await elicitationStarted; + expect(callOptions?.timeout).toBeGreaterThan(MCP_ACTIVE_WORK_TIMEOUT_MS); + vi.advanceTimersByTime(MCP_ACTIVE_WORK_TIMEOUT_MS); + expect(callOptions?.signal.aborted).toBe(false); + + finishElicitation!(); + await Promise.resolve(); + await Promise.resolve(); + vi.advanceTimersByTime(MCP_ACTIVE_WORK_TIMEOUT_MS); + expect(callOptions?.signal.aborted).toBe(true); + expect(await invocation).toBe("failed"); + }); + + it("interrupts an elicitation when the MCP connection closes", async () => { + let requestHandler: + | ((request: { params: unknown }, context: ClientContext) => Promise) + | undefined; + const connectionAbort = new AbortController(); + const client = { + setRequestHandler: (_method: string, handler: unknown) => { + requestHandler = handler as typeof requestHandler; + }, + callTool: async () => { + await requestHandler!( + { + params: { mode: "form", message: "Approve?", requestedSchema: {} }, + }, + { mcpReq: { signal: connectionAbort.signal } } as ClientContext, + ); + return { content: [] }; + }, + }; + + const invocation = Effect.runPromise( + invokeMcpTool({ + toolId: "closed", + toolName: "closed", + args: {}, + transport: "streamable-http", + connector: Effect.succeed({ + // oxlint-disable-next-line executor/no-double-cast -- boundary: minimal fake MCP client implements only invokeMcpTool's surface + client: client as unknown as McpConnection["client"], + close: () => Promise.resolve(), + }), + elicit: () => Effect.callback(() => undefined), + }), + ).then( + () => "completed" as const, + () => "failed" as const, + ); + + await Promise.resolve(); + connectionAbort.abort(); + expect(await invocation).toBe("failed"); + }); + for (const testCase of invocationRejectionCases) { it.effect(testCase.name, () => Effect.gen(function* () { diff --git a/packages/plugins/mcp/src/sdk/invoke.ts b/packages/plugins/mcp/src/sdk/invoke.ts index 4b7a433c7c..a7596fc3ae 100644 --- a/packages/plugins/mcp/src/sdk/invoke.ts +++ b/packages/plugins/mcp/src/sdk/invoke.ts @@ -16,7 +16,7 @@ import { Cause, Effect, Exit, Option, Predicate, Schema } from "effect"; -import type { ProtocolError } from "@modelcontextprotocol/client"; +import type { ClientContext, ProtocolError } from "@modelcontextprotocol/client"; // SDK error classes come through the lazy loader; by the time a tool call can // fail, the connect path has always loaded the module (see client-module.ts). @@ -39,6 +39,95 @@ import { httpStatusFromCause, insufficientScopeFromCause } from "./http-status"; // Helpers // --------------------------------------------------------------------------- +/** + * The MCP SDK's default request timer measures wall-clock time. An elicitation + * is user work, so it must not consume the tool's active-work budget. The SDK + * still gets a long timer as a transport-level backstop; this controller owns + * the normal deadline and is paused while one or more elicitation handlers are + * waiting for input. + */ +export const MCP_ACTIVE_WORK_TIMEOUT_MS = 60_000; +const MCP_SDK_TIMEOUT_BACKSTOP_MS = 2_147_483_647; + +export type ActiveWorkDeadline = { + readonly signal: AbortSignal; + readonly pause: () => void; + readonly resume: () => void; + readonly dispose: () => void; +}; + +export const makeActiveWorkDeadline = ( + timeoutMs: number = MCP_ACTIVE_WORK_TIMEOUT_MS, +): ActiveWorkDeadline => { + const controller = new AbortController(); + let remainingMs = timeoutMs; + let pendingElicitations = 0; + let startedAt: number | undefined; + let timer: ReturnType | undefined; + + const stopTimer = (): void => { + if (timer === undefined || startedAt === undefined) return; + clearTimeout(timer); + timer = undefined; + remainingMs = Math.max(0, remainingMs - (Date.now() - startedAt)); + startedAt = undefined; + }; + + const abortForTimeout = (): void => { + timer = undefined; + startedAt = undefined; + // oxlint-disable-next-line executor/no-error-constructor -- boundary: AbortSignal consumers need a stable timeout reason + controller.abort(new Error("MCP tool invocation exceeded its active-work deadline")); + }; + + const startTimer = (): void => { + if (controller.signal.aborted || pendingElicitations > 0) return; + if (remainingMs <= 0) { + abortForTimeout(); + return; + } + startedAt = Date.now(); + timer = setTimeout(() => { + remainingMs = 0; + abortForTimeout(); + }, remainingMs); + }; + + startTimer(); + + return { + signal: controller.signal, + pause: () => { + pendingElicitations += 1; + if (pendingElicitations === 1) stopTimer(); + }, + resume: () => { + if (pendingElicitations === 0) return; + pendingElicitations -= 1; + if (pendingElicitations === 0) startTimer(); + }, + dispose: () => { + stopTimer(); + // oxlint-disable-next-line executor/no-error-constructor -- boundary: disposing the scoped signal must interrupt SDK work + controller.abort(new Error("MCP tool invocation was disposed")); + }, + }; +}; + +const abortOnSignals = (signals: readonly AbortSignal[]): Effect.Effect => + Effect.callback((resume) => { + // oxlint-disable-next-line executor/no-error-constructor -- boundary: an aborted MCP handler must reject its JSON-RPC response + const abort = () => resume(Effect.fail(new Error("MCP elicitation was cancelled"))); + if (signals.some((signal) => signal.aborted)) { + abort(); + return; + } + for (const signal of signals) signal.addEventListener("abort", abort, { once: true }); + return Effect.sync(() => { + for (const signal of signals) signal.removeEventListener("abort", abort); + }); + }); + const ArgsRecord = Schema.Record(Schema.String, Schema.Unknown); const decodeArgsRecord = Schema.decodeUnknownOption(ArgsRecord); @@ -156,36 +245,53 @@ const toElicitationRequest = (params: McpElicitParams): ElicitationRequest => { }); }; -const installElicitationHandler = (client: McpConnection["client"], elicit: Elicit): void => { - client.setRequestHandler("elicitation/create", async (request: { params: unknown }) => { - const params = decodeElicitParams(request.params); - const req = toElicitationRequest(params); - // Use runPromiseExit so we can inspect typed failures — `elicit` - // fails with `ElicitationDeclinedError` on decline/cancel, which - // we translate into the equivalent MCP elicit response instead of - // surfacing as a JSON-RPC error. - const exit = await Effect.runPromiseExit(elicit(req)); - if (Exit.isSuccess(exit)) { - const response = exit.value; - return { - action: response.action, - ...(response.action === "accept" && response.content - ? { content: decodeElicitContent(response.content) } - : {}), - }; - } - const failure = exit.cause.reasons.find(Cause.isFailReason); - if (failure) { - const err = failure.error; - if (Predicate.isTagged(err, "ElicitationDeclinedError")) { - const action = - Predicate.hasProperty(err, "action") && err.action === "cancel" ? "cancel" : "decline"; - return { action }; +const installElicitationHandler = ( + client: McpConnection["client"], + elicit: Elicit, + deadline: ActiveWorkDeadline, +): void => { + client.setRequestHandler( + "elicitation/create", + async (request: { params: unknown }, ctx: ClientContext) => { + const params = decodeElicitParams(request.params); + const req = toElicitationRequest(params); + deadline.pause(); + // Use runPromiseExit so we can inspect typed failures — `elicit` + // fails with `ElicitationDeclinedError` on decline/cancel, which + // we translate into the equivalent MCP elicit response instead of + // surfacing as a JSON-RPC error. + // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: MCP SDK request handlers are promise callbacks and must release the active-work lease + try { + const exit = await Effect.runPromiseExit( + Effect.raceFirst(elicit(req), abortOnSignals([ctx.mcpReq.signal, deadline.signal])), + ); + if (Exit.isSuccess(exit)) { + const response = exit.value; + return { + action: response.action, + ...(response.action === "accept" && response.content + ? { content: decodeElicitContent(response.content) } + : {}), + }; + } + const failure = exit.cause.reasons.find(Cause.isFailReason); + if (failure) { + const err = failure.error; + if (Predicate.isTagged(err, "ElicitationDeclinedError")) { + const action = + Predicate.hasProperty(err, "action") && err.action === "cancel" + ? "cancel" + : "decline"; + return { action }; + } + } + // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: MCP SDK async request handlers signal unexpected failures by rejecting + throw Cause.squash(exit.cause); + } finally { + deadline.resume(); } - } - // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: MCP SDK async request handlers signal unexpected failures by rejecting - throw Cause.squash(exit.cause); - }); + }, + ); }; // --------------------------------------------------------------------------- @@ -218,10 +324,18 @@ const useConnection = ( onToolListChanged: (() => void) | undefined, ): Effect.Effect => Effect.gen(function* () { - installElicitationHandler(connection.client, elicit); + const deadline = yield* Effect.acquireRelease( + Effect.sync(() => makeActiveWorkDeadline()), + (activeWork) => Effect.sync(activeWork.dispose), + ); + installElicitationHandler(connection.client, elicit, deadline); installToolListChangedHandler(connection.client, onToolListChanged); return yield* Effect.tryPromise({ - try: () => connection.client.callTool({ name: toolName, arguments: args }), + try: () => + connection.client.callTool( + { name: toolName, arguments: args }, + { signal: deadline.signal, timeout: MCP_SDK_TIMEOUT_BACKSTOP_MS }, + ), catch: (cause) => { if (Predicate.isTagged(cause, "McpOAuthReauthorizationRequired")) { return new McpOAuthReauthorizationRequired({ @@ -258,7 +372,7 @@ const useConnection = ( attributes: { "mcp.tool.name": toolName }, }), ); - }); + }).pipe(Effect.scoped); // --------------------------------------------------------------------------- // Public API