feat: relay local Pi task completion

src/connectLocalController.test.ts

@@ -0,0 +1,127 @@
+import { mkdtemp, rm } from 'node:fs/promises';
+import os from 'node:os';
+import path from 'node:path';
+import { afterEach, describe, expect, it, vi } from 'vitest';
+import type { ExtensionAPI, ExtensionContext } from '@earendil-works/pi-coding-agent';
+import { ConnectLocalController } from './connectLocalController.js';
+import type { LocalAgentServerMessage } from './localAgentProtocol.js';
+import type { LocalAgentWsClient } from './wsClient.js';
+
+class FakeWsClient {
+  readonly sent: unknown[] = [];
+  readonly isOpen = true;
+
+  constructor(private readonly onServerMessage: (message: LocalAgentServerMessage) => void) {}
+
+  async connect(): Promise<void> {}
+
+  close(): void {}
+
+  send(message: unknown): void {
+    this.sent.push(message);
+  }
+
+  receive(message: LocalAgentServerMessage): void {
+    this.onServerMessage(message);
+  }
+}
+
+function createHarness() {
+  const handlers = new Map<string, (event: never, ctx: ExtensionContext) => Promise<void> | void>();
+  const sendUserMessage = vi.fn();
+  const pi = {
+    on: vi.fn((event: string, handler: (event: never, ctx: ExtensionContext) => Promise<void> | void) => {
+      handlers.set(event, handler);
+    }),
+    sendUserMessage,
+  } as unknown as ExtensionAPI;
+  let client: FakeWsClient | undefined;
+
+  return {
+    client: () => {
+      if (!client) throw new Error('WebSocket client was not created');
+      return client;
+    },
+    createClient: (options: { onServerMessage: (message: LocalAgentServerMessage) => void }) => {
+      client = new FakeWsClient(options.onServerMessage);
+      return client as never as LocalAgentWsClient;
+    },
+    emit: async (event: string, payload: never) => {
+      const handler = handlers.get(event);
+      if (!handler) throw new Error(`Missing ${event} handler`);
+      await handler(payload, { isIdle: () => true } as never);
+    },
+    pi,
+    sendUserMessage,
+  };
+}
+
+describe('ConnectLocalController remote tasks', () => {
+  const credentialDirs: string[] = [];
+
+  afterEach(async () => {
+    await Promise.all(credentialDirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true })));
+  });
+
+  async function connect(controller: ConnectLocalController, client: () => FakeWsClient): Promise<void> {
+    const connected = controller.connectWithPairingToken('pair-token', 'my-project');
+    await Promise.resolve();
+    await Promise.resolve();
+    client().receive({
+      type: 'pair.ok',
+      projectName: 'my-project',
+      reconnectCredential: 'credential',
+      connectionEpoch: 1,
+    });
+    await connected;
+  }
+
+  it('sends the final assistant response for a dispatched task', async () => {
+    const harness = createHarness();
+    const credentialBaseDir = await mkdtemp(path.join(os.tmpdir(), 'connect-local-test-'));
+    credentialDirs.push(credentialBaseDir);
+    const controller = new ConnectLocalController({
+      pi: harness.pi,
+      credentialBaseDir,
+      createWsClient: harness.createClient,
+    });
+    controller.bindLifecycle();
+    await connect(controller, harness.client);
+
+    harness.client().receive({ type: 'task.dispatch', taskId: 'task-1', prompt: 'Create CODE_INDEX.md' });
+    await harness.emit('message_end', {
+      message: { role: 'assistant', content: [{ type: 'text', text: 'Created CODE_INDEX.md.' }] },
+    } as never);
+    await harness.emit('agent_settled', {} as never);
+
+    expect(harness.sendUserMessage).toHaveBeenCalledWith('Create CODE_INDEX.md');
+    expect(harness.client().sent).toContainEqual({ type: 'task.result', taskId: 'task-1', status: 'accepted' });
+    expect(harness.client().sent).toContainEqual({
+      type: 'task.completed',
+      taskId: 'task-1',
+      text: 'Created CODE_INDEX.md.',
+    });
+  });
+
+  it('reports a failure when a dispatched task settles without an assistant response', async () => {
+    const harness = createHarness();
+    const credentialBaseDir = await mkdtemp(path.join(os.tmpdir(), 'connect-local-test-'));
+    credentialDirs.push(credentialBaseDir);
+    const controller = new ConnectLocalController({
+      pi: harness.pi,
+      credentialBaseDir,
+      createWsClient: harness.createClient,
+    });
+    controller.bindLifecycle();
+    await connect(controller, harness.client);
+
+    harness.client().receive({ type: 'task.dispatch', taskId: 'task-1', prompt: 'Create CODE_INDEX.md' });
+    await harness.emit('agent_settled', {} as never);
+
+    expect(harness.client().sent).toContainEqual({
+      type: 'task.failed',
+      taskId: 'task-1',
+      message: 'Local Pi task ended without an assistant response.',
+    });
+  });
+});

