From c4e89b033737a0bf3d26227d8a48c9a6b58520e7 Mon Sep 17 00:00:00 2001 From: Cristina Poncela Cubeiro <140309543+cristinaponcela@users.noreply.github.com> Date: Thu, 18 Jun 2026 15:55:35 +0200 Subject: [PATCH] feat: ui attach rpc support --- packages/orchestrator/src/attach-ui.ts | 207 ++++++++++++++++++++++ packages/orchestrator/src/cli.ts | 8 +- packages/orchestrator/src/handler.ts | 25 ++- packages/orchestrator/src/ipc/protocol.ts | 17 +- packages/orchestrator/src/ipc/server.ts | 4 + packages/orchestrator/src/radius.ts | 108 ++++++++--- packages/orchestrator/src/supervisor.ts | 34 +++- 7 files changed, 369 insertions(+), 34 deletions(-) create mode 100644 packages/orchestrator/src/attach-ui.ts diff --git a/packages/orchestrator/src/attach-ui.ts b/packages/orchestrator/src/attach-ui.ts new file mode 100644 index 00000000..05c8a9d7 --- /dev/null +++ b/packages/orchestrator/src/attach-ui.ts @@ -0,0 +1,207 @@ +import { randomUUID } from "node:crypto"; +import type { + AgentSessionRuntime, + ExtensionUIContext, + RpcExtensionUIRequest, + RpcExtensionUIResponse, +} from "@earendil-works/pi-coding-agent"; + +import { theme } from "../../coding-agent/src/modes/interactive/theme/theme.ts"; + +type DialogRequest = + | Extract + | Extract + | Extract + | Extract; + +type FireAndForgetRequest = Exclude; + +interface PendingExtensionRequest { + resolve(response: RpcExtensionUIResponse): void; + cancel(): void; +} + +export class AttachUiBridge { + private readonly pendingRequests = new Map(); + private onRequest?: (request: RpcExtensionUIRequest) => void; + + attach(onRequest: (request: RpcExtensionUIRequest) => void): () => void { + this.onRequest = onRequest; + return () => { + if (this.onRequest === onRequest) { + this.onRequest = undefined; + } + }; + } + + handleResponse(response: RpcExtensionUIResponse): void { + const pending = this.pendingRequests.get(response.id); + if (!pending) { + return; + } + this.pendingRequests.delete(response.id); + pending.resolve(response); + } + + createUiContext(): ExtensionUIContext { + const requestDialog = ( + request: DialogRequest, + fallbackValue: T, + parseResponse: (response: RpcExtensionUIResponse) => T, + ): Promise => { + if (!this.onRequest) { + return Promise.resolve(fallbackValue); + } + + return new Promise((resolve) => { + this.pendingRequests.set(request.id, { + resolve: (response) => resolve(parseResponse(response)), + cancel: () => resolve(fallbackValue), + }); + this.onRequest?.(request); + }); + }; + + const emit = (request: FireAndForgetRequest): void => { + if (!this.onRequest) { + return; + } + this.onRequest(request); + }; + + return { + select: async (title, options, opts) => + requestDialog( + { + type: "extension_ui_request", + id: randomUUID(), + method: "select", + title, + options, + timeout: opts?.timeout, + }, + undefined, + (response) => ("value" in response ? response.value : undefined), + ), + confirm: async (title, message, opts) => + requestDialog( + { + type: "extension_ui_request", + id: randomUUID(), + method: "confirm", + title, + message, + timeout: opts?.timeout, + }, + false, + (response) => ("confirmed" in response ? response.confirmed : false), + ), + input: async (title, placeholder, opts) => + requestDialog( + { + type: "extension_ui_request", + id: randomUUID(), + method: "input", + title, + placeholder, + timeout: opts?.timeout, + }, + undefined, + (response) => ("value" in response ? response.value : undefined), + ), + notify: (message, notifyType) => { + emit({ type: "extension_ui_request", id: randomUUID(), method: "notify", message, notifyType }); + }, + onTerminalInput: () => () => {}, + setStatus: (statusKey, statusText) => { + emit({ type: "extension_ui_request", id: randomUUID(), method: "setStatus", statusKey, statusText }); + }, + setWorkingMessage: () => {}, + setWorkingVisible: () => {}, + setWorkingIndicator: () => {}, + setHiddenThinkingLabel: () => {}, + setWidget: (widgetKey, widgetLines, options) => { + if (widgetLines === undefined || Array.isArray(widgetLines)) { + emit({ + type: "extension_ui_request", + id: randomUUID(), + method: "setWidget", + widgetKey, + widgetLines, + widgetPlacement: options?.placement, + }); + } + }, + setFooter: () => {}, + setHeader: () => {}, + setTitle: (title) => { + emit({ type: "extension_ui_request", id: randomUUID(), method: "setTitle", title }); + }, + custom: async () => Promise.reject(new Error("Custom UI not supported in orchestrator attach mode")), + pasteToEditor: (text) => { + emit({ type: "extension_ui_request", id: randomUUID(), method: "set_editor_text", text }); + }, + setEditorText: (text) => { + emit({ type: "extension_ui_request", id: randomUUID(), method: "set_editor_text", text }); + }, + getEditorText: () => "", + editor: async (title, prefill) => + requestDialog( + { type: "extension_ui_request", id: randomUUID(), method: "editor", title, prefill }, + undefined, + (response) => ("value" in response ? response.value : undefined), + ), + addAutocompleteProvider: () => {}, + setEditorComponent: () => {}, + getEditorComponent: () => undefined, + get theme() { + return theme; + }, + getAllThemes: () => [], + getTheme: () => undefined, + setTheme: () => ({ success: false, error: "Theme switching not supported in orchestrator attach mode" }), + getToolsExpanded: () => false, + setToolsExpanded: () => {}, + }; + } + + cancelPendingRequests(): void { + for (const [id, pending] of this.pendingRequests) { + this.pendingRequests.delete(id); + pending.cancel(); + } + } +} + +export async function bindAttachExtensions(runtime: AgentSessionRuntime, uiBridge: AttachUiBridge): Promise { + const session = runtime.session; + await session.bindExtensions({ + uiContext: uiBridge.createUiContext(), + mode: "rpc", + commandContextActions: { + waitForIdle: () => session.agent.waitForIdle(), + newSession: async (options) => runtime.newSession(options), + fork: async (entryId, forkOptions) => { + const result = await runtime.fork(entryId, forkOptions); + return { cancelled: result.cancelled }; + }, + navigateTree: async (targetId, options) => { + const result = await session.navigateTree(targetId, { + summarize: options?.summarize, + customInstructions: options?.customInstructions, + replaceInstructions: options?.replaceInstructions, + label: options?.label, + }); + return { cancelled: result.cancelled }; + }, + switchSession: async (sessionPath, options) => runtime.switchSession(sessionPath, options), + reload: async () => { + await session.reload(); + }, + }, + shutdownHandler: () => {}, + onError: (error) => { + console.error("Extension error in orchestrator attach mode", error); + }, + }); +} diff --git a/packages/orchestrator/src/cli.ts b/packages/orchestrator/src/cli.ts index 02fb3caa..bb763730 100644 --- a/packages/orchestrator/src/cli.ts +++ b/packages/orchestrator/src/cli.ts @@ -4,6 +4,7 @@ import { createConnection } from "node:net"; import { dirname, join } from "node:path"; import { cwd } from "node:process"; import { fileURLToPath } from "node:url"; +import type { RpcCommand, RpcExtensionUIResponse } from "@earendil-works/pi-coding-agent"; import { getSocketPath } from "./config.ts"; import { sendIpcRequest } from "./ipc/client.ts"; import { encodeMessage } from "./ipc/protocol.ts"; @@ -61,7 +62,12 @@ async function attach(instanceId: string): Promise { .map((line) => line.trim()) .filter((line) => line.length > 0); for (const line of lines) { - socket.write(encodeMessage({ type: "attach_rpc", command: JSON.parse(line) })); + const parsed = JSON.parse(line) as RpcCommand | RpcExtensionUIResponse; + if (parsed.type === "extension_ui_response") { + socket.write(encodeMessage(parsed)); + continue; + } + socket.write(encodeMessage({ type: "attach_rpc", command: parsed })); } }); } diff --git a/packages/orchestrator/src/handler.ts b/packages/orchestrator/src/handler.ts index 382030c5..efcfb104 100644 --- a/packages/orchestrator/src/handler.ts +++ b/packages/orchestrator/src/handler.ts @@ -1,4 +1,9 @@ -import type { AgentSessionEvent, RpcCommand } from "@earendil-works/pi-coding-agent"; +import type { + AgentSessionEvent, + RpcCommand, + RpcExtensionUIRequest, + RpcExtensionUIResponse, +} from "@earendil-works/pi-coding-agent"; import type { AttachReadyResponse, AttachRequest, @@ -128,16 +133,26 @@ export function attachIpcInstance( instanceId: string, onResponse: (response: AttachRpcResponse) => void, onSessionEvent: (event: AgentSessionEvent) => void, -): { handleRequest(request: { type: "attach_rpc"; command: RpcCommand }): Promise; close(): void } | undefined { - const handle = supervisor.attachInstance(instanceId, onSessionEvent); + onUiRequest: (request: RpcExtensionUIRequest) => void, +): + | { + handleRequest(request: { type: "attach_rpc"; command: RpcCommand } | RpcExtensionUIResponse): Promise; + close(): void; + } + | undefined { + const handle = supervisor.attachInstance(instanceId, onSessionEvent, onUiRequest); if (!handle) { return undefined; } return { async handleRequest(request): Promise { - const response = await handle.handleRpc(request.command); - onResponse({ type: "attach_rpc_result", response }); + if (request.type === "attach_rpc") { + const response = await handle.handleRpc(request.command); + onResponse({ type: "attach_rpc_result", response }); + return; + } + handle.handleUiResponse(request); }, close(): void { handle.close(); diff --git a/packages/orchestrator/src/ipc/protocol.ts b/packages/orchestrator/src/ipc/protocol.ts index 93d7e597..43929a71 100644 --- a/packages/orchestrator/src/ipc/protocol.ts +++ b/packages/orchestrator/src/ipc/protocol.ts @@ -1,4 +1,10 @@ -import type { AgentSessionEvent, RpcCommand, RpcResponse } from "@earendil-works/pi-coding-agent"; +import type { + AgentSessionEvent, + RpcCommand, + RpcExtensionUIRequest, + RpcExtensionUIResponse, + RpcResponse, +} from "@earendil-works/pi-coding-agent"; import type { InstanceStatus } from "../types.ts"; export interface SpawnRequest { @@ -121,8 +127,13 @@ export interface ResponseMap { } export type OrchestratorResponse = ResponseMap[keyof ResponseMap] | ErrorResponse; -export type AttachClientRequest = AttachRpcRequest; -export type AttachServerResponse = AttachReadyResponse | AttachEventResponse | AttachRpcResponse | ErrorResponse; +export type AttachClientRequest = AttachRpcRequest | RpcExtensionUIResponse; +export type AttachServerResponse = + | AttachReadyResponse + | AttachEventResponse + | AttachRpcResponse + | RpcExtensionUIRequest + | ErrorResponse; export type ProtocolMessage = OrchestratorRequest | OrchestratorResponse | AttachClientRequest | AttachServerResponse; export type ResponseFor = T extends { type: infer K } diff --git a/packages/orchestrator/src/ipc/server.ts b/packages/orchestrator/src/ipc/server.ts index 2dce45d0..bb43a0a1 100644 --- a/packages/orchestrator/src/ipc/server.ts +++ b/packages/orchestrator/src/ipc/server.ts @@ -35,6 +35,7 @@ export interface IpcRequestHandler { instanceId: string, onEvent: (response: AttachRpcResponse) => void, onSessionEvent: (event: import("@earendil-works/pi-coding-agent").AgentSessionEvent) => void, + onUiRequest: (request: import("@earendil-works/pi-coding-agent").RpcExtensionUIRequest) => void, ): | { handleRequest(request: AttachClientRequest): Promise; @@ -80,6 +81,9 @@ export async function startIpcServer(handler: IpcRequestHandler): Promise { socket.write(encodeMessage({ type: "attach_event", event })); }, + (request) => { + socket.write(encodeMessage(request)); + }, ); if (!attachment) { socket.end( diff --git a/packages/orchestrator/src/radius.ts b/packages/orchestrator/src/radius.ts index a6dcad7f..71d39922 100644 --- a/packages/orchestrator/src/radius.ts +++ b/packages/orchestrator/src/radius.ts @@ -8,6 +8,8 @@ const DEFAULT_RADIUS_URL = "https://radius.pi.dev/"; const DEFAULT_ORCHESTRATOR_BASE_PATH = "/v1/"; const ORCHESTRATOR_VERSION = "0.79.6"; const NOT_FOUND_RETRY_THRESHOLD = 3; +const HEARTBEAT_BACKOFF_BASE_MS = 1_000; +const HEARTBEAT_BACKOFF_MAX_MS = 30_000; const RADIUS_PROVIDER = "radius"; interface RegisterMachineResponse extends RadiusRegistration { @@ -25,9 +27,11 @@ interface RadiusPresenceCoordinator { } interface PiHeartbeatState { - timer: NodeJS.Timeout; + timer?: NodeJS.Timeout; + intervalMs: number; radiusPiId: string; consecutiveNotFoundCount: number; + transientFailureCount: number; } class RadiusHttpError extends Error { @@ -75,6 +79,15 @@ function isNotFoundError(error: unknown): error is RadiusHttpError { return error instanceof RadiusHttpError && error.status === 404; } +function computeBackoffDelayMs(failureCount: number): number { + const exponentialDelay = Math.min( + HEARTBEAT_BACKOFF_MAX_MS, + HEARTBEAT_BACKOFF_BASE_MS * 2 ** Math.max(0, failureCount - 1), + ); + const jitterMs = Math.floor(Math.random() * Math.max(250, exponentialDelay / 4)); + return Math.min(HEARTBEAT_BACKOFF_MAX_MS, exponentialDelay + jitterMs); +} + export function getRadiusUrl(): string { return process.env.PI_RADIUS_URL || DEFAULT_RADIUS_URL; } @@ -119,7 +132,9 @@ export function isRadiusEnabled(): boolean { export class RadiusPresence { private machineHeartbeatTimer?: NodeJS.Timeout; + private machineHeartbeatIntervalMs = 0; private machineConsecutiveNotFoundCount = 0; + private machineTransientFailureCount = 0; private readonly piHeartbeatStates = new Map(); private machine?: MachineRecord; private coordinator?: RadiusPresenceCoordinator; @@ -140,11 +155,13 @@ export class RadiusPresence { async stop(): Promise { if (this.machineHeartbeatTimer) { - clearInterval(this.machineHeartbeatTimer); + clearTimeout(this.machineHeartbeatTimer); this.machineHeartbeatTimer = undefined; } for (const [instanceId, state] of this.piHeartbeatStates) { - clearInterval(state.timer); + if (state.timer) { + clearTimeout(state.timer); + } this.piHeartbeatStates.delete(instanceId); } if (!this.machine || !isRadiusEnabled()) { @@ -179,7 +196,9 @@ export class RadiusPresence { async disconnectPi(instance: InstanceRecord): Promise { const state = this.piHeartbeatStates.get(instance.id); if (state) { - clearInterval(state.timer); + if (state.timer) { + clearTimeout(state.timer); + } this.piHeartbeatStates.delete(instance.id); } if (!isRadiusEnabled() || !instance.radiusPiId) { @@ -209,31 +228,54 @@ export class RadiusPresence { }; saveMachine(this.machine); this.machineConsecutiveNotFoundCount = 0; + this.machineTransientFailureCount = 0; return registered; } private startMachineHeartbeat(intervalMs: number): void { + this.machineHeartbeatIntervalMs = intervalMs; + this.scheduleMachineHeartbeat(intervalMs); + } + + private scheduleMachineHeartbeat(delayMs: number): void { if (this.machineHeartbeatTimer) { - clearInterval(this.machineHeartbeatTimer); + clearTimeout(this.machineHeartbeatTimer); } - this.machineHeartbeatTimer = setInterval(() => { + this.machineHeartbeatTimer = setTimeout(() => { void this.heartbeatMachine(); - }, intervalMs); + }, delayMs); } private startPiHeartbeat(instanceId: string, intervalMs: number, radiusPiId: string): void { const existingState = this.piHeartbeatStates.get(instanceId); - if (existingState) { - clearInterval(existingState.timer); + if (existingState?.timer) { + clearTimeout(existingState.timer); } - const timer = setInterval(() => { - void this.heartbeatPi(instanceId); - }, intervalMs); - this.piHeartbeatStates.set(instanceId, { - timer, + const state: PiHeartbeatState = existingState ?? { + intervalMs, radiusPiId, consecutiveNotFoundCount: 0, - }); + transientFailureCount: 0, + }; + state.intervalMs = intervalMs; + state.radiusPiId = radiusPiId; + state.consecutiveNotFoundCount = 0; + state.transientFailureCount = 0; + this.piHeartbeatStates.set(instanceId, state); + this.schedulePiHeartbeat(instanceId, intervalMs); + } + + private schedulePiHeartbeat(instanceId: string, delayMs: number): void { + const state = this.piHeartbeatStates.get(instanceId); + if (!state) { + return; + } + if (state.timer) { + clearTimeout(state.timer); + } + state.timer = setTimeout(() => { + void this.heartbeatPi(instanceId); + }, delayMs); } private async heartbeatMachine(): Promise { @@ -247,21 +289,31 @@ export class RadiusPresence { socketPath: getSocketPath(), }); this.machineConsecutiveNotFoundCount = 0; + this.machineTransientFailureCount = 0; + this.scheduleMachineHeartbeat(this.machineHeartbeatIntervalMs); } catch (error) { if (!isNotFoundError(error)) { - console.error("Radius machine heartbeat failed", error); + this.machineTransientFailureCount += 1; + const delayMs = computeBackoffDelayMs(this.machineTransientFailureCount); + console.error(`Radius machine heartbeat failed; retrying in ${delayMs}ms`, error); + this.scheduleMachineHeartbeat(delayMs); return; } + this.machineTransientFailureCount = 0; this.machineConsecutiveNotFoundCount += 1; if (this.machineConsecutiveNotFoundCount < NOT_FOUND_RETRY_THRESHOLD) { + this.scheduleMachineHeartbeat(this.machineHeartbeatIntervalMs); return; } try { await this.reRegisterMachineAndPis(); } catch (recoveryError) { - console.error("Radius machine re-registration failed", recoveryError); + this.machineTransientFailureCount += 1; + const delayMs = computeBackoffDelayMs(this.machineTransientFailureCount); + console.error(`Radius machine re-registration failed; retrying in ${delayMs}ms`, recoveryError); + this.scheduleMachineHeartbeat(delayMs); } } } @@ -279,14 +331,21 @@ export class RadiusPresence { try { await maybePost(`pis/${state.radiusPiId}/heartbeat`, {}); state.consecutiveNotFoundCount = 0; + state.transientFailureCount = 0; + this.schedulePiHeartbeat(instanceId, state.intervalMs); } catch (error) { if (!isNotFoundError(error)) { - console.error(`Radius Pi heartbeat failed for instance ${instanceId}`, error); + state.transientFailureCount += 1; + const delayMs = computeBackoffDelayMs(state.transientFailureCount); + console.error(`Radius Pi heartbeat failed for instance ${instanceId}; retrying in ${delayMs}ms`, error); + this.schedulePiHeartbeat(instanceId, delayMs); return; } + state.transientFailureCount = 0; state.consecutiveNotFoundCount += 1; if (state.consecutiveNotFoundCount < NOT_FOUND_RETRY_THRESHOLD) { + this.schedulePiHeartbeat(instanceId, state.intervalMs); return; } @@ -294,9 +353,16 @@ export class RadiusPresence { const recovered = await this.reRegisterPi(instanceId); if (!recovered) { console.error(`Radius Pi re-registration skipped for instance ${instanceId}`); + this.schedulePiHeartbeat(instanceId, computeBackoffDelayMs(1)); } } catch (recoveryError) { - console.error(`Radius Pi re-registration failed for instance ${instanceId}`, recoveryError); + state.transientFailureCount += 1; + const delayMs = computeBackoffDelayMs(state.transientFailureCount); + console.error( + `Radius Pi re-registration failed for instance ${instanceId}; retrying in ${delayMs}ms`, + recoveryError, + ); + this.schedulePiHeartbeat(instanceId, delayMs); } } } @@ -320,7 +386,9 @@ export class RadiusPresence { if (!instance) { const state = this.piHeartbeatStates.get(instanceId); if (state) { - clearInterval(state.timer); + if (state.timer) { + clearTimeout(state.timer); + } this.piHeartbeatStates.delete(instanceId); } return false; diff --git a/packages/orchestrator/src/supervisor.ts b/packages/orchestrator/src/supervisor.ts index 9c2ef4b5..886cbf26 100644 --- a/packages/orchestrator/src/supervisor.ts +++ b/packages/orchestrator/src/supervisor.ts @@ -9,9 +9,12 @@ import { createAgentSessionServices, getAgentDir, type RpcCommand, + type RpcExtensionUIRequest, + type RpcExtensionUIResponse, type RpcResponse, SessionManager, } from "@earendil-works/pi-coding-agent"; +import { AttachUiBridge, bindAttachExtensions } from "./attach-ui.ts"; import { radiusPresence } from "./radius.ts"; import { handleRpcCommand } from "./rpc-bridge.ts"; import { getInstance, loadInstances, removeInstance, saveInstances, upsertInstance } from "./storage.ts"; @@ -21,6 +24,7 @@ interface LiveInstance { runtime: AgentSessionRuntime; record: InstanceRecord; subscribers: Set; + uiBridge: AttachUiBridge; unsubscribeSession?: () => void; } @@ -70,7 +74,8 @@ export class OrchestratorSupervisor { upsertInstance(live.record); } - private bindLiveInstance(live: LiveInstance): void { + private async bindLiveInstance(live: LiveInstance): Promise { + await bindAttachExtensions(live.runtime, live.uiBridge); live.unsubscribeSession?.(); live.unsubscribeSession = live.runtime.session.subscribe((event) => { for (const subscriber of live.subscribers) { @@ -79,7 +84,7 @@ export class OrchestratorSupervisor { }); live.runtime.setRebindSession(async () => { this.syncInstanceRecord(live); - this.bindLiveInstance(live); + await this.bindLiveInstance(live); }); } @@ -94,20 +99,33 @@ export class OrchestratorSupervisor { attachInstance( instanceId: string, onEvent: (event: AgentSessionEvent) => void, - ): { handleRpc(command: RpcCommand): Promise; close(): void } | undefined { + onUiRequest: (request: RpcExtensionUIRequest) => void, + ): + | { + handleRpc(command: RpcCommand): Promise; + handleUiResponse(response: RpcExtensionUIResponse): void; + close(): void; + } + | undefined { const live = this.liveInstances.get(instanceId); if (!live) { return undefined; } live.subscribers.add(onEvent); + const detachUi = live.uiBridge.attach(onUiRequest); return { handleRpc: async (command) => { const response = await handleRpcCommand(live.runtime, command); this.syncInstanceRecord(live); return response; }, + handleUiResponse: (response) => { + live.uiBridge.handleResponse(response); + }, close: () => { + detachUi(); live.subscribers.delete(onEvent); + live.uiBridge.cancelPendingRequests(); }, }; } @@ -162,8 +180,13 @@ export class OrchestratorSupervisor { }; const registeredRecord = await radiusPresence.registerPi(record); - const live: LiveInstance = { runtime, record: registeredRecord, subscribers: new Set() }; - this.bindLiveInstance(live); + const live: LiveInstance = { + runtime, + record: registeredRecord, + subscribers: new Set(), + uiBridge: new AttachUiBridge(), + }; + await this.bindLiveInstance(live); this.liveInstances.set(registeredRecord.id, live); upsertInstance(registeredRecord); return cloneInstance(registeredRecord); @@ -177,6 +200,7 @@ export class OrchestratorSupervisor { await radiusPresence.disconnectPi(live.record); live.unsubscribeSession?.(); + live.uiBridge.cancelPendingRequests(); live.runtime.setRebindSession(undefined); await live.runtime.dispose(); this.liveInstances.delete(instanceId);