@vellumai/assistant 0.12.1-staging.1 → 0.12.1

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

@@ -1,395 +0,0 @@
-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

@@ -1,214 +0,0 @@
-/**
- * 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);
-    }
-  }
-}