fix(everything): keep async task tools within their client task (#4986, #4987)

.changeset/everything-async-task-tools.md

@@ -0,0 +1,5 @@
+---
+"@modelcontextprotocol/server-everything": patch
+---
+
+Fix the async task tools: `trigger-elicitation-request-async` now stops polling before the 10-minute TTL it requests for the client's task runs out, instead of polling past it and failing on a task the client has expired (#4986); and `trigger-sampling-request-async` reports the status message of a client task that is already finished when it is created, instead of "No message" (#4987). `trigger-elicitation-request-async` had the same flaw and gets the same fix.

src/everything/__tests__/tasks.test.ts

@@ -422,18 +422,17 @@ describe("trigger-sampling-request-async", () => {
     );
   });
 
-  // KNOWN BUG #4987: trigger-sampling-request-async drops the status message of a task that is already finished when created; the fix changes this assertion.
-  it("does not poll a client task that is already finished when it is created", async () => {
-    // Characterization: the status message is only read from polls, so a
-    // task that fails before it is returned is reported with "No message".
+  it("reports the status message of a client task that is already finished when it is created", async () => {
+    // A task that fails before it is returned is never polled, so its status
+    // message comes from the CreateTaskResult.
     const { taskStore } = await connectTaskClient((_m, taskId) => {
       void taskStore.updateTaskStatus(taskId!, "failed", "instant failure");
     });
     const [text] = await callAsync("trigger-sampling-request-async", {
       prompt: "p",
     });
     expect(text).toMatch(
-      /^\[FAILED\] No message\n\nProgress:\nTask created: \S+$/,
+      /^\[FAILED\] instant failure\n\nProgress:\nTask created: \S+$/,
     );
   });
 
@@ -569,6 +568,20 @@ describe("trigger-elicitation-request-async", () => {
     expect(polls).toHaveLength(4);
   });
 
+  it("reports the status message of a client task that is already finished when it is created", async () => {
+    // A task that fails before it is returned is never polled, so its status
+    // message comes from the CreateTaskResult.
+    const { taskStore } = await connectTaskClient((_m, taskId) => {
+      void taskStore.updateTaskStatus(taskId!, "failed", "instant failure");
+    });
+    const texts = await callAsync("trigger-elicitation-request-async");
+    expect(texts).toEqual([
+      expect.stringMatching(
+        /^\[FAILED\] instant failure\n\nProgress:\nTask created: \S+$/,
+      ),
+    ]);
+  });
+
   it("reports a failed client task", async () => {
     const { taskStore } = await connectTaskClient((_m, taskId) => {
       later(() => taskStore.updateTaskStatus(taskId!, "failed"));
@@ -577,25 +590,54 @@ describe("trigger-elicitation-request-async", () => {
     expect(texts[0]).toMatch(/^\[FAILED\] No message\n\nProgress:\n/);
   });
 
-  it("gives up after 600 polls of a task that never finishes", async () => {
+  /** Also fake `performance`, which the tool's polling deadline is measured with. */
+  function fakeClock() {
+    vi.useFakeTimers({
+      toFake: ["setTimeout", "clearTimeout", "performance"],
+    });
+  }
+
+  it("gives up 5 seconds before the 10-minute TTL of a task that never finishes", async () => {
+    fakeClock();
     // A client that keeps its task longer than the 10-minute TTL asked for.
     await connectTaskClient(() => {}, 60_000);
     const texts = await callAsync("trigger-elicitation-request-async");
     expect(texts[0]).toMatch(
-      /^\[TIMEOUT\] Task timed out after 600 poll attempts\n\nProgress:\n/,
+      /^\[TIMEOUT\] Task timed out after 595 poll attempts, before its 10-minute TTL expired\n\nProgress:\n/,
     );
   });
 
-  // KNOWN BUG #4986: trigger-elicitation-request-async polls for longer than the 10-minute TTL it requests for its task; the fix changes this assertion.
-  it("fails on the last poll instead when the client honors the 10-minute TTL", async () => {
-    // Characterization: 600 one-second polls outlast the 600000 ms TTL the
-    // tool asks for, so a client that expires the task on time answers the
-    // last tasks/get with "not found" and the call fails before the timeout
-    // report.
+  it("times out before a client that honors the 10-minute TTL expires the task", async () => {
+    fakeClock();
+    let created = "";
+    const { taskStore } = await connectTaskClient((_m, taskId) => {
+      created = taskId!;
+    });
+    const texts = await callAsync("trigger-elicitation-request-async");
+    // The last poll reached the task while the client still held it.
+    expect(texts).toHaveLength(1);
+    expect(texts[0]).toMatch(
+      /^\[TIMEOUT\] Task timed out after 595 poll attempts, before its 10-minute TTL expired\n\nProgress:\nTask created: \S+\nPoll 1: working\n/,
+    );
+    expect(texts[0]).toMatch(/\nPoll 595: working$/);
+    // The client still holds the task when the tool gives up, 5 s before the
+    // 600000 ms TTL it asked for runs out.
+    expect((await taskStore.getTask(created))?.ttl).toBe(600000);
+  });
+
+  it("times out without polling when its wait ends past the deadline", async () => {
+    // A timer that fires late: the first reading sets the deadline, and every
+    // later one is already past it.
+    let readings = 0;
+    vi.spyOn(performance, "now").mockImplementation(() =>
+      readings++ === 0 ? 0 : 600_000,
+    );
     await connectTaskClient(() => {});
     const texts = await callAsync("trigger-elicitation-request-async");
     expect(texts).toEqual([
-      "MCP error -32602: MCP error -32602: Failed to retrieve task: Task not found",
+      expect.stringMatching(
+        /^\[TIMEOUT\] Task timed out after 0 poll attempts, before its 10-minute TTL expired\n\nProgress:\nTask created: \S+$/,
+      ),
     ]);
   });
 

src/everything/tools/trigger-elicitation-request-async.ts

@@ -22,8 +22,12 @@ const config = {
 // Poll interval in milliseconds
 const POLL_INTERVAL = 1000;
 
-// Maximum poll attempts before timeout (10 minutes for user input)
-const MAX_POLL_ATTEMPTS = 600;
+// Time-to-live requested for the client's task (10 minutes for user input)
+const TASK_TTL = 600000;
+
+// Stop polling this long before the TTL runs out, so the last tasks/get
+// reaches the client while it still holds the task
+const TTL_SAFETY_MARGIN = 5000;
 
 /**
  * Registers the 'trigger-elicitation-request-async' tool.
@@ -59,13 +63,19 @@ export const registerTriggerElicitationRequestAsyncTool = (
       name,
       config,
       async (args, extra): Promise<CallToolResult> => {
+        // Polling must end before the client may expire the task. The client
+        // starts the TTL when it creates the task, after this point, so a
+        // deadline measured from here is conservative. performance.now() is
+        // monotonic, so a wall-clock adjustment cannot stretch the deadline.
+        const pollDeadline = performance.now() + TASK_TTL - TTL_SAFETY_MARGIN;
+
         // Create the elicitation request WITH task metadata
         // Using z.any() schema to avoid complex type matching with _meta
         const request = {
           method: "elicitation/create" as const,
           params: {
             task: {
-              ttl: 600000, // 10 minutes (user input may take a while)
+              ttl: TASK_TTL, // 10 minutes (user input may take a while)
             },
             message:
               "Please provide inputs for the following fields (async task demo):",
@@ -143,16 +153,26 @@ export const registerTriggerElicitationRequestAsyncTool = (
         // Poll for task completion
         let attempts = 0;
         let taskStatus = elicitResponse.task.status;
-        let taskStatusMessage: string | undefined;
+        // A task that is already finished when created is never polled, so
+        // its status message has to come from the CreateTaskResult
+        let taskStatusMessage = elicitResponse.task.statusMessage;
+        let timedOut = false;
 
         while (
           taskStatus !== "completed" &&
           taskStatus !== "failed" &&
-          taskStatus !== "cancelled" &&
-          attempts < MAX_POLL_ATTEMPTS
+          taskStatus !== "cancelled"
         ) {
           // Wait before polling
           await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL));
+
+          // Give up rather than poll past the TTL the task was created with.
+          // Checked after the wait, so a timer that fires late cannot slip a
+          // poll in after the deadline.
+          if (performance.now() > pollDeadline) {
+            timedOut = true;
+            break;
+          }
           attempts++;
 
           // Get task status from client
@@ -185,12 +205,14 @@ export const registerTriggerElicitationRequestAsyncTool = (
         }
 
         // Check for timeout
-        if (attempts >= MAX_POLL_ATTEMPTS) {
+        if (timedOut) {
           return {
             content: [
               {
                 type: "text",
-                text: `[TIMEOUT] Task timed out after ${MAX_POLL_ATTEMPTS} poll attempts\n\nProgress:\n${statusMessages.join(
+                text: `[TIMEOUT] Task timed out after ${attempts} poll attempts, before its ${
+                  TASK_TTL / 60000
+                }-minute TTL expired\n\nProgress:\n${statusMessages.join(
                   "\n",
                 )}`,
               },

src/everything/tools/trigger-sampling-request-async.ts

@@ -147,7 +147,9 @@ export const registerTriggerSamplingRequestAsyncTool = (server: McpServer) => {
         // Poll for task completion
         let attempts = 0;
         let taskStatus = samplingResponse.task.status;
-        let taskStatusMessage: string | undefined;
+        // A task that is already finished when created is never polled, so
+        // its status message has to come from the CreateTaskResult
+        let taskStatusMessage = samplingResponse.task.statusMessage;
 
         while (
           taskStatus !== "completed" &&