wehub-resource-sync
363 lines
12 KiB
TypeScript
363 lines
12 KiB
TypeScript
import { describe, expect, test } from "bun:test";
|
|
import {
|
|
dispatchRpcInputFrame,
|
|
type PendingExtensionRequest,
|
|
type RpcInputFrameDeps,
|
|
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";
|
|
|
|
type OutputFrame = RpcResponse | object;
|
|
|
|
const makeDeps = (handleCommand: RpcInputFrameDeps["handleCommand"]) => {
|
|
const outputs: OutputFrame[] = [];
|
|
const deps: RpcInputFrameDeps = {
|
|
handleCommand,
|
|
output: obj => {
|
|
outputs.push(obj as OutputFrame);
|
|
},
|
|
errorResponse: (id, command, message) => ({
|
|
id,
|
|
type: "response",
|
|
command,
|
|
success: false,
|
|
error: message,
|
|
}),
|
|
pendingExtensionRequests: new Map<string, PendingExtensionRequest>(),
|
|
onHostToolResult: () => {},
|
|
onHostToolUpdate: () => {},
|
|
onHostUriResult: () => {},
|
|
};
|
|
return { deps, outputs };
|
|
};
|
|
|
|
const flushMicrotasks = () => new Promise<void>(resolve => setImmediate(resolve));
|
|
|
|
const cancelledBashResponse = (id: string): RpcResponse => ({
|
|
id,
|
|
type: "response",
|
|
command: "bash",
|
|
success: true,
|
|
data: {
|
|
output: "",
|
|
exitCode: -1,
|
|
cancelled: true,
|
|
truncated: false,
|
|
totalLines: 0,
|
|
totalBytes: 0,
|
|
outputLines: 0,
|
|
outputBytes: 0,
|
|
},
|
|
});
|
|
|
|
describe("dispatchRpcInputFrame", () => {
|
|
test("bash is dispatched in the background so abort_bash preempts it (issue #4079 A)", async () => {
|
|
const { promise: bashPending, resolve: resolveBash } = Promise.withResolvers<RpcResponse>();
|
|
let abortBashCalled = false;
|
|
|
|
const handleCommand = async (command: RpcCommand): Promise<RpcResponse> => {
|
|
if (command.type === "bash") {
|
|
// Block until abort_bash resolves the shared promise.
|
|
return await bashPending;
|
|
}
|
|
if (command.type === "abort_bash") {
|
|
abortBashCalled = true;
|
|
// Emulate `session.abortBash()` cancelling the in-flight bash so
|
|
// the queued executeBash promise resolves with cancelled=true.
|
|
resolveBash(cancelledBashResponse("b1"));
|
|
return { id: command.id, type: "response", command: "abort_bash", success: true };
|
|
}
|
|
throw new Error(`unexpected command type: ${command.type}`);
|
|
};
|
|
|
|
const { deps, outputs } = makeDeps(handleCommand);
|
|
|
|
// Kick off bash. If the fix works, dispatchRpcInputFrame returns
|
|
// undefined immediately without waiting for handleCommand.
|
|
const bashAwait = dispatchRpcInputFrame({ id: "b1", type: "bash", command: "sleep 9999" }, deps);
|
|
expect(bashAwait).toBeUndefined();
|
|
await flushMicrotasks();
|
|
expect(outputs).toHaveLength(0);
|
|
|
|
// Now dispatch abort_bash. It must run serially (not backgrounded)
|
|
// and resolve after handleCommand completes.
|
|
const abortAwait = dispatchRpcInputFrame({ id: "a1", type: "abort_bash" }, deps);
|
|
expect(abortAwait).toBeInstanceOf(Promise);
|
|
await abortAwait;
|
|
|
|
expect(abortBashCalled).toBe(true);
|
|
expect(outputs[0]).toEqual({
|
|
id: "a1",
|
|
type: "response",
|
|
command: "abort_bash",
|
|
success: true,
|
|
});
|
|
|
|
// The background bash response arrives after abort_bash.
|
|
await flushMicrotasks();
|
|
expect(outputs).toHaveLength(2);
|
|
const bashFrame = outputs[1] as RpcResponse;
|
|
expect(bashFrame.command).toBe("bash");
|
|
expect(bashFrame.success).toBe(true);
|
|
if (bashFrame.command === "bash" && bashFrame.success) {
|
|
expect(bashFrame.data.cancelled).toBe(true);
|
|
expect(bashFrame.data.exitCode).toBe(-1);
|
|
}
|
|
});
|
|
|
|
test("non-bash commands are dispatched serially (ordering preserved)", async () => {
|
|
const started: string[] = [];
|
|
const finished: string[] = [];
|
|
const handleCommand = async (command: RpcCommand): Promise<RpcResponse> => {
|
|
started.push(command.type);
|
|
await Bun.sleep(5);
|
|
finished.push(command.type);
|
|
if (command.type === "abort_retry") {
|
|
return { id: command.id, type: "response", command: "abort_retry", success: true };
|
|
}
|
|
if (command.type === "set_auto_retry") {
|
|
return { id: command.id, type: "response", command: "set_auto_retry", success: true };
|
|
}
|
|
throw new Error(`unexpected: ${command.type}`);
|
|
};
|
|
|
|
const { deps, outputs } = makeDeps(handleCommand);
|
|
|
|
const first = dispatchRpcInputFrame({ id: "c1", type: "abort_retry" }, deps);
|
|
expect(first).toBeInstanceOf(Promise);
|
|
// The input loop awaits each command's promise before pulling the next
|
|
// frame; simulate that contract by awaiting before the next dispatch.
|
|
await first;
|
|
expect(outputs).toHaveLength(1);
|
|
expect(started).toEqual(["abort_retry"]);
|
|
expect(finished).toEqual(["abort_retry"]);
|
|
|
|
const second = dispatchRpcInputFrame({ id: "c2", type: "set_auto_retry", enabled: true }, deps);
|
|
await second;
|
|
expect(outputs).toHaveLength(2);
|
|
expect(started).toEqual(["abort_retry", "set_auto_retry"]);
|
|
});
|
|
|
|
test("bash handler errors surface as an error response on the background frame", async () => {
|
|
const handleCommand = async (command: RpcCommand): Promise<RpcResponse> => {
|
|
if (command.type === "bash") throw new Error("kaboom");
|
|
throw new Error(`unexpected: ${command.type}`);
|
|
};
|
|
|
|
const { deps, outputs } = makeDeps(handleCommand);
|
|
|
|
const awaited = dispatchRpcInputFrame({ id: "b2", type: "bash", command: "echo hi" }, deps);
|
|
expect(awaited).toBeUndefined();
|
|
|
|
// Give the background dispatch a chance to run its catch.
|
|
await flushMicrotasks();
|
|
await flushMicrotasks();
|
|
|
|
expect(outputs).toHaveLength(1);
|
|
expect(outputs[0]).toEqual({
|
|
id: "b2",
|
|
type: "response",
|
|
command: "bash",
|
|
success: false,
|
|
error: "kaboom",
|
|
});
|
|
});
|
|
|
|
test("background bash task is exposed so EOF cleanup can await its response", async () => {
|
|
const bashResponse: RpcResponse = {
|
|
id: "b3",
|
|
type: "response",
|
|
command: "bash",
|
|
success: true,
|
|
data: {
|
|
output: "done",
|
|
exitCode: 0,
|
|
cancelled: false,
|
|
truncated: false,
|
|
totalLines: 1,
|
|
totalBytes: 4,
|
|
outputLines: 1,
|
|
outputBytes: 4,
|
|
},
|
|
};
|
|
const { promise: bashPending, resolve: resolveBash } = Promise.withResolvers<RpcResponse>();
|
|
const { deps, outputs } = makeDeps(async command => {
|
|
if (command.type === "bash") return await bashPending;
|
|
throw new Error(`unexpected: ${command.type}`);
|
|
});
|
|
let trackedTask: Promise<void> | undefined;
|
|
deps.trackBackgroundTask = task => {
|
|
trackedTask = task;
|
|
};
|
|
|
|
const awaited = dispatchRpcInputFrame({ id: "b3", type: "bash", command: "echo done" }, deps);
|
|
expect(awaited).toBeUndefined();
|
|
expect(trackedTask).toBeInstanceOf(Promise);
|
|
expect(outputs).toHaveLength(0);
|
|
|
|
resolveBash(bashResponse);
|
|
await trackedTask;
|
|
|
|
expect(outputs).toEqual([bashResponse]);
|
|
});
|
|
});
|
|
|
|
describe("RpcShutdownCoordinator", () => {
|
|
/** performShutdown spy that records call count and outputs.length at the moment it ran. */
|
|
const makeShutdownRecorder = (outputs: OutputFrame[]) => {
|
|
const state = { calls: 0, outputsAtShutdown: -1 };
|
|
const performShutdown = async () => {
|
|
state.calls++;
|
|
state.outputsAtShutdown = outputs.length;
|
|
};
|
|
return { state, performShutdown };
|
|
};
|
|
|
|
/**
|
|
* Full production-shaped harness: a background-dispatched bash frame whose
|
|
* handler blocks on a gate, tracked by the coordinator exactly as
|
|
* `runRpcMode` wires it (`trackBackgroundTask: task => coordinator.track(task)`).
|
|
*/
|
|
const makeBashHarness = () => {
|
|
const gate = Promise.withResolvers<RpcResponse>();
|
|
const { deps, outputs } = makeDeps(async command => {
|
|
if (command.type === "bash") return await gate.promise;
|
|
throw new Error(`unexpected: ${command.type}`);
|
|
});
|
|
const shutdown = { requested: false };
|
|
const recorder = makeShutdownRecorder(outputs);
|
|
const coordinator = new RpcShutdownCoordinator({
|
|
isShutdownRequested: () => shutdown.requested,
|
|
performShutdown: recorder.performShutdown,
|
|
});
|
|
deps.trackBackgroundTask = task => coordinator.track(task);
|
|
return { gate, deps, outputs, shutdown, recorder, coordinator };
|
|
};
|
|
|
|
test("deferred shutdown drains an in-flight background bash before performShutdown", async () => {
|
|
const { gate, deps, outputs, shutdown, recorder, coordinator } = makeBashHarness();
|
|
|
|
const awaited = dispatchRpcInputFrame({ id: "s1", type: "bash", command: "sleep 9999" }, deps);
|
|
expect(awaited).toBeUndefined();
|
|
|
|
// Extension calls pi.shutdown() while bash is in flight; the input loop
|
|
// re-checks after its next serially-awaited frame.
|
|
shutdown.requested = true;
|
|
const check = coordinator.checkShutdownRequested();
|
|
|
|
// The check must stay pending while the background bash still owes its
|
|
// response frame. Race it against a flushed sentinel: if the check could
|
|
// resolve, its microtask would win before the setImmediate tick.
|
|
const winner = await Promise.race([check.then(() => "shutdown"), flushMicrotasks().then(() => "pending")]);
|
|
expect(winner).toBe("pending");
|
|
expect(recorder.state.calls).toBe(0);
|
|
expect(outputs).toHaveLength(0);
|
|
|
|
gate.resolve(cancelledBashResponse("s1"));
|
|
await check;
|
|
|
|
expect(outputs).toEqual([cancelledBashResponse("s1")]);
|
|
expect(recorder.state.calls).toBe(1);
|
|
// The bash response frame was already written when performShutdown ran.
|
|
expect(recorder.state.outputsAtShutdown).toBe(1);
|
|
});
|
|
|
|
test("settle hook fires the deferred shutdown when no further client frames arrive", async () => {
|
|
const { gate, deps, outputs, shutdown, recorder } = makeBashHarness();
|
|
|
|
const awaited = dispatchRpcInputFrame({ id: "s2", type: "bash", command: "sleep 9999" }, deps);
|
|
expect(awaited).toBeUndefined();
|
|
|
|
// Shutdown requested mid-bash; the stdin loop is parked with no frames,
|
|
// so the test never calls checkShutdownRequested() — only track()'s
|
|
// settle hook can trigger it.
|
|
shutdown.requested = true;
|
|
await flushMicrotasks();
|
|
expect(recorder.state.calls).toBe(0);
|
|
|
|
gate.resolve(cancelledBashResponse("s2"));
|
|
await flushMicrotasks();
|
|
await flushMicrotasks();
|
|
|
|
expect(recorder.state.calls).toBe(1);
|
|
expect(outputs).toEqual([cancelledBashResponse("s2")]);
|
|
expect(recorder.state.outputsAtShutdown).toBe(1);
|
|
});
|
|
|
|
test("concurrent triggers are latched: performShutdown runs exactly once", async () => {
|
|
const outputs: OutputFrame[] = [];
|
|
const recorder = makeShutdownRecorder(outputs);
|
|
const coordinator = new RpcShutdownCoordinator({
|
|
isShutdownRequested: () => true,
|
|
performShutdown: recorder.performShutdown,
|
|
});
|
|
|
|
const gateA = Promise.withResolvers<void>();
|
|
const gateB = Promise.withResolvers<void>();
|
|
coordinator.track(gateA.promise);
|
|
coordinator.track(gateB.promise);
|
|
|
|
// Explicit trigger (input loop) races the settle hooks of both tasks.
|
|
const check = coordinator.checkShutdownRequested();
|
|
gateA.resolve();
|
|
gateB.resolve();
|
|
await check;
|
|
await flushMicrotasks();
|
|
await flushMicrotasks();
|
|
|
|
expect(recorder.state.calls).toBe(1);
|
|
// A later re-check reuses the latched sequence instead of re-running it.
|
|
await coordinator.checkShutdownRequested();
|
|
expect(recorder.state.calls).toBe(1);
|
|
});
|
|
|
|
test("no-op when shutdown was not requested", async () => {
|
|
const outputs: OutputFrame[] = [];
|
|
const recorder = makeShutdownRecorder(outputs);
|
|
const coordinator = new RpcShutdownCoordinator({
|
|
isShutdownRequested: () => false,
|
|
performShutdown: recorder.performShutdown,
|
|
});
|
|
|
|
await coordinator.checkShutdownRequested();
|
|
expect(recorder.state.calls).toBe(0);
|
|
|
|
// A tracked task settling with the flag false never triggers shutdown.
|
|
const gate = Promise.withResolvers<void>();
|
|
coordinator.track(gate.promise);
|
|
gate.resolve();
|
|
await flushMicrotasks();
|
|
await flushMicrotasks();
|
|
expect(recorder.state.calls).toBe(0);
|
|
});
|
|
|
|
test("drain() waits for tasks tracked while draining", async () => {
|
|
const coordinator = new RpcShutdownCoordinator({
|
|
isShutdownRequested: () => false,
|
|
performShutdown: async () => {},
|
|
});
|
|
|
|
const gateA = Promise.withResolvers<void>();
|
|
const gateB = Promise.withResolvers<void>();
|
|
coordinator.track(gateA.promise);
|
|
// When A settles, a new task B enters the set mid-drain.
|
|
void gateA.promise.then(() => {
|
|
coordinator.track(gateB.promise);
|
|
});
|
|
|
|
let drained = false;
|
|
const drain = coordinator.drain().then(() => {
|
|
drained = true;
|
|
});
|
|
|
|
gateA.resolve();
|
|
await flushMicrotasks();
|
|
// A settled and B was tracked mid-drain; drain must keep waiting on B.
|
|
expect(drained).toBe(false);
|
|
|
|
gateB.resolve();
|
|
await drain;
|
|
expect(drained).toBe(true);
|
|
});
|
|
});
|