Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
20 kB · 672 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673/** * RPC Client for programmatic access to the coding agent. * * Spawns the agent in RPC mode and provides a typed API for all operations. */
import type { ChildProcess } from "node:child_process";import type { AgentEvent, AgentMessage, ThinkingLevel } from "@earendil-works/pi-agent-core";import type { ImageContent } from "@earendil-works/pi-ai";import type { AgentSessionMessageReceipt, AgentSessionMessageSafetyStatus } from "../../core/agent-messages.js";import type { BashResult } from "../../core/bash-executor.js";import type { CompactionResult } from "../../core/compaction/index.js";import type { AgentCronJob, AgentHeartbeatDeliveryMode, AgentHeartbeatManagementAction, AgentHeartbeatUpdateAction,} from "../../core/cron-jobs.js";import type { RefinementResult } from "../../core/refinement/index.js";import type { SessionStats } from "../../core/session-stats.js";import { spawnHidden } from "../../utils/child-process.js";import type { AgentConnectionHeartbeat } from "../agent-connection/types.js";import { attachJsonlLineReader, serializeJsonLine } from "./jsonl.js";import type { RpcCommand, RpcObservedSessionEvent, RpcResponse, RpcSessionState, RpcSlashCommand,} from "./rpc-types.js";
// ============================================================================// Types// ============================================================================
/** Extended response timeout for refine requests, which run an LLM pass. */export const REFINE_REQUEST_TIMEOUT_MS = 10 * 60 * 1000;
/** Distributive Omit that works with union types */type DistributiveOmit<T, K extends keyof T> = T extends unknown ? Omit<T, K> : never;
/** RpcCommand without the id field (for internal send) */type RpcCommandBody = DistributiveOmit<RpcCommand, "id">;
export interface RpcClientOptions { /** Path to the CLI entry point (default: searches for dist/cli.js) */ cliPath?: string; /** Working directory for the agent */ cwd?: string; /** Environment variables */ env?: Record<string, string>; /** Provider to use */ provider?: string; /** Model ID to use */ model?: string; /** Additional CLI arguments */ args?: string[];}
export interface ModelInfo { provider: string; id: string; contextWindow: number; reasoning: boolean;}
export type RpcEventListener = (event: AgentEvent) => void;export type RpcObservedSessionListener = (event: RpcObservedSessionEvent) => void;
// ============================================================================// RPC Client// ============================================================================
export class RpcClient { private process: ChildProcess | null = null; private stopReadingStdout: (() => void) | null = null; private eventListeners: RpcEventListener[] = []; private observedSessionListeners: RpcObservedSessionListener[] = []; private pendingRequests: Map<string, { resolve: (response: RpcResponse) => void; reject: (error: Error) => void }> = new Map(); private requestId = 0; private stderr = "";
constructor(private options: RpcClientOptions = {}) {}
/** * Start the RPC agent process. */ async start(): Promise<void> { if (this.process) { throw new Error("Client already started"); }
const cliPath = this.options.cliPath ?? "dist/cli.js"; const args = ["--mode", "rpc"];
if (this.options.provider) { args.push("--provider", this.options.provider); } if (this.options.model) { args.push("--model", this.options.model); } if (this.options.args) { args.push(...this.options.args); }
this.process = spawnHidden("node", [cliPath, ...args], { cwd: this.options.cwd, env: { ...process.env, ...this.options.env }, stdio: ["pipe", "pipe", "pipe"], });
// Collect stderr for debugging this.process.stderr?.on("data", (data) => { this.stderr += data.toString(); process.stderr.write(data); });
// Set up strict JSONL reader for stdout. this.stopReadingStdout = attachJsonlLineReader(this.process.stdout!, (line) => { this.handleLine(line); });
// Wait a moment for process to initialize await new Promise((resolve) => setTimeout(resolve, 100));
if (this.process.exitCode !== null) { throw new Error(`Agent process exited immediately with code ${this.process.exitCode}. Stderr: ${this.stderr}`); } }
/** * Stop the RPC agent process. */ async stop(): Promise<void> { if (!this.process) return;
this.stopReadingStdout?.(); this.stopReadingStdout = null; this.process.kill("SIGTERM");
// Wait for process to exit await new Promise<void>((resolve) => { const timeout = setTimeout(() => { this.process?.kill("SIGKILL"); resolve(); }, 1000);
this.process?.on("exit", () => { clearTimeout(timeout); resolve(); }); });
this.process = null; this.pendingRequests.clear(); }
/** * Subscribe to agent events. */ onEvent(listener: RpcEventListener): () => void { this.eventListeners.push(listener); return () => { const index = this.eventListeners.indexOf(listener); if (index !== -1) { this.eventListeners.splice(index, 1); } }; }
onObservedSessionEvent(listener: RpcObservedSessionListener): () => void { this.observedSessionListeners.push(listener); return () => { const index = this.observedSessionListeners.indexOf(listener); if (index !== -1) { this.observedSessionListeners.splice(index, 1); } }; }
/** * Get collected stderr output (useful for debugging). */ getStderr(): string { return this.stderr; }
// ========================================================================= // Command Methods // =========================================================================
/** * Send a prompt to the agent. * Returns immediately after sending; use onEvent() to receive streaming events. * Use waitForIdle() to wait for completion. */ async prompt(message: string, images?: ImageContent[]): Promise<void> { await this.send({ type: "prompt", message, images }); }
/** * Queue a steering message to interrupt the agent mid-run. */ async steer(message: string, images?: ImageContent[]): Promise<void> { await this.send({ type: "steer", message, images }); }
/** * Queue a follow-up message to be processed after the agent finishes. */ async followUp(message: string, images?: ImageContent[]): Promise<void> { await this.send({ type: "follow_up", message, images }); }
/** * Abort current operation. */ async abort(): Promise<void> { await this.send({ type: "abort" }); }
/** * Start a new session, optionally with parent tracking. * @param parentSession - Optional parent session path for lineage tracking * @returns Object with `cancelled: true` if an extension cancelled the new session */ async newSession(parentSession?: string): Promise<{ cancelled: boolean }> { const response = await this.send({ type: "new_session", parentSession }); return this.getData(response); }
/** * Get current session state. */ async getState(): Promise<RpcSessionState> { const response = await this.send({ type: "get_state" }); return this.getData(response); }
/** * Set model by provider and ID. */ async setModel(provider: string, modelId: string): Promise<{ provider: string; id: string }> { const response = await this.send({ type: "set_model", provider, modelId }); return this.getData(response); }
/** * Cycle to next model. */ async cycleModel(): Promise<{ model: { provider: string; id: string }; thinkingLevel: ThinkingLevel; isScoped: boolean; } | null> { const response = await this.send({ type: "cycle_model" }); return this.getData(response); }
/** * Get list of available models. */ async getAvailableModels(): Promise<ModelInfo[]> { const response = await this.send({ type: "get_available_models" }); return this.getData<{ models: ModelInfo[] }>(response).models; }
/** * Set thinking level. */ async setThinkingLevel(level: ThinkingLevel): Promise<void> { await this.send({ type: "set_thinking_level", level }); }
/** * Cycle thinking level. */ async cycleThinkingLevel(): Promise<{ level: ThinkingLevel } | null> { const response = await this.send({ type: "cycle_thinking_level" }); return this.getData(response); }
/** * Set steering mode. */ async setSteeringMode(mode: "all" | "one-at-a-time"): Promise<void> { await this.send({ type: "set_steering_mode", mode }); }
/** * Set follow-up mode. */ async setFollowUpMode(mode: "all" | "one-at-a-time"): Promise<void> { await this.send({ type: "set_follow_up_mode", mode }); }
/** * Compact session context. */ async compact(customInstructions?: string): Promise<CompactionResult> { const response = await this.send({ type: "compact", customInstructions }); return this.getData(response); }
/** * Refine editable continual harness state. */ async refine( options: { instructions?: string; rollbackId?: string; global?: boolean } = {}, ): Promise<RefinementResult> { // Refinement runs an LLM pass that routinely exceeds the default 30s response // timeout, so use the same extended window as the daemon refine path. const command = { type: "refine", instructions: options.instructions, rollbackId: options.rollbackId } as { type: "refine"; instructions?: string; rollbackId?: string; global?: boolean; }; if (options.global !== undefined) { command.global = options.global; } const response = await this.send(command, REFINE_REQUEST_TIMEOUT_MS); return this.getData(response); }
/** * Set auto-compaction enabled/disabled. */ async setAutoCompaction(enabled: boolean): Promise<void> { await this.send({ type: "set_auto_compaction", enabled }); }
/** * Set auto-retry enabled/disabled. */ async setAutoRetry(enabled: boolean): Promise<void> { await this.send({ type: "set_auto_retry", enabled }); }
/** * Abort in-progress retry. */ async abortRetry(): Promise<void> { await this.send({ type: "abort_retry" }); }
/** * Execute a bash command. */ async bash(command: string): Promise<BashResult> { const response = await this.send({ type: "bash", command }); return this.getData(response); }
/** * Abort running bash command. */ async abortBash(): Promise<void> { await this.send({ type: "abort_bash" }); }
/** * Get session statistics. */ async getSessionStats(): Promise<SessionStats> { const response = await this.send({ type: "get_session_stats" }); return this.getData(response); }
/** * Export session to HTML. */ async exportHtml(outputPath?: string): Promise<{ path: string }> { const response = await this.send({ type: "export_html", outputPath }); return this.getData(response); }
/** * Switch to a different session file. * @returns Object with `cancelled: true` if an extension cancelled the switch */ async switchSession(sessionPath: string): Promise<{ cancelled: boolean }> { const response = await this.send({ type: "switch_session", sessionPath }); return this.getData(response); }
/** * Fork from a specific message. * @returns Object with `text` (the message text) and `cancelled` (if extension cancelled) */ async fork(entryId: string): Promise<{ text: string; cancelled: boolean }> { const response = await this.send({ type: "fork", entryId }); return this.getData(response); }
/** * Clone the current active branch into a new session. * @returns Object with `cancelled: true` if an extension cancelled the clone */ async clone(): Promise<{ cancelled: boolean }> { const response = await this.send({ type: "clone" }); return this.getData(response); }
/** * Get messages available for forking. */ async getForkMessages(): Promise<Array<{ entryId: string; text: string }>> { const response = await this.send({ type: "get_fork_messages" }); return this.getData<{ messages: Array<{ entryId: string; text: string }> }>(response).messages; }
/** * Get text of last assistant message. */ async getLastAssistantText(): Promise<string | null> { const response = await this.send({ type: "get_last_assistant_text" }); return this.getData<{ text: string | null }>(response).text; }
/** * Set the session display name. */ async setSessionName(name: string): Promise<void> { await this.send({ type: "set_session_name", name }); }
/** * Get all messages in the session. */ async getMessages(): Promise<AgentMessage[]> { const response = await this.send({ type: "get_messages" }); return this.getData<{ messages: AgentMessage[] }>(response).messages; }
async sendAgentMessage(targetActiveSessionId: string, message: string): Promise<AgentSessionMessageReceipt> { const response = await this.send({ type: "send_message", targetActiveSessionId, message, }); return this.getData(response); }
async getAgentMessageStatus(): Promise<AgentSessionMessageSafetyStatus> { return this.getData(await this.send({ type: "agent_messages_status" })); }
async pauseAgentMessages(): Promise<AgentSessionMessageSafetyStatus> { return this.getData(await this.send({ type: "agent_messages_pause" })); }
async resumeAgentMessages(): Promise<AgentSessionMessageSafetyStatus> { return this.getData(await this.send({ type: "agent_messages_resume" })); }
async clearAgentMessages(): Promise<number> { const response = await this.send({ type: "agent_messages_clear" }); return this.getData<{ cleared: number }>(response).cleared; }
async listSchedules(includeInactive?: boolean): Promise<AgentCronJob[]> { const response = await this.send({ type: "list_schedules", includeInactive }); return this.getData<{ jobs: AgentCronJob[] }>(response).jobs; }
async addSchedule(schedule: string, prompt: string): Promise<AgentCronJob> { const response = await this.send({ type: "add_schedule", schedule, prompt }); return this.getData<{ job: AgentCronJob }>(response).job; }
async cancelSchedule(jobId: string): Promise<AgentCronJob> { const response = await this.send({ type: "cancel_schedule", jobId }); return this.getData<{ job: AgentCronJob }>(response).job; }
async listHeartbeats(): Promise<AgentConnectionHeartbeat[]> { const response = await this.send({ type: "list_heartbeats" }); return this.getData<{ heartbeats: AgentConnectionHeartbeat[] }>(response).heartbeats; }
async getHeartbeat(): Promise<AgentCronJob | null> { const response = await this.send({ type: "get_heartbeat" }); return this.getData<{ heartbeat: AgentCronJob | null }>(response).heartbeat; }
async setHeartbeat( schedule: string, prompt: string, deliveryMode?: AgentHeartbeatDeliveryMode, ): Promise<AgentCronJob> { const response = await this.send({ type: "set_heartbeat", schedule, prompt, deliveryMode }); const heartbeat = this.getData<{ heartbeat: AgentCronJob | null }>(response).heartbeat; if (!heartbeat) { throw new Error("Daemon did not return the created heartbeat"); } return heartbeat; }
async updateHeartbeat(action: AgentHeartbeatUpdateAction): Promise<AgentCronJob | null> { const response = await this.send({ type: "update_heartbeat", action }); return this.getData<{ heartbeat: AgentCronJob | null }>(response).heartbeat; }
async manageHeartbeat( activeSessionId: string, jobId: string, action: AgentHeartbeatManagementAction, ): Promise<AgentCronJob> { const response = await this.send({ type: "manage_heartbeat", activeSessionId, jobId, action }); const heartbeat = this.getData<{ heartbeat: AgentCronJob | null }>(response).heartbeat; if (!heartbeat) { throw new Error("Daemon did not return the managed heartbeat"); } return heartbeat; }
async observe(activeSessionId: string): Promise<AgentMessage[]> { const response = await this.send({ type: "observe", activeSessionId }); return this.getData<{ messages: AgentMessage[] }>(response).messages; }
async unobserve(activeSessionId: string): Promise<void> { await this.send({ type: "unobserve", activeSessionId }); }
/** * Get available commands (extension commands, prompt templates, skills). */ async getCommands(): Promise<RpcSlashCommand[]> { const response = await this.send({ type: "get_commands" }); return this.getData<{ commands: RpcSlashCommand[] }>(response).commands; }
// ========================================================================= // Helpers // =========================================================================
/** * Wait for agent to become idle (no streaming). * Resolves when agent_end event is received. */ waitForIdle(timeout = 60000): Promise<void> { return new Promise((resolve, reject) => { const timer = setTimeout(() => { unsubscribe(); reject(new Error(`Timeout waiting for agent to become idle. Stderr: ${this.stderr}`)); }, timeout);
const unsubscribe = this.onEvent((event) => { if (event.type === "agent_end") { clearTimeout(timer); unsubscribe(); resolve(); } }); }); }
/** * Collect events until agent becomes idle. */ collectEvents(timeout = 60000): Promise<AgentEvent[]> { return new Promise((resolve, reject) => { const events: AgentEvent[] = []; const timer = setTimeout(() => { unsubscribe(); reject(new Error(`Timeout collecting events. Stderr: ${this.stderr}`)); }, timeout);
const unsubscribe = this.onEvent((event) => { events.push(event); if (event.type === "agent_end") { clearTimeout(timer); unsubscribe(); resolve(events); } }); }); }
/** * Send prompt and wait for completion, returning all events. */ async promptAndWait(message: string, images?: ImageContent[], timeout = 60000): Promise<AgentEvent[]> { const eventsPromise = this.collectEvents(timeout); await this.prompt(message, images); return eventsPromise; }
// ========================================================================= // Internal // =========================================================================
private handleLine(line: string): void { try { const data = JSON.parse(line);
// Check if it's a response to a pending request if (data.type === "response" && data.id && this.pendingRequests.has(data.id)) { const pending = this.pendingRequests.get(data.id)!; this.pendingRequests.delete(data.id); pending.resolve(data as RpcResponse); return; } if (data.type === "observed_session_event" || data.type === "observed_session_closed") { for (const listener of [...this.observedSessionListeners]) { try { listener(data as RpcObservedSessionEvent); } catch { // Listener failures must not block other RPC subscribers. } } return; }
// Otherwise it's an event for (const listener of [...this.eventListeners]) { try { listener(data as AgentEvent); } catch { // Listener failures must not block other RPC subscribers. } } } catch { // Ignore non-JSON lines } }
private async send(command: RpcCommandBody, timeoutMs = 30000): Promise<RpcResponse> { if (!this.process?.stdin) { throw new Error("Client not started"); }
const id = `req_${++this.requestId}`; const fullCommand = { ...command, id } as RpcCommand;
return new Promise((resolve, reject) => { this.pendingRequests.set(id, { resolve, reject });
const timeout = setTimeout(() => { this.pendingRequests.delete(id); reject(new Error(`Timeout waiting for response to ${command.type}. Stderr: ${this.stderr}`)); }, timeoutMs);
this.pendingRequests.set(id, { resolve: (response) => { clearTimeout(timeout); resolve(response); }, reject: (error) => { clearTimeout(timeout); reject(error); }, });
this.process!.stdin!.write(serializeJsonLine(fullCommand)); }); }
private getData<T>(response: RpcResponse): T { if (!response.success) { const errorResponse = response as Extract<RpcResponse, { success: false }>; throw new Error(errorResponse.error); } // Type assertion: we trust response.data matches T based on the command sent. // This is safe because each public method specifies the correct T for its command. const successResponse = response as Extract<RpcResponse, { success: true; data: unknown }>; return successResponse.data as T; }}