fix(ui): preserve device token state across reconnects (#126362)

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

@@ -1,5 +1,5 @@
 /** @vitest-environment node */
-import { createHash } from "node:crypto";
+import { createHash, webcrypto } from "node:crypto";
 import {
   ConnectErrorDetailCodes,
   GATEWAY_CLIENT_CAPS,
@@ -94,6 +94,37 @@ function storeDeviceAuthToken(params: {
   return storeScopedDeviceAuthToken({ ...params, gatewayUrl: DEFAULT_GATEWAY_URL });
 }
 
+function storeDeviceIdentity(deviceId: string) {
+  localStorage.setItem(
+    "openclaw-device-identity-v1",
+    JSON.stringify({
+      version: 1,
+      deviceId,
+      publicKey: "AA",
+      privateKey: "AA",
+      createdAtMs: 1,
+    }),
+  );
+}
+
+function deferDeviceIdentityDigest() {
+  const digest = createDeferred<ArrayBuffer>();
+  const digestMock = vi.fn(() => digest.promise);
+  vi.stubGlobal("crypto", { subtle: { digest: digestMock } });
+  return { digest, digestMock };
+}
+
+function createDeviceTokenState(request: (method: string) => Promise<unknown>) {
+  const state = nodes.createInitialDevicesState({
+    client: {
+      request: request as <T = unknown>(method: string, params?: unknown) => Promise<T>,
+    },
+    connected: true,
+  });
+  state.requestGeneration = 1;
+  return state;
+}
+
 type HandlerMap = {
   close: MockWebSocketHandler[];
   error: MockWebSocketHandler[];
@@ -1705,6 +1736,81 @@ describe("GatewayBrowserClient", () => {
     });
   });
 
+  it("selects the replacement token after a successful self rotation retires the page epoch", async () => {
+    localStorage.clear();
+    storeDeviceIdentity("00");
+    loadOrCreateDeviceIdentityMock.mockResolvedValue({
+      deviceId: "00",
+      privateKey: "private-key", // pragma: allowlist secret
+      publicKey: "public-key", // pragma: allowlist secret
+    });
+    const { digest, digestMock } = deferDeviceIdentityDigest();
+    const state = createDeviceTokenState(async () => ({
+      deviceId: "00",
+      role: "operator",
+      token: "replacement-device-token",
+      scopes: ["operator.read"],
+      rotatedAtMs: 1_800_000_000_000,
+      tokenDelivery: "in-band",
+    }));
+
+    const operation = nodes.rotateDeviceToken(state, {
+      deviceId: "00",
+      gatewayUrl: DEFAULT_GATEWAY_URL,
+      role: "operator",
+    });
+    await vi.waitFor(() => expect(digestMock).toHaveBeenCalledOnce());
+    state.requestGeneration += 1;
+    digest.resolve(new Uint8Array([0]).buffer);
+    await expect(operation).resolves.toEqual({
+      delivery: "in-band",
+      token: "replacement-device-token",
+    });
+
+    vi.stubGlobal("crypto", webcrypto);
+    const nextClient = new GatewayBrowserClient({ url: DEFAULT_GATEWAY_URL });
+    const { connectFrame } = await startConnect(nextClient);
+    expect(connectFrame.params?.auth).toMatchObject({
+      token: "replacement-device-token",
+      deviceToken: "replacement-device-token",
+    });
+    nextClient.stop();
+  });
+
+  it("selects no revoked token after a successful self revocation retires the page epoch", async () => {
+    localStorage.clear();
+    storeDeviceIdentity("00");
+    storeDeviceAuthToken({
+      deviceId: "00",
+      role: "operator",
+      token: "revoked-device-token",
+      scopes: ["operator.read"],
+    });
+    loadOrCreateDeviceIdentityMock.mockResolvedValue({
+      deviceId: "00",
+      privateKey: "private-key", // pragma: allowlist secret
+      publicKey: "public-key", // pragma: allowlist secret
+    });
+    const { digest, digestMock } = deferDeviceIdentityDigest();
+    const state = createDeviceTokenState(async () => ({}));
+
+    const operation = nodes.revokeDeviceToken(state, {
+      deviceId: "00",
+      gatewayUrl: DEFAULT_GATEWAY_URL,
+      role: "operator",
+    });
+    await vi.waitFor(() => expect(digestMock).toHaveBeenCalledOnce());
+    state.requestGeneration += 1;
+    digest.resolve(new Uint8Array([0]).buffer);
+    await operation;
+
+    vi.stubGlobal("crypto", webcrypto);
+    const nextClient = new GatewayBrowserClient({ url: DEFAULT_GATEWAY_URL });
+    const { connectFrame } = await startConnect(nextClient);
+    expect(connectFrame.params?.auth).toBeUndefined();
+    nextClient.stop();
+  });
+
   it("uses a scoped device token when legacy cleanup fails", async () => {
     vi.spyOn(localStorage, "removeItem").mockImplementation(() => {
       throw new Error("storage cleanup blocked");

ui/src/lib/nodes/device-token.test.ts

@@ -5,7 +5,6 @@ import { createStorageMock } from "../../test-helpers/storage.ts";
 import {
   clearDeviceAuthToken,
   loadDeviceAuthToken,
-  revokeDeviceToken,
   rotateDeviceToken,
   storeDeviceAuthToken,
 } from "./index.ts";
@@ -91,34 +90,20 @@ afterEach(() => {
 describe("device token request lifecycle", () => {
   // A retired epoch is a reconnect, not a reason to destroy the credential: the previous
   // token is already dead on the server, so the caller still needs this one to recover.
-  it("returns a rotate response from a retired request epoch without persisting it", async () => {
+  it("persists a rotate response after its request epoch retires before success", async () => {
+    storeIdentity();
+    const { digest, digestMock } = deferIdentityFingerprint();
     const response = deferred<unknown>();
     const state = createState(() => response.promise);
 
     const operation = rotateDeviceToken(state, tokenParams);
     state.requestGeneration += 1;
     response.resolve({ token: "rotated-token", tokenDelivery: "in-band", ...rotationResult });
-
-    expect(await operation).toEqual({ delivery: "in-band", token: "rotated-token" });
-    expect(loadDeviceAuthToken(tokenParams)).toBeNull();
-  });
-
-  it("rechecks rotate ownership after loading the local identity", async () => {
-    storeIdentity();
-    const { digest, digestMock } = deferIdentityFingerprint();
-    const state = createState(async () => ({
-      token: "rotated-token",
-      tokenDelivery: "in-band",
-      ...rotationResult,
-    }));
-
-    const operation = rotateDeviceToken(state, tokenParams);
     await vi.waitFor(() => expect(digestMock).toHaveBeenCalledOnce());
-    state.requestGeneration += 1;
     digest.resolve(new Uint8Array([0]).buffer);
 
     expect(await operation).toEqual({ delivery: "in-band", token: "rotated-token" });
-    expect(loadDeviceAuthToken(tokenParams)).toBeNull();
+    expect(loadDeviceAuthToken(tokenParams)?.token).toBe("rotated-token");
   });
 
   it("reports a cross-device rotation the Gateway withheld the token for", async () => {
@@ -228,21 +213,6 @@ describe("device token request lifecycle", () => {
     expect(loadDeviceAuthToken(tokenParams)).toBeNull();
   });
 
-  it("does not clear a current token when a revoke request retires during identity loading", async () => {
-    storeIdentity();
-    storeDeviceAuthToken({ ...tokenParams, token: "current-token", scopes: ["operator.read"] });
-    const { digest, digestMock } = deferIdentityFingerprint();
-    const state = createState(async () => ({}));
-
-    const operation = revokeDeviceToken(state, tokenParams);
-    await vi.waitFor(() => expect(digestMock).toHaveBeenCalledOnce());
-    state.requestGeneration += 1;
-    digest.resolve(new Uint8Array([0]).buffer);
-    await operation;
-
-    expect(loadDeviceAuthToken(tokenParams)?.token).toBe("current-token");
-  });
-
   it("normalizes malformed persisted scopes without breaking token loading", () => {
     storeDeviceAuthToken({ ...tokenParams, token: "current-token", scopes: [] });
     const key = storedTokenKey();

ui/src/lib/nodes/index.ts

@@ -549,29 +549,23 @@ export async function rotateDeviceToken(
       tokenDelivery?: string;
     }>("device.token.rotate", requestParams);
     const outcome = classifyRotationOutcome(res, requestParams);
-    // A retired epoch stops every state write below, but never the return: the previous
-    // credential is already dead on the server, so discarding this response would leave
-    // the operator locked out with no way to ask for the replacement again.
-    if (!isCurrentNodesRequest(state, client, generation)) {
-      return outcome;
-    }
     if (outcome.delivery === "in-band") {
       const identity = await loadOrCreateDeviceIdentity();
-      if (!isCurrentNodesRequest(state, client, generation)) {
-        return outcome;
-      }
-      const role = res.role ?? params.role;
-      if (res.deviceId === identity.deviceId || params.deviceId === identity.deviceId) {
+      // RPC success retires the old bearer and may immediately reconnect the page.
+      // Commit the exact captured credential scope before fencing render projections.
+      if (res.deviceId === identity.deviceId || requestParams.deviceId === identity.deviceId) {
         storeDeviceAuthToken({
           deviceId: identity.deviceId,
           gatewayUrl,
-          role,
+          role: requestParams.role,
           token: outcome.token,
-          scopes: res.scopes ?? params.scopes ?? [],
+          scopes: res.scopes ?? requestParams.scopes ?? [],
         });
       }
     }
-    await loadDevices(state);
+    if (isCurrentNodesRequest(state, client, generation)) {
+      await loadDevices(state);
+    }
     return outcome;
   } catch (err) {
     if (isCurrentNodesRequest(state, client, generation)) {
@@ -593,18 +587,14 @@ export async function revokeDeviceToken(
   try {
     const { gatewayUrl, ...requestParams } = params;
     await client.request("device.token.revoke", requestParams);
-    if (!isCurrentNodesRequest(state, client, generation)) {
-      return;
-    }
     const identity = await loadOrCreateDeviceIdentity();
-    if (!isCurrentNodesRequest(state, client, generation)) {
-      return;
-    }
-    if (params.deviceId === identity.deviceId) {
+    // Clearing the successfully revoked credential belongs to this captured scope,
+    // not to the page generation invalidated by the resulting reconnect.
+    if (requestParams.deviceId === identity.deviceId) {
       clearDeviceAuthToken({
         deviceId: identity.deviceId,
         gatewayUrl,
-        role: params.role,
+        role: requestParams.role,
       });
     }
     if (isCurrentNodesRequest(state, client, generation)) {