diff --git a/client/src/App.tsx b/client/src/App.tsx index 7cf6d751a..bd66b73b9 100644 --- a/client/src/App.tsx +++ b/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(null); + const latestTaskStatusRef = useRef>(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((resolve) => { + taskPollingWakeRef.current = { taskId, resolve }; + }), + ]); + } + if (taskPollingWakeRef.current?.taskId === taskId) { + taskPollingWakeRef.current = null; + } const taskStatus = await sendMCPRequest( { diff --git a/client/src/__tests__/App.taskPolling.test.tsx b/client/src/__tests__/App.taskPolling.test.tsx index 1eff4951c..65977dff3 100644 --- a/client/src/__tests__/App.taskPolling.test.tsx +++ b/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; + }); + + render(); + + 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 }, + ); + }); });