shumengya/sproutclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(coding-agent): emit rpc prompt response after preflight closes #3049

Browse Source
This commit is contained in:
Mario Zechner
2026-04-15 21:59:45 +02:00
parent b920622110
commit 441c12956a
6 changed files with 467 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
packages/coding-agent/CHANGELOG.md
View File

@@ -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))

4
packages/coding-agent/docs/rpc.md
View File

@@ -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

23
packages/coding-agent/docs/sdk.md
View File

@@ -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:

259
packages/coding-agent/src/core/agent-session.ts
View File

@@ -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<void> {
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();
}

28
packages/coding-agent/src/modes/rpc/rpc-mode.ts
View File

@@ -357,7 +357,7 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise<neve
registerSignalHandlers();
// Handle a single command
const handleCommand = async (command: RpcCommand): Promise<RpcResponse> => {
const handleCommand = async (command: RpcCommand): Promise<RpcResponse | undefined> => {
const id = command.id;
switch (command.type) {
@@ -366,17 +366,27 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise<neve
// =================================================================
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
// Start prompt handling immediately, but emit the authoritative response only after
// prompt preflight succeeds. Queued and immediately handled prompts also count as success.
let preflightSucceeded = false;
void session
.prompt(command.message, {
images: command.images,
streamingBehavior: command.streamingBehavior,
source: "rpc",
preflightResult: (didSucceed) => {
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<neve
const command = parsed as RpcCommand;
try {
const response = await handleCommand(command);
output(response);
if (response) {
output(response);
}
await checkShutdownRequested();
} catch (commandError: unknown) {
output(

283
packages/coding-agent/test/rpc-prompt-response-semantics.test.ts Normal file
View File

@@ -0,0 +1,283 @@
import { existsSync, mkdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { Agent } from "@mariozechner/pi-agent-core";
import {
type AssistantMessage,
type AssistantMessageEvent,
EventStream,
getModel,
type Model,
} from "@mariozechner/pi-ai";
import { afterEach, describe, expect, it, vi } from "vitest";
import { AgentSession } from "../src/core/agent-session.js";
import type { AgentSessionRuntime } from "../src/core/agent-session-runtime.js";
import { AuthStorage } from "../src/core/auth-storage.js";
import { ModelRegistry } from "../src/core/model-registry.js";
import { SessionManager } from "../src/core/session-manager.js";
import { SettingsManager } from "../src/core/settings-manager.js";
import { runRpcMode } from "../src/modes/rpc/rpc-mode.js";
import { createTestResourceLoader } from "./utilities.js";
const rpcIo = vi.hoisted(() => ({
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<AssistantMessageEvent, AssistantMessage> {
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<string, unknown>;
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<void> {
return new Promise((resolve) => setTimeout(resolve, ms));
}
function createRuntimeHost(options: { withAuth: boolean; responseDelayMs: number; model?: Model<any> }): {
runtimeHost: AgentSessionRuntime;
cleanup: () => Promise<void>;
} {
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<any> }): Promise<{
lineHandler: (line: string) => void;
cleanup: () => Promise<void>;
}> {
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();
}
});
});