test(agent): cover run cancellation through the harnesses

src/agent/runner/harness/pi/__tests__/backend-cancellation.test.ts

@@ -0,0 +1,85 @@
+import { piBackend } from '..';
+import type { TaskRunInputs } from '../../types';
+import { HostResolution } from '@shared/host-resolution';
+import { AgentErrorType } from '@agent/signals';
+
+vi.mock('@utils/analytics');
+vi.mock('@utils/debug');
+vi.mock('@agent/yara-hooks', async (original) => ({
+  ...(await original<typeof import('@agent/yara-hooks')>()),
+  prewarmYaraScanner: vi.fn(),
+}));
+vi.mock('@agent/gateway-session', async (original) => ({
+  ...(await original<typeof import('@agent/gateway-session')>()),
+  gatewayAuth: () =>
+    Promise.resolve({ gatewayUrl: 'https://gateway.test', token: 'phe_test' }),
+}));
+vi.mock('../mcp', () => ({
+  fetchInstructions: vi.fn(),
+  setupPostHogMcp: vi.fn().mockRejectedValue(new Error('offline')),
+}));
+vi.mock('../orchestrator-tools', () => ({
+  createPiOrchestratorTools: () => [],
+}));
+
+let host: AbortController;
+const session = vi.hoisted(() => {
+  let finish!: () => void;
+  return {
+    bindExtensions: vi.fn(),
+    subscribe: vi.fn(() => vi.fn()),
+    // The host cancels while the first turn is live.
+    prompt: vi.fn(
+      () =>
+        new Promise<void>((resolve) => {
+          finish = resolve;
+          host.abort();
+        }),
+    ),
+    abort: vi.fn(() => Promise.resolve(finish())),
+  };
+});
+vi.mock('@earendil-works/pi-coding-agent', async (original) => ({
+  ...(await original<typeof import('@earendil-works/pi-coding-agent')>()),
+  createAgentSession: () => Promise.resolve({ session }),
+  DefaultResourceLoader: class {
+    reload = vi.fn();
+  },
+  AuthStorage: { create: vi.fn() },
+  ModelRegistry: {
+    inMemory: () => ({ registerProvider: vi.fn(), find: () => ({}) }),
+  },
+}));
+
+afterEach(() => vi.clearAllMocks());
+
+it.each(['run', 'runTask'] as const)(
+  'aborts the live pi session when the host cancels %s',
+  async (entry) => {
+    host = new AbortController();
+    const inputs = {
+      config: { programId: 'metrics', run: {} },
+      input: { installDir: '/tmp/pi-cancel-test', flags: {} },
+      boot: {
+        credentials: { host: HostResolution.fromRegion('us') },
+        wizardMetadata: {},
+        wizardFlags: {},
+      },
+      emit: vi.fn(),
+      prompt: 'Do the work',
+      spinner: { start: vi.fn(), stop: vi.fn(), message: vi.fn() },
+      model: 'claude-test',
+      signal: host.signal,
+      orchestrator: { currentTaskId: 'task-1' },
+    } as unknown as TaskRunInputs;
+    const running =
+      entry === 'run' ? piBackend.run(inputs) : piBackend.runTask?.(inputs);
+
+    await expect(running).resolves.toEqual({
+      kind: 'abort',
+      classification: AgentErrorType.ABORT,
+      message: 'Agent run cancelled',
+    });
+    expect(session.abort).toHaveBeenCalledOnce();
+  },
+);

src/agent/runner/harness/pi/__tests__/cancellation.test.ts

@@ -1,28 +1,32 @@
 import { bindPiCancellation } from '../cancellation';
 
