Skip to content

fix: allow cancelling task polling #1399

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
qybaihe wants to merge 1 commit into
base: main
Choose a base branch
from
+312 −69
Open

fix: allow cancelling task polling #1399

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
225 changes: 157 additions & 68 deletions client/src/App.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,50 @@ const cloneToolParams = (
}
};

const createPollingCancelledError = (): Error | DOMException => {
if (typeof DOMException !== "undefined") {
return new DOMException("Polling cancelled", "AbortError");
}

const error = new Error("Polling cancelled");
error.name = "AbortError";
return error;
};

const isAbortError = (error: unknown): boolean => {
return (
typeof error === "object" &&
error !== null &&
"name" in error &&
(error as { name?: string }).name === "AbortError"
);
};

const waitForNextTaskPoll = (
pollInterval: number,
signal: AbortSignal,
): Promise<void> => {
return new Promise<void>((resolve, reject) => {
if (signal.aborted) {
reject(createPollingCancelledError());
return;
}

function onAbort() {
window.clearTimeout(timeoutId);
signal.removeEventListener("abort", onAbort);
reject(createPollingCancelledError());
}

const timeoutId = window.setTimeout(() => {
signal.removeEventListener("abort", onAbort);
resolve();
}, pollInterval);

signal.addEventListener("abort", onAbort, { once: true });
});
};