src/connectLocalController.ts

@@ -16,6 +16,11 @@ export interface ConnectLocalControllerDeps {
   createWsClient?: (options: ConstructorParameters<typeof LocalAgentWsClient>[0]) => LocalAgentWsClient;
 }
 
+type RemoteTask = {
+  readonly taskId: string;
+  assistantText: string | undefined;
+};
+
 export class ConnectLocalController {
   private readonly pi: ExtensionAPI;
   private readonly credentialBaseDir: string;
@@ -24,7 +29,7 @@ export class ConnectLocalController {
   private client: LocalAgentWsClient | undefined;
   private projectName: string | undefined;
   private connectionEpoch: number | undefined;
-  private remoteTaskReserved = false;
+  private remoteTask: RemoteTask | undefined;
   private agentBusy = false;
   private pairWaiters: Array<{
     resolve: (value: LocalAgentServerMessage & { type: 'pair.ok' }) => void;
@@ -68,9 +73,37 @@ export class ConnectLocalController {
       this.agentBusy = true;
       this.sendPresence('busy', ctx);
     });
+    this.pi.on('message_end', async (event) => {
+      if (!this.remoteTask || event.message.role !== 'assistant') {
+        return;
+      }
+      const text = event.message.content
+        .filter((block) => block.type === 'text')
+        .map((block) => block.text)
+        .join('');
+      if (text.trim()) {
+        this.remoteTask.assistantText = text;
+      }
+    });
     this.pi.on('agent_settled', async (_event, ctx) => {
       this.agentBusy = false;
-      this.remoteTaskReserved = false;
+      const remoteTask = this.remoteTask;
+      this.remoteTask = undefined;
+      if (remoteTask && this.client?.isOpen) {
+        if (remoteTask.assistantText) {
+          this.client.send({
+            type: 'task.completed',
+            taskId: remoteTask.taskId,
+            text: remoteTask.assistantText,
+          });
+        } else {
+          this.client.send({
+            type: 'task.failed',
+            taskId: remoteTask.taskId,
+            message: 'Local Pi task ended without an assistant response.',
+          });
+        }
+      }
       this.sendPresence('idle', ctx);
     });
     this.pi.on('session_shutdown', async () => {
@@ -186,7 +219,7 @@ export class ConnectLocalController {
     this.client = undefined;
     this.projectName = undefined;
     this.connectionEpoch = undefined;
-    this.remoteTaskReserved = false;
+    this.remoteTask = undefined;
   }
 
   private async ensureClient(wsUrl = this.wsUrl): Promise<void> {
@@ -233,7 +266,7 @@ export class ConnectLocalController {
     if (!this.client?.isOpen) {
       return;
     }
-    if (this.remoteTaskReserved || this.agentBusy) {
+    if (this.remoteTask || this.agentBusy) {
       this.client.send({
         type: 'task.result',
         taskId: message.taskId,
@@ -243,7 +276,7 @@ export class ConnectLocalController {
       return;
     }
 
-    this.remoteTaskReserved = true;
+    this.remoteTask = { taskId: message.taskId, assistantText: undefined };
     this.client.send({ type: 'presence', status: 'busy' });
     this.pi.sendUserMessage(message.prompt);
     this.client.send({
@@ -257,7 +290,7 @@ export class ConnectLocalController {
     if (!this.client?.isOpen) {
       return;
     }
-    if (status === 'idle' && this.remoteTaskReserved) {
+    if (status === 'idle' && this.remoteTask) {
       return;
     }
     if (!ctx.isIdle() && status === 'idle') {