Files
oh-my-pi/packages/coding-agent/src/modes/rpc/rpc-mode.ts
T
can1357 94ee2f99ac fix(coding-agent): use ready-signal protocol for RPC client startup
Replace flaky time-based heuristic (sleep + 500ms race) with a proper
ready-signal protocol. The RPC server now emits {"type":"ready"} on
stdout when initialized, and the client waits for that signal instead
of guessing based on timing.

Fixes CI failure where the process took longer than 600ms to reach
provider validation (due to auth discovery + model registry refresh),
causing start() to resolve even when the process was about to exit.
2026-02-20 12:28:25 +01:00

701 lines
21 KiB
TypeScript

/**
* RPC mode: Headless operation with JSON stdin/stdout protocol.
*
* Used for embedding the agent in other applications.
* Receives commands as JSON on stdin, outputs events and responses as JSON on stdout.
*
* Protocol:
* - Commands: JSON objects with `type` field, optional `id` for correlation
* - Responses: JSON objects with `type: "response"`, `command`, `success`, and optional `data`/`error`
* - Events: AgentSessionEvent objects streamed as they occur
* - Extension UI: Extension UI requests are emitted, client responds with extension_ui_response
*/
import { readJsonl, Snowflake } from "@oh-my-pi/pi-utils";
import type { ExtensionUIContext, ExtensionUIDialogOptions } from "../../extensibility/extensions";
import { type Theme, theme } from "../../modes/theme/theme";
import type { AgentSession } from "../../session/agent-session";
import type {
RpcCommand,
RpcExtensionUIRequest,
RpcExtensionUIResponse,
RpcResponse,
RpcSessionState,
} from "./rpc-types";
// Re-export types for consumers
export type {
RpcCommand,
RpcExtensionUIRequest,
RpcExtensionUIResponse,
RpcResponse,
RpcSessionState,
} from "./rpc-types";
/**
* Run in RPC mode.
* Listens for JSON commands on stdin, outputs events and responses on stdout.
*/
export async function runRpcMode(session: AgentSession): Promise<never> {
// Signal to RPC clients that the server is ready to accept commands
process.stdout.write(`${JSON.stringify({ type: "ready" })}\n`);
const output = (obj: RpcResponse | RpcExtensionUIRequest | object) => {
process.stdout.write(`${JSON.stringify(obj)}\n`);
};
const success = <T extends RpcCommand["type"]>(
id: string | undefined,
command: T,
data?: object | null,
): RpcResponse => {
if (data === undefined) {
return { id, type: "response", command, success: true } as RpcResponse;
}
return { id, type: "response", command, success: true, data } as RpcResponse;
};
const error = (id: string | undefined, command: string, message: string): RpcResponse => {
return { id, type: "response", command, success: false, error: message };
};
// Pending extension UI requests waiting for response
type PendingExtensionRequest = {
resolve: (response: RpcExtensionUIResponse) => void;
reject: (error: Error) => void;
};
const pendingExtensionRequests = new Map<string, PendingExtensionRequest>();
// Shutdown request flag (wrapped in object to allow mutation with const)
const shutdownState = { requested: false };
/**
* Extension UI context that uses the RPC protocol.
*/
class RpcExtensionUIContext implements ExtensionUIContext {
constructor(
private pendingRequests: Map<string, PendingExtensionRequest>,
private output: (obj: RpcResponse | RpcExtensionUIRequest | object) => void,
) {}
/** Helper for dialog methods with signal/timeout support */
#createDialogPromise<T>(
opts: ExtensionUIDialogOptions | undefined,
defaultValue: T,
request: Record<string, unknown>,
parseResponse: (response: RpcExtensionUIResponse) => T,
): Promise<T> {
if (opts?.signal?.aborted) return Promise.resolve(defaultValue);
const id = Snowflake.next() as string;
const { promise, resolve, reject } = Promise.withResolvers<T>();
let timeoutId: NodeJS.Timeout | undefined;
const cleanup = () => {
if (timeoutId) clearTimeout(timeoutId);
opts?.signal?.removeEventListener("abort", onAbort);
this.pendingRequests.delete(id);
};
const onAbort = () => {
cleanup();
resolve(defaultValue);
};
opts?.signal?.addEventListener("abort", onAbort, { once: true });
if (opts?.timeout !== undefined) {
timeoutId = setTimeout(() => {
cleanup();
resolve(defaultValue);
}, opts.timeout);
}
this.pendingRequests.set(id, {
resolve: (response: RpcExtensionUIResponse) => {
cleanup();
resolve(parseResponse(response));
},
reject,
});
this.output({ type: "extension_ui_request", id, ...request } as RpcExtensionUIRequest);
return promise;
}
select(title: string, options: string[], dialogOptions?: ExtensionUIDialogOptions): Promise<string | undefined> {
return this.#createDialogPromise(
dialogOptions,
undefined,
{ method: "select", title, options, timeout: dialogOptions?.timeout },
response =>
"cancelled" in response && response.cancelled
? undefined
: "value" in response
? response.value
: undefined,
);
}
confirm(title: string, message: string, dialogOptions?: ExtensionUIDialogOptions): Promise<boolean> {
return this.#createDialogPromise(
dialogOptions,
false,
{ method: "confirm", title, message, timeout: dialogOptions?.timeout },
response =>
"cancelled" in response && response.cancelled
? false
: "confirmed" in response
? response.confirmed
: false,
);
}
input(
title: string,
placeholder?: string,
dialogOptions?: ExtensionUIDialogOptions,
): Promise<string | undefined> {
return this.#createDialogPromise(
dialogOptions,
undefined,
{ method: "input", title, placeholder, timeout: dialogOptions?.timeout },
response =>
"cancelled" in response && response.cancelled
? undefined
: "value" in response
? response.value
: undefined,
);
}
onTerminalInput(): () => void {
// Raw terminal input not supported in RPC mode
return () => {};
}
notify(message: string, type?: "info" | "warning" | "error"): void {
// Fire and forget - no response needed
this.output({
type: "extension_ui_request",
id: Snowflake.next() as string,
method: "notify",
message,
notifyType: type,
} as RpcExtensionUIRequest);
}
setStatus(key: string, text: string | undefined): void {
// Fire and forget - no response needed
this.output({
type: "extension_ui_request",
id: Snowflake.next() as string,
method: "setStatus",
statusKey: key,
statusText: text,
} as RpcExtensionUIRequest);
}
setWorkingMessage(_message?: string): void {
// Not supported in RPC mode
}
setWidget(key: string, content: unknown): void {
// Only support string arrays in RPC mode - factory functions are ignored
if (content === undefined || Array.isArray(content)) {
this.output({
type: "extension_ui_request",
id: Snowflake.next() as string,
method: "setWidget",
widgetKey: key,
widgetLines: content as string[] | undefined,
} as RpcExtensionUIRequest);
}
// Component factories are not supported in RPC mode - would need TUI access
}
setFooter(_factory: unknown): void {
// Custom footer not supported in RPC mode - requires TUI access
}
setHeader(_factory: unknown): void {
// Custom header not supported in RPC mode - requires TUI access
}
setTitle(title: string): void {
// Fire and forget - host can implement terminal title control
this.output({
type: "extension_ui_request",
id: Snowflake.next() as string,
method: "setTitle",
title,
} as RpcExtensionUIRequest);
}
async custom(): Promise<never> {
// Custom UI not supported in RPC mode
return undefined as never;
}
pasteToEditor(text: string): void {
// Paste handling not supported in RPC mode - falls back to setEditorText
this.setEditorText(text);
}
setEditorText(text: string): void {
// Fire and forget - host can implement editor control
this.output({
type: "extension_ui_request",
id: Snowflake.next() as string,
method: "set_editor_text",
text,
} as RpcExtensionUIRequest);
}
getEditorText(): string {
// Synchronous method can't wait for RPC response
// Host should track editor state locally if needed
return "";
}
async editor(title: string, prefill?: string): Promise<string | undefined> {
const id = Snowflake.next() as string;
const { promise, resolve, reject } = Promise.withResolvers<string | undefined>();
this.pendingRequests.set(id, {
resolve: (response: RpcExtensionUIResponse) => {
this.pendingRequests.delete(id);
if ("cancelled" in response && response.cancelled) {
resolve(undefined);
} else if ("value" in response) {
resolve(response.value);
} else {
resolve(undefined);
}
},
reject,
});
this.output({
type: "extension_ui_request",
id,
method: "editor",
title,
prefill,
} as RpcExtensionUIRequest);
return promise;
}
get theme(): Theme {
return theme;
}
getAllThemes(): Promise<{ name: string; path: string | undefined }[]> {
return Promise.resolve([]);
}
getTheme(_name: string): Promise<Theme | undefined> {
return Promise.resolve(undefined);
}
setTheme(_theme: string | Theme): Promise<{ success: boolean; error?: string }> {
// Theme switching not supported in RPC mode
return Promise.resolve({ success: false, error: "Theme switching not supported in RPC mode" });
}
getToolsExpanded() {
// Tool expansion not supported in RPC mode - no TUI
return false;
}
setToolsExpanded(_expanded: boolean) {
// Tool expansion not supported in RPC mode - no TUI
}
setEditorComponent(): void {
// Custom editor components not supported in RPC mode
}
}
// Set up extensions with RPC-based UI context
const extensionRunner = session.extensionRunner;
if (extensionRunner) {
extensionRunner.initialize(
// ExtensionActions
{
sendMessage: (message, options) => {
session.sendCustomMessage(message, options).catch(e => {
output(error(undefined, "extension_send", e.message));
});
},
sendUserMessage: (content, options) => {
session.sendUserMessage(content, options).catch(e => {
output(error(undefined, "extension_send_user", e.message));
});
},
appendEntry: (customType, data) => {
session.sessionManager.appendCustomEntry(customType, data);
},
setLabel: (targetId, label) => {
session.sessionManager.appendLabelChange(targetId, label);
},
getActiveTools: () => session.getActiveToolNames(),
getAllTools: () => session.getAllToolNames(),
setActiveTools: (toolNames: string[]) => session.setActiveToolsByName(toolNames),
getCommands: () => [],
setModel: async model => {
const key = await session.modelRegistry.getApiKey(model);
if (!key) return false;
await session.setModel(model);
return true;
},
getThinkingLevel: () => session.thinkingLevel,
setThinkingLevel: level => session.setThinkingLevel(level),
},
// ExtensionContextActions
{
getModel: () => session.agent.state.model,
isIdle: () => !session.isStreaming,
abort: () => session.abort(),
hasPendingMessages: () => session.queuedMessageCount > 0,
shutdown: () => {
shutdownState.requested = true;
},
getContextUsage: () => session.getContextUsage(),
getSystemPrompt: () => session.systemPrompt,
compact: async instructionsOrOptions => {
const instructions = typeof instructionsOrOptions === "string" ? instructionsOrOptions : undefined;
const options =
instructionsOrOptions && typeof instructionsOrOptions === "object"
? instructionsOrOptions
: undefined;
await session.compact(instructions, options);
},
},
// ExtensionCommandContextActions - commands invokable via prompt("/command")
{
getContextUsage: () => session.getContextUsage(),
waitForIdle: () => session.agent.waitForIdle(),
newSession: async options => {
const success = await session.newSession({ parentSession: options?.parentSession });
// Note: setup callback runs but no UI feedback in RPC mode
if (success && options?.setup) {
await options.setup(session.sessionManager);
}
return { cancelled: !success };
},
branch: async entryId => {
const result = await session.branch(entryId);
return { cancelled: result.cancelled };
},
navigateTree: async (targetId, options) => {
const result = await session.navigateTree(targetId, { summarize: options?.summarize });
return { cancelled: result.cancelled };
},
switchSession: async sessionPath => {
const success = await session.switchSession(sessionPath);
return { cancelled: !success };
},
reload: async () => {
await session.reload();
},
compact: async instructionsOrOptions => {
const instructions = typeof instructionsOrOptions === "string" ? instructionsOrOptions : undefined;
const options =
instructionsOrOptions && typeof instructionsOrOptions === "object"
? instructionsOrOptions
: undefined;
await session.compact(instructions, options);
},
},
new RpcExtensionUIContext(pendingExtensionRequests, output),
);
extensionRunner.onError(err => {
output({ type: "extension_error", extensionPath: err.extensionPath, event: err.event, error: err.error });
});
// Emit session_start event
await extensionRunner.emit({
type: "session_start",
});
}
// Output all agent events as JSON
session.subscribe(event => {
output(event);
});
// Handle a single command
const handleCommand = async (command: RpcCommand): Promise<RpcResponse> => {
const id = command.id;
switch (command.type) {
// =================================================================
// Prompting
// =================================================================
case "prompt": {
// Don't await - events will stream
// Extension commands are executed immediately, file prompt templates are expanded
// If streaming and streamingBehavior specified, queues via steer/followUp
session
.prompt(command.message, {
images: command.images,
streamingBehavior: command.streamingBehavior,
})
.catch(e => output(error(id, "prompt", e.message)));
return success(id, "prompt");
}
case "steer": {
await session.steer(command.message, command.images);
return success(id, "steer");
}
case "follow_up": {
await session.followUp(command.message, command.images);
return success(id, "follow_up");
}
case "abort": {
await session.abort();
return success(id, "abort");
}
case "abort_and_prompt": {
await session.abort();
session
.prompt(command.message, { images: command.images })
.catch(e => output(error(id, "abort_and_prompt", e.message)));
return success(id, "abort_and_prompt");
}
case "new_session": {
const options = command.parentSession ? { parentSession: command.parentSession } : undefined;
const cancelled = !(await session.newSession(options));
return success(id, "new_session", { cancelled });
}
// =================================================================
// State
// =================================================================
case "get_state": {
const state: RpcSessionState = {
model: session.model,
thinkingLevel: session.thinkingLevel,
isStreaming: session.isStreaming,
isCompacting: session.isCompacting,
steeringMode: session.steeringMode,
followUpMode: session.followUpMode,
interruptMode: session.interruptMode,
sessionFile: session.sessionFile,
sessionId: session.sessionId,
sessionName: session.sessionName,
autoCompactionEnabled: session.autoCompactionEnabled,
messageCount: session.messages.length,
queuedMessageCount: session.queuedMessageCount,
};
return success(id, "get_state", state);
}
// =================================================================
// Model
// =================================================================
case "set_model": {
const models = session.getAvailableModels();
const model = models.find(m => m.provider === command.provider && m.id === command.modelId);
if (!model) {
return error(id, "set_model", `Model not found: ${command.provider}/${command.modelId}`);
}
await session.setModel(model);
return success(id, "set_model", model);
}
case "cycle_model": {
const result = await session.cycleModel();
if (!result) {
return success(id, "cycle_model", null);
}
return success(id, "cycle_model", result);
}
case "get_available_models": {
const models = session.getAvailableModels();
return success(id, "get_available_models", { models });
}
// =================================================================
// Thinking
// =================================================================
case "set_thinking_level": {
session.setThinkingLevel(command.level);
return success(id, "set_thinking_level");
}
case "cycle_thinking_level": {
const level = session.cycleThinkingLevel();
if (!level) {
return success(id, "cycle_thinking_level", null);
}
return success(id, "cycle_thinking_level", { level });
}
// =================================================================
// Queue Modes
// =================================================================
case "set_steering_mode": {
session.setSteeringMode(command.mode);
return success(id, "set_steering_mode");
}
case "set_follow_up_mode": {
session.setFollowUpMode(command.mode);
return success(id, "set_follow_up_mode");
}
case "set_interrupt_mode": {
session.setInterruptMode(command.mode);
return success(id, "set_interrupt_mode");
}
// =================================================================
// Compaction
// =================================================================
case "compact": {
const result = await session.compact(command.customInstructions);
return success(id, "compact", result);
}
case "set_auto_compaction": {
session.setAutoCompactionEnabled(command.enabled);
return success(id, "set_auto_compaction");
}
// =================================================================
// Retry
// =================================================================
case "set_auto_retry": {
session.setAutoRetryEnabled(command.enabled);
return success(id, "set_auto_retry");
}
case "abort_retry": {
session.abortRetry();
return success(id, "abort_retry");
}
// =================================================================
// Bash
// =================================================================
case "bash": {
const result = await session.executeBash(command.command);
return success(id, "bash", result);
}
case "abort_bash": {
session.abortBash();
return success(id, "abort_bash");
}
// =================================================================
// Session
// =================================================================
case "get_session_stats": {
const stats = session.getSessionStats();
return success(id, "get_session_stats", stats);
}
case "export_html": {
const path = await session.exportToHtml(command.outputPath);
return success(id, "export_html", { path });
}
case "switch_session": {
const cancelled = !(await session.switchSession(command.sessionPath));
return success(id, "switch_session", { cancelled });
}
case "branch": {
const result = await session.branch(command.entryId);
return success(id, "branch", { text: result.selectedText, cancelled: result.cancelled });
}
case "get_branch_messages": {
const messages = session.getUserMessagesForBranching();
return success(id, "get_branch_messages", { messages });
}
case "get_last_assistant_text": {
const text = session.getLastAssistantText();
return success(id, "get_last_assistant_text", { text });
}
case "set_session_name": {
const name = command.name.trim();
if (!name) {
return error(id, "set_session_name", "Session name cannot be empty");
}
session.setSessionName(name);
return success(id, "set_session_name");
}
// =================================================================
// Messages
// =================================================================
case "get_messages": {
return success(id, "get_messages", { messages: session.messages });
}
default: {
const unknownCommand = command as { type: string };
return error(undefined, unknownCommand.type, `Unknown command: ${unknownCommand.type}`);
}
}
};
/**
* Check if shutdown was requested and perform shutdown if so.
* Called after handling each command when waiting for the next command.
*/
async function checkShutdownRequested(): Promise<void> {
if (!shutdownState.requested) return;
if (extensionRunner?.hasHandlers("session_shutdown")) {
await extensionRunner.emit({ type: "session_shutdown" });
}
process.exit(0);
}
// Listen for JSON input using Bun's stdin
for await (const parsed of readJsonl(Bun.stdin.stream())) {
try {
// Handle extension UI responses
if ((parsed as RpcExtensionUIResponse).type === "extension_ui_response") {
const response = parsed as RpcExtensionUIResponse;
const pending = pendingExtensionRequests.get(response.id);
if (pending) {
pending.resolve(response);
}
continue;
}
// Handle regular commands
const command = parsed as RpcCommand;
const response = await handleCommand(command);
output(response);
// Check for deferred shutdown request (idle between commands)
await checkShutdownRequested();
} catch (e: any) {
output(error(undefined, "parse", `Failed to parse command: ${e.message}`));
}
}
// stdin closed — RPC client is gone, exit cleanly
process.exit(0);
}