From e2cd316f33a2698b6d0fbd31b61e5695f0710613 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Korm=C3=A1kur?= Date: Sat, 13 Jun 2026 01:02:41 +0000 Subject: [PATCH] fix(coding-agent): complete local rpc prompts --- packages/coding-agent/CHANGELOG.md | 1 + .../coding-agent/src/modes/rpc/rpc-client.ts | 198 ++++++++++++++++-- .../coding-agent/src/modes/rpc/rpc-mode.ts | 4 +- .../test/rpc-client-prompt-result.test.ts | 64 ++++++ .../test/rpc-skill-command.test.ts | 2 +- 5 files changed, 247 insertions(+), 22 deletions(-) create mode 100644 packages/coding-agent/test/rpc-client-prompt-result.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 3b3a021be..eabf47ff2 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -164,6 +164,7 @@ - Fixed interrupt loader state getting stuck after queued-message aborts by removing the session-layer flush/latch path; empty Enter now aborts the active turn and lets the existing post-unwind queue drain resume normally. - Fixed `/goal ` and `/goal set ` during streaming so goal context is steered immediately but objective submission waits for the active turn to finish instead of spamming `AgentBusyError` ([#2454](https://github.com/can1357/oh-my-pi/issues/2454)). - Fixed RPC local-only prompt completion for extension `pi.sendUserMessage(..., { deliverAs: "followUp" | "steer" })` calls that only queue messages, so hosts no longer wait for an `agent_end` that will not be emitted. +- Fixed RPC prompt completion for local-only or already-finished slash command paths, including `/skill:*` prompts and extension `pi.sendUserMessage(..., { deliverAs: "followUp" | "steer" })` calls that only queue messages, so hosts and `RpcClient` waits no longer wait for an `agent_end` that will not be emitted. ### Fixed diff --git a/packages/coding-agent/src/modes/rpc/rpc-client.ts b/packages/coding-agent/src/modes/rpc/rpc-client.ts index 9ac6a3bb3..262da097b 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-client.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-client.ts @@ -23,6 +23,7 @@ import type { RpcHostToolDefinition, RpcHostToolResult, RpcHostToolUpdate, + RpcPromptResultFrame, RpcResponse, RpcSessionState, RpcSubagentEventFrame, @@ -39,6 +40,12 @@ type DistributiveOmit = T extends unknown ? Omit : n /** RpcCommand without the id field (for internal send) */ type RpcCommandBody = DistributiveOmit; +type PromptCompletion = { + promise: Promise; + resolve: () => void; + settled: boolean; +}; + export interface RpcClientOptions { /** Path to the CLI entry point (default: searches for dist/cli.js) */ cliPath?: string; @@ -169,6 +176,15 @@ function isRpcAvailableCommandsUpdateFrame(value: unknown): value is RpcAvailabl return value.type === "available_commands_update" && Array.isArray(value.commands); } +function isRpcPromptResultFrame(value: unknown): value is RpcPromptResultFrame { + if (!isRecord(value)) return false; + return ( + value.type === "prompt_result" && + (value.id === undefined || typeof value.id === "string") && + typeof value.agentInvoked === "boolean" + ); +} + function isRpcHostToolCallRequest(value: unknown): value is RpcHostToolCallRequest { if (!isRecord(value)) return false; return ( @@ -218,6 +234,9 @@ export class RpcClient { #requestId = 0; #extensionUiListeners: Set<(req: RpcExtensionUIRequest) => void> = new Set(); #abortController = new AbortController(); + #promptCompletions = new Map(); + #lastPromptCompletion: PromptCompletion | undefined; + #promptCompletionListeners = new Set<() => void>(); constructor(private options: RpcClientOptions = {}) { this.#customTools = [...(options.customTools ?? [])]; @@ -417,7 +436,7 @@ export class RpcClient { * Use waitForIdle() to wait for completion. */ async prompt(message: string, images?: ImageContent[]): Promise { - await this.#send({ type: "prompt", message, images }); + await this.#sendPrompt(message, images); } /** @@ -738,18 +757,26 @@ export class RpcClient { /** * Wait for agent to become idle (no streaming). - * Resolves when agent_end event is received. + * Resolves on `agent_end` or a local-only `prompt_result`. */ waitForIdle(timeout = 60000): Promise { + const completion = this.#lastPromptCompletion; + if (completion) { + return this.#withIdleTimeout( + completion.promise, + timeout, + `Timeout waiting for agent to become idle. Stderr: ${this.#process?.peekStderr() ?? ""}`, + ); + } + const { promise, resolve, reject } = Promise.withResolvers(); let settled = false; - const unsubscribe = this.onEvent(event => { - if (event.type === "agent_end") { - settled = true; - unsubscribe(); - clearTimeout(timeoutId); - resolve(); - } + const unsubscribe = this.#onPromptCompletion(() => { + if (settled) return; + settled = true; + unsubscribe(); + clearTimeout(timeoutId); + resolve(); }); const timeoutId = this.#startTimeout(timeout, () => { @@ -768,20 +795,34 @@ export class RpcClient { const { promise, resolve, reject } = Promise.withResolvers(); const events: AgentEvent[] = []; let settled = false; + let unsubscribeCompletion = () => {}; + const finish = () => { + if (settled) return; + settled = true; + unsubscribe(); + unsubscribeCompletion(); + clearTimeout(timeoutId); + resolve(events); + }; const unsubscribe = this.onEvent(event => { events.push(event); if (event.type === "agent_end") { - settled = true; - unsubscribe(); - clearTimeout(timeoutId); - resolve(events); + finish(); } }); + const completion = this.#lastPromptCompletion; + if (completion) { + void completion.promise.then(finish); + } else { + unsubscribeCompletion = this.#onPromptCompletion(finish); + } + const timeoutId = this.#startTimeout(timeout, () => { if (settled) return; settled = true; unsubscribe(); + unsubscribeCompletion(); reject(new Error(`Timeout collecting events. Stderr: ${this.#process?.peekStderr() ?? ""}`)); }); return promise; @@ -791,12 +832,122 @@ export class RpcClient { * Send prompt and wait for completion, returning all events. */ async promptAndWait(message: string, images?: ImageContent[], timeout = 60000): Promise { - const eventsPromise = this.collectEvents(timeout); - await this.prompt(message, images); - return eventsPromise; + const events: AgentEvent[] = []; + const unsubscribe = this.onEvent(event => { + events.push(event); + }); + try { + await this.prompt(message, images); + await this.waitForIdle(timeout); + return events; + } finally { + unsubscribe(); + } } // ========================================================================= + async #sendPrompt(message: string, images?: ImageContent[]): Promise { + const id = this.#nextRequestId(); + this.#createPromptCompletion(id); + try { + const response = await this.#send({ type: "prompt", message, images }, 30_000, id); + if (!response.success) { + this.#resolvePromptCompletion(id); + return; + } + if (response.command === "prompt" && response.data?.agentInvoked === false) { + this.#resolvePromptCompletion(id); + } + } catch (error) { + this.#resolvePromptCompletion(id); + throw error; + } + } + + #nextRequestId(): string { + return `req_${++this.#requestId}`; + } + + #createPromptCompletion(id: string): PromptCompletion { + const { promise, resolve } = Promise.withResolvers(); + const completion: PromptCompletion = { + promise, + resolve: () => { + if (completion.settled) return; + completion.settled = true; + resolve(); + this.#notifyPromptCompletion(); + }, + settled: false, + }; + this.#promptCompletions.set(id, completion); + this.#lastPromptCompletion = completion; + return completion; + } + + #resolvePromptCompletion(id: string | undefined): void { + if (!id) { + this.#resolveAllPromptCompletions(); + return; + } + const completion = this.#promptCompletions.get(id); + if (!completion) { + this.#notifyPromptCompletion(); + return; + } + this.#promptCompletions.delete(id); + completion.resolve(); + } + + #resolveAllPromptCompletions(): void { + if (this.#promptCompletions.size === 0) { + this.#notifyPromptCompletion(); + return; + } + const completions = Array.from(this.#promptCompletions.values()); + this.#promptCompletions.clear(); + for (const completion of completions) { + completion.resolve(); + } + } + + #onPromptCompletion(listener: () => void): () => void { + this.#promptCompletionListeners.add(listener); + return () => { + this.#promptCompletionListeners.delete(listener); + }; + } + + #notifyPromptCompletion(): void { + for (const listener of [...this.#promptCompletionListeners]) { + listener(); + } + } + + #withIdleTimeout(source: Promise, timeoutMs: number, timeoutMessage: string): Promise { + const { promise, resolve, reject } = Promise.withResolvers(); + let settled = false; + const timeoutId = this.#startTimeout(timeoutMs, () => { + if (settled) return; + settled = true; + reject(new Error(timeoutMessage)); + }); + void source.then( + value => { + if (settled) return; + settled = true; + clearTimeout(timeoutId); + resolve(value); + }, + error => { + if (settled) return; + settled = true; + clearTimeout(timeoutId); + reject(error); + }, + ); + return promise; + } // Internal // ========================================================================= @@ -812,6 +963,13 @@ export class RpcClient { } } + if (isRpcPromptResultFrame(data)) { + if (!data.agentInvoked) { + this.#resolvePromptCompletion(data.id); + } + return; + } + if (isRpcHostToolCallRequest(data)) { void this.#handleHostToolCall(data); return; @@ -868,14 +1026,16 @@ export class RpcClient { for (const listener of this.#eventListeners) { listener(data); } + + if (data.type === "agent_end") { + this.#resolveAllPromptCompletions(); + } } - #send(command: RpcCommandBody, timeoutMs = 30_000): Promise { + #send(command: RpcCommandBody, timeoutMs = 30_000, id = this.#nextRequestId()): Promise { if (!this.#process?.stdin) { throw new Error("Client not started"); } - - const id = `req_${++this.#requestId}`; const fullCommand = { ...command, id } as RpcCommand; const { promise, resolve, reject } = Promise.withResolvers(); let settled = false; diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index 52d92a3a3..dd9df6f2d 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -80,7 +80,7 @@ export type RpcSessionChangeResult = export type RpcSessionChangeSession = Pick; export type RpcSkillCommandSession = Pick; -export type RpcSkillCommandResult = { agentInvoked: true }; +export type RpcSkillCommandResult = { agentInvoked: false }; export async function tryRunRpcSkillCommand( session: RpcSkillCommandSession, @@ -102,7 +102,7 @@ export async function tryRunRpcSkillCommand( details: built.details, attribution: "user", }); - return { agentInvoked: true }; + return { agentInvoked: false }; } export function reportLocalOnlyPromptResult(input: { diff --git a/packages/coding-agent/test/rpc-client-prompt-result.test.ts b/packages/coding-agent/test/rpc-client-prompt-result.test.ts new file mode 100644 index 000000000..92a116de2 --- /dev/null +++ b/packages/coding-agent/test/rpc-client-prompt-result.test.ts @@ -0,0 +1,64 @@ +import { describe, expect, test } from "bun:test"; +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import { RpcClient } from "@oh-my-pi/pi-coding-agent/modes/rpc/rpc-client"; + +const FAKE_RPC_SERVER = ` +function writeFrame(frame: unknown): void { + process.stdout.write(JSON.stringify(frame) + "\\n"); +} + +writeFrame({ type: "ready" }); +const decoder = new TextDecoder(); +let buffer = ""; +for await (const chunk of Bun.stdin.stream()) { + buffer += decoder.decode(chunk, { stream: true }); + let newline = buffer.indexOf("\\n"); + while (newline !== -1) { + const line = buffer.slice(0, newline).trim(); + buffer = buffer.slice(newline + 1); + if (line.length > 0) { + const command = JSON.parse(line) as { type?: string; id?: string; message?: string }; + if (command.type === "prompt") { + if (command.message === "immediate") { + writeFrame({ + type: "response", + command: "prompt", + id: command.id, + success: true, + data: { agentInvoked: false }, + }); + } else { + writeFrame({ type: "response", command: "prompt", id: command.id, success: true }); + setTimeout(() => { + writeFrame({ type: "prompt_result", id: command.id, agentInvoked: false }); + }, 10); + } + } + } + newline = buffer.indexOf("\\n"); + } +} +`; + +describe("RpcClient prompt completion", () => { + test("waits for local-only prompt completions without agent_end", async () => { + const dir = await fs.mkdtemp(path.join(os.tmpdir(), "omp-rpc-client-prompt-result-")); + const serverPath = path.join(dir, "fake-rpc-server.ts"); + await Bun.write(serverPath, FAKE_RPC_SERVER); + const client = new RpcClient({ cliPath: serverPath, cwd: dir }); + try { + await client.start(); + + await client.prompt("immediate"); + await client.waitForIdle(1000); + + const events = await client.promptAndWait("deferred", undefined, 1000); + expect(events).toEqual([]); + } finally { + client.stop(); + await fs.rm(dir, { recursive: true, force: true }); + } + }); +}); diff --git a/packages/coding-agent/test/rpc-skill-command.test.ts b/packages/coding-agent/test/rpc-skill-command.test.ts index 7a9f56866..a18ab4c78 100644 --- a/packages/coding-agent/test/rpc-skill-command.test.ts +++ b/packages/coding-agent/test/rpc-skill-command.test.ts @@ -30,7 +30,7 @@ describe("tryRunRpcSkillCommand", () => { "/skill:reviewer focus on risks", ); - expect(handled).toEqual({ agentInvoked: true }); + expect(handled).toEqual({ agentInvoked: false }); expect(message?.customType).toBe(SKILL_PROMPT_MESSAGE_TYPE); expect(message?.content).toContain("Review the supplied code carefully."); expect(message?.content).toContain("User: focus on risks");