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 <cursoragent@cursor.com>
This commit is contained in:
Chengdong Zhang
239
src/connectLocalController.ts
Normal file
239
src/connectLocalController.ts
Normal file
@@ -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<typeof LocalAgentWsClient>[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<void> {
|
||||
const parts = args.trim().split(/\s+/).filter(Boolean);
|
||||
const token = parts[0];
|
||||
if (!token) {
|
||||
ctx.ui.notify('Usage: /connect <pairing-token> [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<void> {
|
||||
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<void> {
|
||||
await this.ensureClient();
|
||||
const pairOk = await new Promise<LocalAgentServerMessage & { type: 'pair.ok' }>((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<void> {
|
||||
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<void> {
|
||||
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<LocalAgentServerMessage, { type: 'task.dispatch' }>): 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 });
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user