diff --git a/packages/coding-agent/src/collab/host.ts b/packages/coding-agent/src/collab/host.ts index 5fc22f3c1..74a021a80 100644 --- a/packages/coding-agent/src/collab/host.ts +++ b/packages/coding-agent/src/collab/host.ts @@ -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(); #uiReqSeq = 0; - #pendingUi = new Map(); + #pendingUi = new Map(); #lastStateJson = ""; #stateDebounce: Timer | null = null; #streamingInterval: Timer | null = null; @@ -163,24 +170,24 @@ export class CollabHost { return list; } - requestGuestUi(request: CollabUiRequestDraft, signal?: AbortSignal): Promise | null { + requestGuestUi(request: CollabUiRequestDraft, signal?: AbortSignal): Promise | null { if (!this.#socket || !this.#hasWritablePeers()) return null; const reqId = ++this.#uiReqSeq; const fullRequest: CollabUiRequest = { ...request, reqId }; - const { promise, resolve } = Promise.withResolvers(); + const { promise, resolve } = Promise.withResolvers(); 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 { diff --git a/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts b/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts index bd6d515f9..1c9172544 100644 --- a/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts +++ b/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts @@ -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(); @@ -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, + local: (signal: AbortSignal | undefined) => Promise, ): Promise { 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 = 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(); diff --git a/packages/coding-agent/test/collab/chunked-welcome.test.ts b/packages/coding-agent/test/collab/chunked-welcome.test.ts index 61e54fc85..53d786845 100644 --- a/packages/coding-agent/test/collab/chunked-welcome.test.ts +++ b/packages/coding-agent/test/collab/chunked-welcome.test.ts @@ -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(); - #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"); }); diff --git a/packages/coding-agent/test/collab/guest-subagent-badge.test.ts b/packages/coding-agent/test/collab/guest-subagent-badge.test.ts index efa921a31..6a35f87ca 100644 --- a/packages/coding-agent/test/collab/guest-subagent-badge.test.ts +++ b/packages/coding-agent/test/collab/guest-subagent-badge.test.ts @@ -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(); - #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["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(); }); diff --git a/packages/coding-agent/test/collab/guest-ui-request.test.ts b/packages/coding-agent/test/collab/guest-ui-request.test.ts index 13838859d..59806adf3 100644 --- a/packages/coding-agent/test/collab/guest-ui-request.test.ts +++ b/packages/coding-agent/test/collab/guest-ui-request.test.ts @@ -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(); - #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)[] = []; 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 { @@ -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 { + const { promise, resolve } = Promise.withResolvers(); + 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 }; + result: Promise; + dialog: LocalDialogStub; + requestFrame: CollabFrame & { t: "ui-request" }; + cleanup(): Promise; + }> { + 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(); + } + }); +}); diff --git a/packages/coding-agent/test/collab/helpers/in-memory-relay.ts b/packages/coding-agent/test/collab/helpers/in-memory-relay.ts new file mode 100644 index 000000000..dc693bcee --- /dev/null +++ b/packages/coding-agent/test/collab/helpers/in-memory-relay.ts @@ -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(); + #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; +} diff --git a/packages/coding-agent/test/collab/read-only.test.ts b/packages/coding-agent/test/collab/read-only.test.ts index 26eb1a452..84a7dc006 100644 --- a/packages/coding-agent/test/collab/read-only.test.ts +++ b/packages/coding-agent/test/collab/read-only.test.ts @@ -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(); - #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 () => {