fix: address Copilot comments

Author: LucaButBoring

packages/ext-tasks/src/client/execution.ts

@@ -578,7 +578,16 @@ export class TaskExecution<
 
   cancel(signal?: AbortSignal): Promise<void> {
     throwIfAborted(signal);
-    this.cancelPromise ??= this.cancelTask(this.cancellationController.signal);
+    // The server's cancel ack ends the local lifetime because V2 cancellation
+    // is cooperative: the server may never reach `cancelled`, so waiting for
+    // a terminal snapshot could leave result() pending forever.
+    this.cancelPromise ??= this.cancelTask(
+      this.cancellationController.signal,
+    ).then(() => {
+      this.releaseLifecycleListener?.();
+      this.inputController.abort(this.cancelledError);
+      this.controller.abort(this.cancelledError);
+    });
     return signal === undefined
       ? this.cancelPromise
       : withAbort(this.cancelPromise, signal);

packages/ext-tasks/src/client/task-lifecycle-races.test.ts

@@ -3,8 +3,8 @@ import { describe, expect, it, vi } from "vitest";
 import {
   DispatchError,
   JsonRpcResponseError,
+  TaskCancelledError,
   TaskExecutionClosedError,
-  TaskFailedError,
   TaskUpdatesAlreadyAcquiredError,
   toolDeclaration,
   withTasks,
@@ -178,8 +178,9 @@ describe("task lifecycle and races", () => {
       },
     });
     await execution.close();
+    // Cancelled, not closed: the awaited cancel settled the outcome first.
     await expect(legacyResult(execution)).rejects.toBeInstanceOf(
-      TaskExecutionClosedError,
+      TaskCancelledError,
     );
     await session.close();
   });
