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

test(coding-agent): unified collab test infrastructure and refined guest ui dialog fallback

Browse Source 
- Introduced the `CollabGuestUiResult` type to distinguish between answers and unavailable states during guest UI requests.
- Updated the remote dialog race logic to ignore unavailable guest results and fallback to the local host dialog.
- Centralized the `FakeWebSocket` and `InMemoryRelay` test infrastructure into a shared test helpers file.
- Cleaned up duplicate mock implementations and updated existing test suites to utilize the shared in-memory relay helpers.
This commit is contained in:
can1357
2026-07-02 02:40:08 +02:00
parent 0e2feab742
commit ad4f82eb8f
7 changed files with 327 additions and 466 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/coding-agent/src/collab/host.ts
+17 -10
View File
@@ -110,6 +110,13 @@ const TRANSCRIPT_ENTRY_TOO_LARGE_ERROR = `transcript entry exceeds transcript fe
* ship in a chunk of their own.
*/
const SNAPSHOT_CHUNK_BYTES = 512 * 1024;
/**
* Outcome of {@link CollabHost.requestGuestUi}. `answered` carries the guest's
* response (an `undefined` value is a genuine guest cancel); `unavailable`
* means the collab channel went away (teardown, relay drop) or the request was
* aborted before any guest answered — callers MUST NOT treat it as a cancel.
*/
export type CollabGuestUiResult = { kind: "answered"; value: CollabUiResponseValue } | { kind: "unavailable" };
export class CollabHost {
#ctx: InteractiveModeContext;
@@ -123,7 +130,7 @@ export class CollabHost {
#unsubscribe?: () => void;
#peers = new Map<number, { name: string; canWrite: boolean }>();
#uiReqSeq = 0;
#pendingUi = new Map<number, { request: CollabUiRequest; resolve(value: CollabUiResponseValue): void }>();
#pendingUi = new Map<number, { request: CollabUiRequest; settle(result: CollabGuestUiResult): void }>();
#lastStateJson = "";
#stateDebounce: Timer | null = null;
#streamingInterval: Timer | null = null;
@@ -163,24 +170,24 @@ export class CollabHost {
return list;
}
requestGuestUi(request: CollabUiRequestDraft, signal?: AbortSignal): Promise<CollabUiResponseValue> | null {
requestGuestUi(request: CollabUiRequestDraft, signal?: AbortSignal): Promise<CollabGuestUiResult> | null {
if (!this.#socket || !this.#hasWritablePeers()) return null;
const reqId = ++this.#uiReqSeq;
const fullRequest: CollabUiRequest = { ...request, reqId };
const { promise, resolve } = Promise.withResolvers<CollabUiResponseValue>();
const { promise, resolve } = Promise.withResolvers<CollabGuestUiResult>();
let settled = false;
const settle = (value: CollabUiResponseValue): void => {
const settle = (result: CollabGuestUiResult): void => {
if (settled) return;
settled = true;
signal?.removeEventListener("abort", onAbort);
this.#pendingUi.delete(reqId);
this.#sendWritablePeers({ t: "ui-request-end", reqId });
resolve(value);
resolve(result);
};
const onAbort = (): void => settle(undefined);
if (signal?.aborted) return Promise.resolve(undefined);
const onAbort = (): void => settle({ kind: "unavailable" });
if (signal?.aborted) return Promise.resolve({ kind: "unavailable" });
signal?.addEventListener("abort", onAbort, { once: true });
this.#pendingUi.set(reqId, { request: fullRequest, resolve: settle });
this.#pendingUi.set(reqId, { request: fullRequest, settle });
this.#sendWritablePeers({ t: "ui-request", request: fullRequest });
return promise;
}
@@ -302,7 +309,7 @@ export class CollabHost {
this.#agentsDebounce = null;
clearInterval(this.#streamingInterval ?? undefined);
this.#streamingInterval = null;
for (const pending of this.#pendingUi.values()) pending.resolve(undefined);
for (const pending of this.#pendingUi.values()) pending.settle({ kind: "unavailable" });
this.#pendingUi.clear();
this.#peers.clear();
this.#socket?.close();
@@ -455,7 +462,7 @@ export class CollabHost {
this.#rejectReadOnly("responding to ask", fromPeer);
return;
}
this.#pendingUi.get(reqId)?.resolve(value);
this.#pendingUi.get(reqId)?.settle({ kind: "answered", value });
}
#handlePrompt(text: string, images: ImageContent[] | undefined, fromPeer: number): void {
packages/coding-agent/src/modes/controllers/extension-ui-controller.ts
+20 -17
View File
@@ -44,16 +44,6 @@ function toWireSelectOptions(options: ExtensionUISelectItem[]): CollabUiSelectIt
);
}
function mergeAbortSignals(first: AbortSignal | undefined, second: AbortSignal): AbortSignal {
if (!first) return second;
if (first.aborted) return first;
const controller = new AbortController();
const abort = (): void => controller.abort();
first.addEventListener("abort", abort, { once: true });
second.addEventListener("abort", abort, { once: true });
return controller.signal;
}
export class ExtensionUiController {
#extensionTerminalInputUnsubscribers = new Set<() => void>();
#hookWidgetsAbove = new Map<string, ExtensionUiComponent>();
@@ -589,20 +579,33 @@ export class ExtensionUiController {
);
}
/**
* Race the local hook dialog against a mirrored guest ask. First *answer*
* wins and cancels the other side. A remote `unavailable` settlement
* (collab teardown, relay drop, abort) is NOT an answer: the local dialog
* keeps running — the host user may be mid-keystroke in it — and its
* eventual result is returned.
*/
async #raceCollabDialog(
request: CollabUiRequestDraft,
signal: AbortSignal | undefined,
local: (signal: AbortSignal) => Promise<string | undefined>,
local: (signal: AbortSignal | undefined) => Promise<string | undefined>,
): Promise<string | undefined> {
const host = this.ctx.collabHost;
if (!host) return local(signal ?? new AbortController().signal);
if (!host) return local(signal);
const localAbort = new AbortController();
const remoteAbort = new AbortController();
const remote = host.requestGuestUi(request, mergeAbortSignals(signal, remoteAbort.signal));
if (!remote) return local(signal ?? new AbortController().signal);
const localSignal = mergeAbortSignals(signal, localAbort.signal);
const localWinner = local(localSignal).then((value): CollabDialogWinner => ({ source: "local", value }));
const remoteWinner = remote.then((value): CollabDialogWinner => ({ source: "remote", value }));
const remote = host.requestGuestUi(
request,
signal ? AbortSignal.any([signal, remoteAbort.signal]) : remoteAbort.signal,
);
if (!remote) return local(signal);
const localWinner = local(signal ? AbortSignal.any([signal, localAbort.signal]) : localAbort.signal).then(
(value): CollabDialogWinner => ({ source: "local", value }),
);
const remoteWinner: Promise<CollabDialogWinner> = remote.then(result =>
result.kind === "answered" ? { source: "remote", value: result.value } : localWinner,
);
const winner = await Promise.race([localWinner, remoteWinner]);
if (winner.source === "remote") localAbort.abort();
else remoteAbort.abort();
packages/coding-agent/test/collab/chunked-welcome.test.ts
+6 -109
View File
@@ -14,114 +14,14 @@ import { afterAll, afterEach, beforeAll, describe, expect, it, spyOn } from "bun
import { importRoomKey } from "@oh-my-pi/pi-coding-agent/collab/crypto";
import { CollabGuestLink } from "@oh-my-pi/pi-coding-agent/collab/guest";
import { CollabHost } from "@oh-my-pi/pi-coding-agent/collab/host";
import {
COLLAB_PROTO,
type CollabFrame,
parseCollabLink,
rewriteEnvelopePeer,
unpackEnvelope,
} from "@oh-my-pi/pi-coding-agent/collab/protocol";
import { COLLAB_PROTO, type CollabFrame, parseCollabLink } from "@oh-my-pi/pi-coding-agent/collab/protocol";
import { CollabSocket } from "@oh-my-pi/pi-coding-agent/collab/relay-client";
import type { InteractiveModeContext } from "@oh-my-pi/pi-coding-agent/modes/types";
import type { SessionEntry } from "@oh-my-pi/pi-coding-agent/session/session-entries";
import { installInMemoryRelay, uninstallInMemoryRelay } from "./helpers/in-memory-relay";
// ── In-memory transport (verbatim copy of the relay used in read-only.test.ts) ──
let activeRelay: InMemoryRelay | null = null;
class FakeWebSocket {
static readonly CONNECTING = 0;
static readonly OPEN = 1;
static readonly CLOSING = 2;
static readonly CLOSED = 3;
binaryType = "blob";
readyState: number = FakeWebSocket.CONNECTING;
readonly role: "host" | "guest";
peerId = 0;
onopen: (() => void) | null = null;
onmessage: ((event: { data: unknown }) => void) | null = null;
onerror: (() => void) | null = null;
onclose: ((event: { code: number; reason: string }) => void) | null = null;
readonly #relay: InMemoryRelay;
constructor(url: string) {
const relay = activeRelay;
if (!relay) throw new Error("FakeWebSocket: no active in-memory relay");
this.#relay = relay;
this.role = new URL(url).searchParams.get("role") === "host" ? "host" : "guest";
queueMicrotask(() => {
if (this.readyState !== FakeWebSocket.CONNECTING) return;
this.readyState = FakeWebSocket.OPEN;
relay.connect(this);
this.onopen?.();
});
}
send(data: Uint8Array): void {
if (this.readyState !== FakeWebSocket.OPEN) return;
const bytes = new Uint8Array(data);
queueMicrotask(() => this.#relay.forward(this, bytes));
}
close(_code?: number): void {
if (this.readyState === FakeWebSocket.CLOSED) return;
this.readyState = FakeWebSocket.CLOSED;
this.#relay.disconnect(this);
queueMicrotask(() => this.onclose?.({ code: 1000, reason: "closed" }));
}
deliver(bytes: Uint8Array): void {
if (this.readyState !== FakeWebSocket.OPEN) return;
const copy = new Uint8Array(bytes);
queueMicrotask(() => this.onmessage?.({ data: copy.buffer }));
}
deliverControl(json: string): void {
if (this.readyState !== FakeWebSocket.OPEN) return;
queueMicrotask(() => this.onmessage?.({ data: json }));
}
}
class InMemoryRelay {
#host: FakeWebSocket | null = null;
readonly #guests = new Map<number, FakeWebSocket>();
#nextPeerId = 1;
connect(ws: FakeWebSocket): void {
if (ws.role === "host") {
this.#host = ws;
return;
}
ws.peerId = this.#nextPeerId++;
this.#guests.set(ws.peerId, ws);
this.#host?.deliverControl(JSON.stringify({ t: "peer-joined", peer: ws.peerId }));
}
forward(from: FakeWebSocket, bytes: Uint8Array): void {
if (from.role === "host") {
const envelope = unpackEnvelope(bytes);
if (!envelope) return;
if (envelope.peerId === 0) {
for (const guest of this.#guests.values()) guest.deliver(bytes);
} else {
this.#guests.get(envelope.peerId)?.deliver(bytes);
}
return;
}
rewriteEnvelopePeer(bytes, from.peerId);
this.#host?.deliver(bytes);
}
disconnect(ws: FakeWebSocket): void {
if (ws.role === "host") {
if (this.#host === ws) this.#host = null;
return;
}
this.#guests.delete(ws.peerId);
this.#host?.deliverControl(JSON.stringify({ t: "peer-left", peer: ws.peerId }));
}
}
// In-memory transport: shared FakeWebSocket + InMemoryRelay harness (see
// ./helpers/in-memory-relay), mirroring the relay's forwarding contract.
// ── Host harness with a configurable transcript ────────────────────────────
@@ -225,20 +125,17 @@ function makeFailingGuestContext(failure: Error): InteractiveModeContext {
// ── Shared host/relay ───────────────────────────────────────────────────────
const RealWebSocket = globalThis.WebSocket;
const snapshot = makeLargeSnapshot();
let host: CollabHost;
beforeAll(async () => {
globalThis.WebSocket = FakeWebSocket as unknown as typeof WebSocket;
activeRelay = new InMemoryRelay();
installInMemoryRelay();
host = new CollabHost(makeHostContext(snapshot));
await host.start("ws://localhost:8788");
});
afterAll(async () => {
globalThis.WebSocket = RealWebSocket;
activeRelay = null;
uninstallInMemoryRelay();
await host.stop("test done");
});
packages/coding-agent/test/collab/guest-subagent-badge.test.ts
+5 -102
View File
@@ -6,8 +6,6 @@ import {
COLLAB_PROTO,
type CollabFrame,
formatCollabLink,
rewriteEnvelopePeer,
unpackEnvelope,
} from "@oh-my-pi/pi-coding-agent/collab/protocol";
import { CollabSocket } from "@oh-my-pi/pi-coding-agent/collab/relay-client";
import {
@@ -16,103 +14,10 @@ import {
} from "@oh-my-pi/pi-coding-agent/modes/running-subagent-badge";
import type { InteractiveModeContext } from "@oh-my-pi/pi-coding-agent/modes/types";
import { AgentRegistry } from "@oh-my-pi/pi-coding-agent/registry/agent-registry";
import { installInMemoryRelay, uninstallInMemoryRelay } from "./helpers/in-memory-relay";
let activeRelay: InMemoryRelay | null = null;
const RealWebSocket = globalThis.WebSocket;
class FakeWebSocket {
static readonly CONNECTING = 0;
static readonly OPEN = 1;
static readonly CLOSING = 2;
static readonly CLOSED = 3;
binaryType = "arraybuffer";
readyState: number = FakeWebSocket.CONNECTING;
readonly role: "host" | "guest";
peerId = 0;
onopen: (() => void) | null = null;
onmessage: ((event: { data: unknown }) => void) | null = null;
onerror: (() => void) | null = null;
onclose: ((event: { code: number; reason: string }) => void) | null = null;
readonly #relay: InMemoryRelay;
constructor(url: string) {
const relay = activeRelay;
if (!relay) throw new Error("FakeWebSocket: no active in-memory relay");
this.#relay = relay;
this.role = new URL(url).searchParams.get("role") === "host" ? "host" : "guest";
queueMicrotask(() => {
if (this.readyState !== FakeWebSocket.CONNECTING) return;
this.readyState = FakeWebSocket.OPEN;
relay.connect(this);
this.onopen?.();
});
}
send(data: Uint8Array): void {
if (this.readyState !== FakeWebSocket.OPEN) return;
const bytes = new Uint8Array(data);
queueMicrotask(() => this.#relay.forward(this, bytes));
}
close(_code?: number): void {
if (this.readyState === FakeWebSocket.CLOSED) return;
this.readyState = FakeWebSocket.CLOSED;
this.#relay.disconnect(this);
queueMicrotask(() => this.onclose?.({ code: 1000, reason: "closed" }));
}
deliver(bytes: Uint8Array): void {
if (this.readyState !== FakeWebSocket.OPEN) return;
const copy = new Uint8Array(bytes);
queueMicrotask(() => this.onmessage?.({ data: copy.buffer }));
}
deliverControl(json: string): void {
if (this.readyState !== FakeWebSocket.OPEN) return;
queueMicrotask(() => this.onmessage?.({ data: json }));
}
}
class InMemoryRelay {
#host: FakeWebSocket | null = null;
readonly #guests = new Map<number, FakeWebSocket>();
#nextPeerId = 1;
connect(ws: FakeWebSocket): void {
if (ws.role === "host") {
this.#host = ws;
return;
}
ws.peerId = this.#nextPeerId++;
this.#guests.set(ws.peerId, ws);
this.#host?.deliverControl(JSON.stringify({ t: "peer-joined", peer: ws.peerId }));
}
forward(from: FakeWebSocket, bytes: Uint8Array): void {
if (from.role === "host") {
const envelope = unpackEnvelope(bytes);
if (!envelope) return;
if (envelope.peerId === 0) {
for (const guest of this.#guests.values()) guest.deliver(bytes);
} else {
this.#guests.get(envelope.peerId)?.deliver(bytes);
}
return;
}
rewriteEnvelopePeer(bytes, from.peerId);
this.#host?.deliver(bytes);
}
disconnect(ws: FakeWebSocket): void {
if (ws.role === "host") {
if (this.#host === ws) this.#host = null;
return;
}
this.#guests.delete(ws.peerId);
this.#host?.deliverControl(JSON.stringify({ t: "peer-left", peer: ws.peerId }));
}
}
// In-memory transport: shared FakeWebSocket + InMemoryRelay harness (see
// ./helpers/in-memory-relay), mirroring the relay's forwarding contract.
function makeState(): Extract<CollabFrame, { t: "welcome" }>["state"] {
return {
@@ -200,13 +105,11 @@ function makeGuestContext(counts: number[]): InteractiveModeContext {
beforeEach(() => {
AgentRegistry.resetGlobalForTests();
activeRelay = new InMemoryRelay();
globalThis.WebSocket = FakeWebSocket as unknown as typeof WebSocket;
installInMemoryRelay();
});
afterEach(() => {
globalThis.WebSocket = RealWebSocket;
activeRelay = null;
uninstallInMemoryRelay();
AgentRegistry.resetGlobalForTests();
});
packages/coding-agent/test/collab/guest-ui-request.test.ts
+128 -105
View File
@@ -21,114 +21,18 @@ import {
type CollabSessionState,
formatCollabLink,
parseCollabLink,
rewriteEnvelopePeer,
unpackEnvelope,
} from "@oh-my-pi/pi-coding-agent/collab/protocol";
import { CollabSocket } from "@oh-my-pi/pi-coding-agent/collab/relay-client";
import type {
ExtensionUIDialogOptions,
ExtensionUISelectItem,
} from "@oh-my-pi/pi-coding-agent/extensibility/extensions/types";
import { ExtensionUiController } from "@oh-my-pi/pi-coding-agent/modes/controllers/extension-ui-controller";
import type { InteractiveModeContext, InteractiveSelectorDialogOptions } from "@oh-my-pi/pi-coding-agent/modes/types";
import { installInMemoryRelay, uninstallInMemoryRelay } from "./helpers/in-memory-relay";
// ── In-memory transport (same contract as the other collab tests) ──────────
let activeRelay: InMemoryRelay | null = null;
const RealWebSocket = globalThis.WebSocket;
class FakeWebSocket {
static readonly CONNECTING = 0;
static readonly OPEN = 1;
static readonly CLOSING = 2;
static readonly CLOSED = 3;
binaryType = "arraybuffer";
readyState: number = FakeWebSocket.CONNECTING;
readonly role: "host" | "guest";
peerId = 0;
onopen: (() => void) | null = null;
onmessage: ((event: { data: unknown }) => void) | null = null;
onerror: (() => void) | null = null;
onclose: ((event: { code: number; reason: string }) => void) | null = null;
readonly #relay: InMemoryRelay;
constructor(url: string) {
const relay = activeRelay;
if (!relay) throw new Error("FakeWebSocket: no active in-memory relay");
this.#relay = relay;
this.role = new URL(url).searchParams.get("role") === "host" ? "host" : "guest";
queueMicrotask(() => {
if (this.readyState !== FakeWebSocket.CONNECTING) return;
this.readyState = FakeWebSocket.OPEN;
relay.connect(this);
this.onopen?.();
});
}
send(data: Uint8Array): void {
if (this.readyState !== FakeWebSocket.OPEN) return;
const bytes = new Uint8Array(data);
queueMicrotask(() => this.#relay.forward(this, bytes));
}
close(_code?: number): void {
if (this.readyState === FakeWebSocket.CLOSED) return;
this.readyState = FakeWebSocket.CLOSED;
this.#relay.disconnect(this);
queueMicrotask(() => this.onclose?.({ code: 1000, reason: "closed" }));
}
deliver(bytes: Uint8Array): void {
if (this.readyState !== FakeWebSocket.OPEN) return;
const copy = new Uint8Array(bytes);
queueMicrotask(() => this.onmessage?.({ data: copy.buffer }));
}
deliverControl(json: string): void {
if (this.readyState !== FakeWebSocket.OPEN) return;
queueMicrotask(() => this.onmessage?.({ data: json }));
}
}
class InMemoryRelay {
#host: FakeWebSocket | null = null;
readonly #guests = new Map<number, FakeWebSocket>();
#nextPeerId = 1;
connect(ws: FakeWebSocket): void {
if (ws.role === "host") {
this.#host = ws;
return;
}
ws.peerId = this.#nextPeerId++;
this.#guests.set(ws.peerId, ws);
this.#host?.deliverControl(JSON.stringify({ t: "peer-joined", peer: ws.peerId }));
}
forward(from: FakeWebSocket, bytes: Uint8Array): void {
if (from.role === "host") {
const envelope = unpackEnvelope(bytes);
if (!envelope) return;
if (envelope.peerId === 0) {
for (const guest of this.#guests.values()) guest.deliver(bytes);
} else {
this.#guests.get(envelope.peerId)?.deliver(bytes);
}
return;
}
rewriteEnvelopePeer(bytes, from.peerId);
this.#host?.deliver(bytes);
}
disconnect(ws: FakeWebSocket): void {
if (ws.role === "host") {
if (this.#host === ws) this.#host = null;
return;
}
this.#guests.delete(ws.peerId);
this.#host?.deliverControl(JSON.stringify({ t: "peer-left", peer: ws.peerId }));
}
}
// In-memory transport: shared FakeWebSocket + InMemoryRelay harness (see
// ./helpers/in-memory-relay), same contract as the other collab tests.
// ── Guest harness ───────────────────────────────────────────────────────────
@@ -373,8 +277,7 @@ const harnessCleanups: (() => Promise<void>)[] = [];
let writeSpy: { mockRestore(): void } | null = null;
beforeEach(() => {
activeRelay = new InMemoryRelay();
globalThis.WebSocket = FakeWebSocket as unknown as typeof WebSocket;
installInMemoryRelay();
writeSpy = spyOn(Bun, "write").mockResolvedValue(0);
});
@@ -382,8 +285,7 @@ afterEach(async () => {
for (const cleanup of harnessCleanups.splice(0).reverse()) await cleanup();
writeSpy?.mockRestore();
writeSpy = null;
globalThis.WebSocket = RealWebSocket;
activeRelay = null;
uninstallInMemoryRelay();
});
async function openHarness(opts?: { readOnly?: boolean }): Promise<GuestUiHarness> {
@@ -650,7 +552,7 @@ describe("collab proto handshake (#4049)", () => {
const request = await guest.nextFrame();
if (request.t !== "ui-request") throw new Error(`expected ui-request, got ${request.t}`);
guest.socket.send({ t: "ui-response", reqId: request.request.reqId, value: "Yes" });
expect(await pending).toBe("Yes");
expect(await pending).toEqual({ kind: "answered", value: "Yes" });
} finally {
guest.socket.close();
await host.stop("test done");
@@ -690,3 +592,124 @@ describe("collab proto handshake (#4049)", () => {
}
});
});
// ── Host dialog vs collab teardown (#4049 follow-up) ────────────────────────
//
// `ExtensionUiController.#raceCollabDialog` mirrors a hook dialog to writable
// guests and races the two surfaces. Teardown (/collab stop, non-reconnectable
// relay drop) settles every pending guest ask as `unavailable`; that is NOT a
// guest answer, so the local dialog the host user may be typing in must keep
// running and win with its eventual value. Only a genuine guest settlement —
// answer or explicit cancel (`answered` with undefined) — dismisses it.
/** One local hook-dialog presentation captured from the stub controller. */
interface LocalDialogStub {
title: string;
signal: AbortSignal | undefined;
/** Simulate the host user submitting (string) or cancelling (undefined). */
settle(value: string | undefined): void;
}
/**
* ExtensionUiController with the TUI dialog seam stubbed out: presentations
* are recorded instead of mounted, and abort mirrors `#presentDialog`
* (settles the dialog with undefined).
*/
class StubDialogController extends ExtensionUiController {
readonly localDialogs: LocalDialogStub[] = [];
override showHookSelector(
title: string,
_options: ExtensionUISelectItem[],
dialogOptions?: InteractiveSelectorDialogOptions,
): Promise<string | undefined> {
const { promise, resolve } = Promise.withResolvers<string | undefined>();
let settled = false;
const settle = (value: string | undefined): void => {
if (settled) return;
settled = true;
resolve(value);
};
dialogOptions?.signal?.addEventListener("abort", () => settle(undefined), { once: true });
this.localDialogs.push({ title, signal: dialogOptions?.signal, settle });
return promise;
}
}
describe("collab host dialog vs teardown (#4049 follow-up)", () => {
async function openRace(): Promise<{
host: CollabHost;
controller: StubDialogController;
guest: { socket: CollabSocket; nextFrame(): Promise<CollabFrame> };
result: Promise<string | undefined>;
dialog: LocalDialogStub;
requestFrame: CollabFrame & { t: "ui-request" };
cleanup(): Promise<void>;
}> {
const ctx = makeHostContext();
const host = new CollabHost(ctx);
await host.start("ws://localhost:8787");
ctx.collabHost = host;
const controller = new StubDialogController(ctx);
const guest = await joinRawGuest(host.link, COLLAB_PROTO);
const welcome = await guest.nextFrame();
if (welcome.t !== "welcome") throw new Error(`expected welcome, got ${welcome.t}`);
const result = controller.showCollabAwareSelector("Deploy?", ["Yes", "No"]);
const requestFrame = await guest.nextFrame();
if (requestFrame.t !== "ui-request") throw new Error(`expected ui-request, got ${requestFrame.t}`);
const dialog = controller.localDialogs[0];
if (!dialog) throw new Error("expected the local dialog to be presented alongside the guest ask");
return {
host,
controller,
guest,
result,
dialog,
requestFrame,
cleanup: async () => {
guest.socket.close();
await host.stop("test done");
},
};
}
it("keeps the local dialog running through collab teardown and returns its eventual answer", async () => {
const race = await openRace();
try {
await race.host.stop("host stopped collab");
// Deterministic bug discriminator, no clock: buggy code treated
// teardown's settlement as a remote win — the race resolved
// undefined before the local dialog could answer. Fixed code keeps
// the local dialog live, so `result` stays pending until it
// settles and wins with its value.
race.dialog.settle("stay-local");
expect(await race.result).toBe("stay-local");
expect(race.dialog.signal?.aborted).toBe(false);
} finally {
await race.cleanup();
}
});
it("dismisses the local dialog and returns undefined on a genuine guest cancel", async () => {
const race = await openRace();
try {
race.guest.socket.send({ t: "ui-response", reqId: race.requestFrame.request.reqId, value: undefined });
expect(await race.result).toBeUndefined();
expect(race.dialog.signal?.aborted).toBe(true);
} finally {
await race.cleanup();
}
});
it("dismisses the local dialog and returns the guest's value when the guest answers first", async () => {
const race = await openRace();
try {
race.guest.socket.send({ t: "ui-response", reqId: race.requestFrame.request.reqId, value: "No" });
expect(await race.result).toBe("No");
expect(race.dialog.signal?.aborted).toBe(true);
} finally {
await race.cleanup();
}
});
});
packages/coding-agent/test/collab/helpers/in-memory-relay.ts
+140
View File
@@ -0,0 +1,140 @@
/**
* In-memory collab transport shared by the collab test suites.
*
* `FakeWebSocket` + `InMemoryRelay` replace the real Bun.serve relay and
* loopback WebSocket. They mirror the production relay's forwarding contract
* exactly (4-byte peerId envelope routing, peer-joined/peer-left control
* frames) but deliver every frame on a microtask with zero network or timer
* latency. Real `CollabSocket` / `CollabHost` / `CollabGuestLink` run
* unchanged on top, so sealing, enveloping, the hello→welcome handshake, and
* permission enforcement are all exercised.
*
* Usage: `installInMemoryRelay()` in `beforeAll`/`beforeEach`,
* `uninstallInMemoryRelay()` in the matching `afterAll`/`afterEach`.
*/
import { rewriteEnvelopePeer, unpackEnvelope } from "@oh-my-pi/pi-coding-agent/collab/protocol";
/** Active relay the fake transport routes through; set between install/uninstall. */
let activeRelay: InMemoryRelay | null = null;
/** Pristine constructor captured before any test swaps it. */
const RealWebSocket = globalThis.WebSocket;
export class FakeWebSocket {
static readonly CONNECTING = 0;
static readonly OPEN = 1;
static readonly CLOSING = 2;
static readonly CLOSED = 3;
binaryType = "blob";
readyState: number = FakeWebSocket.CONNECTING;
readonly role: "host" | "guest";
peerId = 0;
onopen: (() => void) | null = null;
onmessage: ((event: { data: unknown }) => void) | null = null;
onerror: (() => void) | null = null;
onclose: ((event: { code: number; reason: string }) => void) | null = null;
readonly #relay: InMemoryRelay;
constructor(url: string) {
const relay = activeRelay;
if (!relay) throw new Error("FakeWebSocket: no active in-memory relay");
this.#relay = relay;
this.role = new URL(url).searchParams.get("role") === "host" ? "host" : "guest";
queueMicrotask(() => {
if (this.readyState !== FakeWebSocket.CONNECTING) return;
this.readyState = FakeWebSocket.OPEN;
relay.connect(this);
this.onopen?.();
});
}
send(data: Uint8Array): void {
if (this.readyState !== FakeWebSocket.OPEN) return;
// Snapshot: the relay rewrites the peerId in place, and the sender may
// reuse the buffer once send() returns.
const bytes = new Uint8Array(data);
queueMicrotask(() => this.#relay.forward(this, bytes));
}
close(_code?: number): void {
if (this.readyState === FakeWebSocket.CLOSED) return;
this.readyState = FakeWebSocket.CLOSED;
this.#relay.disconnect(this);
queueMicrotask(() => this.onclose?.({ code: 1000, reason: "closed" }));
}
/** Relay → this socket: a binary frame, delivered as ArrayBuffer (binaryType "arraybuffer"). */
deliver(bytes: Uint8Array): void {
if (this.readyState !== FakeWebSocket.OPEN) return;
const copy = new Uint8Array(bytes);
queueMicrotask(() => this.onmessage?.({ data: copy.buffer }));
}
/** Relay → this socket: a JSON control message. */
deliverControl(json: string): void {
if (this.readyState !== FakeWebSocket.OPEN) return;
queueMicrotask(() => this.onmessage?.({ data: json }));
}
}
/** Single-room in-memory relay mirroring the production forwarding contract. */
export class InMemoryRelay {
#host: FakeWebSocket | null = null;
readonly #guests = new Map<number, FakeWebSocket>();
#nextPeerId = 1;
connect(ws: FakeWebSocket): void {
if (ws.role === "host") {
this.#host = ws;
return;
}
ws.peerId = this.#nextPeerId++;
this.#guests.set(ws.peerId, ws);
this.#host?.deliverControl(JSON.stringify({ t: "peer-joined", peer: ws.peerId }));
}
forward(from: FakeWebSocket, bytes: Uint8Array): void {
if (from.role === "host") {
const envelope = unpackEnvelope(bytes);
if (!envelope) return;
if (envelope.peerId === 0) {
for (const guest of this.#guests.values()) guest.deliver(bytes);
} else {
this.#guests.get(envelope.peerId)?.deliver(bytes);
}
return;
}
rewriteEnvelopePeer(bytes, from.peerId);
this.#host?.deliver(bytes);
}
disconnect(ws: FakeWebSocket): void {
if (ws.role === "host") {
if (this.#host === ws) this.#host = null;
return;
}
this.#guests.delete(ws.peerId);
this.#host?.deliverControl(JSON.stringify({ t: "peer-left", peer: ws.peerId }));
}
}
/**
* Create a fresh relay and route `new WebSocket(...)` through it.
* Pair with {@link uninstallInMemoryRelay} in the matching after-hook.
*/
export function installInMemoryRelay(): InMemoryRelay {
activeRelay = new InMemoryRelay();
globalThis.WebSocket = FakeWebSocket as unknown as typeof WebSocket;
return activeRelay;
}
/**
* Restore the real WebSocket constructor and drop the active relay. Sockets
* already constructed keep their own relay reference, so in-flight teardown
* (e.g. `host.stop()`) still works after uninstall.
*/
export function uninstallInMemoryRelay(): void {
globalThis.WebSocket = RealWebSocket;
activeRelay = null;
}
packages/coding-agent/test/collab/read-only.test.ts
+11 -123
View File
@@ -11,125 +11,16 @@
import { afterAll, afterEach, beforeAll, describe, expect, it } from "bun:test";
import { importRoomKey } from "@oh-my-pi/pi-coding-agent/collab/crypto";
import { CollabHost } from "@oh-my-pi/pi-coding-agent/collab/host";
import {
COLLAB_PROTO,
type CollabFrame,
parseCollabLink,
rewriteEnvelopePeer,
unpackEnvelope,
} from "@oh-my-pi/pi-coding-agent/collab/protocol";
import { COLLAB_PROTO, type CollabFrame, parseCollabLink } from "@oh-my-pi/pi-coding-agent/collab/protocol";
import { CollabSocket } from "@oh-my-pi/pi-coding-agent/collab/relay-client";
import type { InteractiveModeContext } from "@oh-my-pi/pi-coding-agent/modes/types";
import { installInMemoryRelay, uninstallInMemoryRelay } from "./helpers/in-memory-relay";
// ── In-memory transport ────────────────────────────────────────────────────
// FakeWebSocket + InMemoryRelay replace the real Bun.serve relay and loopback
// WebSocket. They mirror the production relay's forwarding contract exactly
// (4-byte peerId envelope routing, peer-joined/peer-left control frames) but
// deliver every frame on a microtask with zero network or timer latency. Real
// CollabSocket / CollabHost run unchanged on top, so sealing, enveloping, the
// hello→welcome handshake, and read-only enforcement are all exercised.
/** Active relay the fake transport routes through; set for the lifetime of this file. */
let activeRelay: InMemoryRelay | null = null;
class FakeWebSocket {
static readonly CONNECTING = 0;
static readonly OPEN = 1;
static readonly CLOSING = 2;
static readonly CLOSED = 3;
binaryType = "blob";
readyState: number = FakeWebSocket.CONNECTING;
readonly role: "host" | "guest";
peerId = 0;
onopen: (() => void) | null = null;
onmessage: ((event: { data: unknown }) => void) | null = null;
onerror: (() => void) | null = null;
onclose: ((event: { code: number; reason: string }) => void) | null = null;
readonly #relay: InMemoryRelay;
constructor(url: string) {
const relay = activeRelay;
if (!relay) throw new Error("FakeWebSocket: no active in-memory relay");
this.#relay = relay;
this.role = new URL(url).searchParams.get("role") === "host" ? "host" : "guest";
queueMicrotask(() => {
if (this.readyState !== FakeWebSocket.CONNECTING) return;
this.readyState = FakeWebSocket.OPEN;
relay.connect(this);
this.onopen?.();
});
}
send(data: Uint8Array): void {
if (this.readyState !== FakeWebSocket.OPEN) return;
// Snapshot: the relay rewrites the peerId in place, and the sender may
// reuse the buffer once send() returns.
const bytes = new Uint8Array(data);
queueMicrotask(() => this.#relay.forward(this, bytes));
}
close(_code?: number): void {
if (this.readyState === FakeWebSocket.CLOSED) return;
this.readyState = FakeWebSocket.CLOSED;
this.#relay.disconnect(this);
queueMicrotask(() => this.onclose?.({ code: 1000, reason: "closed" }));
}
/** Relay → this socket: a binary frame, delivered as ArrayBuffer (binaryType "arraybuffer"). */
deliver(bytes: Uint8Array): void {
if (this.readyState !== FakeWebSocket.OPEN) return;
const copy = new Uint8Array(bytes);
queueMicrotask(() => this.onmessage?.({ data: copy.buffer }));
}
/** Relay → this socket: a JSON control message. */
deliverControl(json: string): void {
if (this.readyState !== FakeWebSocket.OPEN) return;
queueMicrotask(() => this.onmessage?.({ data: json }));
}
}
/** Single-room in-memory relay mirroring the production forwarding contract. */
class InMemoryRelay {
#host: FakeWebSocket | null = null;
readonly #guests = new Map<number, FakeWebSocket>();
#nextPeerId = 1;
connect(ws: FakeWebSocket): void {
if (ws.role === "host") {
this.#host = ws;
return;
}
ws.peerId = this.#nextPeerId++;
this.#guests.set(ws.peerId, ws);
this.#host?.deliverControl(JSON.stringify({ t: "peer-joined", peer: ws.peerId }));
}
forward(from: FakeWebSocket, bytes: Uint8Array): void {
if (from.role === "host") {
const envelope = unpackEnvelope(bytes);
if (!envelope) return;
if (envelope.peerId === 0) {
for (const guest of this.#guests.values()) guest.deliver(bytes);
} else {
this.#guests.get(envelope.peerId)?.deliver(bytes);
}
return;
}
rewriteEnvelopePeer(bytes, from.peerId);
this.#host?.deliver(bytes);
}
disconnect(ws: FakeWebSocket): void {
if (ws.role === "host") {
if (this.#host === ws) this.#host = null;
return;
}
this.#guests.delete(ws.peerId);
this.#host?.deliverControl(JSON.stringify({ t: "peer-left", peer: ws.peerId }));
}
}
// In-memory transport: FakeWebSocket + InMemoryRelay (see ./helpers/in-memory-relay)
// replace the real Bun.serve relay and loopback WebSocket with a zero-latency
// microtask transport. Real CollabSocket / CollabHost run unchanged on top, so
// sealing, enveloping, the hello→welcome handshake, and read-only enforcement
// are all exercised.
interface HostHarness {
ctx: InteractiveModeContext;
@@ -250,14 +141,12 @@ async function joinAsGuest(link: string, name: string, writeTokenOverride?: stri
// step; it is identical across all three tests (none mutate host config), so it
// runs once. Per-test guest state is reset in afterEach.
const RealWebSocket = globalThis.WebSocket;
const guestCleanups: (() => void)[] = [];
let harness: HostHarness;
let host: CollabHost;
beforeAll(async () => {
globalThis.WebSocket = FakeWebSocket as unknown as typeof WebSocket;
activeRelay = new InMemoryRelay();
installInMemoryRelay();
harness = makeHostContext();
host = new CollabHost(harness.ctx);
// Port is irrelevant: the fake transport routes by the `role` query param.
@@ -273,8 +162,7 @@ afterEach(() => {
afterAll(async () => {
// Restore the real transport first so the global is clean even if stop() throws;
// the host's socket holds its own FakeWebSocket/relay refs, so teardown still works.
globalThis.WebSocket = RealWebSocket;
activeRelay = null;
uninstallInMemoryRelay();
await host.stop("test done");
});
@@ -336,7 +224,7 @@ describe("collab read-only links", () => {
expect(request.request).toMatchObject({ kind: "select", title: "Continue?", options: ["Yes"] });
guest.socket.send({ t: "ui-response", reqId: request.request.reqId, value: "Yes" });
expect(await pending).toBe("Yes");
expect(await pending).toEqual({ kind: "answered", value: "Yes" });
const end = await guest.nextFrame();
expect(end).toEqual({ t: "ui-request-end", reqId: request.request.reqId });
});
@@ -360,7 +248,7 @@ describe("collab read-only links", () => {
expect(replayed).toEqual(firstRequest);
secondGuest.socket.send({ t: "ui-response", reqId: firstRequest.request.reqId, value: "late" });
expect(await pending).toBe("late");
expect(await pending).toEqual({ kind: "answered", value: "late" });
});
it("treats a forged write token as read-only", async () => {