test: share ui reconnect and storage helpers

ui/src/ui/app-gateway.node.test.ts

@@ -22,7 +22,9 @@ type GatewayClientMock = {
 
 const gatewayClientInstances: GatewayClientMock[] = [];
 
-vi.mock("./gateway.ts", () => {
+vi.mock("./gateway.ts", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("./gateway.ts")>();
+
   function resolveGatewayErrorDetailCode(
     error: { details?: unknown } | null | undefined,
   ): string | null {
@@ -81,7 +83,7 @@ vi.mock("./gateway.ts", () => {
     }
   }
 
-  return { GatewayBrowserClient, resolveGatewayErrorDetailCode };
+  return { ...actual, GatewayBrowserClient, resolveGatewayErrorDetailCode };
 });
 
 vi.mock("./controllers/chat.ts", async (importOriginal) => {
@@ -145,6 +147,33 @@ function createHost() {
   } as unknown as Parameters<typeof connectGateway>[0];
 }
 
+function connectHostGateway() {
+  const host = createHost();
+  connectGateway(host);
+  const client = gatewayClientInstances[0];
+  expect(client).toBeDefined();
+  return { host, client };
+}
+
+function emitToolResultEvent(client: GatewayClientMock) {
+  client.emitEvent({
+    event: "agent",
+    payload: {
+      runId: "engine-run-1",
+      seq: 1,
+      stream: "tool",
+      ts: 1,
+      sessionKey: "main",
+      data: {
+        toolCallId: "tool-1",
+        name: "fetch",
+        phase: "result",
+        result: { text: "ok" },
+      },
+    },
+  });
+}
+
 describe("connectGateway", () => {
   beforeEach(() => {
     gatewayClientInstances.length = 0;
@@ -457,55 +486,15 @@ describe("connectGateway", () => {
   });
 
   it("does not reload chat history for each live tool result event", () => {
-    const host = createHost();
+    const { client } = connectHostGateway();
-
+    emitToolResultEvent(client);
-    connectGateway(host);
-    const client = gatewayClientInstances[0];
-    expect(client).toBeDefined();
-
-    client.emitEvent({
-      event: "agent",
-      payload: {
-        runId: "engine-run-1",
-        seq: 1,
-        stream: "tool",
-        ts: 1,
-        sessionKey: "main",
-        data: {
-          toolCallId: "tool-1",
-          name: "fetch",
-          phase: "result",
-          result: { text: "ok" },
-        },
-      },
-    });
 
     expect(loadChatHistoryMock).not.toHaveBeenCalled();
   });
 
   it("reloads chat history once after the final chat event when tool output was used", () => {
-    const host = createHost();
+    const { client } = connectHostGateway();
-
+    emitToolResultEvent(client);
-    connectGateway(host);
-    const client = gatewayClientInstances[0];
-    expect(client).toBeDefined();
-
-    client.emitEvent({
-      event: "agent",
-      payload: {
-        runId: "engine-run-1",
-        seq: 1,
-        stream: "tool",
-        ts: 1,
-        sessionKey: "main",
-        data: {
-          toolCallId: "tool-1",
-          name: "fetch",
-          phase: "result",
-          result: { text: "ok" },
-        },
-      },
-    });
 
     client.emitEvent({
       event: "chat",

ui/src/ui/app-settings.test.ts

@@ -7,6 +7,7 @@ import {
   setTabFromRoute,
   syncThemeWithSettings,
 } from "./app-settings.ts";
+import { createStorageMock } from "./test-helpers/storage.ts";
 import type { ThemeMode, ThemeName } from "./theme.ts";
 
 type Tab =
@@ -66,30 +67,6 @@ type SettingsHost = {
   pendingGatewayToken?: string | null;
 };
 
-function createStorageMock(): Storage {
-  const store = new Map<string, string>();
-  return {
-    get length() {
-      return store.size;
-    },
-    clear() {
-      store.clear();
-    },
-    getItem(key: string) {
-      return store.get(key) ?? null;
-    },
-    key(index: number) {
-      return Array.from(store.keys())[index] ?? null;
-    },
-    removeItem(key: string) {
-      store.delete(key);
-    },
-    setItem(key: string, value: string) {
-      store.set(key, String(value));
-    },
-  };
-}
-
 function setTestWindowUrl(urlString: string) {
   const current = new URL(urlString);
   const history = {

ui/src/ui/gateway.node.test.ts

@@ -1,6 +1,7 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { loadDeviceAuthToken, storeDeviceAuthToken } from "./device-auth.ts";
 import type { DeviceIdentity } from "./device-identity.ts";
+import { createStorageMock } from "./test-helpers/storage.ts";
 
 const wsInstances = vi.hoisted((): MockWebSocket[] => []);
 const loadOrCreateDeviceIdentityMock = vi.hoisted(() =>
@@ -91,30 +92,6 @@ type ConnectFrame = {
   };
 };
 
-function createStorageMock(): Storage {
-  const store = new Map<string, string>();
-  return {
-    get length() {
-      return store.size;
-    },
-    clear() {
-      store.clear();
-    },
-    getItem(key: string) {
-      return store.get(key) ?? null;
-    },
-    key(index: number) {
-      return Array.from(store.keys())[index] ?? null;
-    },
-    removeItem(key: string) {
-      store.delete(key);
-    },
-    setItem(key: string, value: string) {
-      store.set(key, String(value));
-    },
-  };
-}
-
 function getLatestWebSocket(): MockWebSocket {
   const ws = wsInstances.at(-1);
   if (!ws) {
@@ -133,9 +110,7 @@ function parseLatestConnectFrame(ws: MockWebSocket): ConnectFrame {
   return JSON.parse(ws.sent.at(-1) ?? "{}") as ConnectFrame;
 }
 
-async function startConnect(client: InstanceType<typeof GatewayBrowserClient>, nonce = "nonce-1") {
+async function continueConnect(ws: MockWebSocket, nonce = "nonce-1") {
-  client.start();
-  const ws = getLatestWebSocket();
   ws.emitOpen();
   ws.emitMessage({
     type: "event",
@@ -146,6 +121,54 @@ async function startConnect(client: InstanceType<typeof GatewayBrowserClient>, n
   return { ws, connectFrame: parseLatestConnectFrame(ws) };
 }
 
+async function startConnect(client: InstanceType<typeof GatewayBrowserClient>, nonce = "nonce-1") {
+  client.start();
+  return await continueConnect(getLatestWebSocket(), nonce);
+}
+
+function emitRetryableTokenMismatch(ws: MockWebSocket, connectId: string | undefined) {
+  ws.emitMessage({
+    type: "res",
+    id: connectId,
+    ok: false,
+    error: {
+      code: "INVALID_REQUEST",
+      message: "unauthorized",
+      details: { code: "AUTH_TOKEN_MISMATCH", canRetryWithDeviceToken: true },
+    },
+  });
+}
+
+async function startRetriedDeviceTokenConnect(params: {
+  url: string;
+  token: string;
+  retryNonce?: string;
+}) {
+  const client = new GatewayBrowserClient({
+    url: params.url,
+    token: params.token,
+  });
+  const { ws: firstWs, connectFrame: firstConnect } = await startConnect(client);
+  expect(firstConnect.params?.auth?.token).toBe(params.token);
+  expect(firstConnect.params?.auth?.deviceToken).toBeUndefined();
+
+  emitRetryableTokenMismatch(firstWs, firstConnect.id);
+  await vi.waitFor(() => expect(firstWs.readyState).toBe(3));
+  firstWs.emitClose(4008, "connect failed");
+
+  await vi.advanceTimersByTimeAsync(800);
+  const secondWs = getLatestWebSocket();
+  expect(secondWs).not.toBe(firstWs);
+  const { connectFrame: secondConnect } = await continueConnect(
+    secondWs,
+    params.retryNonce ?? "nonce-2",
+  );
+  expect(secondConnect.params?.auth?.token).toBe(params.token);
+  expect(secondConnect.params?.auth?.deviceToken).toBe("stored-device-token");
+
+  return { client, firstWs, secondWs, firstConnect, secondConnect };
+}
+
 describe("GatewayBrowserClient", () => {
   beforeEach(() => {
     const storage = createStorageMock();
@@ -286,43 +309,12 @@ describe("GatewayBrowserClient", () => {
 
   it("retries once with device token after token mismatch when shared token is explicit", async () => {
     vi.useFakeTimers();
-    const client = new GatewayBrowserClient({
+    const { secondWs, secondConnect } = await startRetriedDeviceTokenConnect({
       url: "ws://127.0.0.1:18789",
       token: "shared-auth-token",
     });
 
-    const { ws: ws1, connectFrame: firstConnect } = await startConnect(client);
+    secondWs.emitMessage({
-    expect(firstConnect.params?.auth?.token).toBe("shared-auth-token");
-    expect(firstConnect.params?.auth?.deviceToken).toBeUndefined();
-
-    ws1.emitMessage({
-      type: "res",
-      id: firstConnect.id,
-      ok: false,
-      error: {
-        code: "INVALID_REQUEST",
-        message: "unauthorized",
-        details: { code: "AUTH_TOKEN_MISMATCH", canRetryWithDeviceToken: true },
-      },
-    });
-    await vi.waitFor(() => expect(ws1.readyState).toBe(3));
-    ws1.emitClose(4008, "connect failed");
-
-    await vi.advanceTimersByTimeAsync(800);
-    const ws2 = getLatestWebSocket();
-    expect(ws2).not.toBe(ws1);
-    ws2.emitOpen();
-    ws2.emitMessage({
-      type: "event",
-      event: "connect.challenge",
-      payload: { nonce: "nonce-2" },
-    });
-    await vi.waitFor(() => expect(ws2.sent.length).toBeGreaterThan(0));
-    const secondConnect = parseLatestConnectFrame(ws2);
-    expect(secondConnect.params?.auth?.token).toBe("shared-auth-token");
-    expect(secondConnect.params?.auth?.deviceToken).toBe("stored-device-token");
-
-    ws2.emitMessage({
       type: "res",
       id: secondConnect.id,
       ok: false,
@@ -332,8 +324,8 @@ describe("GatewayBrowserClient", () => {
         details: { code: "AUTH_TOKEN_MISMATCH" },
       },
     });
-    await vi.waitFor(() => expect(ws2.readyState).toBe(3));
+    await vi.waitFor(() => expect(secondWs.readyState).toBe(3));
-    ws2.emitClose(4008, "connect failed");
+    secondWs.emitClose(4008, "connect failed");
     expect(loadDeviceAuthToken({ deviceId: "device-1", role: "operator" })?.token).toBe(
       "stored-device-token",
     );
@@ -345,42 +337,11 @@ describe("GatewayBrowserClient", () => {
 
   it("treats IPv6 loopback as trusted for bounded device-token retry", async () => {
     vi.useFakeTimers();
-    const client = new GatewayBrowserClient({
+    const { client } = await startRetriedDeviceTokenConnect({
       url: "ws://[::1]:18789",
       token: "shared-auth-token",
     });
 
-    const { ws: ws1, connectFrame: firstConnect } = await startConnect(client);
-    expect(firstConnect.params?.auth?.token).toBe("shared-auth-token");
-    expect(firstConnect.params?.auth?.deviceToken).toBeUndefined();
-
-    ws1.emitMessage({
-      type: "res",
-      id: firstConnect.id,
-      ok: false,
-      error: {
-        code: "INVALID_REQUEST",
-        message: "unauthorized",
-        details: { code: "AUTH_TOKEN_MISMATCH", canRetryWithDeviceToken: true },
-      },
-    });
-    await vi.waitFor(() => expect(ws1.readyState).toBe(3));
-    ws1.emitClose(4008, "connect failed");
-
-    await vi.advanceTimersByTimeAsync(800);
-    const ws2 = getLatestWebSocket();
-    expect(ws2).not.toBe(ws1);
-    ws2.emitOpen();
-    ws2.emitMessage({
-      type: "event",
-      event: "connect.challenge",
-      payload: { nonce: "nonce-2" },
-    });
-    await vi.waitFor(() => expect(ws2.sent.length).toBeGreaterThan(0));
-    const secondConnect = parseLatestConnectFrame(ws2);
-    expect(secondConnect.params?.auth?.token).toBe("shared-auth-token");
-    expect(secondConnect.params?.auth?.deviceToken).toBe("stored-device-token");
-
     client.stop();
     vi.useRealTimers();
   });

ui/src/ui/test-helpers/storage.ts

@@ -0,0 +1,23 @@
+export function createStorageMock(): Storage {
+  const store = new Map<string, string>();
+  return {
+    get length() {
+      return store.size;
+    },
+    clear() {
+      store.clear();
+    },
+    getItem(key: string) {
+      return store.get(key) ?? null;
+    },
+    key(index: number) {
+      return Array.from(store.keys())[index] ?? null;
+    },
+    removeItem(key: string) {
+      store.delete(key);
+    },
+    setItem(key: string, value: string) {
+      store.set(key, String(value));
+    },
+  };
+}