From d5acd7dee567745c19285eb0e0417329fe7a79b1 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Thu, 26 Mar 2026 16:41:51 +0000 Subject: [PATCH] test: share ui reconnect and storage helpers --- ui/src/ui/app-gateway.node.test.ts | 81 +++++++--------- ui/src/ui/app-settings.test.ts | 25 +---- ui/src/ui/gateway.node.test.ts | 149 +++++++++++------------------ ui/src/ui/test-helpers/storage.ts | 23 +++++ 4 files changed, 114 insertions(+), 164 deletions(-) create mode 100644 ui/src/ui/test-helpers/storage.ts diff --git a/ui/src/ui/app-gateway.node.test.ts b/ui/src/ui/app-gateway.node.test.ts index d830206444e..d0a878c714f 100644 --- a/ui/src/ui/app-gateway.node.test.ts +++ b/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(); + 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[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(); - - 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" }, - }, - }, - }); + const { client } = connectHostGateway(); + emitToolResultEvent(client); expect(loadChatHistoryMock).not.toHaveBeenCalled(); }); it("reloads chat history once after the final chat event when tool output was used", () => { - const host = createHost(); - - 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" }, - }, - }, - }); + const { client } = connectHostGateway(); + emitToolResultEvent(client); client.emitEvent({ event: "chat", diff --git a/ui/src/ui/app-settings.test.ts b/ui/src/ui/app-settings.test.ts index c119bca8630..089cf8c638c 100644 --- a/ui/src/ui/app-settings.test.ts +++ b/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(); - 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 = { diff --git a/ui/src/ui/gateway.node.test.ts b/ui/src/ui/gateway.node.test.ts index 4872d889173..0558b195885 100644 --- a/ui/src/ui/gateway.node.test.ts +++ b/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(); - 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, nonce = "nonce-1") { - client.start(); - const ws = getLatestWebSocket(); +async function continueConnect(ws: MockWebSocket, nonce = "nonce-1") { ws.emitOpen(); ws.emitMessage({ type: "event", @@ -146,6 +121,54 @@ async function startConnect(client: InstanceType, n return { ws, connectFrame: parseLatestConnectFrame(ws) }; } +async function startConnect(client: InstanceType, 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); - 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({ + secondWs.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)); - ws2.emitClose(4008, "connect failed"); + await vi.waitFor(() => expect(secondWs.readyState).toBe(3)); + 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(); }); diff --git a/ui/src/ui/test-helpers/storage.ts b/ui/src/ui/test-helpers/storage.ts new file mode 100644 index 00000000000..60682f74908 --- /dev/null +++ b/ui/src/ui/test-helpers/storage.ts @@ -0,0 +1,23 @@ +export function createStorageMock(): Storage { + const store = new Map(); + 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)); + }, + }; +}