ton/pi_harness
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: ui attach rpc support

Browse Source
This commit is contained in:
Cristina Poncela Cubeiro
2026-06-18 15:55:35 +02:00
parent 4806b8f9f4
commit c4e89b0337
7 changed files with 369 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/orchestrator/src/attach-ui.ts
+207
View File
@@ -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<RpcExtensionUIRequest, { method: "select" }>
| Extract<RpcExtensionUIRequest, { method: "confirm" }>
| Extract<RpcExtensionUIRequest, { method: "input" }>
| Extract<RpcExtensionUIRequest, { method: "editor" }>;
type FireAndForgetRequest = Exclude<RpcExtensionUIRequest, DialogRequest>;
interface PendingExtensionRequest {
resolve(response: RpcExtensionUIResponse): void;
cancel(): void;
}
export class AttachUiBridge {
private readonly pendingRequests = new Map<string, PendingExtensionRequest>();
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 = <T>(
request: DialogRequest,
fallbackValue: T,
parseResponse: (response: RpcExtensionUIResponse) => T,
): Promise<T> => {
if (!this.onRequest) {
return Promise.resolve(fallbackValue);
}
return new Promise<T>((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<void> {
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);
},
});
}
packages/orchestrator/src/cli.ts
+7 -1
View File
@@ -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<void> {
.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 }));
}
});
}
packages/orchestrator/src/handler.ts
+20 -5
View File
@@ -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<void>; close(): void } | undefined {
const handle = supervisor.attachInstance(instanceId, onSessionEvent);
onUiRequest: (request: RpcExtensionUIRequest) => void,
):
| {
handleRequest(request: { type: "attach_rpc"; command: RpcCommand } | RpcExtensionUIResponse): Promise<void>;
close(): void;
}
| undefined {
const handle = supervisor.attachInstance(instanceId, onSessionEvent, onUiRequest);
if (!handle) {
return undefined;
}
return {
async handleRequest(request): Promise<void> {
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();
packages/orchestrator/src/ipc/protocol.ts
+14 -3
View File
@@ -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 OrchestratorRequest> = T extends { type: infer K }
packages/orchestrator/src/ipc/server.ts
+4
View File
@@ -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<void>;
@@ -80,6 +81,9 @@ export async function startIpcServer(handler: IpcRequestHandler): Promise<Server
(event) => {
socket.write(encodeMessage({ type: "attach_event", event }));
},
(request) => {
socket.write(encodeMessage(request));
},
);
if (!attachment) {
socket.end(
packages/orchestrator/src/radius.ts
+88 -20
View File
@@ -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<string, PiHeartbeatState>();
private machine?: MachineRecord;
private coordinator?: RadiusPresenceCoordinator;
@@ -140,11 +155,13 @@ export class RadiusPresence {
async stop(): Promise<void> {
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<void> {
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<void> {
@@ -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;
packages/orchestrator/src/supervisor.ts
+29 -5
View File
@@ -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<AgentSessionEventListener>;
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<void> {
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<RpcResponse>; close(): void } | undefined {
onUiRequest: (request: RpcExtensionUIRequest) => void,
):
| {
handleRpc(command: RpcCommand): Promise<RpcResponse>;
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);