@@ -237,11 +238,9 @@ describe("task lifecycle and races", () => {
       await new Promise<void>((resolve) => setImmediate(resolve));
       expect(unhandled).toEqual([]);
       const outcome = await result;
-      expect(outcome.status).toBe("failed");
-      if (outcome.status !== "failed")
-        throw new Error("Expected failed outcome");
-      expect(outcome.error).toBeInstanceOf(TaskFailedError);
-      expect(outcome.error.cause).toBeInstanceOf(TaskExecutionClosedError);
+      // Cancelled, not failed-as-closed: the cancel ack settled the outcome
+      // before close ran.
+      expect(outcome.status).toBe("cancelled");
     } finally {
       process.off("unhandledRejection", onUnhandledRejection);
     }
@@ -406,8 +405,10 @@ describe("task lifecycle and races", () => {
         else await expect(execution.cancel()).rejects.toThrow("cancel failed");
         expect(cancelCalls).toBe(retryable ? 2 : 1);
         await execution.close();
+        // A successful cancel settles the outcome as cancelled; only a
+        // failed one leaves close to decide it.
         await expect(legacyResult(execution)).rejects.toBeInstanceOf(
-          TaskExecutionClosedError,
+          retryable ? TaskCancelledError : TaskExecutionClosedError,
         );
         await session.close();
       }),

packages/ext-tasks/src/client/task-protocol-v2.ts

@@ -306,10 +306,11 @@ async function resolveInputRequest<TApplicationContext>(args: {
     request,
   );
   if (acquisition.kind !== "new") {
+    // Throwing fails the execution because reusing a key with a different
+    // request is a protocol violation: the client already answered this key,
+    // and polling on would act on input the application never saw.
     if (acquisition.kind === "incompatible")
-      inputContext.reportError(
-        new Error(`V2 task input key ${inputKey} was reused incompatibly`),
-      );
+      throw new Error(`V2 task input key ${inputKey} was reused incompatibly`);
     return undefined;
   }
 

packages/ext-tasks/src/client/v2-input-task.test.ts

@@ -441,6 +441,73 @@ describe("V2 input and task behavior", () => {
     await session.close();
   });
 
+  it("fails the execution when a committed V2 input key is reused incompatibly", async () => {
+    const port = new FakePort({ generation: "v2", capabilities: {} });
+    let getCalls = 0;
+    port.dispatchHandler = async (request) => {
+      await Promise.resolve();
+      const method = expectRecord(request).method;
+      if (method === "tools/call")
+        return {
+          kind: "result",
+          result: asJson({
+            resultType: "task",
+            taskId: "reuse-incompatible",
+            status: "working",
+            createdAt: "a",
+            lastUpdatedAt: "a",
+            ttlMs: null,
+          }),
+        };
+      if (method === "tasks/get") {
+        getCalls += 1;
+        return {
+          kind: "result",
+          result: asJson({
+            resultType: "complete",
+            taskId: "reuse-incompatible",
+            status: "input_required",
+            createdAt: "a",
+            lastUpdatedAt: String(getCalls),
+            ttlMs: null,
+            inputRequests: {
+              same:
+                getCalls === 1
+                  ? { method: "roots/list" }
+                  : { method: "elicitation/create", params: {} },
+            },
+          }),
+        };
+      }
+      if (method === "tasks/update")
+        return { kind: "result", result: { resultType: "complete" } };
+      throw new Error(`unexpected method ${formatJson(method)}`);
+    };
+    const session = withTasks(port, {
+      tools: {
+        currentTool: () =>
+          toolDeclaration({ name: "x", inputSchema: { type: "object" } }),
+      },
+      onInputRequest: async (request) => {
+        await Promise.resolve();
+        // `as never`, because the generic signature cannot relate a runtime
+        // kind branch to TRequest (the suite-wide fake-handler pattern).
+        return (
+          request.kind === "roots"
+            ? { roots: [{ uri: "file:///workspace" }] }
+            : { action: "cancel" }
+        ) as never;
+      },
+    });
+    const execution = await session.callTool("x");
+    // Key "same" returns with a different request shape — a protocol
+    // violation that must fail the execution, not keep it polling.
+    await expect(legacyResult(execution)).rejects.toThrow(
+      "reused incompatibly",
+    );
+    await session.close();
+  });
+
   it("declines keyed V2 elicitation while withholding sampling and roots", async () => {
     const port = new FakePort({ generation: "v2", capabilities: {} });
     let getCalls = 0;

packages/ext-tasks/src/core/v2/index.test.ts

@@ -902,4 +902,26 @@ describe("V2 runtime wire contracts", () => {
     ).toBe(true);
     expect(hasTaskServerCapabilityV2({ extensions: {} })).toBe(false);
   });
+
+  it("merges the tasks capability into existing client capabilities", () => {
+    const wire = withTaskCapabilityV2({
+      _meta: {
+        "io.modelcontextprotocol/clientCapabilities": {
+          sampling: {},
+          extensions: { "example.com/other": { enabled: true } },
+        },
+      },
+    });
+    expect(wire).toEqual({
+      _meta: {
+        "io.modelcontextprotocol/clientCapabilities": {
+          sampling: {},
+          extensions: {
+            "example.com/other": { enabled: true },
+            "io.modelcontextprotocol/tasks": {},
+          },
+        },
+      },
+    });
+  });
 });

packages/ext-tasks/src/core/v2/integration.ts

@@ -117,7 +117,16 @@ export function withTaskCapabilityV2<
   T extends Readonly<Record<string, JsonValue>>,
 >(params: T): T & Readonly<Record<string, JsonValue>> {
   const existingMetadata = asObjectRecord(params._meta) ?? {};
-  const capability = { extensions: { [TASKS_EXTENSION_ID_V2]: {} } };
+  // Merging (like frameV2TaskRequest) because replacing the entry would drop
+  // any capabilities and extensions the caller already declared.
+  const existingCapability =
+    asObjectRecord(existingMetadata[CLIENT_CAPABILITIES_META_KEY_V2]) ?? {};
+  const existingExtensions =
+    asObjectRecord(existingCapability.extensions) ?? {};
+  const capability = {
+    ...existingCapability,
+    extensions: { ...existingExtensions, [TASKS_EXTENSION_ID_V2]: {} },
+  };
   return {
     ...params,
     _meta: {

packages/ext-tasks/src/receiver/index.ts

@@ -160,6 +160,30 @@ function hasTaskAugmentation(request: unknown): boolean {
   );
 }
 
+function asTaskAugmentationRecord(
+  value: JsonValue,
+): Record<string, JsonValue> | undefined {
+  return value !== null && typeof value === "object" && !Array.isArray(value)
+    ? (value as Record<string, JsonValue>)
+    : undefined;
+}
+
+function validateTaskAugmentation(request: unknown): void {
+  const params = paramsOf(request);
+  // Absent is fine: without a prior handler, a plain request runs as a task.
+  // Object.hasOwn, because indexing types absent keys as JsonValue.
+  if (!Object.hasOwn(params, "task")) return;
+  // A present augmentation must be an object (the 2025-11-25 shape);
+  // `task: true` would leave us guessing at what the requester meant.
+  const augmentation = asTaskAugmentationRecord(params.task);
+  if (augmentation === undefined)
+    throw new Error("Task augmentation must be a JSON object");
+  if (!Object.hasOwn(augmentation, "ttl") || augmentation.ttl === null) return;
+  const ttl = augmentation.ttl;
+  if (typeof ttl !== "number" || !Number.isInteger(ttl) || ttl < 0)
+    throw new RangeError("task.ttl must be a non-negative integer or null");
+}
+
 function taskIdOf(request: unknown): string {
   const taskId = paramsOf(request).taskId;
   if (typeof taskId !== "string")
@@ -328,6 +352,7 @@ export function bindTaskReceiver(
       method,
       (raw) => {
         expire();
+        validateTaskAugmentation(raw);
         const params = paramsOf(raw);
         if (tasks.size >= maxTasks)
           throw new Error(
@@ -439,15 +464,11 @@ export function bindTaskReceiver(
   install("tasks/get", (request) =>
     Promise.resolve(snapshot(get(taskIdOf(request)))),
   );
-  install("tasks/result", async (request) => {
+  install("tasks/result", (request) => {
     const record = get(taskIdOf(request));
-    if (
-      record.task.status === "working" ||
-      record.task.status === "input_required"
-    )
-      throw new Error("Task is not terminal");
-    if (record.task.status === "cancelled")
-      throw new Error("Task was cancelled");
+    // Returning the still-pending promise blocks the response until the task
+    // settles, because tasks/result is a blocking call per 2025-11-25; the
+    // promise already rejects on failure, cancellation, expiry, and close.
     return record.result;
   });
   install("tasks/cancel", (request) => {

packages/ext-tasks/src/receiver/receiver.test.ts

@@ -161,6 +161,58 @@ describe("bindTaskReceiver", () => {
     ).rejects.toThrow("expired");
   });
 
+  it("rejects malformed task augmentations and invalid augmentation TTLs", async () => {
+    const host = new Host();
+    bindTaskReceiver(asClient(host), {
+      methods: { "sampling/createMessage": true },
+      sampling: () => Promise.resolve({ ok: true }),
+      createTaskId: () => "augmented",
+    });
+    // `task: true` is not a spec shape: the augmentation is an object.
+    await expect(
+      host.call("sampling/createMessage", { task: true }),
+    ).rejects.toThrow("Task augmentation must be a JSON object");
+    await expect(
+      host.call("sampling/createMessage", { task: [1] }),
+    ).rejects.toThrow("Task augmentation must be a JSON object");
+    await expect(
+      host.call("sampling/createMessage", { task: { ttl: -1 } }),
+    ).rejects.toThrow("task.ttl must be a non-negative integer or null");
+    await expect(
+      host.call("sampling/createMessage", { task: { ttl: 1.5 } }),
+    ).rejects.toThrow("task.ttl must be a non-negative integer or null");
+    // A well-formed augmentation still creates the task.
+    await expect(
+      host.call("sampling/createMessage", { task: { ttl: null } }),
+    ).resolves.toMatchObject({ task: { taskId: "augmented" } });
+  });
+
+  it("blocks tasks/result until the task settles instead of rejecting while pending", async () => {
+    const host = new Host();
+    const work = deferred<Record<string, never>>();
+    bindTaskReceiver(asClient(host), {
+      methods: { "sampling/createMessage": true },
+      sampling: () => work.promise,
+      createTaskId: () => "blocking",
+    });
+    await host.call("sampling/createMessage");
+    // Issued mid-work: it must block (per 2025-11-25), not reject.
+    let settled = false;
+    const pending = host.call("tasks/result", { taskId: "blocking" });
+    void pending.then(
+      () => {
+        settled = true;
+      },
+      () => {
+        settled = true;
+      },
+    );
+    await flush();
+    expect(settled).toBe(false);
+    work.resolve({});
+    await expect(pending).resolves.toEqual({});
+  });
+
   it("advertises and installs only enabled request methods", () => {
     const host = new Host();
     const binding = bindTaskReceiver(asClient(host), {