diff --git a/packages/coding-agent/src/modes/index.ts b/packages/coding-agent/src/modes/index.ts index 8de01c219..e1dea209b 100644 --- a/packages/coding-agent/src/modes/index.ts +++ b/packages/coding-agent/src/modes/index.ts @@ -8,39 +8,9 @@ import { postmortem } from "@oh-my-pi/pi-utils"; * barrel does not pull print, RPC server, or ACP server mode into the normal * TUI graph. */ -export { InteractiveMode, type InteractiveModeOptions } from "./interactive-mode"; -export { - defineRpcClientTool, - type ModelInfo, - RpcClient, - type RpcClientCustomTool, - type RpcClientOptions, - type RpcClientToolContext, - type RpcClientToolResult, - type RpcEventListener, - type RpcSessionEventListener, - type RpcSubagentEventListener, - type RpcSubagentLifecycleListener, - type RpcSubagentProgressListener, -} from "./rpc/rpc-client"; -export type { - RpcCommand, - RpcHostToolCallRequest, - RpcHostToolCancelRequest, - RpcHostToolDefinition, - RpcHostToolResult, - RpcHostToolUpdate, - RpcResponse, - RpcSessionEventFrame, - RpcSessionState, - RpcSubagentEventFrame, - RpcSubagentFrame, - RpcSubagentLifecycleFrame, - RpcSubagentMessagesResult, - RpcSubagentProgressFrame, - RpcSubagentSnapshot, - RpcSubagentSubscriptionLevel, -} from "./rpc/rpc-types"; +export * from "./interactive-mode"; +export * from "./rpc/rpc-client"; +export * from "./rpc/rpc-types"; postmortem.register("terminal-restore", () => { emergencyTerminalRestore(); diff --git a/packages/coding-agent/src/modes/rpc/rpc-client.ts b/packages/coding-agent/src/modes/rpc/rpc-client.ts index ff7b88740..e847ec735 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-client.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-client.ts @@ -9,6 +9,7 @@ import type { AgentEvent, AgentMessage, AgentToolResult, ThinkingLevel } from "@ import type { CompactionResult } from "@oh-my-pi/pi-agent-core/compaction"; import type { ImageContent, Model } from "@oh-my-pi/pi-ai"; import { isRecord, ptree, readJsonl } from "@oh-my-pi/pi-utils"; +import type { FileSink } from "bun"; import type { BashResult } from "../../exec/bash-executor"; import type { AgentSessionEvent, SessionStats } from "../../session/agent-session"; import type { @@ -353,7 +354,7 @@ export class RpcClient { } /** - * Subscribe to subagent lifecycle frames emitted by the task tool. + * Subscribe to subagent lifecycle frames after setSubagentSubscription("progress" | "events"). */ onSubagentLifecycle(listener: RpcSubagentLifecycleListener): () => void { this.#subagentLifecycleListeners.add(listener); @@ -361,7 +362,7 @@ export class RpcClient { } /** - * Subscribe to aggregated subagent progress frames emitted by the task tool. + * Subscribe to aggregated subagent progress frames after setSubagentSubscription("progress" | "events"). */ onSubagentProgress(listener: RpcSubagentProgressListener): () => void { this.#subagentProgressListeners.add(listener); @@ -449,8 +450,8 @@ export class RpcClient { } /** - * Configure subagent frames emitted by the RPC server. - * Progress emits lifecycle/progress frames; events additionally emits raw subagent session events. + * Configure subagent frames emitted by the RPC server. Servers default to "off". + * "progress" emits lifecycle/progress frames; "events" additionally emits raw subagent session events. */ async setSubagentSubscription(level: RpcSubagentSubscriptionLevel): Promise { const response = await this.#send({ type: "set_subagent_subscription", level }); @@ -939,7 +940,7 @@ export class RpcClient { if (!this.#process?.stdin) { throw new Error("Client not started"); } - const stdin = this.#process.stdin as import("bun").FileSink; + const stdin = this.#process.stdin as FileSink; stdin.write(`${JSON.stringify(frame)}\n`); const flushResult = stdin.flush(); if (isPromise(flushResult)) { diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index a74880cd4..ceeb809cf 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -59,6 +59,47 @@ type RpcOutput = ( | object, ) => void; +export type RpcSessionChangeCommand = Extract< + RpcCommand, + { type: "new_session" } | { type: "switch_session" } | { type: "branch" } +>; + +export type RpcSessionChangeResult = + | { type: "new_session"; data: { cancelled: boolean } } + | { type: "switch_session"; data: { cancelled: boolean } } + | { type: "branch"; data: { text: string; cancelled: boolean } }; + +export type RpcSessionChangeSession = Pick; +export type RpcSubagentResetRegistry = Pick; + +export async function handleRpcSessionChange( + session: RpcSessionChangeSession, + command: RpcSessionChangeCommand, + subagentRegistry?: RpcSubagentResetRegistry, +): Promise { + switch (command.type) { + case "new_session": { + const options = command.parentSession ? { parentSession: command.parentSession } : undefined; + const cancelled = !(await session.newSession(options)); + if (!cancelled) subagentRegistry?.clear(); + return { type: "new_session", data: { cancelled } }; + } + + case "switch_session": { + const cancelled = !(await session.switchSession(command.sessionPath)); + if (!cancelled) subagentRegistry?.clear(); + return { type: "switch_session", data: { cancelled } }; + } + + case "branch": { + const result = await session.branch(command.entryId); + if (!result.cancelled) subagentRegistry?.clear(); + return { type: "branch", data: { text: result.selectedText, cancelled: result.cancelled } }; + } + } + throw new Error("Unsupported RPC session change command"); +} + function normalizeHostToolDefinitions(tools: RpcHostToolDefinition[]): RpcHostToolDefinition[] { return tools.map((tool, index) => { const name = typeof tool.name === "string" ? tool.name.trim() : ""; @@ -516,9 +557,8 @@ export async function runRpcMode( } case "new_session": { - const options = command.parentSession ? { parentSession: command.parentSession } : undefined; - const cancelled = !(await session.newSession(options)); - return success(id, "new_session", { cancelled }); + const result = await handleRpcSessionChange(session, command, subagentRegistry); + return success(id, result.type, result.data); } // ================================================================= @@ -589,7 +629,10 @@ export async function runRpcMode( } case "get_subagents": { - return success(id, "get_subagents", { subagents: subagentRegistry?.getSubagents() ?? [] }); + if (!subagentRegistry) { + return error(id, "get_subagents", "Subagent event bus is unavailable"); + } + return success(id, "get_subagents", { subagents: subagentRegistry.getSubagents() }); } case "get_subagent_messages": { @@ -727,14 +770,10 @@ export async function runRpcMode( return success(id, "export_html", { path }); } - case "switch_session": { - const cancelled = !(await session.switchSession(command.sessionPath)); - return success(id, "switch_session", { cancelled }); - } - + case "switch_session": case "branch": { - const result = await session.branch(command.entryId); - return success(id, "branch", { text: result.selectedText, cancelled: result.cancelled }); + const result = await handleRpcSessionChange(session, command, subagentRegistry); + return success(id, result.type, result.data); } case "get_branch_messages": { @@ -894,8 +933,9 @@ export async function runRpcMode( // Check for deferred shutdown request (idle between commands) await checkShutdownRequested(); - } catch (e: any) { - output(error(undefined, "parse", `Failed to parse command: ${e.message}`)); + } catch (e: unknown) { + const message = e instanceof Error ? e.message : String(e); + output(error(undefined, "parse", `Failed to parse command: ${message}`)); } } diff --git a/packages/coding-agent/src/modes/rpc/rpc-subagents.ts b/packages/coding-agent/src/modes/rpc/rpc-subagents.ts index 614c696bb..6eeef3521 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-subagents.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-subagents.ts @@ -27,6 +27,8 @@ export interface RpcSubagentTranscriptSelector { type RpcSubagentOutput = (frame: RpcSubagentFrame) => void; +const MAX_RETAINED_TRANSCRIPT_REFERENCES = 256; + function isSessionMessageEntry(entry: FileEntry): entry is SessionMessageEntry { return entry.type === "message"; } @@ -35,6 +37,10 @@ function statusFromLifecycle(status: SubagentLifecyclePayload["status"]): AgentP return status === "started" ? "running" : status; } +function isTerminalLifecycleStatus(status: SubagentLifecyclePayload["status"]): boolean { + return status !== "started"; +} + export async function readRpcSubagentTranscript(sessionFile: string, fromByte = 0): Promise { let startByte = Number.isFinite(fromByte) ? Math.max(0, Math.trunc(fromByte)) : 0; const file = Bun.file(sessionFile); @@ -63,9 +69,10 @@ export async function readRpcSubagentTranscript(sessionFile: string, fromByte = export class RpcSubagentRegistry { #subagents = new Map(); + #transcriptSessionFilesBySubagentId = new Map(); #unsubscribers: Array<() => void> = []; #output: RpcSubagentOutput; - #subscriptionLevel: RpcSubagentSubscriptionLevel = "progress"; + #subscriptionLevel: RpcSubagentSubscriptionLevel = "off"; constructor(eventBus: EventBus, output: RpcSubagentOutput) { this.#output = output; @@ -86,6 +93,12 @@ export class RpcSubagentRegistry { for (const unsubscribe of this.#unsubscribers) unsubscribe(); this.#unsubscribers = []; this.#subagents.clear(); + this.#transcriptSessionFilesBySubagentId.clear(); + } + + clear(): void { + this.#subagents.clear(); + this.#transcriptSessionFilesBySubagentId.clear(); } setSubscriptionLevel(level: RpcSubagentSubscriptionLevel): void { @@ -100,8 +113,30 @@ export class RpcSubagentRegistry { return [...this.#subagents.values()].sort((a, b) => a.index - b.index || a.id.localeCompare(b.id)); } + #rememberTranscriptSession(subagentId: string, sessionFile: string | undefined): void { + if (!sessionFile) return; + this.#transcriptSessionFilesBySubagentId.delete(subagentId); + this.#transcriptSessionFilesBySubagentId.set(subagentId, sessionFile); + while (this.#transcriptSessionFilesBySubagentId.size > MAX_RETAINED_TRANSCRIPT_REFERENCES) { + const oldest = this.#transcriptSessionFilesBySubagentId.keys().next(); + if (oldest.done) break; + this.#transcriptSessionFilesBySubagentId.delete(oldest.value); + } + } + + #hasTranscriptSessionFile(sessionFile: string): boolean { + for (const snapshot of this.#subagents.values()) { + if (snapshot.sessionFile === sessionFile) return true; + } + for (const transcriptSessionFile of this.#transcriptSessionFilesBySubagentId.values()) { + if (transcriptSessionFile === sessionFile) return true; + } + return false; + } + handleLifecycle(payload: SubagentLifecyclePayload): void { const existing = this.#subagents.get(payload.id); + const sessionFile = payload.sessionFile ?? existing?.sessionFile; const snapshot: RpcSubagentSnapshot = { id: payload.id, index: payload.index, @@ -111,12 +146,17 @@ export class RpcSubagentRegistry { status: statusFromLifecycle(payload.status), task: existing?.task, assignment: existing?.assignment, - sessionFile: payload.sessionFile ?? existing?.sessionFile, + sessionFile, parentToolCallId: payload.parentToolCallId ?? existing?.parentToolCallId, lastUpdate: Date.now(), progress: existing?.progress, }; - this.#subagents.set(payload.id, snapshot); + this.#rememberTranscriptSession(payload.id, sessionFile); + if (isTerminalLifecycleStatus(payload.status)) { + this.#subagents.delete(payload.id); + } else { + this.#subagents.set(payload.id, snapshot); + } if (this.#subscriptionLevel !== "off") { this.#output({ type: "subagent_lifecycle", payload }); } @@ -125,6 +165,8 @@ export class RpcSubagentRegistry { handleProgress(payload: SubagentProgressPayload): void { const progress = payload.progress; const existing = this.#subagents.get(progress.id); + const sessionFile = payload.sessionFile ?? existing?.sessionFile; + this.#rememberTranscriptSession(progress.id, sessionFile); this.#subagents.set(progress.id, { id: progress.id, index: payload.index, @@ -134,7 +176,7 @@ export class RpcSubagentRegistry { status: progress.status, task: payload.task, assignment: payload.assignment, - sessionFile: payload.sessionFile ?? existing?.sessionFile, + sessionFile, lastUpdate: Date.now(), parentToolCallId: payload.parentToolCallId ?? existing?.parentToolCallId, progress, @@ -152,16 +194,15 @@ export class RpcSubagentRegistry { resolveSessionFile(selector: RpcSubagentTranscriptSelector): string { if (selector.subagentId) { const snapshot = this.#subagents.get(selector.subagentId); - if (!snapshot?.sessionFile) { + const sessionFile = snapshot?.sessionFile ?? this.#transcriptSessionFilesBySubagentId.get(selector.subagentId); + if (!sessionFile) { throw new Error(`Unknown subagent or session file unavailable: ${selector.subagentId}`); } - return snapshot.sessionFile; + return sessionFile; } if (selector.sessionFile) { - for (const snapshot of this.#subagents.values()) { - if (snapshot.sessionFile === selector.sessionFile) return selector.sessionFile; - } + if (this.#hasTranscriptSessionFile(selector.sessionFile)) return selector.sessionFile; throw new Error("Unknown subagent session file"); } diff --git a/packages/coding-agent/src/task/executor.ts b/packages/coding-agent/src/task/executor.ts index 2f1fae54d..a0ebd42c1 100644 --- a/packages/coding-agent/src/task/executor.ts +++ b/packages/coding-agent/src/task/executor.ts @@ -928,14 +928,7 @@ export async function runSubprocess(options: ExecutorOptions): Promise, @@ -428,7 +428,7 @@ export class TaskTool implements AgentTool agent.name === params.agent); if (!asyncEnabled || selectedAgent?.blocking === true) { - return this.#executeSync(_toolCallId, params, signal, onUpdate); + return this.#executeSync(toolCallId, params, signal, onUpdate); } const manager = this.session.asyncJobManager; @@ -438,12 +438,12 @@ export class TaskTool implements AgentTool, ); try { - const result = await this.#executeSync(_toolCallId, singleParams, runSignal, undefined, [ - uniqueId, - ]); + const result = await this.#executeSync(toolCallId, singleParams, runSignal, undefined, [uniqueId]); const finalText = result.content.find(part => part.type === "text")?.text ?? "(no output)"; const singleResult = result.details?.results[0]; // A missing per-task result means #executeSync failed at the @@ -708,7 +706,7 @@ export class TaskTool implements AgentTool, @@ -1036,7 +1034,7 @@ export class TaskTool implements AgentTool = {}): AgentProgress { }; } +function createRegistryWithSnapshot(): RpcSubagentRegistry { + const eventBus = new EventBus(); + const registry = new RpcSubagentRegistry(eventBus, () => {}); + eventBus.emit(TASK_SUBAGENT_LIFECYCLE_CHANNEL, { + id: "SubagentA", + index: 0, + agent: "task", + agentSource: "bundled", + status: "started", + sessionFile: "/tmp/subagent.jsonl", + } satisfies SubagentLifecyclePayload); + expect(registry.getSubagents()).toHaveLength(1); + return registry; +} + +type SessionChangeStubOptions = { + newSession?: boolean; + switchSession?: boolean; + branch?: { selectedText: string; cancelled: boolean }; +}; + +function createSessionChangeSession(options: SessionChangeStubOptions): RpcSessionChangeSession { + return { + newSession: async (_options?: unknown) => options.newSession ?? true, + switchSession: async (_sessionPath: string) => options.switchSession ?? true, + branch: async (_entryId: string) => options.branch ?? { selectedText: "branched text", cancelled: false }, + }; +} + describe("RPC subagent registry", () => { - test("emits progress frames and snapshots tracked subagents", () => { + test("defaults subagent frame emission to off while tracking snapshots", () => { const eventBus = new EventBus(); const frames: RpcSubagentFrame[] = []; const registry = new RpcSubagentRegistry(eventBus, frame => frames.push(frame)); @@ -69,6 +104,52 @@ describe("RPC subagent registry", () => { sessionFile: "/tmp/subagent.jsonl", progress: createProgress(), }; + const eventPayload: SubagentEventPayload = { + id: "SubagentA", + event: { type: "agent_start" }, + }; + + expect(registry.getSubscriptionLevel()).toBe("off"); + eventBus.emit(TASK_SUBAGENT_LIFECYCLE_CHANNEL, lifecycle); + eventBus.emit(TASK_SUBAGENT_PROGRESS_CHANNEL, progressPayload); + eventBus.emit(TASK_SUBAGENT_EVENT_CHANNEL, eventPayload); + + expect(frames).toHaveLength(0); + expect(registry.getSubagents()).toMatchObject([ + { + id: "SubagentA", + status: "running", + sessionFile: "/tmp/subagent.jsonl", + }, + ]); + registry.dispose(); + }); + + test("emits progress frames after explicit progress subscription and snapshots tracked subagents", () => { + const eventBus = new EventBus(); + const frames: RpcSubagentFrame[] = []; + const registry = new RpcSubagentRegistry(eventBus, frame => frames.push(frame)); + registry.setSubscriptionLevel("progress"); + const lifecycle: SubagentLifecyclePayload = { + id: "SubagentA", + index: 0, + agent: "task", + agentSource: "bundled", + description: "Worker", + status: "started", + sessionFile: "/tmp/subagent.jsonl", + parentToolCallId: "toolu_parent", + }; + const progressPayload: SubagentProgressPayload = { + index: 0, + agent: "task", + agentSource: "bundled", + task: "Do work", + assignment: "Implement work", + parentToolCallId: "toolu_parent", + sessionFile: "/tmp/subagent.jsonl", + progress: createProgress(), + }; eventBus.emit(TASK_SUBAGENT_LIFECYCLE_CHANNEL, lifecycle); eventBus.emit(TASK_SUBAGENT_PROGRESS_CHANNEL, progressPayload); @@ -88,19 +169,136 @@ describe("RPC subagent registry", () => { registry.dispose(); }); + test("clears stale snapshots when the active RPC session changes", () => { + const eventBus = new EventBus(); + const registry = new RpcSubagentRegistry(eventBus, () => {}); + eventBus.emit(TASK_SUBAGENT_LIFECYCLE_CHANNEL, { + id: "SubagentA", + index: 0, + agent: "task", + agentSource: "bundled", + status: "started", + sessionFile: "/tmp/subagent.jsonl", + } satisfies SubagentLifecyclePayload); + + expect(registry.getSubagents()).toHaveLength(1); + registry.clear(); + + expect(registry.getSubagents()).toHaveLength(0); + registry.dispose(); + }); + + test("clears stale snapshots after successful RPC session changes", async () => { + const cases: Array<{ + command: RpcSessionChangeCommand; + session: RpcSessionChangeSession; + expected: RpcSessionChangeResult; + }> = [ + { + command: { type: "new_session", parentSession: "/tmp/parent.jsonl" }, + session: createSessionChangeSession({ newSession: true }), + expected: { type: "new_session", data: { cancelled: false } }, + }, + { + command: { type: "switch_session", sessionPath: "/tmp/next.jsonl" }, + session: createSessionChangeSession({ switchSession: true }), + expected: { type: "switch_session", data: { cancelled: false } }, + }, + { + command: { type: "branch", entryId: "entry-1" }, + session: createSessionChangeSession({ branch: { selectedText: "Branch text", cancelled: false } }), + expected: { type: "branch", data: { text: "Branch text", cancelled: false } }, + }, + ]; + + for (const testCase of cases) { + const registry = createRegistryWithSnapshot(); + try { + const result = await handleRpcSessionChange(testCase.session, testCase.command, registry); + + expect(result).toEqual(testCase.expected); + expect(registry.getSubagents()).toHaveLength(0); + expect(() => registry.resolveSessionFile({ subagentId: "SubagentA" })).toThrow( + /Unknown subagent or session file unavailable/, + ); + } finally { + registry.dispose(); + } + } + }); + + test("keeps stale snapshots when RPC session changes are cancelled", async () => { + const cases: Array<{ + command: RpcSessionChangeCommand; + session: RpcSessionChangeSession; + expected: RpcSessionChangeResult; + }> = [ + { + command: { type: "new_session", parentSession: "/tmp/parent.jsonl" }, + session: createSessionChangeSession({ newSession: false }), + expected: { type: "new_session", data: { cancelled: true } }, + }, + { + command: { type: "switch_session", sessionPath: "/tmp/next.jsonl" }, + session: createSessionChangeSession({ switchSession: false }), + expected: { type: "switch_session", data: { cancelled: true } }, + }, + { + command: { type: "branch", entryId: "entry-1" }, + session: createSessionChangeSession({ branch: { selectedText: "", cancelled: true } }), + expected: { type: "branch", data: { text: "", cancelled: true } }, + }, + ]; + + for (const testCase of cases) { + const registry = createRegistryWithSnapshot(); + try { + const result = await handleRpcSessionChange(testCase.session, testCase.command, registry); + + expect(result).toEqual(testCase.expected); + expect(registry.getSubagents()).toMatchObject([{ id: "SubagentA" }]); + expect(registry.resolveSessionFile({ subagentId: "SubagentA" })).toBe("/tmp/subagent.jsonl"); + } finally { + registry.dispose(); + } + } + }); + + test("prunes terminal lifecycle snapshots while retaining transcript selectors", () => { + const eventBus = new EventBus(); + const registry = new RpcSubagentRegistry(eventBus, () => {}); + const sessionFile = "/tmp/subagent.jsonl"; + eventBus.emit(TASK_SUBAGENT_LIFECYCLE_CHANNEL, { + id: "SubagentA", + index: 0, + agent: "task", + agentSource: "bundled", + status: "started", + sessionFile, + } satisfies SubagentLifecyclePayload); + + expect(registry.getSubagents()).toHaveLength(1); + eventBus.emit(TASK_SUBAGENT_LIFECYCLE_CHANNEL, { + id: "SubagentA", + index: 0, + agent: "task", + agentSource: "bundled", + status: "completed", + sessionFile, + } satisfies SubagentLifecyclePayload); + + expect(registry.getSubagents()).toHaveLength(0); + expect(registry.resolveSessionFile({ subagentId: "SubagentA" })).toBe(sessionFile); + expect(registry.resolveSessionFile({ sessionFile })).toBe(sessionFile); + registry.dispose(); + }); + test("gates raw subagent events behind the events subscription level", () => { const eventBus = new EventBus(); const frames: RpcSubagentFrame[] = []; const registry = new RpcSubagentRegistry(eventBus, frame => frames.push(frame)); const eventPayload: SubagentEventPayload = { id: "SubagentA", - index: 0, - agent: "task", - agentSource: "bundled", - task: "Do work", - assignment: "Implement work", - parentToolCallId: "toolu_parent", - sessionFile: "/tmp/subagent.jsonl", event: { type: "agent_start" }, }; @@ -111,7 +309,7 @@ describe("RPC subagent registry", () => { eventBus.emit(TASK_SUBAGENT_EVENT_CHANNEL, eventPayload); expect(frames).toHaveLength(1); - expect(frames[0]).toMatchObject({ type: "subagent_event", payload: { id: "SubagentA" } }); + expect(frames[0]).toEqual({ type: "subagent_event", payload: eventPayload }); registry.dispose(); }); }); @@ -195,7 +393,7 @@ function handle(frame) { write({ type: "notice", level: "info", message: "subagent test" }); write({ type: "subagent_lifecycle", payload: { id: "SubagentA", index: 0, agent: "task", agentSource: "bundled", status: "started", sessionFile: "/tmp/subagent.jsonl" } }); write({ type: "subagent_progress", payload: { index: 0, agent: "task", agentSource: "bundled", task: "Do work", assignment: "Implement work", sessionFile: "/tmp/subagent.jsonl", progress } }); - write({ type: "subagent_event", payload: { id: "SubagentA", index: 0, agent: "task", agentSource: "bundled", task: "Do work", assignment: "Implement work", sessionFile: "/tmp/subagent.jsonl", event: { type: "agent_start" } } }); + write({ type: "subagent_event", payload: { id: "SubagentA", event: { type: "agent_start" } } }); write({ type: "agent_end", messages: [] }); } }