From 441c12956a6c67b4d3edf077f72033b26926094e Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Wed, 15 Apr 2026 21:59:45 +0200 Subject: [PATCH] fix(coding-agent): emit rpc prompt response after preflight closes #3049 --- packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/docs/rpc.md | 4 +- packages/coding-agent/docs/sdk.md | 23 +- .../coding-agent/src/core/agent-session.ts | 259 ++++++++-------- .../coding-agent/src/modes/rpc/rpc-mode.ts | 28 +- .../rpc-prompt-response-semantics.test.ts | 283 ++++++++++++++++++ 6 files changed, 467 insertions(+), 131 deletions(-) create mode 100644 packages/coding-agent/test/rpc-prompt-response-semantics.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index a5639932..fc5418ad 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -5,6 +5,7 @@ ### Fixed - Fixed the `plan-mode` example extension to allow `eza` in the read-only bash allowlist instead of the deprecated `exa` command ([#3160](https://github.com/badlogic/pi-mono/issues/3160)) +- Fixed RPC `prompt` to wait for prompt preflight success before emitting its single authoritative response, while still treating handled and queued prompts as success ([#3049](https://github.com/badlogic/pi-mono/issues/3049)) - Fixed Alt keybindings inside Zellij by skipping the Kitty keyboard protocol query there and enabling xterm `modifyOtherKeys` mode 2 directly ([#3163](https://github.com/badlogic/pi-mono/issues/3163)) - Fixed `/scoped-models` reordering to propagate into the `/model` scoped tab, preserving the user-defined scoped model order instead of re-sorting it ([#3217](https://github.com/badlogic/pi-mono/issues/3217)) - Fixed `session_shutdown` to fire on `SIGHUP` and `SIGTERM` in interactive, print, and RPC modes so extensions can run shutdown cleanup on those signal-driven exits ([#3212](https://github.com/badlogic/pi-mono/issues/3212)) diff --git a/packages/coding-agent/docs/rpc.md b/packages/coding-agent/docs/rpc.md index 6572fa98..54defea9 100644 --- a/packages/coding-agent/docs/rpc.md +++ b/packages/coding-agent/docs/rpc.md @@ -41,7 +41,7 @@ In particular, Node `readline` is not protocol-compliant for RPC mode because it #### prompt -Send a user prompt to the agent. Returns immediately; events stream asynchronously. +Send a user prompt to the agent. The command response is emitted after the prompt is accepted, queued, or handled. Events continue streaming asynchronously after acceptance. ```json {"id": "req-1", "type": "prompt", "message": "Hello, world!"} @@ -72,6 +72,8 @@ Response: {"id": "req-1", "type": "response", "command": "prompt", "success": true} ``` +`success: true` means the prompt was accepted, queued, or handled immediately. `success: false` means the prompt was rejected before acceptance. Failures after acceptance are reported through the normal event and message stream, not as a second `response` for the same request id. + The `images` field is optional. Each image uses `ImageContent` format: `{"type": "image", "data": "base64-encoded-data", "mimeType": "image/png"}`. #### steer diff --git a/packages/coding-agent/docs/sdk.md b/packages/coding-agent/docs/sdk.md index 5b5e0f67..8b70077f 100644 --- a/packages/coding-agent/docs/sdk.md +++ b/packages/coding-agent/docs/sdk.md @@ -182,6 +182,25 @@ unsubscribe = session.subscribe(() => {}); ### Prompting and Message Queueing +`PromptOptions` controls prompt expansion, queueing behavior while streaming, and prompt preflight notifications: + +```typescript +interface PromptOptions { + expandPromptTemplates?: boolean; + images?: ImageContent[]; + streamingBehavior?: "steer" | "followUp"; + source?: InputSource; + preflightResult?: (success: boolean) => void; +} +``` + +`preflightResult` is called once per `prompt()` invocation: + +- `true` when the prompt was accepted, queued, or handled immediately +- `false` when prompt preflight rejected before acceptance + +It fires before `prompt()` resolves. `prompt()` still resolves only after the full accepted run finishes, including retries. Failures after acceptance are reported through the normal event and message stream, not through `preflightResult(false)`. + The `prompt()` method handles prompt templates, extension commands, and message sending: ```typescript @@ -200,8 +219,10 @@ await session.prompt("After you're done, also check X", { streamingBehavior: "fo **Behavior:** - **Extension commands** (e.g., `/mycommand`): Execute immediately, even during streaming. They manage their own LLM interaction via `pi.sendMessage()`. -- **File-based prompt templates** (from `.md` files): Expanded to their content before sending/queueing. +- **File-based prompt templates** (from `.md` files): Expanded to their content before sending or queueing. - **During streaming without `streamingBehavior`**: Throws an error. Use `steer()` or `followUp()` directly, or specify the option. +- **`preflightResult(true)`**: Means the prompt was accepted, queued, or handled immediately. +- **`preflightResult(false)`**: Means preflight rejected before acceptance. For explicit queueing during streaming: diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 21364666..bccc5b8a 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -180,6 +180,8 @@ export interface PromptOptions { streamingBehavior?: "steer" | "followUp"; /** Source of input for extension input event handlers. Defaults to "interactive". */ source?: InputSource; + /** Internal hook used by RPC mode to observe prompt preflight acceptance or rejection. */ + preflightResult?: (success: boolean) => void; } /** Result from cycleModel() */ @@ -928,139 +930,154 @@ export class AgentSession { */ async prompt(text: string, options?: PromptOptions): Promise { const expandPromptTemplates = options?.expandPromptTemplates ?? true; + const preflightResult = options?.preflightResult; + let messages: AgentMessage[] | undefined; - // Handle extension commands first (execute immediately, even during streaming) - // Extension commands manage their own LLM interaction via pi.sendMessage() - if (expandPromptTemplates && text.startsWith("/")) { - const handled = await this._tryExecuteExtensionCommand(text); - if (handled) { - // Extension command executed, no prompt to send + try { + // Handle extension commands first (execute immediately, even during streaming) + // Extension commands manage their own LLM interaction via pi.sendMessage() + if (expandPromptTemplates && text.startsWith("/")) { + const handled = await this._tryExecuteExtensionCommand(text); + if (handled) { + // Extension command executed, no prompt to send + preflightResult?.(true); + return; + } + } + + // Emit input event for extension interception (before skill/template expansion) + let currentText = text; + let currentImages = options?.images; + if (this._extensionRunner?.hasHandlers("input")) { + const inputResult = await this._extensionRunner.emitInput( + currentText, + currentImages, + options?.source ?? "interactive", + ); + if (inputResult.action === "handled") { + preflightResult?.(true); + return; + } + if (inputResult.action === "transform") { + currentText = inputResult.text; + currentImages = inputResult.images ?? currentImages; + } + } + + // Expand skill commands (/skill:name args) and prompt templates (/template args) + let expandedText = currentText; + if (expandPromptTemplates) { + expandedText = this._expandSkillCommand(expandedText); + expandedText = expandPromptTemplate(expandedText, [...this.promptTemplates]); + } + + // If streaming, queue via steer() or followUp() based on option + if (this.isStreaming) { + if (!options?.streamingBehavior) { + throw new Error( + "Agent is already processing. Specify streamingBehavior ('steer' or 'followUp') to queue the message.", + ); + } + if (options.streamingBehavior === "followUp") { + await this._queueFollowUp(expandedText, currentImages); + } else { + await this._queueSteer(expandedText, currentImages); + } + preflightResult?.(true); return; } - } - // Emit input event for extension interception (before skill/template expansion) - let currentText = text; - let currentImages = options?.images; - if (this._extensionRunner?.hasHandlers("input")) { - const inputResult = await this._extensionRunner.emitInput( - currentText, - currentImages, - options?.source ?? "interactive", - ); - if (inputResult.action === "handled") { - return; - } - if (inputResult.action === "transform") { - currentText = inputResult.text; - currentImages = inputResult.images ?? currentImages; - } - } + // Flush any pending bash messages before the new prompt + this._flushPendingBashMessages(); - // Expand skill commands (/skill:name args) and prompt templates (/template args) - let expandedText = currentText; - if (expandPromptTemplates) { - expandedText = this._expandSkillCommand(expandedText); - expandedText = expandPromptTemplate(expandedText, [...this.promptTemplates]); - } - - // If streaming, queue via steer() or followUp() based on option - if (this.isStreaming) { - if (!options?.streamingBehavior) { + // Validate model + if (!this.model) { throw new Error( - "Agent is already processing. Specify streamingBehavior ('steer' or 'followUp') to queue the message.", + "No model selected.\n\n" + + `Use /login or set an API key environment variable. See ${join(getDocsPath(), "providers.md")}\n\n` + + "Then use /model to select a model.", ); } - if (options.streamingBehavior === "followUp") { - await this._queueFollowUp(expandedText, currentImages); - } else { - await this._queueSteer(expandedText, currentImages); + + if (!this._modelRegistry.hasConfiguredAuth(this.model)) { + const isOAuth = this._modelRegistry.isUsingOAuth(this.model); + if (isOAuth) { + throw new Error( + `Authentication failed for "${this.model.provider}". ` + + `Credentials may have expired or network is unavailable. ` + + `Run '/login ${this.model.provider}' to re-authenticate.`, + ); + } + throw new Error( + `No API key found for ${this.model.provider}.\n\n` + + `Use /login or set an API key environment variable. See ${join(getDocsPath(), "providers.md")}`, + ); } + + // Check if we need to compact before sending (catches aborted responses) + const lastAssistant = this._findLastAssistantMessage(); + if (lastAssistant) { + await this._checkCompaction(lastAssistant, false); + } + + // Build messages array (custom message if any, then user message) + messages = []; + + // Add user message + const userContent: (TextContent | ImageContent)[] = [{ type: "text", text: expandedText }]; + if (currentImages) { + userContent.push(...currentImages); + } + messages.push({ + role: "user", + content: userContent, + timestamp: Date.now(), + }); + + // Inject any pending "nextTurn" messages as context alongside the user message + for (const msg of this._pendingNextTurnMessages) { + messages.push(msg); + } + this._pendingNextTurnMessages = []; + + // Emit before_agent_start extension event + if (this._extensionRunner) { + const result = await this._extensionRunner.emitBeforeAgentStart( + expandedText, + currentImages, + this._baseSystemPrompt, + ); + // Add all custom messages from extensions + if (result?.messages) { + for (const msg of result.messages) { + messages.push({ + role: "custom", + customType: msg.customType, + content: msg.content, + display: msg.display, + details: msg.details, + timestamp: Date.now(), + }); + } + } + // Apply extension-modified system prompt, or reset to base + if (result?.systemPrompt) { + this.agent.state.systemPrompt = result.systemPrompt; + } else { + // Ensure we're using the base prompt (in case previous turn had modifications) + this.agent.state.systemPrompt = this._baseSystemPrompt; + } + } + } catch (error) { + preflightResult?.(false); + throw error; + } + + if (!messages) { return; } - // Flush any pending bash messages before the new prompt - this._flushPendingBashMessages(); - - // Validate model - if (!this.model) { - throw new Error( - "No model selected.\n\n" + - `Use /login or set an API key environment variable. See ${join(getDocsPath(), "providers.md")}\n\n` + - "Then use /model to select a model.", - ); - } - - if (!this._modelRegistry.hasConfiguredAuth(this.model)) { - const isOAuth = this._modelRegistry.isUsingOAuth(this.model); - if (isOAuth) { - throw new Error( - `Authentication failed for "${this.model.provider}". ` + - `Credentials may have expired or network is unavailable. ` + - `Run '/login ${this.model.provider}' to re-authenticate.`, - ); - } - throw new Error( - `No API key found for ${this.model.provider}.\n\n` + - `Use /login or set an API key environment variable. See ${join(getDocsPath(), "providers.md")}`, - ); - } - - // Check if we need to compact before sending (catches aborted responses) - const lastAssistant = this._findLastAssistantMessage(); - if (lastAssistant) { - await this._checkCompaction(lastAssistant, false); - } - - // Build messages array (custom message if any, then user message) - const messages: AgentMessage[] = []; - - // Add user message - const userContent: (TextContent | ImageContent)[] = [{ type: "text", text: expandedText }]; - if (currentImages) { - userContent.push(...currentImages); - } - messages.push({ - role: "user", - content: userContent, - timestamp: Date.now(), - }); - - // Inject any pending "nextTurn" messages as context alongside the user message - for (const msg of this._pendingNextTurnMessages) { - messages.push(msg); - } - this._pendingNextTurnMessages = []; - - // Emit before_agent_start extension event - if (this._extensionRunner) { - const result = await this._extensionRunner.emitBeforeAgentStart( - expandedText, - currentImages, - this._baseSystemPrompt, - ); - // Add all custom messages from extensions - if (result?.messages) { - for (const msg of result.messages) { - messages.push({ - role: "custom", - customType: msg.customType, - content: msg.content, - display: msg.display, - details: msg.details, - timestamp: Date.now(), - }); - } - } - // Apply extension-modified system prompt, or reset to base - if (result?.systemPrompt) { - this.agent.state.systemPrompt = result.systemPrompt; - } else { - // Ensure we're using the base prompt (in case previous turn had modifications) - this.agent.state.systemPrompt = this._baseSystemPrompt; - } - } - + preflightResult?.(true); await this.agent.prompt(messages); await this.waitForRetry(); } diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index 2ca456b8..0f17515b 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -357,7 +357,7 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise => { + const handleCommand = async (command: RpcCommand): Promise => { const id = command.id; switch (command.type) { @@ -366,17 +366,27 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise { + if (didSucceed) { + preflightSucceeded = true; + output(success(id, "prompt")); + } + }, }) - .catch((e) => output(error(id, "prompt", e.message))); - return success(id, "prompt"); + .catch((e) => { + if (!preflightSucceeded) { + output(error(id, "prompt", e.message)); + } + }); + return undefined; } case "steer": { @@ -686,7 +696,9 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise ({ + outputLines: [] as string[], + lineHandler: undefined as ((line: string) => void) | undefined, +})); + +vi.mock("../src/core/output-guard.js", () => ({ + takeOverStdout: vi.fn(), + writeRawStdout: (line: string) => { + rpcIo.outputLines.push(line); + }, +})); + +vi.mock("../src/modes/interactive/theme/theme.js", () => ({ theme: {} })); + +vi.mock("../src/modes/rpc/jsonl.js", () => ({ + attachJsonlLineReader: vi.fn((_stream: NodeJS.ReadableStream, onLine: (line: string) => void) => { + rpcIo.lineHandler = onLine; + return () => {}; + }), + serializeJsonLine: (value: unknown) => `${JSON.stringify(value)}\n`, +})); + +class MockAssistantStream extends EventStream { + constructor() { + super( + (event) => event.type === "done" || event.type === "error", + (event) => { + if (event.type === "done") return event.message; + if (event.type === "error") return event.error; + throw new Error("Unexpected event type"); + }, + ); + } +} + +function createAssistantMessage(text: string): AssistantMessage { + return { + role: "assistant", + content: [{ type: "text", text }], + api: "anthropic-messages", + provider: "anthropic", + model: "claude-sonnet-4-5", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: Date.now(), + }; +} + +type ParsedOutputLine = Record; + +function parseOutputLines(outputLines: string[]): ParsedOutputLine[] { + return outputLines + .flatMap((line) => line.split("\n")) + .filter((line) => line.trim().length > 0) + .map((line) => JSON.parse(line) as ParsedOutputLine); +} + +function getPromptResponses(outputLines: string[], id: string): ParsedOutputLine[] { + return parseOutputLines(outputLines).filter( + (record) => record.id === id && record.type === "response" && record.command === "prompt", + ); +} + +function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +function createRuntimeHost(options: { withAuth: boolean; responseDelayMs: number; model?: Model }): { + runtimeHost: AgentSessionRuntime; + cleanup: () => Promise; +} { + const tempDir = join(tmpdir(), `pi-rpc-prompt-${Date.now()}-${Math.random().toString(36).slice(2)}`); + mkdirSync(tempDir, { recursive: true }); + + const model = options.model ?? getModel("anthropic", "claude-sonnet-4-5"); + if (!model) { + throw new Error("Test model not found"); + } + + const agent = new Agent({ + getApiKey: () => "test-key", + initialState: { + model, + systemPrompt: "Test", + tools: [], + }, + streamFn: (_model, _context, _options) => { + const stream = new MockAssistantStream(); + queueMicrotask(() => { + stream.push({ type: "start", partial: createAssistantMessage("") }); + setTimeout(() => { + stream.push({ type: "done", reason: "stop", message: createAssistantMessage("done") }); + }, options.responseDelayMs); + }); + return stream; + }, + }); + + const sessionManager = SessionManager.inMemory(); + const settingsManager = SettingsManager.create(tempDir, tempDir); + const authStorage = AuthStorage.create(join(tempDir, "auth.json")); + const modelRegistry = ModelRegistry.create(authStorage, tempDir); + if (options.withAuth) { + authStorage.setRuntimeApiKey("anthropic", "test-key"); + } + + const session = new AgentSession({ + agent, + sessionManager, + settingsManager, + cwd: tempDir, + modelRegistry, + resourceLoader: createTestResourceLoader(), + }); + + const runtimeHost = { + session, + newSession: vi.fn(async () => ({ cancelled: true })), + switchSession: vi.fn(async () => ({ cancelled: true })), + fork: vi.fn(async () => ({ cancelled: true, selectedText: "" })), + dispose: vi.fn(async () => {}), + } as unknown as AgentSessionRuntime; + + return { + runtimeHost, + cleanup: async () => { + try { + if (session.isStreaming) { + await session.abort(); + } + } catch { + // ignore test cleanup failures + } + session.dispose(); + if (existsSync(tempDir)) { + rmSync(tempDir, { recursive: true }); + } + }, + }; +} + +async function startRpcMode(options: { withAuth: boolean; responseDelayMs: number; model?: Model }): Promise<{ + lineHandler: (line: string) => void; + cleanup: () => Promise; +}> { + rpcIo.outputLines = []; + rpcIo.lineHandler = undefined; + + const { runtimeHost, cleanup } = createRuntimeHost(options); + void runRpcMode(runtimeHost); + await vi.waitFor(() => expect(rpcIo.lineHandler).toBeDefined()); + + return { lineHandler: rpcIo.lineHandler!, cleanup }; +} + +describe("RPC prompt response semantics", () => { + afterEach(() => { + rpcIo.outputLines = []; + rpcIo.lineHandler = undefined; + }); + + it("emits one failure response when prompt preflight rejects", async () => { + const { lineHandler, cleanup } = await startRpcMode({ + withAuth: false, + responseDelayMs: 0, + model: { + id: "fake-model", + name: "Fake Model", + api: "openai-completions", + provider: "fake-provider", + baseUrl: "https://example.invalid", + reasoning: false, + input: [], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 0, + maxTokens: 0, + }, + }); + + try { + lineHandler(JSON.stringify({ id: "b1", type: "prompt", message: "Hello" })); + + await vi.waitFor(() => { + const responses = getPromptResponses(rpcIo.outputLines, "b1"); + expect(responses).toHaveLength(1); + expect(responses[0]).toMatchObject({ + id: "b1", + type: "response", + command: "prompt", + success: false, + error: "No API key found for fake-provider.\n\nUse /login or set an API key environment variable. See /Users/badlogic/workspaces/pi-mono-2/packages/coding-agent/docs/providers.md", + }); + }); + } finally { + await cleanup(); + } + }); + + it("emits one success response when prompt preflight succeeds", async () => { + const { lineHandler, cleanup } = await startRpcMode({ withAuth: true, responseDelayMs: 0 }); + + try { + lineHandler(JSON.stringify({ id: "b2", type: "prompt", message: "Hello" })); + + await vi.waitFor(() => { + const responses = getPromptResponses(rpcIo.outputLines, "b2"); + expect(responses).toHaveLength(1); + expect(responses[0]).toMatchObject({ + id: "b2", + type: "response", + command: "prompt", + success: true, + }); + }); + } finally { + await cleanup(); + } + }); + + it("emits one success response when prompt is queued during streaming", async () => { + const { lineHandler, cleanup } = await startRpcMode({ withAuth: true, responseDelayMs: 100 }); + + try { + lineHandler(JSON.stringify({ id: "b3-start", type: "prompt", message: "Start" })); + await vi.waitFor(() => { + expect(getPromptResponses(rpcIo.outputLines, "b3-start")).toHaveLength(1); + }); + + rpcIo.outputLines = []; + lineHandler( + JSON.stringify({ + id: "b3", + type: "prompt", + message: "Queue this", + streamingBehavior: "followUp", + }), + ); + + await vi.waitFor(() => { + const responses = getPromptResponses(rpcIo.outputLines, "b3"); + expect(responses).toHaveLength(1); + expect(responses[0]).toMatchObject({ + id: "b3", + type: "response", + command: "prompt", + success: true, + }); + }); + + await sleep(150); + } finally { + await cleanup(); + } + }); +});