Add Tasks support (#1013)

* ### client/src/App.tsx

- Imports:
  - Added: `Task` and `GetTaskResultSchema` to the same import block.
  - UI Icons: Added `ListTodo` from `lucide-react` to the icon imports.
  - Components: Added `TasksTab` import beside `ToolsTab`.
  - Config utils: Added `getMCPTaskTtl` to the config utils import block.

- State additions:
  - `const [tasks, setTasks] = useState<Task[]>([]);`
  - Extended `errors` state to include a `tasks` key: `tasks: null,`
  - `const [selectedTask, setSelectedTask] = useState<Task | null>(null);`
  - `const [isPollingTask, setIsPollingTask] = useState(false);`
  - `const [nextTaskCursor, setNextTaskCursor] = useState<string | undefined>();`

- Hook: `useConnection({...})` return value usage extended
  - Added destructured functions: `cancelTask: cancelMcpTask` and `listTasks: listMcpTasks` from the custom hook.

- Notification handling:
  - In `onNotification`, added:
    - If `method === "notifications/tasks/list_changed"`, call `listTasks()` (voided).
    - If `method === "notifications/tasks/status"`, treat `notification.params` as a `Task` and update `tasks` state: replace if exists by `taskId`, otherwise prepend. Also update `selectedTask` if it’s the same `taskId`.

- Tab routing:
  - When computing valid tabs, added `tasks` when server declares capability: `...(serverCapabilities?.tasks ? ["tasks"] : []),`
  - When choosing a default tab, added a branch to fall back to `"tasks"` if neither resources/prompts/tools are present but tasks are.

- Effect for Tasks tab:
  - When `mcpClient` is connected and `activeTab === "tasks"`, invoke `listTasks()`.

- Tools → task-augmented calls integration in `callTool`:
  - Parameter signature supports `runAsTask?: boolean` (already present in this file), but now:
    - If `runAsTask` is true, augment the `tools/call` request’s `params` with a `task` object: `{ task: { ttl: getMCPTaskTtl(config) } }`.
  - Use a permissive result schema for tool call: `sendMCPRequest(request, z.any(), "tools")` to avoid version-mismatch schema issues.
  - Task reference detection introduced:
    - `isTaskResult` helper checks for a nested `task` object with `taskId` (i.e., `response.task.taskId`).
  - When task is detected:
    - Set `isPollingTask(true)`.
    - Immediately set a temporary `toolResult` that includes text content “Task created: … Polling for status…” and `_meta` with `"io.modelcontextprotocol/related-task": { taskId }`.
    - Start a polling loop:
      - Delay 1s between polls.
      - Call `tasks/get` with `GetTaskResultSchema` for status.
      - If status is `completed`: call `tasks/result` with `z.any()` to retrieve the final result and set it as `toolResult`; call `listTasks()`.
      - If status is `failed` or `cancelled`: set an error `toolResult` content that includes the status + `statusMessage`; call `listTasks()`.
      - Else (still running): update `toolResult` content with current `status`/`statusMessage` and preserve `_meta` related-task.
    - After loop, set `isPollingTask(false)`.
  - When not a task response, set `toolResult` directly from response (cast to `CompatibilityCallToolResult`).

- Tasks list + cancel helpers in App:
  - `listTasks`: uses `listMcpTasks(nextTaskCursor)` from the hook, updates `tasks`, `nextTaskCursor`, and clears `errors.tasks`.
  - `cancelTask`: calls `cancelMcpTask(taskId)`, updates `tasks` array by `taskId`, updates `selectedTask` if it matches, and clears `errors.tasks`.

- UI integration:
  - Added a `TabsTrigger` for “Tasks” with `<ListTodo />` icon, disabled unless server supports tasks.
  - Added `<TasksTab />` to the main `TabsContent` block, passing: `tasks`, `listTasks`, `clearTasks`, `cancelTask`, `selectedTask`, `setSelectedTask`, `error={errors.tasks}`, `nextCursor={nextTaskCursor}`.
  - Passed `isPollingTask={isPollingTask}` and `toolResult` into `ToolsTab` so the Tools tab can show the live “Polling Task…” state and block reruns while polling.

Note: The raw diff is long; the key hunks align with the above bullet points (imports, state, notifications, tab wiring, request augmentation, polling loop, UI additions).

---

### client/src/components/ToolsTab.tsx

- Props shape changed:
  - Added `isPollingTask?: boolean` prop in the destructured props and in the prop types.
  - The `callTool` callback signature is now `(name, params, metadata?, runAsTask?) => Promise<void>` (runAsTask added earlier; test updates elsewhere reflect this).

- Local state additions:
  - `const [runAsTask, setRunAsTask] = useState(false);`

- Reset behavior:
  - When switching tools (`useEffect` on `selectedTool`), reset `runAsTask(false)`.
  - When clearing the list in `ListPane.clearItems`, also call `setRunAsTask(false)`.

- UI additions:
  - New checkbox control block to toggle “Run as task”:
    - Checkbox `id="run-as-task"`, bound to `runAsTask`, with `onCheckedChange` → `setRunAsTask(checked)`.
    - Label “Run as task”.
  - Run button disabling conditions expanded to include `isPollingTask`.
  - Run button text shows spinner with conditional label:
    - If `isToolRunning || isPollingTask` → show spinner and text `isPollingTask ? "Polling Task..." : "Running..."`.

- Call invocation change:
  - When clicking “Run Tool”, the `callTool` is invoked with `(selectedTool.name, params, metadata?, runAsTask)`.

- ToolResults relay:
  - Passes `isPollingTask` to `<ToolResults />`.

---

### client/src/components/ToolResults.tsx

- Props shape changed:
  - Added optional prop: `isPollingTask?: boolean`.

- Task-running banner logic:
  - Extracts related task from the tool result’s `_meta["io.modelcontextprotocol/related-task"]` if present.
  - Computes `isTaskRunning` as `isPollingTask ||` a text-heuristic against `structuredResult.content` entries that contain text like “Polling” or “Task status”.
  - Header “Tool Result:” now conditionally shows:
    - `Error` (red) if `isError` is true, else
    - `Task Running` (yellow) if `isTaskRunning`, else
    - `Success` (green).

No other changes to validation or rendering of content blocks.

---

### client/src/components/TasksTab.tsx (new file)

- A brand new tab to list and inspect tasks.
- Key elements:
  - Imports `Task` type and multiple status icons.
  - `TaskStatusIcon` component maps task `status` to an icon and color.
  - Main `TasksTab` props: `tasks`, `listTasks`, `clearTasks`, `cancelTask`, `selectedTask`, `setSelectedTask`, `error`, `nextCursor`.
  - Left column (`ListPane`): lists tasks, shows status icon, `taskId`, `status`, and last update time; button text changes to “List More Tasks” if `nextCursor` present; disables button if no cursor and list non-empty.
  - Right column:
    - Shows error `Alert` if `error` prop provided.
    - If a task is selected: header with `Task Details`, a Cancel button when `status === "working"` (shows a spinner while cancelling), and a grid of task fields: Status (with colored label and icon), Last Updated, Created At, TTL (shows “Infinite” if `ttl === null`, otherwise shows numeric with `s` suffix), optional Status Message, and full task JSON via `JsonView`.
    - If no task is selected: centered empty state with a “Refresh Tasks” button.

---

### client/src/lib/hooks/useConnection.ts

- Imports added from `@modelcontextprotocol/sdk/types.js`:
  - `ListTasksResultSchema`, `CancelTaskResultSchema`, `TaskStatusNotificationSchema`.

- Client capabilities on `connect`:
  - Added `tasks: { list: {}, cancel: {} }` into the `clientCapabilities` passed to `new Client(...)`.

- Notification handling setup:
  - The hook’s notification schema registration now includes the `TaskStatusNotificationSchema` in the `setNotificationHandler` list so the app receives `notifications/tasks/status`.

- New hook functions:
  - `cancelTask(taskId: string)` sends `tasks/cancel` with `CancelTaskResultSchema`.
  - `listTasks(cursor?: string)` sends `tasks/list` with `ListTasksResultSchema`.

- Exports:
  - Returned object now includes `cancelTask` and `listTasks`.

---

### client/src/utils/configUtils.ts

- Added a new getter:
  - `export const getMCPTaskTtl = (config: InspectorConfig): number => { return config.MCP_TASK_TTL.value as number; };`

---

### client/src/lib/configurationTypes.ts

- `InspectorConfig` type extended with a new item:
  - `MCP_TASK_TTL: ConfigItem;`
  - Includes descriptive JSDoc about default TTL in milliseconds for newly created tasks.

---

### client/src/lib/constants.ts

- `DEFAULT_INSPECTOR_CONFIG` extended with a default for task TTL:
  - Key: `MCP_TASK_TTL`
  - Label: `"Task TTL"`
  - Description: `"Default Time-to-Live (TTL) in milliseconds for newly created tasks"`
  - Default `value: 60000`
  - `is_session_item: false`

---

### client/src/components/__tests__/ToolsTab.test.tsx

- Expectations updated due to new `callTool` signature (4th arg `runAsTask`). Everywhere the test asserts a `callTool` invocation, an additional trailing `false` argument was added to reflect the default state when the box isn’t checked.
  - Examples of added trailing `false` at various assertion points (line offsets from diff): after calls around prior lines 132, 157, 193, 236, 257, 279, 297, 818, 1082 (now passing 4 arguments: name, params, metadata-or-undefined, false).

---

### Additional notes

- The Tasks feature is wired end-to-end:
  - Client capability declaration (list/cancel)
  - Notification handler for `notifications/tasks/status`
  - Tools call augmentation with `task` `{ ttl }`
  - Polling loop using `tasks/get` and `tasks/result`
  - UI feedback in Tools and dedicated Tasks tab
  - Configurable TTL via new config item and getter

* Fix issue where displayed task on task tab was not updated when task completes, leaving the cancel task button available.

* Handle button the same when input_requred or working state

* ### What changed
- Added **receiver-side (client-side) Tasks support** to the Inspector client for task-augmented **server → client** requests.
- When the server sends `sampling/createMessage` or `elicitation/create` with `params.task`, the client now:
  - Creates a local task and immediately returns a `CreateTaskResult` (`{ task: ... }`).
  - Runs the underlying user-driven flow asynchronously.
  - Updates task state to `completed` / `failed` and makes the final payload available via `tasks/result`.
- Implemented server polling handlers backed by the local task store:
  - `tasks/get`
  - `tasks/result` (blocks until payload is ready)
  - `tasks/list`
  - `tasks/cancel`
- Advertises task-augmented request capabilities via `capabilities.tasks.requests.sampling.createMessage` and `capabilities.tasks.requests.elicitation.create` when the corresponding callbacks are provided.
- Emits best-effort `notifications/tasks/status` updates as local task status changes.

* Fix review requests.

* In the ToolsTab, when a tool call results in an elicitation request, the app automatically switches to the elicitation tab, and once submitted, switches back to the ToolsTab to view the result of the tool call.

However, when a tool call results in a sampling request, it does not switch automatically to the SamplingTab and when the sampling request form is submitted it does not automatically switch back to the ToolsTab to view the result of the tool call.

These changes make the sampling use case work like the elicitation one. Tests were added.
This commit is contained in:
Cliff Hall 2026-01-22 23:40:59 -05:00 • committed by GitHub
commit 7ef0971c81
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 1479 additions and 108 deletions

View file

@ -40,4 +40,9 @@ export type InspectorConfig = {
* Session token for authenticating with the MCP Proxy Server. This token is displayed in the proxy server console on startup.
*/
MCP_PROXY_AUTH_TOKEN: ConfigItem;
/**
* Default Time-to-Live (TTL) in milliseconds for newly created tasks.
*/
MCP_TASK_TTL: ConfigItem;
};

View file

@ -76,4 +76,11 @@ export const DEFAULT_INSPECTOR_CONFIG: InspectorConfig = {
value: "",
is_session_item: true,
},
MCP_TASK_TTL: {
label: "Task TTL",
description:
"Default Time-to-Live (TTL) in milliseconds for newly created tasks",
value: 60000,
is_session_item: false,
},
} as const;

