fix: refresh the runtime JWT on demand instead of waiting for the next GetProject poll

web-admin/src/features/edit-session/updateEditSessionDevJWT.spec.ts

@@ -0,0 +1,184 @@
+import { adminServiceGetDeployment } from "@rilldata/web-admin/client";
+import {
+  createMockUserJwtRefresher,
+  createUpdateEditSessionDevJWT,
+  editSessionRuntimeAuth,
+} from "@rilldata/web-admin/features/edit-session/updateEditSessionDevJWT";
+import {
+  selectedMockUserJWT,
+  selectedMockUserStore,
+} from "@rilldata/web-common/features/dashboards/granular-access-policies/stores";
+import { createQueryClient } from "@rilldata/web-common/lib/svelte-query/globalQueryClient";
+import { RUNTIME_ACCESS_TOKEN_DEFAULT_TTL } from "@rilldata/web-common/runtime-client/constants";
+import {
+  RuntimeClient,
+  type JwtRefresher,
+} from "@rilldata/web-common/runtime-client/v2";
+import { get } from "svelte/store";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+vi.mock("@rilldata/web-admin/client", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@rilldata/web-admin/client")>()),
+  adminServiceGetDeployment: vi.fn(),
+}));
+
+describe("edit session View As", () => {
+  const fetchMock = vi.fn<typeof fetch>();
+  const getDeployment = vi.mocked(adminServiceGetDeployment);
+  const mockUser = { email: "viewer@example.com" };
+
+  beforeEach(() => {
+    vi.useFakeTimers();
+    fetchMock.mockReset().mockImplementation(() =>
+      Promise.resolve(
+        new Response(new Uint8Array(0), {
+          status: 200,
+          headers: { "content-type": "application/proto" },
+        }),
+      ),
+    );
+    vi.stubGlobal("fetch", fetchMock);
+    getDeployment.mockReset();
+    selectedMockUserStore.set(null);
+    selectedMockUserJWT.set(null);
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+    vi.unstubAllGlobals();
+  });
+
+  function createClient() {
+    return new RuntimeClient({
+      host: "http://localhost:9009",
+      instanceId: "test-instance",
+      jwt: "edit-session-jwt",
+    });
+  }
+
+  /** Sends a request; ConnectRPC may reject the minimal mock response, which does not matter here. */
+  async function ping(client: RuntimeClient) {
+    try {
+      await client.runtimeService.ping({});
+    } catch {
+      // no-op
+    }
+  }
+
+  function sentAuthHeader(index: number) {
+    const headers = fetchMock.mock.calls[index]?.[1]?.headers as
+      | Headers
+      | undefined;
+    return headers?.get("Authorization") ?? null;
+  }
+
+  it("installs the mock user's JWT and restores the user's when View As ends", async () => {
+    getDeployment.mockResolvedValueOnce({ accessToken: "mock-token" });
+    const client = createClient();
+    const queryClient = createQueryClient();
+    const updateDevJWT = createUpdateEditSessionDevJWT(
+      "deployment-1",
+      "edit-session-jwt",
+    );
+
+    await updateDevJWT(queryClient, client, mockUser);
+    expect(client.getJwt()).toBe("mock-token");
+    expect(client.authContext).toBe("mock");
+    expect(get(selectedMockUserStore)).toBe(mockUser);
+    expect(get(selectedMockUserJWT)).toBe("mock-token");
+
+    await updateDevJWT(queryClient, client, null);
+    expect(client.getJwt()).toBe("edit-session-jwt");
+    expect(client.authContext).toBe("user");
+    expect(get(selectedMockUserStore)).toBeNull();
+    expect(get(selectedMockUserJWT)).toBeNull();
+  });
+
+  it("reissues the mock user's JWT with the same attributes", async () => {
+    getDeployment.mockResolvedValueOnce({ accessToken: "mock-token-2" });
+    const abort = new AbortController();
+
+    const refresh = createMockUserJwtRefresher("deployment-1", mockUser);
+    await expect(refresh(abort.signal)).resolves.toBe("mock-token-2");
+
+    expect(getDeployment).toHaveBeenCalledTimes(1);
+    const [deploymentId, body, signal] = getDeployment.mock.calls[0];
+    expect(deploymentId).toBe("deployment-1");
+    expect(body.attributes).toMatchObject({
+      email: "viewer@example.com",
+      domain: "example.com",
+    });
+    expect(body.accessTokenTtlSeconds).toBe(
+      RUNTIME_ACCESS_TOKEN_DEFAULT_TTL / 1000,
+    );
+    expect(signal).toBe(abort.signal);
+    expect(get(selectedMockUserJWT)).toBe("mock-token-2");
+  });
+
+  it("keeps the mock credentials when a user JWT refresh lands after View As started", async () => {
+    getDeployment
+      .mockResolvedValueOnce({ accessToken: "mock-token" })
+      .mockResolvedValueOnce({ accessToken: "mock-token-2" });
+    const client = createClient();
+    const queryClient = createQueryClient();
+    let resolveUserRefresh!: (jwt: string) => void;
+    const userRefreshJwt: JwtRefresher = vi.fn(
+      () =>
+        new Promise<string>((resolve) => {
+          resolveUserRefresh = resolve;
+        }),
+    );
+    const mockRefreshJwt = createMockUserJwtRefresher("deployment-1", mockUser);
+    // Plays the edit layout and RuntimeProvider: the client gets whatever the layout
+    // derives from the GetProject JWT and the View As selection
+    function forwardToProvider(userJwt: string) {
+      const auth = editSessionRuntimeAuth({
+        userJwt,
+        userRefreshJwt,
+        mockUser: get(selectedMockUserStore),
+        mockJwt: get(selectedMockUserJWT),
+        mockRefreshJwt,
+      });
+      client.updateJwt(auth.jwt, auth.authContext);
+      client.setJwtRefresher(auth.refreshJwt);
+    }
+    forwardToProvider("edit-session-jwt");
+
+    // The user's JWT expires and a request starts an on-demand refresh
+    vi.advanceTimersByTime(RUNTIME_ACCESS_TOKEN_DEFAULT_TTL);
+    const request = ping(client);
+    await vi.advanceTimersByTimeAsync(1);
+    expect(userRefreshJwt).toHaveBeenCalledTimes(1);
+
+    // View As starts while that refresh is still in flight
+    const updateDevJWT = createUpdateEditSessionDevJWT(
+      "deployment-1",
+      "edit-session-jwt",
+    );
+    await updateDevJWT(queryClient, client, mockUser);
+    forwardToProvider("edit-session-jwt");
+    await request;
+    expect(sentAuthHeader(0)).toBe("Bearer mock-token");
+
+    // The user refresh completes and the layout forwards the new user JWT; the mock credentials stay
+    resolveUserRefresh("edit-session-jwt-2");
+    await vi.advanceTimersByTimeAsync(1);
+    forwardToProvider("edit-session-jwt-2");
+    expect(client.getJwt()).toBe("mock-token");
+    expect(client.authContext).toBe("mock");
+
+    // When the mock token expires, it is the mock issuer that is asked, not the user's
+    vi.advanceTimersByTime(RUNTIME_ACCESS_TOKEN_DEFAULT_TTL);
+    await ping(client);
+    expect(getDeployment).toHaveBeenCalledTimes(2);
+    expect(userRefreshJwt).toHaveBeenCalledTimes(1);
+    expect(client.getJwt()).toBe("mock-token-2");
+    expect(sentAuthHeader(1)).toBe("Bearer mock-token-2");
+
+    // View As ends: the latest user JWT applies
+    await updateDevJWT(queryClient, client, null);
+    forwardToProvider("edit-session-jwt-2");
+    expect(client.getJwt()).toBe("edit-session-jwt-2");
+    expect(client.authContext).toBe("user");
+  });
+});

