fix(coding-agent): emit rpc prompt response after preflight closes #3049
This commit is contained in:
@@ -5,6 +5,7 @@
|
|||||||
### Fixed
|
### 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 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 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 `/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))
|
- 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))
|
||||||
|
|||||||
@@ -41,7 +41,7 @@ In particular, Node `readline` is not protocol-compliant for RPC mode because it
|
|||||||
|
|
||||||
#### prompt
|
#### 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
|
```json
|
||||||
{"id": "req-1", "type": "prompt", "message": "Hello, world!"}
|
{"id": "req-1", "type": "prompt", "message": "Hello, world!"}
|
||||||
@@ -72,6 +72,8 @@ Response:
|
|||||||
{"id": "req-1", "type": "response", "command": "prompt", "success": true}
|
{"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"}`.
|
The `images` field is optional. Each image uses `ImageContent` format: `{"type": "image", "data": "base64-encoded-data", "mimeType": "image/png"}`.
|
||||||
|
|
||||||
#### steer
|
#### steer
|
||||||
|
|||||||
@@ -182,6 +182,25 @@ unsubscribe = session.subscribe(() => {});
|
|||||||
|
|
||||||
### Prompting and Message Queueing
|
### 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:
|
The `prompt()` method handles prompt templates, extension commands, and message sending:
|
||||||
|
|
||||||
```typescript
|
```typescript
|
||||||
@@ -200,8 +219,10 @@ await session.prompt("After you're done, also check X", { streamingBehavior: "fo
|
|||||||
|
|
||||||
**Behavior:**
|
**Behavior:**
|
||||||
- **Extension commands** (e.g., `/mycommand`): Execute immediately, even during streaming. They manage their own LLM interaction via `pi.sendMessage()`.
|
- **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.
|
- **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:
|
For explicit queueing during streaming:
|
||||||
|
|
||||||
|
|||||||
@@ -180,6 +180,8 @@ export interface PromptOptions {
|
|||||||
streamingBehavior?: "steer" | "followUp";
|
streamingBehavior?: "steer" | "followUp";
|
||||||
/** Source of input for extension input event handlers. Defaults to "interactive". */
|
/** Source of input for extension input event handlers. Defaults to "interactive". */
|
||||||
source?: InputSource;
|
source?: InputSource;
|
||||||
|
/** Internal hook used by RPC mode to observe prompt preflight acceptance or rejection. */
|
||||||
|
preflightResult?: (success: boolean) => void;
|
||||||
}
|
}
|
||||||
|
|
||||||
/** Result from cycleModel() */
|
/** Result from cycleModel() */
|
||||||
@@ -928,139 +930,154 @@ export class AgentSession {
|
|||||||
*/
|
*/
|
||||||
async prompt(text: string, options?: PromptOptions): Promise<void> {
|
async prompt(text: string, options?: PromptOptions): Promise<void> {
|
||||||
const expandPromptTemplates = options?.expandPromptTemplates ?? true;
|
const expandPromptTemplates = options?.expandPromptTemplates ?? true;
|
||||||
|
const preflightResult = options?.preflightResult;
|
||||||
|
let messages: AgentMessage[] | undefined;
|
||||||
|
|
||||||
// Handle extension commands first (execute immediately, even during streaming)
|
try {
|
||||||
// Extension commands manage their own LLM interaction via pi.sendMessage()
|
// Handle extension commands first (execute immediately, even during streaming)
|
||||||
if (expandPromptTemplates && text.startsWith("/")) {
|
// Extension commands manage their own LLM interaction via pi.sendMessage()
|
||||||
const handled = await this._tryExecuteExtensionCommand(text);
|
if (expandPromptTemplates && text.startsWith("/")) {
|
||||||
if (handled) {
|
const handled = await this._tryExecuteExtensionCommand(text);
|
||||||
// Extension command executed, no prompt to send
|
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;
|
return;
|
||||||
}
|
}
|
||||||
}
|
|
||||||
|
|
||||||
// Emit input event for extension interception (before skill/template expansion)
|
// Flush any pending bash messages before the new prompt
|
||||||
let currentText = text;
|
this._flushPendingBashMessages();
|
||||||
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;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Expand skill commands (/skill:name args) and prompt templates (/template args)
|
// Validate model
|
||||||
let expandedText = currentText;
|
if (!this.model) {
|
||||||
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(
|
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);
|
if (!this._modelRegistry.hasConfiguredAuth(this.model)) {
|
||||||
} else {
|
const isOAuth = this._modelRegistry.isUsingOAuth(this.model);
|
||||||
await this._queueSteer(expandedText, currentImages);
|
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;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Flush any pending bash messages before the new prompt
|
preflightResult?.(true);
|
||||||
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;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
await this.agent.prompt(messages);
|
await this.agent.prompt(messages);
|
||||||
await this.waitForRetry();
|
await this.waitForRetry();
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -357,7 +357,7 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise<neve
|
|||||||
registerSignalHandlers();
|
registerSignalHandlers();
|
||||||
|
|
||||||
// Handle a single command
|
// Handle a single command
|
||||||
const handleCommand = async (command: RpcCommand): Promise<RpcResponse> => {
|
const handleCommand = async (command: RpcCommand): Promise<RpcResponse | undefined> => {
|
||||||
const id = command.id;
|
const id = command.id;
|
||||||
|
|
||||||
switch (command.type) {
|
switch (command.type) {
|
||||||
@@ -366,17 +366,27 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise<neve
|
|||||||
// =================================================================
|
// =================================================================
|
||||||
|
|
||||||
case "prompt": {
|
case "prompt": {
|
||||||
// Don't await - events will stream
|
// Start prompt handling immediately, but emit the authoritative response only after
|
||||||
// Extension commands are executed immediately, file prompt templates are expanded
|
// prompt preflight succeeds. Queued and immediately handled prompts also count as success.
|
||||||
// If streaming and streamingBehavior specified, queues via steer/followUp
|
let preflightSucceeded = false;
|
||||||
session
|
void session
|
||||||
.prompt(command.message, {
|
.prompt(command.message, {
|
||||||
images: command.images,
|
images: command.images,
|
||||||
streamingBehavior: command.streamingBehavior,
|
streamingBehavior: command.streamingBehavior,
|
||||||
source: "rpc",
|
source: "rpc",
|
||||||
|
preflightResult: (didSucceed) => {
|
||||||
|
if (didSucceed) {
|
||||||
|
preflightSucceeded = true;
|
||||||
|
output(success(id, "prompt"));
|
||||||
|
}
|
||||||
|
},
|
||||||
})
|
})
|
||||||
.catch((e) => output(error(id, "prompt", e.message)));
|
.catch((e) => {
|
||||||
return success(id, "prompt");
|
if (!preflightSucceeded) {
|
||||||
|
output(error(id, "prompt", e.message));
|
||||||
|
}
|
||||||
|
});
|
||||||
|
return undefined;
|
||||||
}
|
}
|
||||||
|
|
||||||
case "steer": {
|
case "steer": {
|
||||||
@@ -686,7 +696,9 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise<neve
|
|||||||
const command = parsed as RpcCommand;
|
const command = parsed as RpcCommand;
|
||||||
try {
|
try {
|
||||||
const response = await handleCommand(command);
|
const response = await handleCommand(command);
|
||||||
output(response);
|
if (response) {
|
||||||
|
output(response);
|
||||||
|
}
|
||||||
await checkShutdownRequested();
|
await checkShutdownRequested();
|
||||||
} catch (commandError: unknown) {
|
} catch (commandError: unknown) {
|
||||||
output(
|
output(
|
||||||
|
|||||||
283
packages/coding-agent/test/rpc-prompt-response-semantics.test.ts
Normal file
283
packages/coding-agent/test/rpc-prompt-response-semantics.test.ts
Normal 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();
|
||||||
|
}
|
||||||
|
});
|
||||||
|
});
|
||||||
Reference in New Issue
Block a user