fix(coding-runtime): close Pi extension lifecycle races

tests/unit/pi-conversation-runtime.test.ts

@@ -14,6 +14,7 @@ import {
 import { PiConversationRuntime } from '../../electron/coding-runtime/pi/runtime';
 import { PiSessionProjectionError } from '../../electron/coding-runtime/pi/session-projector';
 import { PiSessionRegistry } from '../../electron/coding-runtime/pi/session-registry';
+import { PiManagedExtensionHost } from '../../electron/coding-runtime/pi/extension-host';
 import {
   PiWorkerPool,
   type PiConversationWorker,
@@ -114,6 +115,18 @@ class RuntimeFakeWorker implements PiConversationWorker {
   }
 }
 
+class TrackingExtensionHost extends PiManagedExtensionHost {
+  readonly runs = new Map<string, { generation: number; runId: string }>();
+
+  override async bindRun(conversationId: string, generation: number, runId: string): Promise<void> {
+    this.runs.set(conversationId, { generation, runId });
+  }
+
+  override async clearRun(conversationId: string, _generation: number, runId?: string): Promise<void> {
+    if (!runId || this.runs.get(conversationId)?.runId === runId) this.runs.delete(conversationId);
+  }
+}
+
 afterEach(async () => {
   await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true })));
 });
@@ -189,10 +202,12 @@ describe('Pi Conversation runtime', () => {
         };
       },
     });
+    const trackingHost = new TrackingExtensionHost();
     const runtime = new PiConversationRuntime({
       pool,
       registry: new PiSessionRegistry({ projectStore }),
       createId: (kind) => `${kind}-fixed`,
+      extensionHost: trackingHost,
       resolveModel: async (candidate) => {
         if (candidate.accountId !== 'account-b' || candidate.modelId !== 'model-b') {
           throw new Error('model unavailable');
@@ -483,6 +498,17 @@ describe('Pi Conversation runtime', () => {
     await expect(runtime.getSnapshot(forkTarget.id)).rejects.toMatchObject({
       publicError: { code: 'CODING_CONVERSATION_NOT_FOUND' },
     });
+    await runtime.prompt({
+      clientRequestId: 'request-active-recover',
+      conversationId: left.id,
+      mode: 'prompt',
+      text: 'Recover this active run',
+      attachments: [],
+    });
+    expect(trackingHost.runs.get(left.id)?.runId).toBe('run-fixed');
+    await runtime.recover(left.id);
+    expect(trackingHost.runs.has(left.id)).toBe(false);
+    expect((await runtime.getSnapshot(left.id)).run).toEqual({ status: 'idle' });
     unsubscribe();
   });
 });

tests/unit/pi-extension-host.test.ts

@@ -5,6 +5,7 @@ import { tmpdir } from 'node:os';
 import path from 'node:path';
 import { afterEach, describe, expect, it } from 'vitest';
 import { PiManagedExtensionHost } from '../../electron/coding-runtime/pi/extension-host';
+import { PiProjectWriteLeaseCoordinator } from '../../electron/coding-runtime/pi/write-lease';
 
 const roots: string[] = [];
 const hosts: PiManagedExtensionHost[] = [];
@@ -29,6 +30,40 @@ async function post(
 }
 
 describe('managed Pi extension bridge', () => {
+  it('finishes a queued HTTP request while closing a holder and waiter', async () => {
+    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-close-'));
+    roots.push(root);
+    const leases = new PiProjectWriteLeaseCoordinator();
+    const host = new PiManagedExtensionHost(leases);
+    hosts.push(host);
+    const holder = await host.registerWorker({
+      conversationId: 'conversation-holder', generation: 1, projectId: 'project-a', extensionsDir: root,
+    });
+    const waiter = await host.registerWorker({
+      conversationId: 'conversation-waiter', generation: 1, projectId: 'project-a', extensionsDir: root,
+    });
+    await Promise.all([
+      host.bindRun('conversation-holder', 1, 'run-holder'),
+      host.bindRun('conversation-waiter', 1, 'run-waiter'),
+    ]);
+    const held = await post(holder, {
+      action: 'lease.acquire', conversationId: 'conversation-holder', workerGeneration: 1,
+      runId: 'run-holder', resourceId: 'held-tool',
+    });
+    expect(held.status).toBe(200);
+    const waiting = post(waiter, {
+      action: 'lease.acquire', conversationId: 'conversation-waiter', workerGeneration: 1,
+      runId: 'run-waiter', resourceId: 'waiting-tool',
+    });
+    await expect.poll(() => leases.waitingCount('project-a')).toBe(1);
+
+    await expect(Promise.race([
+      host.close().then(() => 'closed'),
+      new Promise<string>((resolve) => setTimeout(() => resolve('timeout'), 500)),
+    ])).resolves.toBe('closed');
+    await expect(waiting).resolves.toMatchObject({ status: 409 });
+  });
+
   it('materializes the active run into a replacement generation before spawn', async () => {
     const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-rebuild-'));
     roots.push(root);

tests/unit/pi-extension-ui-projector.test.ts

@@ -40,5 +40,15 @@ describe('Pi extension UI projector', () => {
     expect(projector.getDiagnostics()).toEqual([{
       method: 'setWidget', reason: 'unsupported-ui-method',
     }]);
+    for (let index = 0; index < 1_000; index += 1) {
+      projector.project('conversation-a', 1, 'run-1', {
+        type: 'extension_ui_request', id: `unknown-${index}`, method: `unknown-${index}`,
+      });
+    }
+    const diagnostics = projector.getDiagnostics();
+    expect(diagnostics).toHaveLength(256);
+    expect(diagnostics.at(-1)).toEqual({
+      method: 'unknown-999', reason: 'invalid-ui-payload',
+    });
   });
 });

tests/unit/pi-interaction.test.ts

@@ -75,4 +75,43 @@ describe('Pi interaction store', () => {
     expect(changes.map(({ status }) => status)).toEqual(['cancelled', 'cancelled']);
     expect(store.list()).toEqual([]);
   });
+
+  it('makes response ownership atomic across double submit and generation invalidation', async () => {
+    const sent: PiRpcCommand[] = [];
+    const changes: ConversationInteraction[] = [];
+    let cancelResource = () => undefined;
+    let releaseSend = () => undefined;
+    const sendGate = new Promise<void>((resolve) => { releaseSend = resolve; });
+    const store = new PiInteractionStore({
+      getState: () => ({
+        conversationId: 'conversation-a', workerId: 'worker-a', state: 'running', generation: 1,
+        session: { piSessionId: 'session-a', sessionKey: 'key-a' },
+      }),
+      getActiveRun: () => ({ generation: 1, runId: 'run-1' }),
+      send: async (_conversationId, command) => {
+        sent.push(command);
+        await sendGate;
+      },
+      trackGenerationResource: (input) => {
+        cancelResource = input.cancel;
+        return () => undefined;
+      },
+    }, (interaction) => changes.push(interaction));
+    store.open('conversation-a', 1, 'run-1', {
+      type: 'extension_ui_request', id: 'question-race', method: 'confirm', title: 'Continue?',
+    });
+
+    const first = store.respond('conversation-a', {
+      interactionId: 'question-race', confirmed: true,
+    });
+    await expect(store.respond('conversation-a', {
+      interactionId: 'question-race', confirmed: false,
+    })).rejects.toThrow('already in progress');
+    cancelResource();
+    releaseSend();
+
+    await expect(first).rejects.toThrow('stale');
+    expect(sent).toHaveLength(1);
+    expect(changes.map(({ status }) => status)).toEqual(['cancelled']);
+  });
 });