commit 94721227c0b359ff6ad6b49fc2296a97f4fd84d0 Author: Chengdong Zhang Date: Thu Jul 23 20:55:36 2026 +0800 feat(connect-local): add Pi extension for local agent bridge Monorepo Pi package with outbound WebSocket, pairing/reconnect flow, presence reporting, and single-flight task.dispatch handling. Co-authored-by: Cursor diff --git a/README.md b/README.md new file mode 100644 index 0000000..04de0b0 --- /dev/null +++ b/README.md @@ -0,0 +1,32 @@ +# Connect Local Pi Extension + +Install with official Pi: + +```bash +pi install git:git@github.com:your-org/agentic-chat-core.git#packages/connect-local +``` + +Prerequisite: [Pi coding agent](https://www.npmjs.com/package/@earendil-works/pi-coding-agent) installed globally or in your environment. + +## Commands + +- `/connect [project-name]` — pair this interactive session with the cloud platform +- `/disconnect` — close the bridge and clear stored reconnect credentials +- `/connected` — show whether the bridge socket is connected + +Obtain a pairing token from Slack: `/pi local new`. + +## Configuration + +- `PI_PLATFORM_LOCAL_AGENT_WS_URL` — WebSocket URL (default `ws://127.0.0.1:3000/api/local-agent/ws`) +- Reconnect credentials are stored under `~/.pi-platform/connect-local/` (outside your repo) + +## Development + +From monorepo root: + +```bash +npm run build -w packages/shared +npm run build -w packages/connect-local +npm run test -w packages/connect-local +``` diff --git a/package.json b/package.json new file mode 100644 index 0000000..8fb5ebe --- /dev/null +++ b/package.json @@ -0,0 +1,26 @@ +{ + "name": "@pi-platform/connect-local", + "version": "0.1.18", + "private": true, + "type": "module", + "keywords": ["pi-package"], + "main": "./dist/index.js", + "types": "./dist/index.d.ts", + "pi": { + "extensions": ["./dist/index.js"] + }, + "scripts": { + "build": "tsc", + "typecheck": "tsc --noEmit", + "test": "vitest run" + }, + "dependencies": { + "@earendil-works/pi-coding-agent": "^0.80.10", + "@pi-platform/shared": "*", + "ws": "^8.21.1" + }, + "devDependencies": { + "@types/ws": "^8.18.1", + "vitest": "^4.1.10" + } +} diff --git a/src/connectLocalController.ts b/src/connectLocalController.ts new file mode 100644 index 0000000..77931e1 --- /dev/null +++ b/src/connectLocalController.ts @@ -0,0 +1,239 @@ +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 }); + } +} diff --git a/src/credentialStore.test.ts b/src/credentialStore.test.ts new file mode 100644 index 0000000..d2db365 --- /dev/null +++ b/src/credentialStore.test.ts @@ -0,0 +1,30 @@ +import { describe, expect, it } from 'vitest'; +import { mkdtemp, rm } from 'node:fs/promises'; +import os from 'node:os'; +import path from 'node:path'; +import { + credentialFilePath, + loadReconnectCredential, + saveReconnectCredential, +} from './credentialStore.js'; + +describe('credentialStore', () => { + it('persists and loads reconnect credentials outside the repo cwd', async () => { + const dir = await mkdtemp(path.join(os.tmpdir(), 'pi-local-agent-')); + try { + await saveReconnectCredential(dir, { + projectName: 'my-app', + credential: 'secret-cred', + wsUrl: 'ws://localhost:3000/api/local-agent/ws', + updatedAt: '2026-07-23T00:00:00.000Z', + }); + + expect(credentialFilePath(dir)).toContain('local-agent-reconnect.json'); + const loaded = await loadReconnectCredential(dir); + expect(loaded?.projectName).toBe('my-app'); + expect(loaded?.credential).toBe('secret-cred'); + } finally { + await rm(dir, { recursive: true, force: true }); + } + }); +}); diff --git a/src/credentialStore.ts b/src/credentialStore.ts new file mode 100644 index 0000000..cfe265c --- /dev/null +++ b/src/credentialStore.ts @@ -0,0 +1,42 @@ +import { mkdir, readFile, rm, writeFile } from 'node:fs/promises'; +import path from 'node:path'; + +export interface StoredReconnectCredential { + readonly projectName: string; + readonly credential: string; + readonly wsUrl: string; + readonly updatedAt: string; +} + +export function credentialFilePath(baseDir: string): string { + return path.join(baseDir, 'local-agent-reconnect.json'); +} + +export async function loadReconnectCredential(baseDir: string): Promise { + try { + const raw = await readFile(credentialFilePath(baseDir), 'utf8'); + const parsed = JSON.parse(raw) as StoredReconnectCredential; + if (!parsed?.projectName || !parsed?.credential || !parsed?.wsUrl) { + return undefined; + } + return parsed; + } catch { + return undefined; + } +} + +export async function saveReconnectCredential( + baseDir: string, + value: StoredReconnectCredential, +): Promise { + await mkdir(baseDir, { recursive: true }); + await writeFile(credentialFilePath(baseDir), JSON.stringify(value, null, 2), 'utf8'); +} + +export async function clearReconnectCredential(baseDir: string): Promise { + try { + await rm(credentialFilePath(baseDir)); + } catch { + /* already absent */ + } +} diff --git a/src/index.test.ts b/src/index.test.ts new file mode 100644 index 0000000..354a78c --- /dev/null +++ b/src/index.test.ts @@ -0,0 +1,24 @@ +import { describe, expect, it, vi } from 'vitest'; +import type { ExtensionAPI, ExtensionCommandContext } from '@earendil-works/pi-coding-agent'; +import connectLocalExtension from './index.js'; + +describe('connectLocalExtension', () => { + it('registers connect, disconnect, and connected commands without blocking', () => { + const commands = new Map Promise }>(); + const pi = { + registerCommand: vi.fn((name: string, options: { handler: (args: string, ctx: ExtensionCommandContext) => Promise }) => { + commands.set(name, options); + }), + on: vi.fn(), + } as unknown as ExtensionAPI; + + connectLocalExtension(pi); + + expect(commands.has('connect')).toBe(true); + expect(commands.has('disconnect')).toBe(true); + expect(commands.has('connected')).toBe(true); + expect(pi.on).toHaveBeenCalledWith('session_shutdown', expect.any(Function)); + expect(pi.on).toHaveBeenCalledWith('agent_start', expect.any(Function)); + expect(pi.on).toHaveBeenCalledWith('agent_settled', expect.any(Function)); + }); +}); diff --git a/src/index.ts b/src/index.ts new file mode 100644 index 0000000..f5c1cb2 --- /dev/null +++ b/src/index.ts @@ -0,0 +1,12 @@ +import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; +import { ConnectLocalController } from './connectLocalController.js'; + +export default function connectLocalExtension(pi: ExtensionAPI): void { + const controller = new ConnectLocalController({ pi }); + controller.registerCommands(); + controller.bindLifecycle(); + controller.scheduleAutoReconnect(); +} + +export { ConnectLocalController } from './connectLocalController.js'; +export { LocalAgentWsClient } from './wsClient.js'; diff --git a/src/wsClient.test.ts b/src/wsClient.test.ts new file mode 100644 index 0000000..2ba507f --- /dev/null +++ b/src/wsClient.test.ts @@ -0,0 +1,57 @@ +import { EventEmitter } from 'node:events'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { DEFAULT_HEARTBEAT_INTERVAL_MS } from '@pi-platform/shared'; +import { LocalAgentWsClient } from './wsClient.js'; + +let latestSocket: { send: ReturnType; close: () => void } | undefined; + +class FakeSocket extends EventEmitter { + static OPEN = 1; + readonly OPEN = 1; + readyState = 0; + send = vi.fn(); + + constructor(public readonly url: string) { + super(); + latestSocket = this; + queueMicrotask(() => { + this.readyState = 1; + this.emit('open'); + }); + } + + close = vi.fn(() => { + this.readyState = 3; + this.emit('close'); + }); +} + +describe('LocalAgentWsClient', () => { + beforeEach(() => { + vi.useFakeTimers(); + latestSocket = undefined; + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it('connects and sends heartbeats on an interval', async () => { + const onServerMessage = vi.fn(); + const client = new LocalAgentWsClient({ + wsUrl: 'ws://127.0.0.1:3000/api/local-agent/ws', + onServerMessage, + WebSocketImpl: FakeSocket as never, + }); + + await client.connect(); + expect(client.isOpen).toBe(true); + + vi.advanceTimersByTime(DEFAULT_HEARTBEAT_INTERVAL_MS); + expect(latestSocket?.send).toHaveBeenCalled(); + const payload = JSON.parse(latestSocket!.send.mock.calls[0][0]); + expect(payload.type).toBe('heartbeat'); + + client.close(); + }); +}); diff --git a/src/wsClient.ts b/src/wsClient.ts new file mode 100644 index 0000000..c1b8961 --- /dev/null +++ b/src/wsClient.ts @@ -0,0 +1,103 @@ +import { + DEFAULT_HEARTBEAT_INTERVAL_MS, + localAgentServerMessageSchema, + type LocalAgentServerMessage, +} from '@pi-platform/shared'; +import { WebSocket } from 'ws'; + +export interface LocalAgentWsClientOptions { + wsUrl: string; + onServerMessage: (message: LocalAgentServerMessage) => void; + onClose?: () => void; + WebSocketImpl?: typeof WebSocket; +} + +export class LocalAgentWsClient { + private readonly wsUrl: string; + private readonly onServerMessage: LocalAgentWsClientOptions['onServerMessage']; + private readonly onClose?: () => void; + private readonly WebSocketImpl: typeof WebSocket; + private ws: WebSocket | undefined; + private heartbeatTimer: ReturnType | undefined; + + constructor(options: LocalAgentWsClientOptions) { + this.wsUrl = options.wsUrl; + this.onServerMessage = options.onServerMessage; + this.onClose = options.onClose; + this.WebSocketImpl = options.WebSocketImpl ?? WebSocket; + } + + get isOpen(): boolean { + return this.ws?.readyState === 1; + } + + connect(): Promise { + if (this.isOpen) { + return Promise.resolve(); + } + + return new Promise((resolve, reject) => { + const ws = new this.WebSocketImpl(this.wsUrl); + this.ws = ws; + + ws.once('open', () => { + this.startHeartbeat(); + resolve(); + }); + ws.once('error', (err) => { + reject(err); + }); + ws.on('message', (data) => { + this.handleMessage(String(data)); + }); + ws.on('close', () => { + this.stopHeartbeat(); + this.onClose?.(); + }); + }); + } + + send(message: unknown): void { + if (!this.ws || this.ws.readyState !== 1) { + throw new Error('WebSocket is not connected'); + } + this.ws.send(JSON.stringify(message)); + } + + close(): void { + this.stopHeartbeat(); + this.ws?.close(); + this.ws = undefined; + } + + private handleMessage(raw: string): void { + let parsed: unknown; + try { + parsed = JSON.parse(raw); + } catch { + return; + } + const result = localAgentServerMessageSchema.safeParse(parsed); + if (result.success) { + this.onServerMessage(result.data); + } + } + + private startHeartbeat(): void { + this.stopHeartbeat(); + this.heartbeatTimer = setInterval(() => { + if (!this.isOpen) { + return; + } + const heartbeat = { type: 'heartbeat' as const, at: new Date().toISOString() }; + this.send(heartbeat); + }, DEFAULT_HEARTBEAT_INTERVAL_MS); + } + + private stopHeartbeat(): void { + if (this.heartbeatTimer) { + clearInterval(this.heartbeatTimer); + this.heartbeatTimer = undefined; + } + } +} diff --git a/tsconfig.json b/tsconfig.json new file mode 100644 index 0000000..9094125 --- /dev/null +++ b/tsconfig.json @@ -0,0 +1,10 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "outDir": "./dist", + "rootDir": "./src" + }, + "include": ["src"], + "exclude": ["src/**/*.test.ts"], + "references": [{ "path": "../shared" }] +}