@vellumai/assistant 0.12.1 → 0.12.2-staging.2

package/src/desktop/desktop-stream-bridge.test.ts

@@ -0,0 +1,395 @@
+import { mkdtempSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterAll, describe, expect, test } from "bun:test";
+
+import { newFakeDesktop, newViewer, settle } from "./__tests__/fake-desktop.js";
+import { DesktopDependencyInstaller } from "./desktop-dependencies.js";
+import {
+  type DesktopSessionManager,
+  type DesktopTcpHandlers,
+  type DesktopTcpSocket,
+  destroyDesktopSessionManager,
+  getDesktopSessionManager,
+} from "./desktop-session-manager.js";
+import { DesktopStreamBridge } from "./desktop-stream-bridge.js";
+
+const profileDir = mkdtempSync(join(tmpdir(), "desktop-bridge-test-"));
+afterAll(() => {
+  rmSync(profileDir, { recursive: true, force: true });
+});
+
+class FakeWs {
+  readonly sent: Uint8Array[] = [];
+  closeCode: number | null = null;
+  closeReason: string | null = null;
+  /** What `send` reports back; Bun returns 0 for a dropped frame. */
+  sendResult: number | null = null;
+
+  send(data: Uint8Array): number {
+    this.sent.push(data);
+    return this.sendResult ?? data.byteLength;
+  }
+
+  close(code?: number, reason?: string): void {
+    this.closeCode = code ?? null;
+    this.closeReason = reason ?? null;
+  }
+}
+
+class FakeTcp implements DesktopTcpSocket {
+  readonly writes: Uint8Array[] = [];
+  ended = false;
+  handlers!: DesktopTcpHandlers;
+
+  write(data: Uint8Array): number {
+    this.writes.push(data);
+    return data.byteLength;
+  }
+
+  end(): void {
+    this.ended = true;
+  }
+}
+
+function newBridge(
+  manager: DesktopSessionManager,
+  ensureInstalled: () => Promise<void> = async () => {},
+  isEnabled: () => boolean = () => true,
+) {
+  const ws = new FakeWs();
+  const tcp = new FakeTcp();
+  let release: (() => void) | null = null;
+  const gate = new Promise<void>((resolve) => {
+    release = resolve;
+  });
+  const bridge = new DesktopStreamBridge(ws, {
+    manager,
+    ensureInstalled,
+    isEnabled,
+    connect: async (_port, handlers) => {
+      tcp.handlers = handlers;
+      await gate;
+      return tcp;
+    },
+  });
+  return { ws, tcp, bridge, connectNow: () => release!() };
+}
+
+/** Whether the manager's viewer slot is free, observed through acquire. */
+function slotIsFree(manager: DesktopSessionManager): boolean {
+  const { viewer } = newViewer();
+  const result = manager.acquireViewerSlot(viewer);
+  if (result.ok) {
+    manager.releaseViewerSlot(viewer);
+  }
+  return result.ok;
+}
+
+describe("DesktopStreamBridge", () => {
+  test("refuses a disabled desktop before installation or process startup", async () => {
+    const h = newFakeDesktop({ profileDir });
+    let installs = 0;
+    const b = newBridge(
+      h.manager,
+      async () => {
+        installs++;
+      },
+      () => false,
+    );
+    await b.bridge.start();
+    expect(b.ws.closeCode).toBe(4008);
+    expect(installs).toBe(0);
+    expect(h.spawned).toEqual([]);
+    expect(slotIsFree(h.manager)).toBe(true);
+  });
+
+  test("does not start the desktop when disabled during installation", async () => {
+    const h = newFakeDesktop({ profileDir });
+    let enabled = true;
+    const install = Promise.withResolvers<void>();
+    const b = newBridge(
+      h.manager,
+      () => install.promise,
+      () => enabled,
+    );
+    const started = b.bridge.start();
+    enabled = false;
+    install.resolve();
+    await started;
+    expect(b.ws.closeCode).toBe(4008);
+    expect(slotIsFree(h.manager)).toBe(true);
+    expect(b.tcp.handlers).toBeUndefined();
+    expect(h.spawned).toEqual([]);
+  });
+
+  for (const direction of ["viewer", "desktop"] as const) {
+    test(`blocks ${direction} frames when the flag is disabled after connection`, async () => {
+      const h = newFakeDesktop({ profileDir });
+      let enabled = true;
+      const b = newBridge(
+        h.manager,
+        async () => {},
+        () => enabled,
+      );
+      b.connectNow();
+      await b.bridge.start();
+      enabled = false;
+      const frame = new Uint8Array([1, 2, 3]);
+      if (direction === "viewer") {
+        b.bridge.handleClientFrame(frame);
+      } else {
+        b.tcp.handlers.onData(frame);
+      }
+      expect(b.ws.closeCode).toBe(4008);
+      expect(b.ws.sent).toHaveLength(0);
+      expect(b.tcp.writes).toHaveLength(0);
+      expect(b.tcp.ended).toBe(true);
+      expect(slotIsFree(h.manager)).toBe(true);
+      await h.manager.destroy();
+    });
+  }
+
+  test("direct viewers share background setup across a timeout and reconnect", async () => {
+    const h = newFakeDesktop({ profileDir });
+    let ready = false;
+    let installs = 0;
+    const install = Promise.withResolvers<void>();
+    const installer = new DesktopDependencyInstaller({
+      supported: () => true,
+      ready: () => ready,
+      install: async () => {
+        installs++;
+        await install.promise;
+        ready = true;
+      },
+      notify: async () => {},
+    });
+    const ensureInstalled = () => installer.ensureReady();
+    const first = newBridge(h.manager, ensureInstalled);
+    const firstStart = first.bridge.start();
+    await settle();
+    expect(installs).toBe(1);
+    expect(h.spawned).toEqual([]);
+
+    const busy = newBridge(h.manager, async () => {
+      throw new Error("A rejected viewer must not start setup");
+    });
+    await busy.bridge.start();
+    expect(busy.ws.closeCode).toBe(4013);
+
+    first.bridge.handleClose();
+    expect(slotIsFree(h.manager)).toBe(true);
+    const retry = newBridge(h.manager, ensureInstalled);
+    const retryStart = retry.bridge.start();
+    await settle();
+    expect(installs).toBe(1);
+    expect(h.spawned).toEqual([]);
+
+    install.resolve();
+    retry.connectNow();
+    await Promise.all([firstStart, retryStart]);
+    expect(retry.ws.closeCode).toBeNull();
+    expect(h.spawned.length).toBeGreaterThan(0);
+    retry.tcp.handlers.onData(new Uint8Array([1, 2, 3]));
+    expect(retry.ws.sent).toEqual([new Uint8Array([1, 2, 3])]);
+    expect(first.ws.sent).toEqual([]);
+    expect(first.tcp.handlers).toBeUndefined();
+    retry.bridge.handleClose();
+  });
+
+  test("setup continues after the last viewer closes without starting a desktop", async () => {
+    const h = newFakeDesktop({ profileDir });
+    const install = Promise.withResolvers<void>();
+    const b = newBridge(h.manager, () => install.promise);
+    const started = b.bridge.start();
+    b.bridge.handleClose();
+    install.resolve();
+    await started;
+    expect(h.spawned).toEqual([]);
+    expect(slotIsFree(h.manager)).toBe(true);
+  });
+
+  test("setup failure releases the slot and permits a successful retry", async () => {
+    const h = newFakeDesktop({ profileDir });
+    let ready = false;
+    let attempts = 0;
+    const installer = new DesktopDependencyInstaller({
+      supported: () => true,
+      ready: () => ready,
+      install: async () => {
+        if (++attempts === 1) {
+          throw new Error("Download failed");
+        }
+        ready = true;
+      },
+      notify: async () => {},
+    });
+    const ensureInstalled = () => installer.ensureReady();
+    const first = newBridge(h.manager, ensureInstalled);
+    await first.bridge.start();
+    expect(first.ws.closeCode).toBe(4011);
+    expect(slotIsFree(h.manager)).toBe(true);
+    expect(h.spawned).toEqual([]);
+
+    const retry = newBridge(h.manager, ensureInstalled);
+    retry.connectNow();
+    await retry.bridge.start();
+    expect(attempts).toBe(2);
+    expect(retry.ws.closeCode).toBeNull();
+    retry.bridge.handleClose();
+  });
+
+  test("shutdown during setup closes the viewer without starting a desktop", async () => {
+    const h = newFakeDesktop({ profileDir });
+    const install = Promise.withResolvers<void>();
+    const b = newBridge(h.manager, () => install.promise);
+    const started = b.bridge.start();
+    await h.manager.destroy();
+    expect(b.ws.closeCode).toBe(1001);
+    install.resolve();
+    await started;
+    expect(h.spawned).toEqual([]);
+    expect(b.tcp.handlers).toBeUndefined();
+  });
+
+  test("pumps bytes both ways and flushes frames that beat the VNC socket", async () => {
+    const h = newFakeDesktop({ profileDir });
+    const b = newBridge(h.manager);
+    const started = b.bridge.start();
+    await settle();
+
+    b.bridge.handleClientFrame(new Uint8Array([1, 2, 3]));
+    b.bridge.handleClientFrame(new Uint8Array([4]).buffer);
+    b.bridge.handleClientFrame("text frames are not RFB");
+    expect(b.tcp.writes).toEqual([]);
+
+    b.connectNow();
+    await started;
+    expect(b.tcp.writes).toEqual([
+      new Uint8Array([1, 2, 3]),
+      new Uint8Array([4]),
+    ]);
+
+    b.tcp.handlers.onData(new Uint8Array([9, 9]));
+    expect(b.ws.sent).toEqual([new Uint8Array([9, 9])]);
+    expect(b.ws.closeCode).toBeNull();
+    expect(slotIsFree(h.manager)).toBe(false);
+
+    b.bridge.handleClose();
+    expect(b.tcp.ended).toBe(true);
+    expect(slotIsFree(h.manager)).toBe(true);
+  });
+
+  test("a second socket closes 4013 and leaves the first untouched", async () => {
+    const h = newFakeDesktop({ profileDir });
+    const first = newBridge(h.manager);
+    first.connectNow();
+    await first.bridge.start();
+
+    const second = newBridge(h.manager);
+    await second.bridge.start();
+
+    expect(second.ws.closeCode).toBe(4013);
+    expect(first.ws.closeCode).toBeNull();
+
+    second.bridge.handleClose();
+    expect(slotIsFree(h.manager)).toBe(false);
+  });
+
+  test("closes 4011 when the desktop fails to start and frees the slot", async () => {
+    const h = newFakeDesktop({ profileDir });
+    h.setVncReady(false);
+    const b = newBridge(h.manager);
+    await b.bridge.start();
+
+    expect(b.ws.closeCode).toBe(4011);
+    expect(slotIsFree(h.manager)).toBe(true);
+  });
+
+  test("closes 4011 and frees the slot when the start rejects with a plain error", async () => {
+    const h = newFakeDesktop({ profileDir });
+    h.manager.ensureDesktopRunning = () =>
+      Promise.reject(new Error("unexpected"));
+    const b = newBridge(h.manager);
+    await b.bridge.start();
+
+    expect(b.ws.closeCode).toBe(4011);
+    expect(b.ws.closeReason).toBe("Desktop failed to start");
+    expect(slotIsFree(h.manager)).toBe(true);
+  });
+
+  test("closes 4011 when the desktop dies under the viewer", async () => {
+    const h = newFakeDesktop({ profileDir });
+    const b = newBridge(h.manager);
+    b.connectNow();
+    await b.bridge.start();
+    await settle();
+
+    h.child("window-manager").exit(1);
+    await settle();
+
+    expect(b.ws.closeCode).toBe(4011);
+    expect(b.ws.closeReason).toContain("window-manager exited");
+    expect(b.tcp.ended).toBe(true);
+  });
+
+  test("closes 4011 when the viewer drops a frame", async () => {
+    const h = newFakeDesktop({ profileDir });
+    const b = newBridge(h.manager);
+    b.connectNow();
+    await b.bridge.start();
+
+    b.ws.sendResult = 0;
+    b.tcp.handlers.onData(new Uint8Array([1]));
+
+    expect(b.ws.closeCode).toBe(4011);
+    expect(b.ws.closeReason).toBe("Viewer too slow");
+    expect(b.tcp.ended).toBe(true);
+    expect(slotIsFree(h.manager)).toBe(true);
+  });
+
+  test("closes 1001 for a socket arriving after shutdown began", async () => {
+    const h = newFakeDesktop({ profileDir });
+    await h.manager.destroy();
+    const b = newBridge(h.manager);
+    await b.bridge.start();
+
+    expect(b.ws.closeCode).toBe(1001);
+    expect(b.ws.closeReason).toBe("The assistant is shutting down");
+    expect(h.spawned).toEqual([]);
+  });
+
+  test("the shared manager is latched by shutdown even when no desktop was ever served", async () => {
+    await destroyDesktopSessionManager();
+
+    const ws = new FakeWs();
+    const bridge = new DesktopStreamBridge(ws, {
+      isEnabled: () => true,
+      connect: async () => {
+        throw new Error("must not dial the VNC port");
+      },
+    });
+    await bridge.start();
+
+    expect(ws.closeCode).toBe(1001);
+    // Nothing is spawned: the start is refused before any binary is looked up.
+    await expect(
+      getDesktopSessionManager().ensureDesktopRunning(),
+    ).rejects.toThrow("Desktop ingress is closed");
+  });
+
+  test("closes a live viewer with 1001 when the runtime shuts down", async () => {
+    const h = newFakeDesktop({ profileDir, exitOnTerm: true });
+    const b = newBridge(h.manager);
+    b.connectNow();
+    await b.bridge.start();
+    await settle();
+
+    await h.manager.destroy();
+
+    expect(b.ws.closeCode).toBe(1001);
+    expect(b.ws.closeReason).toBe("The assistant is shutting down");
+    expect(b.tcp.ended).toBe(true);
+  });
+});

