smashimo/oh-my-pi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge PR #5156: fix(rpc): prevent rpc-ui extension UI response deadlock (@roboomp)

Browse Source
This commit is contained in:
can1357
2026-07-14 18:39:26 +02:00
parent 58b80dac42 3bbeaa400f
commit 3539403e3c
4 changed files with 443 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/coding-agent/CHANGELOG.md
+1
View File
@@ -234,6 +234,7 @@
- Fixed macOS runtime diagnostics (e.g. `MallocStackLogging: can't turn off malloc stack logging because it was not enabled`) written directly to fd 2 by libmalloc painting into the TUI viewport. While the TUI owns the terminal, stderr is now redirected to the omp log file and restored at every ownership handoff (external editor, Ctrl+Z suspend, shutdown, crash restore); fatal crash reports still reach the real terminal.
- Fixed custom model/provider config discovery so `~/.omp/agent/models.yaml` loads when `models.yml` is absent, while preserving `.yml` precedence and only migrating legacy `models.json` when neither YAML file exists. ([#5145](https://github.com/can1357/oh-my-pi/issues/5145))
- Fixed throttled live command output holding a quiet final chunk until the command exited.
- Fixed rpc-ui extension UI and host tool responses deadlocking login, session lifecycle, and queued commands while an active command awaited the same side channel. ([#5153](https://github.com/can1357/oh-my-pi/issues/5153))
## [16.4.2] - 2026-07-10
packages/coding-agent/src/modes/rpc/host-tools.ts
+15
View File
@@ -75,6 +75,7 @@ export class RpcHostToolBridge {
#output: RpcHostToolOutput;
#definitions = new Map<string, RpcHostToolDefinition>();
#pendingCalls = new Map<string, PendingHostToolCall>();
#closedError: Error | undefined;
constructor(output: RpcHostToolOutput) {
this.#output = output;
@@ -126,6 +127,10 @@ export class RpcHostToolBridge {
return Promise.reject(new Error(`Host tool "${definition.name}" was aborted`));
}
if (this.#closedError) {
return Promise.reject(this.#closedError);
}
const id = Snowflake.next() as string;
const { promise, resolve, reject } = Promise.withResolvers<AgentToolResult<unknown>>();
let settled = false;
@@ -183,4 +188,14 @@ export class RpcHostToolBridge {
pending.reject(error);
}
}
/** Reject active and future host tool requests after the RPC client disconnects. */
close(message: string): void {
if (!this.#closedError) this.#closedError = new Error(message);
const pendingCalls = Array.from(this.#pendingCalls.values());
this.#pendingCalls.clear();
for (const pending of pendingCalls) {
pending.reject(this.#closedError);
}
}
}
packages/coding-agent/src/modes/rpc/rpc-mode.ts
+124 -54
View File
@@ -59,6 +59,29 @@ export type PendingExtensionRequest = {
reject: (error: Error) => void;
};
/** Pending extension UI request map that can fail closed when the RPC client disconnects. */
export class RpcPendingExtensionRequests extends Map<string, PendingExtensionRequest> {
#closedError: Error | undefined;
override set(id: string, request: PendingExtensionRequest): this {
if (this.#closedError) {
request.reject(this.#closedError);
return this;
}
return super.set(id, request);
}
/** Reject every active and future extension UI request. */
rejectAll(message: string): void {
if (!this.#closedError) this.#closedError = new Error(message);
const requests = Array.from(this.values());
this.clear();
for (const request of requests) {
request.reject(this.#closedError);
}
}
}
type RpcOutput = (
obj:
| RpcResponse
@@ -246,15 +269,40 @@ function isRpcExtensionUIResponse(value: unknown): value is RpcExtensionUIRespon
return value.type === "extension_ui_response" && typeof value.id === "string";
}
/** Dispatch side-channel frames that must overtake the serialized command queue. */
export function dispatchRpcControlFrame(parsed: unknown, deps: RpcInputFrameDeps): boolean {
if (isRpcExtensionUIResponse(parsed)) {
const pending = deps.pendingExtensionRequests.get(parsed.id);
if (pending) pending.resolve(parsed);
return true;
}
if (isRpcHostToolResult(parsed)) {
deps.onHostToolResult(parsed);
return true;
}
if (isRpcHostToolUpdate(parsed)) {
deps.onHostToolUpdate(parsed);
return true;
}
if (isRpcHostUriResult(parsed)) {
deps.onHostUriResult(parsed);
return true;
}
return false;
}
/**
* Dispatch a single parsed frame from the RPC input stream.
*
* Bash commands are dispatched in the background so the caller (the stdin loop
* in {@link runRpcMode}) can keep reading subsequent frames while a shell
* command is still running. This lets a client send `abort_bash` (or any other
* command) while a long-running `bash` is in flight. Response correlation is
* preserved via each command's `id`; ordering across concurrent commands is
* not guaranteed and clients MUST match on `id`.
* Bash commands are dispatched in the background so the caller can keep reading
* subsequent frames while a shell command is still running. This lets a client
* send `abort_bash` while a long-running `bash` is in flight. Response
* correlation is preserved via each command's `id`; ordering across concurrent
* commands is not guaranteed and clients MUST match on `id`.
*
* @returns `undefined` when the frame was routed to a side-channel handler
* (extension UI response, host tool/URI frames) or dispatched in the
@@ -263,28 +311,7 @@ function isRpcExtensionUIResponse(value: unknown): value is RpcExtensionUIRespon
* on non-`bash` commands propagate; the caller is expected to wrap them.
*/
export function dispatchRpcInputFrame(parsed: unknown, deps: RpcInputFrameDeps): Promise<void> | undefined {
// Side-channel: extension UI responses resolve a pending dialog promise.
if (isRpcExtensionUIResponse(parsed)) {
const pending = deps.pendingExtensionRequests.get(parsed.id);
if (pending) pending.resolve(parsed);
return undefined;
}
if (isRpcHostToolResult(parsed)) {
deps.onHostToolResult(parsed);
return undefined;
}
if (isRpcHostToolUpdate(parsed)) {
deps.onHostToolUpdate(parsed);
return undefined;
}
if (isRpcHostUriResult(parsed)) {
deps.onHostUriResult(parsed);
return undefined;
}
if (dispatchRpcControlFrame(parsed, deps)) return undefined;
// Regular RPC command. The transport contract states each remaining frame
// is an {@link RpcCommand}; `handleCommand`'s `default` arm surfaces
// unknown discriminants as an error response, so we do not shape-check
@@ -313,6 +340,59 @@ export function dispatchRpcInputFrame(parsed: unknown, deps: RpcInputFrameDeps):
})();
}
/** Serializes ordinary RPC commands while allowing control frames to dispatch immediately. */
export class RpcInputDispatcher {
#tail: Promise<void> = Promise.resolve();
#tasks = new Set<Promise<void>>();
readonly #deps: RpcInputFrameDeps;
readonly #afterSerialCommand: (() => Promise<void>) | undefined;
constructor(options: { deps: RpcInputFrameDeps; afterSerialCommand?: () => Promise<void> }) {
this.#deps = options.deps;
this.#afterSerialCommand = options.afterSerialCommand;
}
/** Accept a parsed input frame without blocking the stdin reader. */
dispatch(parsed: unknown): void {
if (dispatchRpcControlFrame(parsed, this.#deps)) return;
const command = parsed as RpcCommand;
if (command.type === "bash") {
dispatchRpcInputFrame(command, this.#deps);
return;
}
const task = this.#tail.then(
() => this.#dispatchSerialCommand(command),
() => this.#dispatchSerialCommand(command),
);
this.#tail = task.catch(() => {});
this.#tasks.add(task);
void task.finally(() => {
this.#tasks.delete(task);
});
}
/** Await every accepted serial command, including commands queued before EOF. */
async drain(): Promise<void> {
while (this.#tasks.size > 0) {
await Promise.allSettled(Array.from(this.#tasks));
}
}
async #dispatchSerialCommand(command: RpcCommand): Promise<void> {
try {
const awaited = dispatchRpcInputFrame(command, this.#deps);
if (awaited) await awaited;
} catch (err: unknown) {
const message = err instanceof Error ? err.message : String(err);
this.#deps.output(this.#deps.errorResponse(command.id, command.type, message));
} finally {
await this.#afterSerialCommand?.();
}
}
}
/**
* Coordinates deferred shutdown with in-flight background input tasks.
*
@@ -551,7 +631,7 @@ export async function runRpcMode(
const extensionUserMessageTracker = new RpcExtensionUserMessageTracker();
const pendingExtensionRequests = new Map<string, PendingExtensionRequest>();
const pendingExtensionRequests = new RpcPendingExtensionRequests();
const hostToolBridge = new RpcHostToolBridge(output);
const hostUriBridge = new RpcHostUriBridge(output);
const subagentRegistry = eventBus ? new RpcSubagentRegistry(eventBus, output) : undefined;
@@ -1279,35 +1359,25 @@ export async function runRpcMode(
onHostUriResult: frame => hostUriBridge.handleResult(frame),
};
// Listen for JSON input using Bun's stdin. Frame dispatch lives in
// dispatchRpcInputFrame so it can be exercised directly by tests; see the
// helper's docstring for the concurrency contract.
const inputDispatcher = new RpcInputDispatcher({
deps: dispatchFrameDeps,
afterSerialCommand: () => shutdownCoordinator.checkShutdownRequested(),
});
// Keep the stdin reader moving: side-channel frames dispatch immediately,
// ordinary commands serialize through inputDispatcher, and bash remains
// background-dispatched so abort_bash can overtake it.
for await (const parsed of readJsonl(Bun.stdin.stream())) {
try {
const awaited = dispatchRpcInputFrame(parsed, dispatchFrameDeps);
if (awaited) {
await awaited;
// Check for deferred shutdown request (idle between commands).
// Background-dispatched bash frames skip this check so a later
// abort_bash can still be read; the coordinator re-checks when
// each tracked task settles, so a shutdown requested mid-bash
// fires once the response frame is written even if no further
// client frames arrive.
await shutdownCoordinator.checkShutdownRequested();
}
} catch (e: unknown) {
const message = e instanceof Error ? e.message : String(e);
output(error(undefined, "parse", `Failed to parse command: ${message}`));
}
inputDispatcher.dispatch(parsed);
}
// Background bash tasks may still owe response frames; drain them before
// tearing down (stdin EOF ends the frame stream, not in-flight work).
await shutdownCoordinator.drain();
// stdin closed — RPC client is gone, exit cleanly
hostToolBridge.rejectAllPending("RPC client disconnected before host tool execution completed");
// stdin closed — RPC client is gone. Fail pending side-channel requests
// first so active/queued commands can settle, then drain accepted work.
pendingExtensionRequests.rejectAll("RPC client disconnected before extension UI response completed");
hostToolBridge.close("RPC client disconnected before host tool execution completed");
hostUriBridge.clear("RPC client disconnected before host URI request completed");
await inputDispatcher.drain();
await shutdownCoordinator.drain();
subagentRegistry?.dispose();
process.exit(0);
}
packages/coding-agent/test/rpc-input-frame.test.ts
+303 -3
View File
@@ -1,15 +1,27 @@
import { describe, expect, test } from "bun:test";
import { RpcHostToolBridge } from "@oh-my-pi/pi-coding-agent/modes/rpc/host-tools";
import {
dispatchRpcInputFrame,
type PendingExtensionRequest,
RpcInputDispatcher,
type RpcInputFrameDeps,
RpcPendingExtensionRequests,
RpcShutdownCoordinator,
} from "@oh-my-pi/pi-coding-agent/modes/rpc/rpc-mode";
import type { RpcCommand, RpcResponse } from "@oh-my-pi/pi-coding-agent/modes/rpc/rpc-types";
import type {
RpcCommand,
RpcExtensionUIResponse,
RpcHostToolCallRequest,
RpcHostToolCancelRequest,
RpcResponse,
} from "@oh-my-pi/pi-coding-agent/modes/rpc/rpc-types";
type OutputFrame = RpcResponse | object;
const makeDeps = (handleCommand: RpcInputFrameDeps["handleCommand"]) => {
const makeDeps = (
handleCommand: RpcInputFrameDeps["handleCommand"],
options?: { pendingExtensionRequests?: Map<string, PendingExtensionRequest> },
) => {
const outputs: OutputFrame[] = [];
const deps: RpcInputFrameDeps = {
handleCommand,
@@ -23,7 +35,7 @@ const makeDeps = (handleCommand: RpcInputFrameDeps["handleCommand"]) => {
success: false,
error: message,
}),
pendingExtensionRequests: new Map<string, PendingExtensionRequest>(),
pendingExtensionRequests: options?.pendingExtensionRequests ?? new Map<string, PendingExtensionRequest>(),
onHostToolResult: () => {},
onHostToolUpdate: () => {},
onHostUriResult: () => {},
@@ -33,6 +45,21 @@ const makeDeps = (handleCommand: RpcInputFrameDeps["handleCommand"]) => {
const flushMicrotasks = () => new Promise<void>(resolve => setImmediate(resolve));
const requestExtensionInput = (deps: RpcInputFrameDeps, id: string, message: string) => {
const response = Promise.withResolvers<RpcExtensionUIResponse>();
deps.pendingExtensionRequests.set(id, {
resolve: response.resolve,
reject: error => response.reject(error),
});
deps.output({
type: "extension_ui_request",
id,
method: "input",
message,
});
return response.promise;
};
const cancelledBashResponse = (id: string): RpcResponse => ({
id,
type: "response",
@@ -202,6 +229,279 @@ describe("dispatchRpcInputFrame", () => {
});
});
describe("RpcInputDispatcher", () => {
test("control frames resolve extension UI requests while an ordinary command is active", async () => {
let depsRef: RpcInputFrameDeps;
const { deps, outputs } = makeDeps(async command => {
if (command.type !== "prompt") throw new Error(`unexpected command type: ${command.type}`);
const response = await requestExtensionInput(depsRef, "ui-active", "Continue?");
return {
id: command.id,
type: "response",
command: "prompt",
success: true,
data: { agentInvoked: "value" in response && response.value === "continue" },
};
});
depsRef = deps;
const dispatcher = new RpcInputDispatcher({ deps });
dispatcher.dispatch({ id: "prompt-1", type: "prompt", message: "ask extension" });
await flushMicrotasks();
expect(outputs).toEqual([
{
type: "extension_ui_request",
id: "ui-active",
method: "input",
message: "Continue?",
},
]);
dispatcher.dispatch({ type: "extension_ui_response", id: "ui-active", value: "continue" });
await dispatcher.drain();
expect(outputs).toEqual([
{
type: "extension_ui_request",
id: "ui-active",
method: "input",
message: "Continue?",
},
{
id: "prompt-1",
type: "response",
command: "prompt",
success: true,
data: { agentInvoked: true },
},
]);
});
test("ordinary commands stay serialized while first command is blocked", async () => {
const releaseFirst = Promise.withResolvers<void>();
const started: string[] = [];
const { deps, outputs } = makeDeps(async command => {
started.push(command.type);
if (command.type === "abort_retry") {
await releaseFirst.promise;
return { id: command.id, type: "response", command: "abort_retry", success: true };
}
if (command.type === "get_state") {
return {
id: command.id,
type: "response",
command: "get_state",
success: true,
data: {
thinkingLevel: undefined,
isStreaming: false,
isCompacting: false,
steeringMode: "all",
followUpMode: "all",
interruptMode: "immediate",
sessionId: "session-1",
autoCompactionEnabled: false,
messageCount: 0,
queuedMessageCount: 0,
todoPhases: [],
},
};
}
throw new Error(`unexpected command type: ${command.type}`);
});
const dispatcher = new RpcInputDispatcher({ deps });
dispatcher.dispatch({ id: "first", type: "abort_retry" });
dispatcher.dispatch({ id: "second", type: "get_state" });
await flushMicrotasks();
expect(started).toEqual(["abort_retry"]);
expect(outputs).toHaveLength(0);
releaseFirst.resolve();
await dispatcher.drain();
expect(started).toEqual(["abort_retry", "get_state"]);
expect((outputs[0] as RpcResponse).id).toBe("first");
expect((outputs[1] as RpcResponse).id).toBe("second");
expect((outputs[1] as RpcResponse).command).toBe("get_state");
});
test("serial command rejection emits an error response and does not poison the queue", async () => {
const started: string[] = [];
const { deps, outputs } = makeDeps(async command => {
started.push(command.type);
if (command.type === "abort_retry") throw new Error("retry controller exploded");
if (command.type === "set_auto_retry") {
return { id: command.id, type: "response", command: "set_auto_retry", success: true };
}
throw new Error(`unexpected command type: ${command.type}`);
});
const dispatcher = new RpcInputDispatcher({ deps });
dispatcher.dispatch({ id: "bad", type: "abort_retry" });
dispatcher.dispatch({ id: "next", type: "set_auto_retry", enabled: true });
await dispatcher.drain();
expect(started).toEqual(["abort_retry", "set_auto_retry"]);
expect(outputs).toEqual([
{
id: "bad",
type: "response",
command: "abort_retry",
success: false,
error: "retry controller exploded",
},
{
id: "next",
type: "response",
command: "set_auto_retry",
success: true,
},
]);
});
test("drain after EOF rejects active and queued host tool requests without emitting new calls", async () => {
const disconnectMessage = "RPC client disconnected before host tool execution completed";
const hostToolFrames: Array<RpcHostToolCallRequest | RpcHostToolCancelRequest> = [];
const bridge = new RpcHostToolBridge(frame => {
hostToolFrames.push(frame);
});
const [tool] = bridge.setTools([
{
name: "host_wait",
description: "Waits for host process",
parameters: {
type: "object",
properties: {},
additionalProperties: false,
},
},
]);
const started: string[] = [];
const { deps, outputs } = makeDeps(async command => {
if (command.type !== "prompt") throw new Error(`unexpected command type: ${command.type}`);
started.push(command.id ?? "");
await tool.execute(`toolu_${command.id}`, {});
return {
id: command.id,
type: "response",
command: "prompt",
success: true,
data: { agentInvoked: true },
};
});
const dispatcher = new RpcInputDispatcher({ deps });
dispatcher.dispatch({ id: "active", type: "prompt", message: "active host tool" });
dispatcher.dispatch({ id: "queued", type: "prompt", message: "queued host tool" });
await flushMicrotasks();
expect(started).toEqual(["active"]);
expect(hostToolFrames).toHaveLength(1);
expect(hostToolFrames[0]).toMatchObject({
type: "host_tool_call",
toolCallId: "toolu_active",
toolName: "host_wait",
arguments: {},
});
bridge.close(disconnectMessage);
await dispatcher.drain();
expect(started).toEqual(["active", "queued"]);
expect(hostToolFrames).toHaveLength(1);
expect(outputs).toEqual([
{
id: "active",
type: "response",
command: "prompt",
success: false,
error: disconnectMessage,
},
{
id: "queued",
type: "response",
command: "prompt",
success: false,
error: disconnectMessage,
},
]);
});
test("drain after EOF rejects active and future extension UI requests", async () => {
const disconnectMessage = "RPC client disconnected before extension UI response completed";
const pendingExtensionRequests = new RpcPendingExtensionRequests();
const started: string[] = [];
let depsRef: RpcInputFrameDeps;
const { deps, outputs } = makeDeps(
async command => {
if (command.type !== "prompt") throw new Error(`unexpected command type: ${command.type}`);
started.push(command.id ?? "");
await requestExtensionInput(depsRef, `${command.id}-dialog`, command.message);
return {
id: command.id,
type: "response",
command: "prompt",
success: true,
data: { agentInvoked: true },
};
},
{ pendingExtensionRequests },
);
depsRef = deps;
const dispatcher = new RpcInputDispatcher({ deps });
dispatcher.dispatch({ id: "active", type: "prompt", message: "active dialog" });
dispatcher.dispatch({ id: "queued", type: "prompt", message: "queued dialog" });
await flushMicrotasks();
expect(started).toEqual(["active"]);
expect(outputs).toEqual([
{
type: "extension_ui_request",
id: "active-dialog",
method: "input",
message: "active dialog",
},
]);
pendingExtensionRequests.rejectAll(disconnectMessage);
await dispatcher.drain();
expect(started).toEqual(["active", "queued"]);
expect(outputs).toEqual([
{
type: "extension_ui_request",
id: "active-dialog",
method: "input",
message: "active dialog",
},
{
id: "active",
type: "response",
command: "prompt",
success: false,
error: disconnectMessage,
},
{
type: "extension_ui_request",
id: "queued-dialog",
method: "input",
message: "queued dialog",
},
{
id: "queued",
type: "response",
command: "prompt",
success: false,
error: disconnectMessage,
},
]);
});
});
describe("RpcShutdownCoordinator", () => {
/** performShutdown spy that records call count and outputs.length at the moment it ran. */
const makeShutdownRecorder = (outputs: OutputFrame[]) => {