import os from 'node:os'; import path from 'node:path'; import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext } from '@earendil-works/pi-coding-agent'; import type { LocalAgentServerMessage } from '@pi-platform/shared'; import { clearReconnectCredential, loadReconnectCredential, saveReconnectCredential, } from './credentialStore.js'; import { LocalAgentWsClient } from './wsClient.js'; export interface ConnectLocalControllerDeps { pi: ExtensionAPI; credentialBaseDir?: string; wsUrl?: string; createWsClient?: (options: ConstructorParameters[0]) => LocalAgentWsClient; } export class ConnectLocalController { private readonly pi: ExtensionAPI; private readonly credentialBaseDir: string; private readonly wsUrl: string; private readonly createWsClient: ConnectLocalControllerDeps['createWsClient']; private client: LocalAgentWsClient | undefined; private projectName: string | undefined; private connectionEpoch: number | undefined; private remoteTaskReserved = false; private agentBusy = false; private pairWaiters: Array<{ resolve: (value: LocalAgentServerMessage & { type: 'pair.ok' }) => void; reject: (error: Error) => void; }> = []; constructor(deps: ConnectLocalControllerDeps) { this.pi = deps.pi; this.credentialBaseDir = deps.credentialBaseDir ?? path.join(os.homedir(), '.pi-platform', 'connect-local'); this.wsUrl = deps.wsUrl ?? process.env.PI_PLATFORM_LOCAL_AGENT_WS_URL ?? 'ws://127.0.0.1:3000/api/local-agent/ws'; this.createWsClient = deps.createWsClient ?? ((options) => new LocalAgentWsClient(options)); } registerCommands(): void { this.pi.registerCommand('connect', { description: 'Pair this interactive Pi session with the cloud platform', handler: async (args, ctx) => { await this.handleConnect(args, ctx); }, }); this.pi.registerCommand('disconnect', { description: 'Disconnect the local agent bridge', handler: async (_args, ctx) => { await this.handleDisconnect(ctx); }, }); this.pi.registerCommand('connected', { description: 'Show local agent bridge connection status', handler: async (_args, ctx) => { this.notifyStatus(ctx); }, }); } bindLifecycle(): void { this.pi.on('agent_start', async (_event, ctx) => { this.agentBusy = true; this.sendPresence('busy', ctx); }); this.pi.on('agent_settled', async (_event, ctx) => { this.agentBusy = false; this.remoteTaskReserved = false; this.sendPresence('idle', ctx); }); this.pi.on('session_shutdown', async () => { this.shutdown(); }); } scheduleAutoReconnect(): void { queueMicrotask(() => { void this.tryAutoReconnect(); }); } async handleConnect(args: string, ctx: ExtensionCommandContext): Promise { const parts = args.trim().split(/\s+/).filter(Boolean); const token = parts[0]; if (!token) { ctx.ui.notify('Usage: /connect [project-name]', 'warning'); return; } const projectName = parts[1] ?? path.basename(ctx.cwd); try { await this.connectWithPairingToken(token, projectName); ctx.ui.notify(`Connected local project "${projectName}"`, 'info'); } catch (error) { const message = error instanceof Error ? error.message : String(error); ctx.ui.notify(`Connect failed: ${message}`, 'error'); } } async handleDisconnect(ctx: ExtensionCommandContext): Promise { this.shutdown(); await clearReconnectCredential(this.credentialBaseDir); ctx.ui.notify('Disconnected from cloud platform', 'info'); } notifyStatus(ctx: ExtensionCommandContext): void { if (!this.client?.isOpen || !this.projectName) { ctx.ui.notify('Local agent: offline', 'info'); return; } ctx.ui.notify(`Local agent: online (${this.projectName})`, 'info'); } async connectWithPairingToken(token: string, projectName: string): Promise { await this.ensureClient(); const pairOk = await new Promise((resolve, reject) => { const timeout = setTimeout(() => reject(new Error('Pairing timed out')), 30_000); this.pairWaiters.push({ resolve: (value) => { clearTimeout(timeout); resolve(value); }, reject: (error) => { clearTimeout(timeout); reject(error); }, }); try { this.client!.send({ type: 'pair', token, projectName }); } catch (error) { clearTimeout(timeout); reject(error); } }); this.projectName = pairOk.projectName; this.connectionEpoch = pairOk.connectionEpoch; await saveReconnectCredential(this.credentialBaseDir, { projectName: pairOk.projectName, credential: pairOk.reconnectCredential, wsUrl: this.wsUrl, updatedAt: new Date().toISOString(), }); } async tryAutoReconnect(): Promise { const stored = await loadReconnectCredential(this.credentialBaseDir); if (!stored?.credential || !stored.projectName) { return; } try { await this.ensureClient(stored.wsUrl || this.wsUrl); this.client!.send({ type: 'reconnect', credential: stored.credential, projectName: stored.projectName, connectionEpoch: this.connectionEpoch, }); this.projectName = stored.projectName; } catch { /* extension must not block Pi; reconnect can be retried manually */ } } shutdown(): void { this.client?.close(); this.client = undefined; this.projectName = undefined; this.connectionEpoch = undefined; this.remoteTaskReserved = false; } private async ensureClient(wsUrl = this.wsUrl): Promise { if (this.client?.isOpen) { return; } this.client = this.createWsClient!({ wsUrl, onServerMessage: (message) => this.handleServerMessage(message), onClose: () => { this.projectName = undefined; }, }); await this.client.connect(); } private handleServerMessage(message: LocalAgentServerMessage): void { if (message.type === 'pair.ok') { const waiter = this.pairWaiters.shift(); waiter?.resolve(message); return; } if (message.type === 'error') { const waiter = this.pairWaiters.shift(); waiter?.reject(new Error(message.message)); return; } if (message.type === 'task.dispatch') { void this.handleTaskDispatch(message); } } private handleTaskDispatch(message: Extract): void { if (!this.client?.isOpen) { return; } if (this.remoteTaskReserved || this.agentBusy) { this.client.send({ type: 'task.result', taskId: message.taskId, status: 'rejected', reason: 'agent_busy', }); return; } this.remoteTaskReserved = true; this.client.send({ type: 'presence', status: 'busy' }); this.pi.sendUserMessage(message.prompt); this.client.send({ type: 'task.result', taskId: message.taskId, status: 'accepted', }); } private sendPresence(status: 'idle' | 'busy', ctx: ExtensionContext): void { if (!this.client?.isOpen) { return; } if (status === 'idle' && this.remoteTaskReserved) { return; } if (!ctx.isIdle() && status === 'idle') { return; } this.client.send({ type: 'presence', status }); } }