diff --git a/src/connectLocalController.test.ts b/src/connectLocalController.test.ts new file mode 100644 index 0000000..d49f974 --- /dev/null +++ b/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 {} + + close(): void {} + + send(message: unknown): void { + this.sent.push(message); + } + + receive(message: LocalAgentServerMessage): void { + this.onServerMessage(message); + } +} + +function createHarness() { + const handlers = new Map Promise | void>(); + const sendUserMessage = vi.fn(); + const pi = { + on: vi.fn((event: string, handler: (event: never, ctx: ExtensionContext) => Promise | 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 { + 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.', + }); + }); +}); diff --git a/src/connectLocalController.ts b/src/connectLocalController.ts index df076aa..ef317c4 100644 --- a/src/connectLocalController.ts +++ b/src/connectLocalController.ts @@ -16,6 +16,11 @@ export interface ConnectLocalControllerDeps { createWsClient?: (options: ConstructorParameters[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 { @@ -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') {