From 829ad3acba67e542795917432d67a2f535e3d3a0 Mon Sep 17 00:00:00 2001 From: Zain Dana Harper Date: Sun, 19 Jul 2026 20:50:08 -0700 Subject: [PATCH] Add cancel control for task polling --- client/src/App.tsx | 143 ++++++++++++- client/src/__tests__/App.taskPolling.test.tsx | 198 ++++++++++++++++++ client/src/components/ToolsTab.tsx | 130 +++++++----- 3 files changed, 417 insertions(+), 54 deletions(-) diff --git a/client/src/App.tsx b/client/src/App.tsx index 7cf6d751a..9cced7248 100644 --- a/client/src/App.tsx +++ b/client/src/App.tsx @@ -110,6 +110,13 @@ type PrefilledAppsToolCall = { result: CompatibilityCallToolResult; }; +type ActivePollingTask = { + taskId: string; + runId: number; + cancelled: boolean; + wakePollingDelay?: () => void; +}; + const hasAppResourceUri = (tool: Tool): boolean => { return Boolean(getToolUiResourceUri(tool)); }; @@ -307,6 +314,7 @@ const App = () => { const [selectedTool, setSelectedTool] = useState(null); const [selectedTask, setSelectedTask] = useState(null); const [isPollingTask, setIsPollingTask] = useState(false); + const [pollingTaskId, setPollingTaskId] = useState(null); const [nextResourceCursor, setNextResourceCursor] = useState< string | undefined >(); @@ -320,6 +328,8 @@ const App = () => { const [nextTaskCursor, setNextTaskCursor] = useState(); const progressTokenRef = useRef(0); const prefilledAppsToolCallIdRef = useRef(0); + const activePollingTaskRef = useRef(null); + const pollingRunIdRef = useRef(0); const [activeTab, setActiveTab] = useState(() => { const hash = window.location.hash.slice(1); @@ -1106,8 +1116,15 @@ const App = () => { if (runAsTask && isTaskResult(response)) { const taskId = response.task.taskId; const pollInterval = response.task.pollInterval; + const pollingRunId = ++pollingRunIdRef.current; + activePollingTaskRef.current = { + taskId, + runId: pollingRunId, + cancelled: false, + }; // Set polling state BEFORE setting tool result for proper UI update setIsPollingTask(true); + setPollingTaskId(taskId); // Safely extract any _meta from the original response (if present) const initialResponseMeta = response && @@ -1129,12 +1146,41 @@ const App = () => { }; setToolResult(latestToolResult); + const isPollingCancelled = () => + activePollingTaskRef.current?.runId !== pollingRunId || + activePollingTaskRef.current.cancelled; + + const waitForNextPoll = () => + new Promise((resolve) => { + const timeoutId = setTimeout(() => { + const activePollingTask = activePollingTaskRef.current; + if (activePollingTask?.runId === pollingRunId) { + activePollingTask.wakePollingDelay = undefined; + } + resolve(); + }, pollInterval); + + const activePollingTask = activePollingTaskRef.current; + if (activePollingTask?.runId === pollingRunId) { + activePollingTask.wakePollingDelay = () => { + clearTimeout(timeoutId); + activePollingTask.wakePollingDelay = undefined; + resolve(); + }; + } + }); + // Polling loop let taskCompleted = false; while (!taskCompleted) { try { - // Wait for 1 second before polling - await new Promise((resolve) => setTimeout(resolve, pollInterval)); + // Wait for the server-provided poll interval before polling. + await waitForNextPoll(); + + if (isPollingCancelled()) { + taskCompleted = true; + break; + } const taskStatus = await sendMCPRequest( { @@ -1144,6 +1190,11 @@ const App = () => { GetTaskResultSchema, ); + if (isPollingCancelled()) { + taskCompleted = true; + break; + } + if ( taskStatus.status === "completed" || taskStatus.status === "failed" || @@ -1163,6 +1214,10 @@ const App = () => { }, CompatibilityCallToolResultSchema, ); + if (isPollingCancelled()) { + taskCompleted = true; + break; + } console.log(`Result received for task ${taskId}:`, result); latestToolResult = result as CompatibilityCallToolResult; setToolResult(latestToolResult); @@ -1219,6 +1274,10 @@ const App = () => { }, CompatibilityCallToolResultSchema, ); + if (isPollingCancelled()) { + taskCompleted = true; + break; + } void listTasks(); } else { latestToolResult = { @@ -1239,6 +1298,10 @@ const App = () => { } } } catch (pollingError) { + if (isPollingCancelled()) { + taskCompleted = true; + break; + } console.error("Error polling task status:", pollingError); latestToolResult = { content: [ @@ -1253,9 +1316,13 @@ const App = () => { taskCompleted = true; } } - setIsPollingTask(false); - // Clear any validation errors since tool execution completed - setErrors((prev) => ({ ...prev, tools: null })); + if (activePollingTaskRef.current?.runId === pollingRunId) { + activePollingTaskRef.current = null; + setIsPollingTask(false); + setPollingTaskId(null); + // Clear any validation errors since tool execution completed + setErrors((prev) => ({ ...prev, tools: null })); + } return latestToolResult; } else { const directResult = response as CompatibilityCallToolResult; @@ -1312,6 +1379,70 @@ const App = () => { } }; + const cancelTaskPolling = useCallback(async () => { + const activePollingTask = activePollingTaskRef.current; + if (!activePollingTask) return; + + try { + const response = await cancelMcpTask(activePollingTask.taskId); + const currentActivePollingTask = activePollingTaskRef.current; + if (currentActivePollingTask?.runId !== activePollingTask.runId) { + return; + } + + currentActivePollingTask.cancelled = true; + currentActivePollingTask.wakePollingDelay?.(); + activePollingTaskRef.current = null; + setIsPollingTask(false); + setPollingTaskId(null); + setTasks((prev) => { + const exists = prev.some((t) => t.taskId === response.taskId); + return exists + ? prev.map((t) => (t.taskId === response.taskId ? response : t)) + : [response, ...prev]; + }); + if (selectedTaskRef.current?.taskId === response.taskId) { + setSelectedTask(response); + } + setToolResult({ + content: [ + { + type: "text", + text: `Task cancelled: ${response.statusMessage || response.taskId}`, + }, + ], + isError: true, + _meta: { + "io.modelcontextprotocol/related-task": { + taskId: response.taskId, + }, + }, + }); + setErrors((prev) => ({ ...prev, tools: null })); + void listTasks(); + } catch (e) { + const message = (e as Error).message ?? String(e); + setErrors((prev) => ({ + ...prev, + tools: message, + })); + setToolResult({ + content: [ + { + type: "text", + text: `Error cancelling task: ${message}`, + }, + ], + isError: true, + _meta: { + "io.modelcontextprotocol/related-task": { + taskId: activePollingTask.taskId, + }, + }, + }); + } + }, [cancelMcpTask, listTasks]); + const handleRootsChange = async () => { await sendNotification({ method: "notifications/roots/list_changed" }); }; @@ -1652,6 +1783,8 @@ const App = () => { }} toolResult={toolResult} isPollingTask={isPollingTask} + pollingTaskId={pollingTaskId} + cancelTaskPolling={cancelTaskPolling} nextCursor={nextToolCursor} error={errors.tools} resourceContent={resourceContentMap} diff --git a/client/src/__tests__/App.taskPolling.test.tsx b/client/src/__tests__/App.taskPolling.test.tsx index 1eff4951c..f39cd51ed 100644 --- a/client/src/__tests__/App.taskPolling.test.tsx +++ b/client/src/__tests__/App.taskPolling.test.tsx @@ -119,6 +119,9 @@ jest.mock("../components/ToolsTab", () => ({ default: ({ callTool, toolResult, + isPollingTask, + pollingTaskId, + cancelTaskPolling, }: { callTool: ( name: string, @@ -127,6 +130,9 @@ jest.mock("../components/ToolsTab", () => ({ runAsTask?: boolean, ) => Promise; toolResult: { content: Array<{ type: string; text: string }> } | null; + isPollingTask?: boolean; + pollingTaskId?: string | null; + cancelTaskPolling?: () => Promise; }) => (
+ {isPollingTask && pollingTaskId && cancelTaskPolling && ( + + )} {toolResult && (
{toolResult.content.map((c, i) => ( @@ -333,4 +349,186 @@ describe("App - task polling with input_required status", () => { ); expect(resultCalls).toHaveLength(1); }); + + it("lets users cancel the active polling task from the Tools tab", async () => { + const taskId = "task-cancel-789"; + const cancelTask = jest.fn(async () => ({ + taskId, + status: "cancelled", + statusMessage: "Cancelled by user", + })); + 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: 50 }, + }; + } + + if (request.method === "tasks/list") { + return { tasks: [] }; + } + + throw new Error( + `Unexpected method after cancellation: ${request.method}`, + ); + }); + + mockUseConnection.mockReturnValue({ + 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, + 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 })); + fireEvent.click( + await screen.findByRole("button", { name: /cancel polling task/i }), + ); + + await waitFor(() => { + expect(cancelTask).toHaveBeenCalledWith(taskId); + }); + await waitFor(() => { + expect(screen.getByTestId("tool-result")).toHaveTextContent( + "Task cancelled: Cancelled by user", + ); + }); + + await new Promise((resolve) => setTimeout(resolve, 75)); + const taskGetCalls = makeRequest.mock.calls.filter( + ([req]) => (req as { method: string }).method === "tasks/get", + ); + expect(taskGetCalls).toHaveLength(0); + }); + + it("does not let an in-flight polling response overwrite a cancelled task", async () => { + const taskId = "task-cancel-in-flight"; + let resolveTaskGet: + ((value: { taskId: string; status: string }) => void) | undefined; + const taskGetResponse = new Promise<{ taskId: string; status: string }>( + (resolve) => { + resolveTaskGet = resolve; + }, + ); + const cancelTask = jest.fn(async () => ({ + taskId, + status: "cancelled", + statusMessage: "Cancelled by user", + })); + 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: 0 }, + }; + } + + if (request.method === "tasks/get") { + return taskGetResponse; + } + + if (request.method === "tasks/result") { + return { + content: [{ type: "text", text: "late final result" }], + }; + } + + if (request.method === "tasks/list") { + return { tasks: [] }; + } + + throw new Error(`Unexpected method: ${request.method}`); + }); + + mockUseConnection.mockReturnValue({ + 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, + 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 })); + const cancelButton = await screen.findByRole("button", { + name: /cancel polling task/i, + }); + await waitFor(() => { + const taskGetCalls = makeRequest.mock.calls.filter( + ([req]) => (req as { method: string }).method === "tasks/get", + ); + expect(taskGetCalls).toHaveLength(1); + }); + fireEvent.click(cancelButton); + + await waitFor(() => { + expect(screen.getByTestId("tool-result")).toHaveTextContent( + "Task cancelled: Cancelled by user", + ); + }); + + resolveTaskGet?.({ taskId, status: "completed" }); + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(screen.getByTestId("tool-result")).toHaveTextContent( + "Task cancelled: Cancelled by user", + ); + const resultCalls = makeRequest.mock.calls.filter( + ([req]) => (req as { method: string }).method === "tasks/result", + ); + expect(resultCalls).toHaveLength(0); + }); }); diff --git a/client/src/components/ToolsTab.tsx b/client/src/components/ToolsTab.tsx index 15b85fd67..c23f09c6f 100644 --- a/client/src/components/ToolsTab.tsx +++ b/client/src/components/ToolsTab.tsx @@ -173,6 +173,8 @@ const ToolsTab = ({ setSelectedTool, toolResult, isPollingTask, + pollingTaskId, + cancelTaskPolling, nextCursor, error, resourceContent, @@ -193,6 +195,8 @@ const ToolsTab = ({ setSelectedTool: (tool: Tool | null) => void; toolResult: CompatibilityCallToolResult | null; isPollingTask?: boolean; + pollingTaskId?: string | null; + cancelTaskPolling?: () => Promise; nextCursor: ListToolsResult["nextCursor"]; error: string | null; resourceContent: Record; @@ -203,6 +207,7 @@ const ToolsTab = ({ const [params, setParams] = useState>({}); const [runAsTask, setRunAsTask] = useState(false); const [isToolRunning, setIsToolRunning] = useState(false); + const [isCancellingPollingTask, setIsCancellingPollingTask] = useState(false); const [isOutputSchemaExpanded, setIsOutputSchemaExpanded] = useState(false); const [isMetadataExpanded, setIsMetadataExpanded] = useState(false); const [metadataEntries, setMetadataEntries] = useState< @@ -805,58 +810,85 @@ const ToolsTab = ({
)} - + {isPollingTask && pollingTaskId && cancelTaskPolling && ( + )} - +