-it('observes session abort and keeps listener failures out of abort dispatch', async () => {
-  const controller = new AbortController();
-  let release: (() => void) | undefined;
-  const abort = vi.fn(
-    () =>
-      new Promise<void>((resolve) => {
-        release = resolve;
-      }),
-  );
-  const binding = bindPiCancellation(controller.signal, { abort });
-  expect(() => controller.abort()).not.toThrow();
-  await Promise.resolve();
-  expect(abort).toHaveBeenCalledOnce();
-  let settled = false;
-  const waiting = binding.settle().then(() => {
-    settled = true;
-  });
-  await Promise.resolve();
-  expect(settled).toBe(false);
-  release?.();
-  await waiting;
-  expect(settled).toBe(true);
-});
+it.each(['after', 'before'])(
+  'aborts the session once and waits for idle when the host cancels %s binding',
+  async (when) => {
+    const controller = new AbortController();
+    if (when === 'before') controller.abort();
+    let release: (() => void) | undefined;
+    const abort = vi.fn(
+      () =>
+        new Promise<void>((resolve) => {
+          release = resolve;
+        }),
+    );
+    const binding = bindPiCancellation(controller.signal, { abort });
+    expect(() => controller.abort()).not.toThrow();
+    await Promise.resolve();
+    expect(abort).toHaveBeenCalledOnce();
+    let settled = false;
+    const waiting = binding.settle().then(() => {
+      settled = true;
+    });
+    await Promise.resolve();
+    expect(settled).toBe(false);
+    release?.();
+    await waiting;
+    expect(settled).toBe(true);
+  },
+);
 
 it('contains a synchronous abort throw and a diagnostic callback throw', async () => {
   const controller = new AbortController();

src/agent/runner/harness/pi/__tests__/gateway.test.ts

@@ -159,6 +159,7 @@ describe('withGatewayRemint', () => {
     turns: unknown[],
     remintedAt = Date.now() + HOUR,
   ) {
+    const controller = new AbortController();
     const prompts: string[] = [];
     const session = {
       prompt: vi.fn((text: string) => {
@@ -185,8 +186,9 @@ describe('withGatewayRemint', () => {
         modelId: 'openai/gpt-5.6-terra',
       }),
       continueText: 'continue',
+      signal: controller.signal,
     });
-    return { wrapped, registry, refreshAuth, prompts };
+    return { wrapped, registry, refreshAuth, prompts, controller };
   }
 
   it.each([
@@ -316,4 +318,20 @@ describe('withGatewayRemint', () => {
     expect(refreshAuth).not.toHaveBeenCalled();
     expect(prompts).toEqual(['do it']);
   });
+
+  it('does not continue when the host cancels during the re-mint', async () => {
+    const { wrapped, registry, refreshAuth, prompts, controller } = harness(
+      gatewayAuth('phe_old', Date.now() - 1),
+      [rejected],
+    );
+    refreshAuth.mockImplementation(() => {
+      controller.abort();
+      return Promise.resolve(gatewayAuth('phe_new', Date.now() + HOUR));
+    });
+
+    await wrapped.prompt('do it');
+
+    expect(registry.registerProvider).not.toHaveBeenCalled();
+    expect(prompts).toEqual(['do it']);
+  });
 });

src/agent/runner/sequence/orchestrator/__tests__/executor.test.ts

@@ -41,38 +41,52 @@ describe('drainQueue', () => {
     return Promise.resolve();
   };
 
-  it('waits for live siblings after a fatal error and starts no dependents', async () => {
-    const fatal = new RunTaskFatal({
-      code: ErrorCodes.AgentOrchestratorTasksFailed,
-      message: 'Authentication failed',
-    });
-    let release!: () => void;
-    const blocked = new Promise<void>((resolve) => {
-      release = resolve;
-    });
-    q.enqueue({ type: 'fatal' });
-    const sibling = q.enqueue({ type: 'sibling' });
-    q.enqueue({ type: 'dependent', dependsOn: [sibling.id] });
-    const started: string[] = [];
-    const drain = drainQueue(q, async (task) => {
-      started.push(task.type);
-      if (task.type === 'fatal') throw fatal;
-      await blocked;
-      q.complete(task.id, HANDOFF);
-    });
-    let settled = false;
-    const result = drain.catch((error: unknown) => {
-      settled = true;
-      return error;
-    });
-    await new Promise<void>((resolve) => setImmediate(resolve));
-    expect(settled).toBe(false);
-    expect(started).toEqual(['fatal', 'sibling']);
-    release();
-    expect(await result).toBe(fatal);
-    expect(q.get(sibling.id)?.status).toBe(TaskStatus.Done);
-    expect(started).toEqual(['fatal', 'sibling']);
-  });
+  it.each(['a fatal error', 'host cancellation'])(
+    'waits for live siblings after %s and starts no dependents',
+    async (ending) => {
+      const controller = new AbortController();
+      const fatal = new RunTaskFatal({
+        code: ErrorCodes.AgentOrchestratorTasksFailed,
+        message: 'Authentication failed',
+      });
+      let release!: () => void;
+      const blocked = new Promise<void>((resolve) => {
+        release = resolve;
+      });
+      q.enqueue({ type: 'first' });
+      const sibling = q.enqueue({ type: 'sibling' });
+      q.enqueue({ type: 'dependent', dependsOn: [sibling.id] });
+      const started: string[] = [];
+      const drain = drainQueue(
+        q,
+        async (task) => {
+          started.push(task.type);
+          if (task.type === 'first') {
+            await Promise.resolve(); // Let the sibling start first.
+            if (ending === 'a fatal error') throw fatal;
+            controller.abort();
+            return;
+          }
+          await blocked;
+          q.complete(task.id, HANDOFF);
+        },
+        { maxStarts: 50, signal: controller.signal },
+      );
+      let settled = false;
+      const result = drain
+        .catch((error: unknown) => error)
+        .finally(() => {
+          settled = true;
+        });
+      await new Promise<void>((resolve) => setImmediate(resolve));
+      expect(settled).toBe(false);
+      expect(started).toEqual(['first', 'sibling']);
+      release();
+      expect(await result).toBe(ending === 'a fatal error' ? fatal : undefined);
+      expect(q.get(sibling.id)?.status).toBe(TaskStatus.Done);
+      expect(started).toEqual(['first', 'sibling']);
+    },
+  );
 
   it('preserves a fatal failure when a sibling completes in the same turn', async () => {
     q.enqueue({ type: 'success' });