web-admin/src/features/edit-session/updateEditSessionDevJWT.ts

@@ -8,7 +8,11 @@ import {
   invalidateCanvasQueries,
   removeInactiveResourceQueries,
 } from "@rilldata/web-common/runtime-client/invalidation";
-import type { RuntimeClient } from "@rilldata/web-common/runtime-client/v2";
+import type {
+  AuthContext,
+  JwtRefresher,
+  RuntimeClient,
+} from "@rilldata/web-common/runtime-client/v2";
 import type { QueryClient } from "@tanstack/svelte-query";
 import { adminServiceGetDeployment } from "@rilldata/web-admin/client";
 import { getDomain } from "@rilldata/web-admin/features/projects/user-management/selectors.ts";
@@ -29,28 +33,14 @@ export function createUpdateEditSessionDevJWT(
     const prevMockUser = get(selectedMockUserStore);
     selectedMockUserStore.set(mockUser);
 
+    // The edit layout derives the provider's credentials from the stores set here (see editSessionRuntimeAuth);
+    // the client is updated directly as well so the invalidations below already run with the new token
     if (mockUser === null) {
       selectedMockUserJWT.set(null);
       client.updateJwt(editSessionJwt, "user");
     } else {
       try {
-        const { name, email, groups, admin, ...customAttributes } = mockUser;
-
-        const { accessToken } = await adminServiceGetDeployment(deploymentId, {
-          attributes: {
-            email,
-            ...(email ? { domain: getDomain(email) } : {}),
-            name: name || "Mock User",
-            admin: !!admin,
-            groups: groups || [],
-            ...customAttributes,
-          },
-          // Make sure to match TTL, GetDeployment defaults to 24-hours
-          accessTokenTtlSeconds: RUNTIME_ACCESS_TOKEN_DEFAULT_TTL / 1000, // Seconds param vs milliseconds constant
-        });
-
-        if (!accessToken) throw new Error("No JWT returned");
-
+        const accessToken = await issueMockUserJwt(deploymentId, mockUser);
         selectedMockUserJWT.set(accessToken);
         client.updateJwt(accessToken, "mock");
       } catch {
@@ -69,3 +59,78 @@ export function createUpdateEditSessionDevJWT(
     return invalidateCanvasQueries(queryClient, client.instanceId);
   };
 }
+
+/**
+ * Reissues the mock user's JWT when a request finds it expired, e.g. after a laptop sleep,
+ * so the preview does not wait for a token that never comes. The mock token has the same TTL as the user's.
+ */
+export function createMockUserJwtRefresher(
+  deploymentId: string,
+  mockUser: MockUser,
+): JwtRefresher {
+  return async (signal) => {
+    const accessToken = await issueMockUserJwt(deploymentId, mockUser, signal);
+    selectedMockUserJWT.set(accessToken);
+    return accessToken;
+  };
+}
+
+/**
+ * Picks the credentials the edit layout's RuntimeProvider runs with: the mock user's JWT and refresher while
+ * View As is active, the user's otherwise. The provider must follow the View As selection: if it kept forwarding
+ * the user's JWT, the next GetProject update (the 15-minute poll, a focus refetch, or an on-demand refresh that was
+ * still in flight when View As started) would push the user's token back onto the client with the "user" context,
+ * and the preview would show the user's data while View As stays selected.
+ */
+export function editSessionRuntimeAuth(opts: {
+  userJwt: string | undefined;
+  userRefreshJwt: JwtRefresher | undefined;
+  mockUser: MockUser | null;
+  mockJwt: string | null;
+  mockRefreshJwt: JwtRefresher | undefined;
+}): {
+  jwt: string | undefined;
+  authContext: AuthContext;
+  refreshJwt: JwtRefresher | undefined;
+} {
+  if (opts.mockUser && opts.mockJwt) {
+    return {
+      jwt: opts.mockJwt,
+      authContext: "mock",
+      refreshJwt: opts.mockRefreshJwt,
+    };
+  }
+  return {
+    jwt: opts.userJwt,
+    authContext: "user",
+    refreshJwt: opts.userRefreshJwt,
+  };
+}
+
+async function issueMockUserJwt(
+  deploymentId: string,
+  mockUser: MockUser,
+  signal?: AbortSignal,
+): Promise<string> {
+  const { name, email, groups, admin, ...customAttributes } = mockUser;
+
+  const { accessToken } = await adminServiceGetDeployment(
+    deploymentId,
+    {
+      attributes: {
+        email,
+        ...(email ? { domain: getDomain(email) } : {}),
+        name: name || "Mock User",
+        admin: !!admin,
+        groups: groups || [],
+        ...customAttributes,
+      },
+      // Make sure to match TTL, GetDeployment defaults to 24-hours
+      accessTokenTtlSeconds: RUNTIME_ACCESS_TOKEN_DEFAULT_TTL / 1000, // Seconds param vs milliseconds constant
+    },
+    signal,
+  );
+
+  if (!accessToken) throw new Error("No JWT returned");
+  return accessToken;
+}