View file

@ -1,10 +1,15 @@
import { renderHook, act } from "@testing-library/react";
import { useConnection } from "../useConnection";
import { z } from "zod";
import { z } from "zod/v3";
import {
ClientRequest,
CreateTaskResultSchema,
JSONRPCMessage,
} from "@modelcontextprotocol/sdk/types.js";
import type {
AnySchema,
SchemaOutput,
} from "@modelcontextprotocol/sdk/server/zod-compat.js";
import { DEFAULT_INSPECTOR_CONFIG, CLIENT_IDENTITY } from "../../constants";
import {
SSEClientTransportOptions,
@ -128,7 +133,7 @@ const mockDiscoverScopes = discoverScopes as jest.MockedFunction<
>;
describe("useConnection", () => {
const defaultProps = {
const defaultProps: Parameters<typeof useConnection>[0] = {
transportType: "sse" as const,
command: "",
args: "",
@ -164,8 +169,10 @@ describe("useConnection", () => {
test: z.string(),
});
const mockSchemaAny: AnySchema = mockSchema as unknown as AnySchema;
await act(async () => {
await result.current.makeRequest(mockRequest, mockSchema);
await result.current.makeRequest(mockRequest, mockSchemaAny);
});
expect(mockClient.request).toHaveBeenCalledWith(
@ -204,8 +211,10 @@ describe("useConnection", () => {
test: z.string(),
});
const mockSchemaAny: AnySchema = mockSchema as unknown as AnySchema;
await act(async () => {
await result.current.makeRequest(mockRequest, mockSchema, {
await result.current.makeRequest(mockRequest, mockSchemaAny, {
timeout: 1000,
maxTotalTimeout: 2000,
resetTimeoutOnProgress: false,
@ -224,8 +233,251 @@ describe("useConnection", () => {
});
});
describe("Receiver-side Tasks (task-augmented incoming requests)", () => {
beforeEach(() => {
jest.clearAllMocks();
});
test("declares tasks.requests.sampling.createMessage when onPendingRequest is provided", async () => {
const Client = jest.requireMock(
"@modelcontextprotocol/sdk/client/index.js",
).Client;
const propsWithPending = {
...defaultProps,
onPendingRequest: jest.fn(),
};
const { result } = renderHook(() => useConnection(propsWithPending));
await act(async () => {
await result.current.connect();
});
expect(Client).toHaveBeenCalledWith(
expect.any(Object),
expect.objectContaining({
capabilities: expect.objectContaining({
tasks: expect.objectContaining({
requests: expect.objectContaining({
sampling: expect.objectContaining({
createMessage: {},
}),
}),
}),
}),
}),
);
});
test("task-augmented sampling/createMessage returns { task } and tasks/result blocks until resolved", async () => {
let pendingResolve: ((value: unknown) => void) | undefined;
let pendingReject: ((reason?: unknown) => void) | undefined;
const mockOnPendingRequest = jest.fn((_request, resolve, reject) => {
pendingResolve = resolve;
pendingReject = reject;
});
const propsWithPending = {
...defaultProps,
onPendingRequest: mockOnPendingRequest,
};
const { result } = renderHook(() => useConnection(propsWithPending));
await act(async () => {
await result.current.connect();
});
const samplingRequest = {
method: "sampling/createMessage",
params: {
task: { ttl: 0 },
messages: [
{
role: "user",
content: { type: "text", text: "hello" },
},
],
maxTokens: 1,
},
};
// Locate the sampling/createMessage handler
const samplingHandlerCall = mockClient.setRequestHandler.mock.calls.find(
(call) => {
try {
const schema = call[0];
const parseResult =
schema.safeParse && schema.safeParse(samplingRequest);
return parseResult?.success;
} catch {
return false;
}
},
);
expect(samplingHandlerCall).toBeDefined();
const [, samplingHandler] = samplingHandlerCall;
// Invoke handler; should return a CreateTaskResult immediately
let createTaskResult: SchemaOutput<typeof CreateTaskResultSchema>;
await act(async () => {
createTaskResult = await samplingHandler(samplingRequest);
});
expect(createTaskResult).toHaveProperty("task");
expect(createTaskResult.task).toEqual(
expect.objectContaining({
taskId: expect.any(String),
status: "input_required",
ttl: 0,
createdAt: expect.any(String),
lastUpdatedAt: expect.any(String),
}),
);
expect(mockOnPendingRequest).toHaveBeenCalledTimes(1);
expect(pendingResolve).toBeDefined();
expect(pendingReject).toBeDefined();
const taskId = createTaskResult.task.taskId as string;
// Locate tasks/get and tasks/result handlers
const taskGetRequest = { method: "tasks/get", params: { taskId } };
const taskResultRequest = { method: "tasks/result", params: { taskId } };
const taskGetHandlerCall = mockClient.setRequestHandler.mock.calls.find(
(call) => {
try {
const schema = call[0];
const parseResult =
schema.safeParse && schema.safeParse(taskGetRequest);
return parseResult?.success;
} catch {
return false;
}
},
);
const taskResultHandlerCall =
mockClient.setRequestHandler.mock.calls.find((call) => {
try {
const schema = call[0];
const parseResult =
schema.safeParse && schema.safeParse(taskResultRequest);
return parseResult?.success;
} catch {
return false;
}
});
expect(taskGetHandlerCall).toBeDefined();
expect(taskResultHandlerCall).toBeDefined();
const [, taskGetHandler] = taskGetHandlerCall;
const [, taskResultHandler] = taskResultHandlerCall;
// Verify tasks/get sees the in-progress task
const getBefore = await taskGetHandler(taskGetRequest);
expect(getBefore.status).toBe("input_required");
// tasks/result should block until user flow resolves
const payloadPromise = taskResultHandler(taskResultRequest);
const race = await Promise.race([
payloadPromise.then(() => "resolved"),
new Promise((r) => setTimeout(() => r("timeout"), 10)),
]);
expect(race).toBe("timeout");
const mockPayload = {
model: "test-model",
role: "assistant",
content: { type: "text", text: "ok" },
};
await act(async () => {
pendingResolve!(mockPayload);
// Let the background updater run
await new Promise((r) => setTimeout(r, 0));
});
await expect(payloadPromise).resolves.toEqual(mockPayload);
const getAfter = await taskGetHandler(taskGetRequest);
expect(getAfter.status).toBe("completed");
});
test("task-augmented elicitation/create returns { task } immediately", async () => {
const mockOnElicitationRequest = jest.fn();
const propsWithElicitation = {
...defaultProps,
onElicitationRequest: mockOnElicitationRequest,
};
const { result } = renderHook(() => useConnection(propsWithElicitation));
await act(async () => {
await result.current.connect();
});
const elicitationRequest = {
method: "elicitation/create",
params: {
task: { ttl: 0 },
message: "Please provide your name",
requestedSchema: {
type: "object",
properties: {
name: { type: "string" },
},
required: ["name"],
},
},
};
const elicitRequestHandlerCall =
mockClient.setRequestHandler.mock.calls.find((call) => {
try {
const schema = call[0];
const parseResult =
schema.safeParse && schema.safeParse(elicitationRequest);
return parseResult?.success;
} catch {
return false;
}
});
expect(elicitRequestHandlerCall).toBeDefined();
const [, handler] = elicitRequestHandlerCall;
mockOnElicitationRequest.mockImplementation((_request, resolve) => {
resolve({ action: "accept", content: { name: "test" } });
});
const resultValue = await handler(elicitationRequest);
expect(resultValue).toHaveProperty("task");
expect(resultValue.task).toEqual(
expect.objectContaining({
taskId: expect.any(String),
status: "input_required",
ttl: 0,
}),
);
});
});
test("throws error when mcpClient is not connected", async () => {
const { result } = renderHook(() => useConnection(defaultProps));
const { result } = renderHook(() => {
const { makeRequest } = useConnection(defaultProps) as unknown as {
makeRequest: (
request: ClientRequest,
schema: AnySchema,
) => Promise<unknown>;
};
return { makeRequest };
});
const mockRequest: ClientRequest = {
method: "ping",
@ -236,8 +488,10 @@ describe("useConnection", () => {
test: z.string(),
});
const mockSchemaAny: AnySchema = mockSchema as unknown as AnySchema;
await expect(
result.current.makeRequest(mockRequest, mockSchema),
result.current.makeRequest(mockRequest, mockSchemaAny),
).rejects.toThrow("MCP client not connected");
});

View file

@ -12,6 +12,7 @@ import {
import {
ClientNotification,
ClientRequest,
ClientResult,
CreateMessageRequestSchema,
ListRootsRequestSchema,
ResourceUpdatedNotificationSchema,
@ -32,13 +33,22 @@ import {
LoggingLevel,
ElicitRequestSchema,
Implementation,
Task,
CreateTaskResultSchema,
GetTaskRequestSchema,
GetTaskPayloadRequestSchema,
ListTasksRequestSchema,
CancelTaskRequestSchema,
ListTasksResultSchema,
CancelTaskResultSchema,
TaskStatusNotificationSchema,
} from "@modelcontextprotocol/sdk/types.js";
import type {
AnySchema,
SchemaOutput,
} from "@modelcontextprotocol/sdk/server/zod-compat.js";
import { RequestOptions } from "@modelcontextprotocol/sdk/shared/protocol.js";
import { useEffect, useState } from "react";
import { useEffect, useRef, useState } from "react";
import { useToast } from "@/lib/hooks/useToast";
import { ConnectionStatus, CLIENT_IDENTITY } from "../constants";
import { Notification } from "../notificationTypes";
@ -56,6 +66,7 @@ import {
} from "../auth";
import {
getMCPProxyAddress,
getMCPTaskTtl,
getMCPServerRequestMaxTotalTimeout,
resetRequestTimeoutOnProgress,
getMCPProxyAuthToken,
@ -131,6 +142,18 @@ export function useConnection({
const [serverImplementation, setServerImplementation] =
useState<Implementation | null>(null);
type ReceiverTaskRecord = {
task: Task;
payloadPromise: Promise<ClientResult>;
resolvePayload: (payload: ClientResult) => void;
rejectPayload: (reason?: unknown) => void;
cleanupTimeoutId?: ReturnType<typeof setTimeout>;
};
// Tasks created locally in response to *incoming* task-augmented requests
// (e.g. `sampling/createMessage` and `elicitation/create` with `params.task`).
const receiverTasksRef = useRef<Map<string, ReceiverTaskRecord>>(new Map());
useEffect(() => {
if (!oauthClientId) {
clearClientInformationFromSessionStorage({
@ -433,6 +456,22 @@ export function useConnection({
roots: {
listChanged: true,
},
tasks: {
list: {},
cancel: {},
...(onPendingRequest || onElicitationRequest
? {
requests: {
...(onPendingRequest
? { sampling: { createMessage: {} } }
: undefined),
...(onElicitationRequest
? { elicitation: { create: {} } }
: undefined),
},
}
: undefined),
},
},
};
@ -714,6 +753,7 @@ export function useConnection({
ResourceListChangedNotificationSchema,
ToolListChangedNotificationSchema,
PromptListChangedNotificationSchema,
TaskStatusNotificationSchema,
].forEach((notificationSchema) => {
client.setNotificationHandler(notificationSchema, onNotification);
});
@ -793,11 +833,218 @@ export function useConnection({
setServerCapabilities(capabilities ?? null);
setCompletionsSupported(capabilities?.completions !== undefined);
const nowIso = () => new Date().toISOString();
const makeTaskId = () => {
// Prefer UUID when available; otherwise fall back to a reasonably unique id.
const cryptoAny = globalThis.crypto as unknown as
| { randomUUID?: () => string }
| undefined;
return (
cryptoAny?.randomUUID?.() ??
`task_${Date.now()}_${Math.random().toString(16).slice(2)}`
);
};
const emitTaskStatus = async (task: Task) => {
// Best-effort; task status notifications are optional.
try {
const notification: ClientNotification = {
method: "notifications/tasks/status",
params: task,
} as unknown as ClientNotification;
await client.notification(notification);
pushHistory(notification);
} catch (e) {
console.warn("Failed to send notifications/tasks/status", e);
}
};
const upsertReceiverTask = async (task: Task) => {
// Update task record and emit status notification.
const record = receiverTasksRef.current.get(task.taskId);
if (record) {
receiverTasksRef.current.set(task.taskId, { ...record, task });
}
await emitTaskStatus(task);
};
const createReceiverTask = (opts: {
ttl?: number;
initialStatus: Task["status"];
statusMessage?: string;
pollInterval?: number;
}): ReceiverTaskRecord => {
const taskId = makeTaskId();
const createdAt = nowIso();
const ttl = opts.ttl ?? getMCPTaskTtl(config);
let resolvePayload: (payload: ClientResult) => void = () => undefined;
let rejectPayload: (reason?: unknown) => void = () => undefined;
const payloadPromise = new Promise<ClientResult>((resolve, reject) => {
resolvePayload = resolve;
rejectPayload = reject;
});
const task: Task = {
taskId,
status: opts.initialStatus,
ttl,
createdAt,
lastUpdatedAt: createdAt,
...(opts.pollInterval !== undefined
? { pollInterval: opts.pollInterval }
: undefined),
...(opts.statusMessage ? { statusMessage: opts.statusMessage } : {}),
};
const record: ReceiverTaskRecord = {
task,
payloadPromise,
resolvePayload,
rejectPayload,
};
// Cleanup after TTL (best-effort).
if (ttl !== null && ttl > 0) {
record.cleanupTimeoutId = setTimeout(() => {
receiverTasksRef.current.delete(taskId);
}, ttl);
}
receiverTasksRef.current.set(taskId, record);
void emitTaskStatus(task);
return record;
};
// Server -> client Tasks handlers (receiver side)
client.setRequestHandler(ListTasksRequestSchema, async () => {
return {
tasks: Array.from(receiverTasksRef.current.values()).map(
(r) => r.task,
),
};
});
client.setRequestHandler(GetTaskRequestSchema, async (request) => {
const record = receiverTasksRef.current.get(request.params.taskId);
if (!record) {
throw new McpError(
ErrorCode.InvalidParams,
`Unknown taskId: ${request.params.taskId}`,
);
}
return record.task;
});
client.setRequestHandler(GetTaskPayloadRequestSchema, async (request) => {
const record = receiverTasksRef.current.get(request.params.taskId);
if (!record) {
throw new McpError(
ErrorCode.InvalidParams,
`Unknown taskId: ${request.params.taskId}`,
);
}
// Block until the task payload is ready.
return await record.payloadPromise;
});
client.setRequestHandler(CancelTaskRequestSchema, async (request) => {
const record = receiverTasksRef.current.get(request.params.taskId);
if (!record) {
throw new McpError(
ErrorCode.InvalidParams,
`Unknown taskId: ${request.params.taskId}`,
);
}
const terminalStatuses: Task["status"][] = [
"completed",
"failed",
"cancelled",
];
if (!terminalStatuses.includes(record.task.status)) {
const updated: Task = {
...record.task,
status: "cancelled",
lastUpdatedAt: nowIso(),
statusMessage: "Cancelled",
};
receiverTasksRef.current.set(request.params.taskId, {
...record,
task: updated,
});
// Unblock any pending `tasks/result`.
record.rejectPayload(
new McpError(ErrorCode.InternalError, "Task was cancelled"),
);
await emitTaskStatus(updated);
}
return receiverTasksRef.current.get(request.params.taskId)!.task;
});
if (onPendingRequest) {
client.setRequestHandler(CreateMessageRequestSchema, (request) => {
return new Promise((resolve, reject) => {
onPendingRequest(request, resolve, reject);
const taskSpec = (request as { params?: { task?: { ttl?: number } } })
.params?.task;
if (!taskSpec) {
return new Promise((resolve, reject) => {
onPendingRequest(request, resolve, reject);
});
}
// Task-augmented sampling request: return a task immediately and
// allow the server to poll via `tasks/get` and `tasks/result`.
const record = createReceiverTask({
ttl: taskSpec.ttl,
initialStatus: "input_required",
statusMessage: "Awaiting user input",
});
// Background runner to complete and resolve this specific task record.
void (async () => {
try {
const payload = await new Promise((resolve, reject) => {
onPendingRequest(request, resolve, reject);
});
record.resolvePayload(payload as ClientResult);
const updated: Task = {
...record.task,
status: "completed",
lastUpdatedAt: nowIso(),
};
receiverTasksRef.current.set(record.task.taskId, {
...record,
task: updated,
});
await upsertReceiverTask(updated);
} catch (e) {
record.rejectPayload(e);
const updated: Task = {
...record.task,
status: "failed",
lastUpdatedAt: nowIso(),
statusMessage: e instanceof Error ? e.message : "Task failed",
};
receiverTasksRef.current.set(record.task.taskId, {
...record,
task: updated,
});
await upsertReceiverTask(updated);
}
})();
const createTaskResult: SchemaOutput<typeof CreateTaskResultSchema> =
{
task: record.task,
};
return createTaskResult;
});
}
@ -807,6 +1054,64 @@ export function useConnection({
});
}
if (onElicitationRequest) {
client.setRequestHandler(ElicitRequestSchema, (request) => {
const taskSpec = (request as { params?: { task?: { ttl?: number } } })
.params?.task;
if (!taskSpec) {
return new Promise((resolve) => {
onElicitationRequest(request, resolve);
});
}
const record = createReceiverTask({
ttl: taskSpec.ttl,
initialStatus: "input_required",
statusMessage: "Awaiting user input",
});
// Run elicitation flow and resolve the task payload.
void (async () => {
try {
const payload = await new Promise((resolve) => {
onElicitationRequest(request, resolve);
});
record.resolvePayload(payload as ClientResult);
const updated: Task = {
...record.task,
status: "completed",
lastUpdatedAt: nowIso(),
};
receiverTasksRef.current.set(record.task.taskId, {
...record,
task: updated,
});
await upsertReceiverTask(updated);
} catch (e) {
record.rejectPayload(e);
const updated: Task = {
...record.task,
status: "failed",
lastUpdatedAt: nowIso(),
statusMessage: e instanceof Error ? e.message : "Task failed",
};
receiverTasksRef.current.set(record.task.taskId, {
...record,
task: updated,
});
await upsertReceiverTask(updated);
}
})();
const createTaskResult: SchemaOutput<typeof CreateTaskResultSchema> =
{
task: record.task,
};
return createTaskResult;
});
}
if (capabilities?.logging && defaultLoggingLevel) {
lastRequest = "logging/setLevel";
await client.setLoggingLevel(defaultLoggingLevel);
@ -822,14 +1127,6 @@ export function useConnection({
lastRequest = "";
}
if (onElicitationRequest) {
client.setRequestHandler(ElicitRequestSchema, async (request) => {
return new Promise((resolve) => {
onElicitationRequest(request, resolve);
});
});
}
setMcpClient(client);
setConnectionStatus("connected");
} catch (e) {
@ -855,7 +1152,35 @@ export function useConnection({
}
};
const cancelTask = async (taskId: string) => {
return makeRequest(
{
method: "tasks/cancel",
params: { taskId },
},
CancelTaskResultSchema,
);
};
const listTasks = async (cursor?: string) => {
return makeRequest(
{
method: "tasks/list",
params: { cursor },
},
ListTasksResultSchema,
);
};
const disconnect = async () => {
// Clear any receiver-side tasks + cleanup timers
receiverTasksRef.current.forEach((record) => {
if (record.cleanupTimeoutId) {
clearTimeout(record.cleanupTimeoutId);
}
});
receiverTasksRef.current.clear();
if (transportType === "streamable-http")
await (
clientTransport as StreamableHTTPClientTransport
@ -885,6 +1210,8 @@ export function useConnection({
requestHistory,
clearRequestHistory,
makeRequest,
cancelTask,
listTasks,
sendNotification,
handleCompletion,
completionsSupported,