package/src/desktop/desktop-stream-bridge.ts

@@ -0,0 +1,214 @@
+/**
+ * Pumps raw RFB bytes between one `/v1/desktop/stream` WebSocket and the
+ * desktop's VNC server on loopback. Every frame is binary RFB data; outcomes
+ * are signaled with close codes only (see `DESKTOP_CLOSE`).
+ */
+
+import { desktopDependencyInstaller } from "./desktop-dependencies.js";
+import {
+  connectLoopback,
+  DESKTOP_CLOSE,
+  DESKTOP_VNC_PORT,
+  type DesktopLoss,
+  type DesktopSessionManager,
+  DesktopStartError,
+  type DesktopTcpHandlers,
+  type DesktopTcpSocket,
+  type DesktopViewer,
+  getDesktopSessionManager,
+  START_FAILED_LOSS,
+} from "./desktop-session-manager.js";
+
+/** Frames a client may send before the VNC socket is up; RFB handshakes are tiny. */
+const MAX_PENDING_BYTES = 1024 * 1024;
+
+interface DesktopStreamClientSocket {
+  /** Bun's `ServerWebSocket.send`: `0` means the frame was dropped. */
+  send(data: Uint8Array): number;
+  close(code?: number, reason?: string): void;
+}
+
+interface DesktopStreamBridgeOptions {
+  readonly isEnabled: () => boolean;
+  readonly manager?: DesktopSessionManager;
+  readonly ensureInstalled?: () => Promise<void>;
+  readonly connect?: (
+    port: number,
+    handlers: DesktopTcpHandlers,
+  ) => Promise<DesktopTcpSocket>;
+}
+
+export class DesktopStreamBridge {
+  private readonly ws: DesktopStreamClientSocket;
+  private readonly isEnabled: () => boolean;
+  private readonly manager: DesktopSessionManager;
+  private readonly ensureInstalled: () => Promise<void>;
+  private readonly connect: NonNullable<DesktopStreamBridgeOptions["connect"]>;
+  private readonly viewer: DesktopViewer;
+
+  private tcp: DesktopTcpSocket | null = null;
+  /** Client bytes waiting on the VNC socket, or on backpressure once it is up. */
+  private pending: Uint8Array[] = [];
+  private pendingBytes = 0;
+  private closed = false;
+  private ownsSlot = false;
+
+  constructor(
+    ws: DesktopStreamClientSocket,
+    options: DesktopStreamBridgeOptions,
+  ) {
+    this.ws = ws;
+    this.isEnabled = options.isEnabled;
+    this.manager = options.manager ?? getDesktopSessionManager();
+    this.ensureInstalled =
+      options.ensureInstalled ??
+      (() => desktopDependencyInstaller.ensureReady());
+    this.connect = options.connect ?? connectLoopback;
+    this.viewer = { onDesktopLost: (loss) => this.lose(loss) };
+  }
+
+  /** Claim the viewer slot, start the desktop, and dial its VNC port. */
+  async start(): Promise<void> {
+    if (!this.checkEnabled()) {
+      return;
+    }
+    const slot = this.manager.acquireViewerSlot(this.viewer);
+    if (!slot.ok) {
+      this.lose(slot.loss);
+      return;
+    }
+    this.ownsSlot = true;
+
+    try {
+      // Direct viewers share setup with the modal, including across reconnects.
+      await this.ensureInstalled();
+      if (!this.checkEnabled()) {
+        return;
+      }
+      await this.manager.ensureDesktopRunning();
+    } catch (err) {
+      // Usually already closed through onDesktopLost; the error covers a
+      // start this viewer joined late and never got the callback for.
+      this.lose(
+        err instanceof DesktopStartError ? err.loss : START_FAILED_LOSS,
+      );
+      return;
+    }
+    if (!this.checkEnabled()) {
+      return;
+    }
+
+    let tcp: DesktopTcpSocket;
+    try {
+      tcp = await this.connect(DESKTOP_VNC_PORT, {
+        onData: (data) => {
+          if (!this.checkEnabled()) {
+            return;
+          }
+          if (this.ws.send(data) === 0) {
+            this.fail(DESKTOP_CLOSE.failed, "Viewer too slow");
+          }
+        },
+        onDrain: () => this.flush(),
+        onClose: () =>
+          this.fail(DESKTOP_CLOSE.failed, "Desktop connection closed"),
+        onError: () =>
+          this.fail(DESKTOP_CLOSE.failed, "Desktop connection failed"),
+      });
+    } catch {
+      this.fail(DESKTOP_CLOSE.failed, "Desktop connection failed");
+      return;
+    }
+    if (!this.checkEnabled()) {
+      tcp.end();
+      return;
+    }
+    this.tcp = tcp;
+    this.flush();
+  }
+
+  handleClientFrame(message: string | Uint8Array | ArrayBuffer): void {
+    if (!this.checkEnabled() || typeof message === "string") {
+      return;
+    }
+    const bytes =
+      message instanceof ArrayBuffer ? new Uint8Array(message) : message;
+    if (this.pendingBytes + bytes.byteLength > MAX_PENDING_BYTES) {
+      this.fail(DESKTOP_CLOSE.failed, "Desktop stream backlog exceeded");
+      return;
+    }
+    this.pending.push(bytes);
+    this.pendingBytes += bytes.byteLength;
+    this.flush();
+  }
+
+  /** The client socket closed; the caller has already seen the close frame. */
+  handleClose(): void {
+    this.release();
+  }
+
+  private flush(): void {
+    if (!this.checkEnabled()) {
+      return;
+    }
+    const tcp = this.tcp;
+    if (!tcp) {
+      return;
+    }
+    while (this.pending.length > 0) {
+      const chunk = this.pending[0]!;
+      const written = tcp.write(chunk);
+      if (written < 0) {
+        this.fail(DESKTOP_CLOSE.failed, "Desktop connection closed");
+        return;
+      }
+      this.pendingBytes -= written;
+      if (written < chunk.byteLength) {
+        this.pending[0] = chunk.subarray(written);
+        return;
+      }
+      this.pending.shift();
+    }
+  }
+
+  private checkEnabled(): boolean {
+    if (this.closed) {
+      return false;
+    }
+    if (!this.isEnabled()) {
+      this.fail(
+        DESKTOP_CLOSE.unavailable,
+        "Desktop is not available on this assistant",
+      );
+      return false;
+    }
+    return true;
+  }
+
+  private lose({ code, reason }: DesktopLoss): void {
+    this.fail(code, reason);
+  }
+
+  private fail(code: number, reason: string): void {
+    if (this.closed) {
+      return;
+    }
+    this.release();
+    this.ws.close(code, reason);
+  }
+
+  private release(): void {
+    if (this.closed) {
+      return;
+    }
+    this.closed = true;
+    this.pending = [];
+    this.pendingBytes = 0;
+    this.tcp?.end();
+    this.tcp = null;
+    if (this.ownsSlot) {
+      this.ownsSlot = false;
+      this.manager.releaseViewerSlot(this.viewer);
+    }
+  }
+}