Wake task polling on terminal notifications

client/src/App.tsx

@@ -102,6 +102,11 @@ import {
 import MetadataTab from "./components/MetadataTab";
 
 const CONFIG_LOCAL_STORAGE_KEY = "inspectorConfig_v1";
+const TERMINAL_TASK_STATUSES: Task["status"][] = [
+  "completed",
+  "failed",
+  "cancelled",
+];
 
 type PrefilledAppsToolCall = {
   id: number;
@@ -370,6 +375,11 @@ const App = () => {
   } = useDraggableSidebar(320);
 
   const selectedTaskRef = useRef<Task | null>(null);
+  const latestTaskStatusRef = useRef<Map<string, Task["status"]>>(new Map());
+  const taskPollingWakeRef = useRef<{
+    taskId: string;
+    resolve: () => void;
+  } | null>(null);
   useEffect(() => {
     selectedTaskRef.current = selectedTask;
   }, [selectedTask]);
@@ -410,6 +420,7 @@ const App = () => {
 
       if (notification.method === "notifications/tasks/status") {
         const task = notification.params as unknown as Task;
+        latestTaskStatusRef.current.set(task.taskId, task.status);
         setTasks((prev) => {
           const exists = prev.some((t) => t.taskId === task.taskId);
           if (exists) {
@@ -421,6 +432,13 @@ const App = () => {
         if (selectedTaskRef.current?.taskId === task.taskId) {
           setSelectedTask(task);
         }
+        if (
+          TERMINAL_TASK_STATUSES.includes(task.status) &&
+          taskPollingWakeRef.current?.taskId === task.taskId
+        ) {
+          taskPollingWakeRef.current.resolve();
+          taskPollingWakeRef.current = null;
+        }
       }
     },
     onPendingRequest: (request, resolve, reject) => {
@@ -1133,8 +1151,22 @@ const App = () => {
         let taskCompleted = false;
         while (!taskCompleted) {
           try {
-            // Wait for 1 second before polling
-            await new Promise((resolve) => setTimeout(resolve, pollInterval));
+            const latestNotifiedStatus =
+              latestTaskStatusRef.current.get(taskId);
+            if (
+              !latestNotifiedStatus ||
+              !TERMINAL_TASK_STATUSES.includes(latestNotifiedStatus)
+            ) {
+              await Promise.race([
+                new Promise((resolve) => setTimeout(resolve, pollInterval)),
+                new Promise<void>((resolve) => {
+                  taskPollingWakeRef.current = { taskId, resolve };
+                }),
+              ]);
+            }
+            if (taskPollingWakeRef.current?.taskId === taskId) {
+              taskPollingWakeRef.current = null;
+            }
 
             const taskStatus = await sendMCPRequest(
               {

client/src/__tests__/App.taskPolling.test.tsx

@@ -1,4 +1,10 @@
-import { fireEvent, render, screen, waitFor } from "@testing-library/react";
+import {
+  act,
+  fireEvent,
+  render,
+  screen,
+  waitFor,
+} from "@testing-library/react";
 import "@testing-library/jest-dom";
 import App from "../App";
 import { useConnection } from "../lib/hooks/useConnection";
@@ -333,4 +339,96 @@ describe("App - task polling with input_required status", () => {
     );
     expect(resultCalls).toHaveLength(1);
   });
+
+  it("wakes task polling immediately when a terminal task status notification arrives", async () => {
+    const taskId = "task-notified-789";
+    let capturedOnNotification:
+      | ((notification: { method: string; params?: unknown }) => void)
+      | undefined;
+
+    const makeRequest = jest.fn(async (request: { method: string }) => {
+      if (request.method === "tools/list") {
+        return {
+          tools: [
+            {
+              name: "myTool",
+              inputSchema: { type: "object", properties: {} },
+            },
+          ],
+          nextCursor: undefined,
+        };
+      }
+
+      if (request.method === "tools/call") {
+        return {
+          task: { taskId, status: "working", pollInterval: 30000 },
+        };
+      }
+
+      if (request.method === "tasks/get") {
+        return { taskId, status: "completed" };
+      }
+
+      if (request.method === "tasks/result") {
+        return {
+          content: [{ type: "text", text: "notified task result" }],
+        };
+      }
+
+      if (request.method === "tasks/list") {
+        return { tasks: [] };
+      }
+
+      throw new Error(`Unexpected method: ${request.method}`);
+    });
+
+    mockUseConnection.mockImplementation((options) => {
+      capturedOnNotification =
+        options.onNotification as typeof capturedOnNotification;
+      return {
+        connectionStatus: "connected",
+        serverCapabilities: { tools: { listChanged: true } },
+        serverImplementation: null,
+        mcpClient: {
+          request: jest.fn(),
+          notification: jest.fn(),
+          close: jest.fn(),
+        } as unknown as Client,
+        requestHistory: [],
+        clearRequestHistory: jest.fn(),
+        makeRequest,
+        cancelTask: jest.fn(),
+        listTasks: jest.fn(),
+        sendNotification: jest.fn(),
+        handleCompletion: jest.fn(),
+        completionsSupported: false,
+        connect: jest.fn(),
+        disconnect: jest.fn(),
+      } as ReturnType<typeof useConnection>;
+    });
+
+    render(<App />);
+
+    fireEvent.click(screen.getByRole("button", { name: /run task tool/i }));
+
+    await waitFor(() => {
+      expect(capturedOnNotification).toBeDefined();
+    });
+
+    act(() => {
+      capturedOnNotification?.({
+        method: "notifications/tasks/status",
+        params: { taskId, status: "completed" },
+      });
+    });
+
+    await waitFor(
+      () => {
+        expect(screen.getByTestId("tool-result")).toHaveTextContent(
+          "notified task result",
+        );
+      },
+      { timeout: 500 },
+    );
+  });
 });