const filterReservedMetadata = (
metadata: Record<string, string>,
): Record<string, string> => {
Expand Down Expand Up @@ -304,6 +348,7 @@ const App = () => {
const [selectedTool, setSelectedTool] = useState<Tool | null>(null);
const [selectedTask, setSelectedTask] = useState<Task | null>(null);
const [isPollingTask, setIsPollingTask] = useState(false);
const pollingAbortControllerRef = useRef<AbortController | null>(null);
const [nextResourceCursor, setNextResourceCursor] = useState<
string | undefined
>();
Expand All @@ -318,6 +363,16 @@ const App = () => {
const progressTokenRef = useRef(0);
const prefilledAppsToolCallIdRef = useRef(0);

const cancelToolPolling = useCallback(() => {
pollingAbortControllerRef.current?.abort();
}, []);

useEffect(() => {
return () => {
pollingAbortControllerRef.current?.abort();
};
}, []);

const [activeTab, setActiveTab] = useState<string>(() => {
const hash = window.location.hash.slice(1);
const initialTab = hash || "resources";
Expand Down Expand Up @@ -839,9 +894,10 @@ const App = () => {
request: ClientRequest,
schema: T,
tabKey?: keyof typeof errors,
options?: Parameters<typeof makeRequest>[2],
): Promise<SchemaOutput<T>> => {
try {
const response = await makeRequest(request, schema);
const response = await makeRequest(request, schema, options);
if (tabKey !== undefined) {
clearError(tabKey);
}
Expand Down Expand Up @@ -1070,6 +1126,8 @@ const App = () => {
if (runAsTask && isTaskResult(response)) {
const taskId = response.task.taskId;
const pollInterval = response.task.pollInterval;
const abortController = new AbortController();
pollingAbortControllerRef.current = abortController;
// Set polling state BEFORE setting tool result for proper UI update
setIsPollingTask(true);
// Safely extract any _meta from the original response (if present)
Expand All @@ -1095,100 +1153,130 @@ const App = () => {

// Polling loop
let taskCompleted = false;
while (!taskCompleted) {
try {
// Wait for 1 second before polling
await new Promise((resolve) => setTimeout(resolve, pollInterval));

const taskStatus = await sendMCPRequest(
{
method: "tasks/get",
params: { taskId },
},
GetTaskResultSchema,
);
try {
while (!taskCompleted && !abortController.signal.aborted) {
try {
await waitForNextTaskPoll(pollInterval, abortController.signal);

if (
taskStatus.status === "completed" ||
taskStatus.status === "failed" ||
taskStatus.status === "cancelled"
) {
taskCompleted = true;
console.log(
`Polling complete for task ${taskId}: ${taskStatus.status}`,
const taskStatus = await sendMCPRequest(
{
method: "tasks/get",
params: { taskId },
},
GetTaskResultSchema,
undefined,
{ signal: abortController.signal },
);

if (taskStatus.status === "completed") {
console.log(`Fetching result for task ${taskId}`);
const result = await sendMCPRequest(
{
method: "tasks/result",
params: { taskId },
},
CompatibilityCallToolResultSchema,
if (
taskStatus.status === "completed" ||
taskStatus.status === "failed" ||
taskStatus.status === "cancelled"
) {
taskCompleted = true;
console.log(
`Polling complete for task ${taskId}: ${taskStatus.status}`,
);
console.log(`Result received for task ${taskId}:`, result);
latestToolResult = result as CompatibilityCallToolResult;
setToolResult(latestToolResult);

// Refresh tasks list to show completed state
void listTasks();
if (taskStatus.status === "completed") {
console.log(`Fetching result for task ${taskId}`);
const result = await sendMCPRequest(
{
method: "tasks/result",
params: { taskId },
},
CompatibilityCallToolResultSchema,
undefined,
{ signal: abortController.signal },
);
console.log(`Result received for task ${taskId}:`, result);
latestToolResult = result as CompatibilityCallToolResult;
setToolResult(latestToolResult);

// Refresh tasks list to show completed state
void listTasks();
} else {
latestToolResult = {
content: [
{
type: "text",
text: `Task ${taskStatus.status}: ${taskStatus.statusMessage || "No additional information"}`,
},
],
isError: true,
};
setToolResult(latestToolResult);
// Refresh tasks list to show failed/cancelled state
void listTasks();
}
} else {
// Update status message while polling
// Safely extract any _meta from the original response (if present)
const pollingResponseMeta =
response &&
typeof response === "object" &&
"_meta" in (response as Record<string, unknown>)
? ((response as { _meta?: Record<string, unknown> })
._meta ?? {})
: undefined;
latestToolResult = {
content: [
{
type: "text",
text: `Task ${taskStatus.status}: ${taskStatus.statusMessage || "No additional information"}`,
text: `Task status: ${taskStatus.status}${taskStatus.statusMessage ? ` - ${taskStatus.statusMessage}` : ""}. Polling...`,
},
],
isError: true,
_meta: {
...(pollingResponseMeta || {}),
"io.modelcontextprotocol/related-task": { taskId },
},
};
setToolResult(latestToolResult);
// Refresh tasks list to show failed/cancelled state
// Refresh tasks list to show progress
void listTasks();
}
} else {
// Update status message while polling
// Safely extract any _meta from the original response (if present)
const pollingResponseMeta =
response &&
typeof response === "object" &&
"_meta" in (response as Record<string, unknown>)
? ((response as { _meta?: Record<string, unknown> })._meta ??
{})
: undefined;
} catch (pollingError) {
if (
abortController.signal.aborted ||
isAbortError(pollingError)
) {
latestToolResult = {
content: [
{
type: "text",
text: `Polling cancelled for task ${taskId}. The task may still be running on the server.`,
},
],
_meta: {
...(initialResponseMeta || {}),
"io.modelcontextprotocol/related-task": { taskId },
},
};
setToolResult(latestToolResult);
taskCompleted = true;
break;
}

console.error("Error polling task status:", pollingError);
latestToolResult = {
content: [
{
type: "text",
text: `Task status: ${taskStatus.status}${taskStatus.statusMessage ? ` - ${taskStatus.statusMessage}` : ""}. Polling...`,
text: `Error polling task status: ${pollingError instanceof Error ? pollingError.message : String(pollingError)}`,
},
],
_meta: {
...(pollingResponseMeta || {}),
"io.modelcontextprotocol/related-task": { taskId },
},
isError: true,
};
setToolResult(latestToolResult);
// Refresh tasks list to show progress
void listTasks();
taskCompleted = true;
}
} catch (pollingError) {
console.error("Error polling task status:", pollingError);
latestToolResult = {
content: [
{
type: "text",
text: `Error polling task status: ${pollingError instanceof Error ? pollingError.message : String(pollingError)}`,
},
],
isError: true,
};
setToolResult(latestToolResult);
taskCompleted = true;
}
} finally {
setIsPollingTask(false);
if (pollingAbortControllerRef.current === abortController) {
pollingAbortControllerRef.current = null;
}
}
setIsPollingTask(false);
// Clear any validation errors since tool execution completed
setErrors((prev) => ({ ...prev, tools: null }));
return latestToolResult;
Expand Down Expand Up @@ -1587,6 +1675,7 @@ const App = () => {
}}
toolResult={toolResult}
isPollingTask={isPollingTask}
cancelPolling={cancelToolPolling}
nextCursor={nextToolCursor}
error={errors.tools}
resourceContent={resourceContentMap}
Expand Down
Loading