From 33d8f6695e337b739e0951eda73edc9df92a48a8 Mon Sep 17 00:00:00 2001 From: Nishant Bangarwa Date: Thu, 5 Nov 2026 01:21:13 +0530 Subject: [PATCH 1/3] fix: refresh the runtime JWT on demand instead of waiting for the next GetProject poll A request that found its JWT expired waited passively for up to 60s for the 15-minute GetProject refetch and then failed with "Timed out waiting for a fresh JWT" whenever a laptop sleep, a network drop or a throttled tab had skipped that refetch. Expiry was also derived from when the client first saw the token string, so a token handed over from the query cache was credited with a full TTL and the runtime rejected it with "jwt is expired". - RuntimeClient takes a JWT refresher; a request that finds the token stale triggers it right away, shares one in-flight refresh, retries with backoff while the issuer is unreachable and fails fast on 401/403 - Expiry comes from the token's exp claim (capped at its lifetime from receipt), falling back to the default TTL for tokens without claims - A runtime "jwt is expired" rejection refreshes once and retries the request - RuntimeProvider refreshes a stale token when the tab becomes visible, gains focus or the network comes back - The project and edit layouts wire a refresher that refetches GetProject, the bearer-token GetProject on public URLs or the View As credentials - GetProject stops polling after a 401/403 instead of repeating it forever --- .../projects/project-query-options.spec.ts | 117 ++++++ .../projects/project-query-options.ts | 121 +++++- .../[organization]/[project]/+layout.svelte | 18 +- .../[project]/-/edit/+layout.svelte | 14 +- .../runtime-client/v2/RuntimeProvider.svelte | 25 +- web-common/src/runtime-client/v2/index.ts | 6 +- .../runtime-client/v2/runtime-client.spec.ts | 345 +++++++++++++++++- .../src/runtime-client/v2/runtime-client.ts | 228 +++++++++++- 8 files changed, 844 insertions(+), 30 deletions(-) create mode 100644 web-admin/src/features/projects/project-query-options.spec.ts diff --git a/web-admin/src/features/projects/project-query-options.spec.ts b/web-admin/src/features/projects/project-query-options.spec.ts new file mode 100644 index 000000000000..50446d13dfdb --- /dev/null +++ b/web-admin/src/features/projects/project-query-options.spec.ts @@ -0,0 +1,117 @@ +import { V1DeploymentStatus, type RpcStatus } from "@rilldata/web-admin/client"; +import { + baseGetProjectQueryOptions, + createRuntimeJwtRefresher, +} from "@rilldata/web-admin/features/projects/project-query-options"; +import { queryClient } from "@rilldata/web-common/lib/svelte-query/globalQueryClient"; +import { afterEach, describe, expect, it, vi } from "vitest"; + +const FifteenMinutes = 15 * 60 * 1000; +const runningDeployment = { + deployment: { status: V1DeploymentStatus.DEPLOYMENT_STATUS_RUNNING }, +}; + +function httpError(status: number): RpcStatus { + return { response: { status, data: { message: "nope" } } } as RpcStatus; +} + +function refetchInterval(state: { data?: unknown; error?: unknown }) { + const compute = baseGetProjectQueryOptions.refetchInterval; + if (typeof compute !== "function") { + throw new Error("expected refetchInterval to be a function"); + } + return compute({ state } as unknown as Parameters[0]); +} + +describe("baseGetProjectQueryOptions.refetchInterval", () => { + it("polls a running deployment at half the JWT TTL", () => { + expect(refetchInterval({ data: runningDeployment })).toBe(FifteenMinutes); + }); + + it("keeps polling after a transient failure so the JWT still gets refreshed", () => { + expect( + refetchInterval({ data: runningDeployment, error: httpError(503) }), + ).toBe(FifteenMinutes); + }); + + it("stops polling once the session or the project access is gone", () => { + for (const status of [401, 403]) { + expect( + refetchInterval({ data: runningDeployment, error: httpError(status) }), + ).toBe(false); + } + }); +}); + +describe("createRuntimeJwtRefresher", () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + + function mockFetchQuery(data: unknown) { + return vi + .spyOn(queryClient, "fetchQuery") + .mockImplementation(() => Promise.resolve(data) as never); + } + + it("refetches GetProject for the signed-in user and returns its JWT", async () => { + const fetchQuery = mockFetchQuery({ jwt: "fresh" }); + const refresh = createRuntimeJwtRefresher("user", { + organization: "org", + project: "proj", + branch: "dev", + }); + + await expect(refresh?.()).resolves.toBe("fresh"); + expect(fetchQuery).toHaveBeenCalledWith( + expect.objectContaining({ + queryKey: ["/v1/orgs/org/projects/proj", { branch: "dev" }], + }), + ); + }); + + it("refetches GetProject with the magic token on public URLs", async () => { + const fetchQuery = mockFetchQuery({ jwt: "fresh" }); + const refresh = createRuntimeJwtRefresher("magic", { + organization: "org", + project: "proj", + token: "magic", + }); + + await expect(refresh?.()).resolves.toBe("fresh"); + expect(fetchQuery).toHaveBeenCalledWith( + expect.objectContaining({ + queryKey: ["/v1/orgs/org/projects/proj", "token/magic"], + }), + ); + }); + + it("refetches the impersonated user's deployment credentials for View As", async () => { + const fetchQuery = mockFetchQuery({ accessToken: "fresh" }); + const refresh = createRuntimeJwtRefresher("mock", { + organization: "org", + project: "proj", + branch: "dev", + mockedUserId: "u1", + }); + + await expect(refresh?.()).resolves.toBe("fresh"); + expect(fetchQuery).toHaveBeenCalledWith( + expect.objectContaining({ + queryKey: [ + "/v1/orgs/org/projects/proj/credentials", + { userId: "u1", branch: "dev" }, + ], + }), + ); + }); + + it("has nothing to refetch for embeds", () => { + expect( + createRuntimeJwtRefresher("embed", { + organization: "org", + project: "proj", + }), + ).toBeUndefined(); + }); +}); diff --git a/web-admin/src/features/projects/project-query-options.ts b/web-admin/src/features/projects/project-query-options.ts index 012c09240c9d..07f8bcab9f3e 100644 --- a/web-admin/src/features/projects/project-query-options.ts +++ b/web-admin/src/features/projects/project-query-options.ts @@ -1,20 +1,44 @@ import { + adminServiceGetDeploymentCredentials, + adminServiceGetProject, + getAdminServiceGetDeploymentCredentialsQueryKey, + getAdminServiceGetProjectQueryKey, V1DeploymentStatus, type RpcStatus, + type V1GetDeploymentCredentialsResponse, type V1GetProjectResponse, } from "@rilldata/web-admin/client"; +import { + adminServiceGetProjectWithBearerToken, + getAdminServiceGetProjectWithBearerTokenQueryKey, +} from "@rilldata/web-admin/features/public-urls/get-project-with-bearer-token"; +import { extractErrorStatusCode } from "@rilldata/web-common/lib/errors"; +import { queryClient } from "@rilldata/web-common/lib/svelte-query/globalQueryClient"; import { RUNTIME_ACCESS_TOKEN_DEFAULT_TTL } from "@rilldata/web-common/runtime-client/constants"; -import type { CreateQueryOptions } from "@tanstack/svelte-query"; +import type { + AuthContext, + JwtRefresher, +} from "@rilldata/web-common/runtime-client/v2"; +import type { + CreateQueryOptions, + FetchQueryOptions, +} from "@tanstack/svelte-query"; const PollTimeWhenProjectDeploymentPending = 1000; const PollTimeWhenProjectDeploymentError = 5000; const PollTimeWhenProjectDeploymentOk = RUNTIME_ACCESS_TOKEN_DEFAULT_TTL / 2; // Proactively refetch the JWT before it expires +// A refetch that completed this recently already carries a fresh JWT; hand it over instead of fetching again. +const FreshJwtStaleTime = 5000; export const baseGetProjectQueryOptions: Partial< CreateQueryOptions > = { gcTime: Math.min(RUNTIME_ACCESS_TOKEN_DEFAULT_TTL, 1000 * 60 * 5), // Make sure we don't keep a stale JWT in the cache refetchInterval: (query) => { + // The session or the project access is gone: polling would repeat the same 401/403 until the tab is closed. + const errorStatus = extractErrorStatusCode(query.state.error); + if (errorStatus === 401 || errorStatus === 403) return false; + const status = query.state.data?.deployment?.status; switch (status) { case V1DeploymentStatus.DEPLOYMENT_STATUS_PENDING: @@ -38,3 +62,98 @@ export const baseGetProjectQueryOptions: Partial< // refetch. Focus/reconnect refetch still kicks in once data exceeds this window. staleTime: PollTimeWhenProjectDeploymentOk, }; + +/** + * Creates the callback `RuntimeProvider` uses to fetch a fresh JWT when a request finds the current one + * expired, which happens when a laptop sleep or a network drop skipped the scheduled refetch above. + * It re-runs the admin query that issued the JWT for the active auth mode, keyed exactly like the + * layout's own query so a refetch here and a scheduled one dedupe. + */ +export function createRuntimeJwtRefresher( + authContext: AuthContext, + opts: { + organization: string; + project: string; + branch?: string; + /** Magic auth token of the public URL being viewed. */ + token?: string; + /** User impersonated with View As. */ + mockedUserId?: string; + }, +): JwtRefresher | undefined { + const { organization, project, branch, token, mockedUserId } = opts; + switch (authContext) { + case "mock": { + if (!mockedUserId) return undefined; + const body = { userId: mockedUserId, ...(branch ? { branch } : {}) }; + return () => + fetchFreshJwt( + { + queryKey: getAdminServiceGetDeploymentCredentialsQueryKey( + organization, + project, + body, + ), + queryFn: ({ signal }) => + adminServiceGetDeploymentCredentials( + organization, + project, + body, + signal, + ), + }, + (data) => data.accessToken, + ); + } + case "magic": { + if (!token) return undefined; + return () => + fetchFreshJwt( + { + queryKey: getAdminServiceGetProjectWithBearerTokenQueryKey( + organization, + project, + token, + ), + queryFn: ({ signal }) => + adminServiceGetProjectWithBearerToken( + organization, + project, + token, + undefined, + signal, + ), + }, + (data) => data.jwt, + ); + } + case "user": { + const params = branch ? { branch } : undefined; + return () => + fetchFreshJwt( + { + queryKey: getAdminServiceGetProjectQueryKey( + organization, + project, + params, + ), + queryFn: ({ signal }) => + adminServiceGetProject(organization, project, params, signal), + }, + (data) => data.jwt, + ); + } + case "embed": + // Embed tokens come from the host page and last 24 hours; there is nothing to refetch. + return undefined; + } +} + +function fetchFreshJwt( + options: FetchQueryOptions, + selectJwt: (data: TData) => string | undefined, +): Promise { + return queryClient + .fetchQuery({ ...options, staleTime: FreshJwtStaleTime }) + .then(selectJwt); +} diff --git a/web-admin/src/routes/[organization]/[project]/+layout.svelte b/web-admin/src/routes/[organization]/[project]/+layout.svelte index d9e4c2cd3eff..854cd48f43c5 100644 --- a/web-admin/src/routes/[organization]/[project]/+layout.svelte +++ b/web-admin/src/routes/[organization]/[project]/+layout.svelte @@ -44,7 +44,10 @@ import ProjectBuilding from "@rilldata/web-admin/features/projects/ProjectBuilding.svelte"; import ProjectHeader from "../../../features/projects/header/ProjectHeader.svelte"; import ProjectTabs from "@rilldata/web-admin/features/projects/ProjectTabs.svelte"; - import { baseGetProjectQueryOptions } from "@rilldata/web-admin/features/projects/project-query-options"; + import { + baseGetProjectQueryOptions, + createRuntimeJwtRefresher, + } from "@rilldata/web-admin/features/projects/project-query-options"; import { resolveRuntimeConnection } from "@rilldata/web-admin/features/projects/project-runtime"; import RedeployProjectCta from "@rilldata/web-admin/features/projects/RedeployProjectCTA.svelte"; import SlimProjectHeader from "@rilldata/web-admin/features/projects/SlimProjectHeader.svelte"; @@ -228,6 +231,18 @@ resolveRuntimeConnection(projectData, mockUser, onPublicURLPage), ); + // Re-requests the runtime JWT from the query that issued it when a request finds it expired, e.g. after a + // laptop sleep or a network drop skipped the scheduled refetch, instead of waiting for the next poll. + let refreshJwt = $derived( + createRuntimeJwtRefresher(runtime.authContext, { + organization, + project, + branch: activeBranch, + token, + mockedUserId, + }), + ); + // --- Side effects (cache invalidation and telemetry) --- // Track previous status so we invalidate only on *transitions*, not on every @@ -349,6 +364,7 @@ instanceId={runtime.instanceId} jwt={runtime.jwt} authContext={runtime.authContext} + {refreshJwt} > {#if !onWelcomePage} {:else if isReady && deployment?.id && instanceId && runtimeHost && jwt} {#key `${runtimeHost}::${instanceId}::${hasPrimaryDeployment}`} - + {#if !inProjectWelcomePage} import { useQueryClient } from "@tanstack/svelte-query"; - import { onDestroy, setContext } from "svelte"; + import { onDestroy, onMount, setContext } from "svelte"; import { featureFlags } from "../../features/feature-flags"; import { invalidateRuntimeQueries } from "../invalidation"; import { @@ -8,7 +8,7 @@ evictRuntimeClient, RUNTIME_CONTEXT_KEY, } from "./context"; - import type { AuthContext } from "./runtime-client"; + import type { AuthContext, JwtRefresher } from "./runtime-client"; import { createEnvFileStore } from "@rilldata/web-common/features/env-management/env-file-store.ts"; const queryClient = useQueryClient(); @@ -17,6 +17,9 @@ export let instanceId: string; export let jwt: string | undefined = undefined; export let authContext: AuthContext = "user"; + // Fetches a fresh JWT when a request finds the current one expired, e.g. after a laptop sleep + // skipped the scheduled refetch. Without it, requests wait for the parent to pass a new `jwt`. + export let refreshJwt: JwtRefresher | undefined = undefined; // Returns a cached instance if a load function already created one for this host+instanceId. // If host/instanceId change, the parent's {#key} re-mounts us. @@ -29,10 +32,28 @@ // Handle JWT-only changes (15-min refresh, View As with same host) $: { const authContextChanged = client.updateJwt(jwt, authContext); + client.setJwtRefresher(refreshJwt); if (authContextChanged) void invalidateRuntimeQueries(queryClient, instanceId); } + // A tab that was asleep, hidden or offline has likely missed the scheduled JWT refetch. + // Refresh as soon as it is back so the first requests don't have to wait for it. + onMount(() => { + const refreshIfStale = () => void client.refreshJwtIfStale(); + const refreshWhenVisible = () => { + if (document.visibilityState !== "hidden") refreshIfStale(); + }; + window.addEventListener("online", refreshIfStale); + window.addEventListener("focus", refreshWhenVisible); + document.addEventListener("visibilitychange", refreshWhenVisible); + return () => { + window.removeEventListener("online", refreshIfStale); + window.removeEventListener("focus", refreshWhenVisible); + document.removeEventListener("visibilitychange", refreshWhenVisible); + }; + }); + onDestroy(() => { featureFlags.clearRuntimeClient(); evictRuntimeClient(client); diff --git a/web-common/src/runtime-client/v2/index.ts b/web-common/src/runtime-client/v2/index.ts index 921cb40acbe9..e6c05b42c678 100644 --- a/web-common/src/runtime-client/v2/index.ts +++ b/web-common/src/runtime-client/v2/index.ts @@ -1,4 +1,8 @@ -export { RuntimeClient, type AuthContext } from "./runtime-client"; +export { + RuntimeClient, + type AuthContext, + type JwtRefresher, +} from "./runtime-client"; export { getRuntimeClient, evictRuntimeClient, diff --git a/web-common/src/runtime-client/v2/runtime-client.spec.ts b/web-common/src/runtime-client/v2/runtime-client.spec.ts index 682069036d8a..00b1a61c9cd3 100644 --- a/web-common/src/runtime-client/v2/runtime-client.spec.ts +++ b/web-common/src/runtime-client/v2/runtime-client.spec.ts @@ -25,12 +25,41 @@ function createFetchMock() { ); } +/** Extract the Authorization header from the given fetch call. */ +function getAuthHeader( + fetchMock: ReturnType, + index: number, +): string | null { + const call = fetchMock.mock.calls[index] as + | [RequestInfo, RequestInit | undefined] + | undefined; + const headers = call?.[1]?.headers as Headers | undefined; + return headers?.get?.("Authorization") ?? null; +} + /** Extract the Authorization header from the most recent fetch call. */ function getLastAuthHeader(fetchMock: ReturnType): string | null { - if (fetchMock.mock.calls.length === 0) return null; - const lastCall = fetchMock.mock.calls[fetchMock.mock.calls.length - 1]; - const headers: Headers | undefined = lastCall[1]?.headers; - return headers?.get?.("Authorization") ?? null; + return getAuthHeader(fetchMock, fetchMock.mock.calls.length - 1); +} + +/** Builds an unsigned JWT whose payload carries the given claims (times in seconds since the epoch). */ +function makeJwt(claims: Record): string { + const payload = btoa(JSON.stringify(claims)) + .replace(/\+/g, "-") + .replace(/\//g, "_") + .replace(/=+$/, ""); + return `eyJhbGciOiJub25lIn0.${payload}.signature`; +} + +/** A unary Connect error response as the runtime returns it for an expired token. */ +function expiredJwtResponse() { + return new Response( + JSON.stringify({ + code: "unauthenticated", + message: "rpc error: code = Unauthenticated desc = jwt is expired", + }), + { status: 401, headers: { "content-type": "application/json" } }, + ); } describe("RuntimeClient JWT refresh", () => { @@ -319,4 +348,312 @@ describe("RuntimeClient JWT refresh", () => { expect(client.authContext).toBe("user"); }); }); + + // ── Active refresh through the registered refresher ──────────────── + + describe("refresher", () => { + it("asks the refresher for a new JWT instead of waiting for updateJwt", async () => { + const client = createClient({ jwt: "expiring-token" }); + const refreshJwt = vi.fn().mockResolvedValue("fresh-token"); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + await makeRequest(client); + + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(getLastAuthHeader(fetchMock)).toBe("Bearer fresh-token"); + }); + + it("shares one in-flight refresh between concurrent requests", async () => { + const client = createClient({ jwt: "expiring-token" }); + let resolveRefresh!: (jwt: string) => void; + const refreshJwt = vi.fn( + () => + new Promise((resolve) => { + resolveRefresh = resolve; + }), + ); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + const requests = Promise.all([ + makeRequest(client), + makeRequest(client), + makeRequest(client), + ]); + await vi.advanceTimersByTimeAsync(1); + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(fetchMock).not.toHaveBeenCalled(); + + resolveRefresh("fresh-token"); + await requests; + + expect(fetchMock).toHaveBeenCalledTimes(3); + for (let i = 0; i < 3; i++) { + expect(getAuthHeader(fetchMock, i)).toBe("Bearer fresh-token"); + } + }); + + it("retries a failed refresh with backoff until it succeeds", async () => { + const client = createClient({ jwt: "expiring-token" }); + const refreshJwt = vi + .fn() + .mockRejectedValueOnce(new Error("Network Error")) + .mockResolvedValueOnce("fresh-token"); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + let resolved = false; + const promise = makeRequest(client).then(() => { + resolved = true; + }); + + // The first attempt fails; the retry waits for the 1s backoff + await vi.advanceTimersByTimeAsync(500); + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(resolved).toBe(false); + + await vi.advanceTimersByTimeAsync(1000); + await promise; + + expect(refreshJwt).toHaveBeenCalledTimes(2); + expect(getLastAuthHeader(fetchMock)).toBe("Bearer fresh-token"); + }); + + it("fails at once when the issuer refuses because the session or access is gone", async () => { + const client = createClient({ jwt: "expiring-token" }); + const refreshJwt = vi.fn().mockRejectedValue({ + response: { + status: 403, + data: { message: "does not have permission to read project" }, + }, + }); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + const error = await client.runtimeService.ping({}).catch((e: Error) => e); + + expect(error).toBeInstanceOf(Error); + expect((error as Error).message).toContain( + "Failed to refresh the runtime JWT: does not have permission to read project", + ); + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("gives up after the timeout and reports why the refresh kept failing", async () => { + const client = createClient({ jwt: "expiring-token" }); + const refreshJwt = vi.fn().mockRejectedValue(new Error("Network Error")); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + const promise = client.runtimeService.ping({}).catch((e: Error) => e); + // 60s deadline plus the longest backoff + await vi.advanceTimersByTimeAsync(75_000); + const error = await promise; + + expect((error as Error).message).toContain( + "Timed out waiting for a fresh JWT: Network Error", + ); + // Several attempts, but spaced out by the backoff rather than a tight loop + expect(refreshJwt.mock.calls.length).toBeGreaterThan(1); + expect(refreshJwt.mock.calls.length).toBeLessThan(20); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("backs off when the issuer hands back an already expired JWT", async () => { + const client = createClient({ jwt: "expiring-token" }); + const now = Date.now() / 1000; + const refreshJwt = vi + .fn() + .mockResolvedValue(makeJwt({ iat: now - 3600, exp: now - 60 })); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + const promise = client.runtimeService.ping({}).catch((e: Error) => e); + await vi.advanceTimersByTimeAsync(75_000); + const error = await promise; + + expect((error as Error).message).toContain( + "Timed out waiting for a fresh JWT: the issuer returned an expired JWT", + ); + expect(refreshJwt.mock.calls.length).toBeLessThan(20); + }); + + it("ignores a refresher registered for another auth context", async () => { + // View As in Rill Developer sets a mock JWT directly on the client; the provider's + // refresher for the "user" context must not replace it with the user's token + const client = createClient({ jwt: "user-token", authContext: "user" }); + const refreshJwt = vi.fn().mockResolvedValue("fresh-user-token"); + client.setJwtRefresher(refreshJwt); + client.updateJwt("mock-token", "mock"); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + const promise = client.runtimeService.ping({}).catch((e: Error) => e); + await vi.advanceTimersByTimeAsync(61_000); + const error = await promise; + + expect(refreshJwt).not.toHaveBeenCalled(); + expect((error as Error).message).toContain( + "Timed out waiting for a fresh JWT", + ); + + // Back in the provider's context, the refresher applies again + client.updateJwt("user-token", "user"); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + await makeRequest(client); + + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(getLastAuthHeader(fetchMock)).toBe("Bearer fresh-user-token"); + }); + + it("refreshes once and retries when the runtime rejects the JWT as expired", async () => { + // The client's clock can disagree with the runtime's; the runtime wins + const client = createClient({ jwt: "token-runtime-rejects" }); + const refreshJwt = vi.fn().mockResolvedValue("fresh-token"); + client.setJwtRefresher(refreshJwt); + // The retry reuses the request's Headers object, so capture the header as each call is made + const sentAuthHeaders: (string | null)[] = []; + fetchMock.mockImplementation((_, init) => { + sentAuthHeaders.push( + (init?.headers as Headers | undefined)?.get("Authorization") ?? null, + ); + return Promise.resolve( + sentAuthHeaders.length === 1 + ? expiredJwtResponse() + : new Response(new Uint8Array(0), { + status: 200, + headers: { "content-type": "application/proto" }, + }), + ); + }); + + await makeRequest(client); + + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(sentAuthHeaders).toEqual([ + "Bearer token-runtime-rejects", + "Bearer fresh-token", + ]); + }); + + it("does not retry on other authentication errors", async () => { + const client = createClient({ jwt: "bad-token" }); + const refreshJwt = vi.fn().mockResolvedValue("fresh-token"); + client.setJwtRefresher(refreshJwt); + fetchMock.mockImplementationOnce(() => + Promise.resolve( + new Response( + JSON.stringify({ code: "unauthenticated", message: "bad token" }), + { status: 401, headers: { "content-type": "application/json" } }, + ), + ), + ); + + const error = await client.runtimeService.ping({}).catch((e: Error) => e); + + expect(error).toBeInstanceOf(Error); + expect(refreshJwt).not.toHaveBeenCalled(); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + + it("does not retry without a refresher, e.g. for embeds", async () => { + const client = createClient({ jwt: "embed-token", authContext: "embed" }); + fetchMock.mockImplementationOnce(() => + Promise.resolve(expiredJwtResponse()), + ); + + const error = await client.runtimeService.ping({}).catch((e: Error) => e); + + expect(error).toBeInstanceOf(Error); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + }); + + // ── Proactive refresh (tab visible again, network back) ──────────── + + describe("refreshJwtIfStale", () => { + it("refreshes a stale JWT and leaves a fresh one alone", async () => { + const client = createClient({ jwt: "token" }); + const refreshJwt = vi.fn().mockResolvedValue("fresh-token"); + client.setJwtRefresher(refreshJwt); + + await client.refreshJwtIfStale(); + expect(refreshJwt).not.toHaveBeenCalled(); + + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + await client.refreshJwtIfStale(); + + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(client.getJwt()).toBe("fresh-token"); + }); + + it("leaves refresh failures to the next request", async () => { + const client = createClient({ jwt: "token" }); + const refreshJwt = vi.fn().mockRejectedValue(new Error("Network Error")); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + await expect(client.refreshJwtIfStale()).resolves.toBeUndefined(); + expect(client.getJwt()).toBe("token"); + }); + }); + + // ── Expiry derived from the token's claims ───────────────────────── + + describe("expiry from the token's claims", () => { + it("treats a token handed over from a cache as expired once its exp claim has passed", async () => { + // The GetProject cache can hand a new RuntimeProvider a token issued long ago. Crediting it + // with a full TTL from receipt made the runtime reject requests with "jwt is expired". + const now = Date.now() / 1000; + const client = createClient({ + jwt: makeJwt({ iat: now - 40 * 60, exp: now - 10 * 60 }), + }); + const refreshJwt = vi.fn().mockResolvedValue("fresh-token"); + client.setJwtRefresher(refreshJwt); + + await makeRequest(client); + + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(getLastAuthHeader(fetchMock)).toBe("Bearer fresh-token"); + }); + + it("trusts an exp claim beyond the default TTL", async () => { + const now = Date.now() / 1000; + const oneHourToken = makeJwt({ iat: now, exp: now + 60 * 60 }); + const client = createClient({ jwt: oneHourToken }); + vi.advanceTimersByTime(RUNTIME_ACCESS_TOKEN_DEFAULT_TTL + 60_000); + + await makeRequest(client); + + expect(fetchMock).toHaveBeenCalledTimes(1); + expect(getLastAuthHeader(fetchMock)).toBe(`Bearer ${oneHourToken}`); + }); + + it("caps a fresh token at its lifetime from receipt when the client clock runs behind the issuer", async () => { + // The issuer's clock is 5 minutes ahead: exp looks 35 minutes away, but the token only lives 30 + const now = Date.now() / 1000; + const skew = 5 * 60; + const client = createClient({ + jwt: makeJwt({ iat: now + skew, exp: now + skew + 30 * 60 }), + }); + const refreshJwt = vi.fn().mockResolvedValue("fresh-token"); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + await makeRequest(client); + + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(getLastAuthHeader(fetchMock)).toBe("Bearer fresh-token"); + }); + + it("falls back to the default TTL for tokens without claims", async () => { + const client = createClient({ jwt: makeJwt({ sub: "dev-user" }) }); + vi.advanceTimersByTime(RUNTIME_ACCESS_TOKEN_DEFAULT_TTL / 2); + + await makeRequest(client); + + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + }); }); diff --git a/web-common/src/runtime-client/v2/runtime-client.ts b/web-common/src/runtime-client/v2/runtime-client.ts index e46e4f7e3212..56408003a321 100644 --- a/web-common/src/runtime-client/v2/runtime-client.ts +++ b/web-common/src/runtime-client/v2/runtime-client.ts @@ -1,8 +1,15 @@ -import { createClient, type Client, type Transport } from "@connectrpc/connect"; +import { + Code, + ConnectError, + createClient, + type Client, + type Transport, +} from "@connectrpc/connect"; import { createConnectTransport } from "@connectrpc/connect-web"; import { QueryService } from "../../proto/gen/rill/runtime/v1/queries_connect"; import { RuntimeService } from "../../proto/gen/rill/runtime/v1/api_connect"; import { ConnectorService } from "../../proto/gen/rill/runtime/v1/connectors_connect"; +import { extractErrorMessage, extractErrorStatusCode } from "../../lib/errors"; import { RUNTIME_ACCESS_TOKEN_DEFAULT_TTL, JWT_EXPIRY_WARNING_WINDOW, @@ -12,6 +19,18 @@ import { RequestQueue, createQueueInterceptor } from "./request-queue"; export type AuthContext = "user" | "mock" | "magic" | "embed"; +/** + * Fetches a fresh JWT from whoever issued the current one, e.g. by refetching GetProject in Rill Cloud. + * Resolves with the new token; rejects when the issuer could not be reached or refused. + */ +export type JwtRefresher = () => Promise; + +// How long a request waits for a fresh JWT before it fails. +const JWT_REFRESH_TIMEOUT = 60_000; +// Delay before retrying a failed refresh: doubles from the base up to the max while the issuer keeps failing. +const JWT_REFRESH_BACKOFF_BASE = 1_000; +const JWT_REFRESH_BACKOFF_MAX = 10_000; + export class RuntimeClient { readonly host: string; readonly instanceId: string; @@ -20,8 +39,15 @@ export class RuntimeClient { // JWT state (mutable; read by the transport interceptor) private currentJwt: string | undefined; - private jwtReceivedAt: number; + private jwtExpiresAt = 0; private _authContext: AuthContext; + // Refresher registered by RuntimeProvider, bound to the auth context it was registered under + private refreshJwt: JwtRefresher | undefined; + private refreshJwtAuthContext: AuthContext | undefined; + // In-flight refresh shared by every request waiting on it, and the backoff after failures + private jwtRefresh: Promise | undefined; + private jwtRefreshBackoff = 0; + private jwtRefreshNotBefore = 0; private disposed = false; // Cached service clients (created once per RuntimeClient) @@ -50,7 +76,7 @@ export class RuntimeClient { this.host = opts.host; this.instanceId = opts.instanceId; this.currentJwt = opts.jwt; - this.jwtReceivedAt = opts.jwt ? Date.now() : 0; + this.jwtExpiresAt = opts.jwt ? jwtExpiry(opts.jwt, Date.now()) : 0; this._authContext = opts.authContext ?? "user"; this.requestQueue = new RequestQueue(); @@ -61,11 +87,32 @@ export class RuntimeClient { createQueueInterceptor(this.requestQueue), // JWT interceptor adds auth header just before the network call (next) => async (req) => { - if (this.currentJwt) { + await this.waitForFreshJwt(); + const jwt = this.currentJwt; + if (jwt) req.header.set("Authorization", `Bearer ${jwt}`); + try { + return await next(req); + } catch (err) { + // The runtime's clock is the source of truth: when it rejects a token this client still + // considered valid, refresh once and retry. The runtime rejects expired tokens before the + // handler runs, so the retry is safe; only unary requests can be re-sent as is. + if ( + !jwt || + req.stream || + !this.canRefreshJwt() || + !isExpiredJwtError(err) + ) { + throw err; + } + if (jwt === this.currentJwt) this.jwtExpiresAt = 0; await this.waitForFreshJwt(); - req.header.set("Authorization", `Bearer ${this.currentJwt}`); + if (this.currentJwt) { + req.header.set("Authorization", `Bearer ${this.currentJwt}`); + } else { + req.header.delete("Authorization"); + } + return next(req); } - return next(req); }, ], }); @@ -81,13 +128,38 @@ export class RuntimeClient { if (jwt !== this.currentJwt) { this.currentJwt = jwt; - this.jwtReceivedAt = Date.now(); + this.jwtExpiresAt = jwt ? jwtExpiry(jwt, Date.now()) : 0; } if (authContext) this._authContext = authContext; return authContextChanged; } + /** + * Registers the callback that fetches a fresh JWT when the current one is about to expire. + * It stays bound to the auth context it was registered under: a mock JWT that View As sets directly + * on the client must not be replaced by the provider's token, so the refresher is ignored until the + * client is back in that context. + */ + setJwtRefresher(refreshJwt: JwtRefresher | undefined): void { + this.refreshJwt = refreshJwt; + this.refreshJwtAuthContext = this._authContext; + } + + /** + * Refreshes the JWT if it is expired or about to expire, e.g. when the tab becomes visible again + * after a laptop sleep skipped the scheduled refresh. Failures are left to the next request, + * which retries with backoff and surfaces the error. + */ + async refreshJwtIfStale(): Promise { + if (this.isJwtFresh() || !this.canRefreshJwt()) return; + try { + await this.runJwtRefresh(); + } catch { + // Reported by the next request that needs the token + } + } + /** Returns the current JWT (used by SSE clients and other non-query consumers). */ getJwt(): string | undefined { return this.currentJwt; @@ -117,28 +189,97 @@ export class RuntimeClient { } /** - * If the JWT is close to expiring, wait in a polling loop for - * RuntimeProvider to call updateJwt() with a fresh token. - * Ported from http-client.ts:maybeWaitForFreshJWT. + * Blocks until the JWT is fresh. With a refresher registered it asks the issuer for a new token + * right away, retrying with backoff while the issuer is unreachable; without one it waits for an + * updateJwt() call. Gives up after JWT_REFRESH_TIMEOUT, or at once when the issuer refuses because + * the session or the project access is gone. */ private async waitForFreshJwt(): Promise { - // Embeds have a 24h backend-issued TTL; skip client-side expiry check + // Embeds have a 24h backend-issued TTL and no refresher; skip client-side expiry check if (this._authContext === "embed") return; - const deadline = Date.now() + 60_000; - let expiresAt = this.jwtReceivedAt + RUNTIME_ACCESS_TOKEN_DEFAULT_TTL; - while (Date.now() + JWT_EXPIRY_WARNING_WINDOW > expiresAt) { + const deadline = Date.now() + JWT_REFRESH_TIMEOUT; + let lastError: unknown; + while (!this.isJwtFresh()) { if (this.disposed) { throw new DOMException("Request cancelled", "AbortError"); } if (Date.now() > deadline) { - throw new Error("Timed out waiting for a fresh JWT"); + throw new Error( + "Timed out waiting for a fresh JWT" + + (lastError ? `: ${extractErrorMessage(lastError)}` : ""), + ); + } + if (!this.canRefreshJwt()) { + // Wait for whoever owns the token (e.g. View As) to call updateJwt() + await sleep(CHECK_RUNTIME_STORE_FOR_JWT_INTERVAL); + continue; + } + try { + await this.runJwtRefresh(); + } catch (err) { + const status = extractErrorStatusCode(err); + if (status === 401 || status === 403) { + throw new Error( + `Failed to refresh the runtime JWT: ${extractErrorMessage(err)}`, + ); + } + lastError = err; + } + } + } + + private isJwtFresh(): boolean { + return ( + !this.currentJwt || + Date.now() + JWT_EXPIRY_WARNING_WINDOW <= this.jwtExpiresAt + ); + } + + private canRefreshJwt(): boolean { + return ( + !!this.refreshJwt && this.refreshJwtAuthContext === this._authContext + ); + } + + /** Runs one refresh, shared by every caller that joins while it is in flight. */ + private runJwtRefresh(): Promise { + this.jwtRefresh ??= this.refreshJwtOnce().finally(() => { + this.jwtRefresh = undefined; + }); + return this.jwtRefresh; + } + + private async refreshJwtOnce(): Promise { + const delay = this.jwtRefreshNotBefore - Date.now(); + if (delay > 0) await sleep(delay); + const refreshJwt = this.refreshJwt; + const authContext = this._authContext; + if (this.disposed || !refreshJwt) return; + try { + const jwt = await refreshJwt(); + // The provider may have switched auth context (View As) meanwhile; its token wins + if ( + this.disposed || + this.refreshJwt !== refreshJwt || + this._authContext !== authContext + ) { + return; } - await new Promise((r) => - setTimeout(r, CHECK_RUNTIME_STORE_FOR_JWT_INTERVAL), + if (!jwt) throw new Error("the issuer returned no JWT"); + this.updateJwt(jwt); + if (!this.isJwtFresh()) { + throw new Error("the issuer returned an expired JWT"); + } + this.jwtRefreshBackoff = 0; + this.jwtRefreshNotBefore = 0; + } catch (err) { + this.jwtRefreshBackoff = Math.min( + Math.max(this.jwtRefreshBackoff * 2, JWT_REFRESH_BACKOFF_BASE), + JWT_REFRESH_BACKOFF_MAX, ); - // Re-check: provider may have called updateJwt() while we waited - expiresAt = this.jwtReceivedAt + RUNTIME_ACCESS_TOKEN_DEFAULT_TTL; + this.jwtRefreshNotBefore = Date.now() + this.jwtRefreshBackoff; + throw err; } } @@ -147,3 +288,52 @@ export class RuntimeClient { this.requestQueue.clear(); } } + +/** + * Matches the runtime's rejection of an expired token. It is raised by the auth interceptor + * before the handler runs, so a request that gets it can be retried with a fresh token. + */ +function isExpiredJwtError(err: unknown): boolean { + return ( + err instanceof ConnectError && + err.code === Code.Unauthenticated && + err.rawMessage.includes("jwt is expired") + ); +} + +/** + * Returns when the JWT expires, in milliseconds since the epoch. + * The exp claim is the source of truth, so a token handed over from a cache is not credited with a + * full TTL again. A fresh token is also capped at receivedAt plus its own lifetime (exp minus iat), + * which keeps a client clock that runs behind the issuer's from presenting tokens the runtime + * already rejects. Tokens without readable claims fall back to the default TTL from receipt. + */ +function jwtExpiry(jwt: string, receivedAt: number): number { + const claims = decodeJwtClaims(jwt); + if (typeof claims?.exp !== "number") { + return receivedAt + RUNTIME_ACCESS_TOKEN_DEFAULT_TTL; + } + const expiresAt = claims.exp * 1000; + if (typeof claims.iat !== "number") return expiresAt; + return Math.min(expiresAt, receivedAt + (claims.exp - claims.iat) * 1000); +} + +/** Decodes the payload of a JWT without verifying it. Returns undefined for anything that is not a JWT. */ +function decodeJwtClaims(jwt: string): Record | undefined { + const parts = jwt.split("."); + if (parts.length !== 3) return undefined; + try { + const base64 = parts[1].replace(/-/g, "+").replace(/_/g, "/"); + const bytes = Uint8Array.from(atob(base64), (c) => c.charCodeAt(0)); + const claims: unknown = JSON.parse(new TextDecoder().decode(bytes)); + return typeof claims === "object" && claims !== null + ? (claims as Record) + : undefined; + } catch { + return undefined; + } +} + +function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} From 9935b526643ca2f1b51c0ce58d52c98a19b0d780 Mon Sep 17 00:00:00 2001 From: Nishant Bangarwa Date: Thu, 5 Nov 2026 01:37:34 +0530 Subject: [PATCH 2/3] fix: bound JWT refresh attempts and keep a refresher per auth context Follow-ups from review of the on-demand JWT refresh: - A refresh attempt pins the refresher and the auth context it serves before the backoff sleep and revalidates both after waking and after the call, so a user token refreshed while View As switched the context is never applied - Each refresher call is abandoned after 20s with its AbortSignal aborted (the admin refresher cancels the GetProject query so the retry is not deduped into the stalled request), waiting requests are bounded by their own 60s deadline, and updateJwt() or dispose() wakes them immediately instead of the 50ms poll - Refreshers are kept per auth context: the provider's stays registered for its context while the editor's View As registers one that reissues its mock token through GetDeployment, and the user's applies again when View As ends --- .../updateEditSessionDevJWT.spec.ts | 94 ++++++++++++ .../edit-session/updateEditSessionDevJWT.ts | 55 ++++--- .../projects/project-query-options.spec.ts | 29 +++- .../projects/project-query-options.ts | 17 ++- .../runtime-client/v2/runtime-client.spec.ts | 112 ++++++++++++++ .../src/runtime-client/v2/runtime-client.ts | 142 +++++++++++++----- 6 files changed, 391 insertions(+), 58 deletions(-) create mode 100644 web-admin/src/features/edit-session/updateEditSessionDevJWT.spec.ts diff --git a/web-admin/src/features/edit-session/updateEditSessionDevJWT.spec.ts b/web-admin/src/features/edit-session/updateEditSessionDevJWT.spec.ts new file mode 100644 index 000000000000..bc65fa0a9fd9 --- /dev/null +++ b/web-admin/src/features/edit-session/updateEditSessionDevJWT.spec.ts @@ -0,0 +1,94 @@ +import { adminServiceGetDeployment } from "@rilldata/web-admin/client"; +import { createUpdateEditSessionDevJWT } from "@rilldata/web-admin/features/edit-session/updateEditSessionDevJWT"; +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 } from "@rilldata/web-common/runtime-client/v2"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("@rilldata/web-admin/client", async (importOriginal) => ({ + ...(await importOriginal()), + adminServiceGetDeployment: vi.fn(), +})); + +describe("createUpdateEditSessionDevJWT", () => { + const fetchMock = vi.fn(); + const getDeployment = vi.mocked(adminServiceGetDeployment); + + 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(); + }); + + afterEach(() => { + vi.useRealTimers(); + vi.unstubAllGlobals(); + }); + + function createClient() { + return new RuntimeClient({ + host: "http://localhost:9009", + instanceId: "test-instance", + jwt: "edit-session-jwt", + }); + } + + it("reissues an expired mock JWT instead of waiting for a refetch that never comes", async () => { + getDeployment + .mockResolvedValueOnce({ accessToken: "mock-token-1" }) + .mockResolvedValueOnce({ accessToken: "mock-token-2" }); + const client = createClient(); + const updateDevJWT = createUpdateEditSessionDevJWT( + "deployment-1", + "edit-session-jwt", + ); + + await updateDevJWT(createQueryClient(), client, { + email: "viewer@example.com", + }); + expect(client.getJwt()).toBe("mock-token-1"); + expect(client.authContext).toBe("mock"); + + // The mock token's TTL passes without a scheduled refetch (e.g. the laptop was asleep) + vi.advanceTimersByTime(RUNTIME_ACCESS_TOKEN_DEFAULT_TTL); + try { + await client.runtimeService.ping({}); + } catch { + // ConnectRPC may reject the minimal mock response; the request still went out + } + + expect(getDeployment).toHaveBeenCalledTimes(2); + expect(getDeployment.mock.calls[1][0]).toBe("deployment-1"); + expect(getDeployment.mock.calls[1][1].attributes).toMatchObject({ + email: "viewer@example.com", + domain: "example.com", + }); + expect(client.getJwt()).toBe("mock-token-2"); + const headers = fetchMock.mock.calls[0][1]?.headers as Headers; + expect(headers.get("Authorization")).toBe("Bearer mock-token-2"); + }); + + it("restores the edit session JWT and the user context 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, { email: "viewer@example.com" }); + await updateDevJWT(queryClient, client, null); + + expect(client.getJwt()).toBe("edit-session-jwt"); + expect(client.authContext).toBe("user"); + }); +}); diff --git a/web-admin/src/features/edit-session/updateEditSessionDevJWT.ts b/web-admin/src/features/edit-session/updateEditSessionDevJWT.ts index c94b182d044a..acfdf60017da 100644 --- a/web-admin/src/features/edit-session/updateEditSessionDevJWT.ts +++ b/web-admin/src/features/edit-session/updateEditSessionDevJWT.ts @@ -31,28 +31,21 @@ export function createUpdateEditSessionDevJWT( if (mockUser === null) { selectedMockUserJWT.set(null); + // Drop the mock token's refresher before leaving its context; the provider's applies again for the user + if (client.authContext === "mock") client.setJwtRefresher(undefined); 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"); + // The mock token expires like the user's; reissue it when a request finds it expired + // (e.g. after a laptop sleep) instead of leaving the preview waiting for a token that never comes + client.setJwtRefresher(async (signal) => { + const jwt = await issueMockUserJwt(deploymentId, mockUser, signal); + selectedMockUserJWT.set(jwt); + return jwt; + }); } catch { // Reset the user and make sure we are not in a errored state. selectedMockUserStore.set(prevMockUser); @@ -69,3 +62,31 @@ export function createUpdateEditSessionDevJWT( return invalidateCanvasQueries(queryClient, client.instanceId); }; } + +async function issueMockUserJwt( + deploymentId: string, + mockUser: MockUser, + signal?: AbortSignal, +): Promise { + 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; +} diff --git a/web-admin/src/features/projects/project-query-options.spec.ts b/web-admin/src/features/projects/project-query-options.spec.ts index 50446d13dfdb..4b800d9dd5e8 100644 --- a/web-admin/src/features/projects/project-query-options.spec.ts +++ b/web-admin/src/features/projects/project-query-options.spec.ts @@ -44,6 +44,8 @@ describe("baseGetProjectQueryOptions.refetchInterval", () => { }); describe("createRuntimeJwtRefresher", () => { + const neverAborted = new AbortController().signal; + afterEach(() => { vi.restoreAllMocks(); }); @@ -62,7 +64,7 @@ describe("createRuntimeJwtRefresher", () => { branch: "dev", }); - await expect(refresh?.()).resolves.toBe("fresh"); + await expect(refresh?.(neverAborted)).resolves.toBe("fresh"); expect(fetchQuery).toHaveBeenCalledWith( expect.objectContaining({ queryKey: ["/v1/orgs/org/projects/proj", { branch: "dev" }], @@ -78,7 +80,7 @@ describe("createRuntimeJwtRefresher", () => { token: "magic", }); - await expect(refresh?.()).resolves.toBe("fresh"); + await expect(refresh?.(neverAborted)).resolves.toBe("fresh"); expect(fetchQuery).toHaveBeenCalledWith( expect.objectContaining({ queryKey: ["/v1/orgs/org/projects/proj", "token/magic"], @@ -95,7 +97,7 @@ describe("createRuntimeJwtRefresher", () => { mockedUserId: "u1", }); - await expect(refresh?.()).resolves.toBe("fresh"); + await expect(refresh?.(neverAborted)).resolves.toBe("fresh"); expect(fetchQuery).toHaveBeenCalledWith( expect.objectContaining({ queryKey: [ @@ -106,6 +108,27 @@ describe("createRuntimeJwtRefresher", () => { ); }); + it("cancels the admin request when the client abandons the refresh", () => { + vi.spyOn(queryClient, "fetchQuery").mockImplementation( + () => new Promise(() => {}) as never, + ); + const cancelQueries = vi + .spyOn(queryClient, "cancelQueries") + .mockResolvedValue(); + const abort = new AbortController(); + const refresh = createRuntimeJwtRefresher("user", { + organization: "org", + project: "proj", + }); + + void refresh?.(abort.signal); + abort.abort(); + + expect(cancelQueries).toHaveBeenCalledWith({ + queryKey: ["/v1/orgs/org/projects/proj"], + }); + }); + it("has nothing to refetch for embeds", () => { expect( createRuntimeJwtRefresher("embed", { diff --git a/web-admin/src/features/projects/project-query-options.ts b/web-admin/src/features/projects/project-query-options.ts index 07f8bcab9f3e..401de92386ef 100644 --- a/web-admin/src/features/projects/project-query-options.ts +++ b/web-admin/src/features/projects/project-query-options.ts @@ -86,7 +86,7 @@ export function createRuntimeJwtRefresher( case "mock": { if (!mockedUserId) return undefined; const body = { userId: mockedUserId, ...(branch ? { branch } : {}) }; - return () => + return (signal) => fetchFreshJwt( { queryKey: getAdminServiceGetDeploymentCredentialsQueryKey( @@ -103,11 +103,12 @@ export function createRuntimeJwtRefresher( ), }, (data) => data.accessToken, + signal, ); } case "magic": { if (!token) return undefined; - return () => + return (signal) => fetchFreshJwt( { queryKey: getAdminServiceGetProjectWithBearerTokenQueryKey( @@ -125,11 +126,12 @@ export function createRuntimeJwtRefresher( ), }, (data) => data.jwt, + signal, ); } case "user": { const params = branch ? { branch } : undefined; - return () => + return (signal) => fetchFreshJwt( { queryKey: getAdminServiceGetProjectQueryKey( @@ -141,6 +143,7 @@ export function createRuntimeJwtRefresher( adminServiceGetProject(organization, project, params, signal), }, (data) => data.jwt, + signal, ); } case "embed": @@ -152,8 +155,14 @@ export function createRuntimeJwtRefresher( function fetchFreshJwt( options: FetchQueryOptions, selectJwt: (data: TData) => string | undefined, + signal: AbortSignal, ): Promise { + // The client abandons a call that stalls; cancel the request so the retry is not deduped into it + const cancel = () => + void queryClient.cancelQueries({ queryKey: options.queryKey }); + signal.addEventListener("abort", cancel, { once: true }); return queryClient .fetchQuery({ ...options, staleTime: FreshJwtStaleTime }) - .then(selectJwt); + .then(selectJwt) + .finally(() => signal.removeEventListener("abort", cancel)); } diff --git a/web-common/src/runtime-client/v2/runtime-client.spec.ts b/web-common/src/runtime-client/v2/runtime-client.spec.ts index 00b1a61c9cd3..87d5f2492990 100644 --- a/web-common/src/runtime-client/v2/runtime-client.spec.ts +++ b/web-common/src/runtime-client/v2/runtime-client.spec.ts @@ -480,6 +480,118 @@ describe("RuntimeClient JWT refresh", () => { expect(refreshJwt.mock.calls.length).toBeLessThan(20); }); + it("abandons a refresh that stalls, cancels it and tries again", async () => { + const client = createClient({ jwt: "expiring-token" }); + const signals: AbortSignal[] = []; + const refreshJwt = vi.fn((signal: AbortSignal) => { + signals.push(signal); + return new Promise(() => {}); + }); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + const promise = client.runtimeService.ping({}).catch((e: Error) => e); + + // An attempt is given 20s, then the retry waits out the 1s backoff + await vi.advanceTimersByTimeAsync(21_001); + expect(signals[0].aborted).toBe(true); + expect(refreshJwt).toHaveBeenCalledTimes(2); + + // The request itself still fails at the 60s deadline + await vi.advanceTimersByTimeAsync(40_000); + const error = await promise; + expect((error as Error).message).toContain( + "Timed out waiting for a fresh JWT: the issuer did not respond within 20s", + ); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("releases requests waiting on a stalled refresh as soon as a JWT arrives from elsewhere", async () => { + // e.g. the scheduled GetProject refetch completes while the refresh hangs + const client = createClient({ jwt: "expiring-token" }); + client.setJwtRefresher(vi.fn(() => new Promise(() => {}))); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + let resolved = false; + const promise = makeRequest(client).then(() => { + resolved = true; + }); + await vi.advanceTimersByTimeAsync(5_000); + expect(resolved).toBe(false); + + client.updateJwt("fresh-token"); + await promise; + + expect(resolved).toBe(true); + expect(getLastAuthHeader(fetchMock)).toBe("Bearer fresh-token"); + }); + + it("does not apply a user token refreshed during the backoff after View As switched the context", async () => { + const client = createClient({ jwt: "user-token", authContext: "user" }); + const refreshJwt = vi + .fn() + .mockRejectedValueOnce(new Error("Network Error")) + .mockResolvedValue("fresh-user-token"); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + const promise = makeRequest(client); + // The first attempt failed; the retry is sleeping through the backoff + await vi.advanceTimersByTimeAsync(500); + expect(refreshJwt).toHaveBeenCalledTimes(1); + + // View As installs a mock token while the retry sleeps + client.updateJwt("mock-token", "mock"); + await promise; + await vi.advanceTimersByTimeAsync(2_000); + + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(client.getJwt()).toBe("mock-token"); + expect(client.authContext).toBe("mock"); + expect(getLastAuthHeader(fetchMock)).toBe("Bearer mock-token"); + }); + + it("keeps a refresher per auth context so View As can register its own", async () => { + const client = createClient({ jwt: "user-token", authContext: "user" }); + const refreshUserJwt = vi.fn().mockResolvedValue("fresh-user-token"); + client.setJwtRefresher(refreshUserJwt); + + // View As in the editor installs a mock token and a refresher that reissues it + client.updateJwt("mock-token", "mock"); + const refreshMockJwt = vi.fn().mockResolvedValue("fresh-mock-token"); + client.setJwtRefresher(refreshMockJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + await makeRequest(client); + expect(refreshMockJwt).toHaveBeenCalledTimes(1); + expect(refreshUserJwt).not.toHaveBeenCalled(); + expect(getLastAuthHeader(fetchMock)).toBe("Bearer fresh-mock-token"); + + // View As ends: the user's refresher applies again without being registered anew + client.updateJwt("user-token", "user"); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + await makeRequest(client); + expect(refreshUserJwt).toHaveBeenCalledTimes(1); + expect(getLastAuthHeader(fetchMock)).toBe("Bearer fresh-user-token"); + }); + + it("removes the current context's refresher when given undefined", async () => { + const client = createClient({ jwt: "token" }); + const refreshJwt = vi.fn().mockResolvedValue("fresh-token"); + client.setJwtRefresher(refreshJwt); + client.setJwtRefresher(undefined); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + const promise = client.runtimeService.ping({}).catch((e: Error) => e); + await vi.advanceTimersByTimeAsync(61_000); + const error = await promise; + + expect(refreshJwt).not.toHaveBeenCalled(); + expect((error as Error).message).toContain( + "Timed out waiting for a fresh JWT", + ); + }); + it("ignores a refresher registered for another auth context", async () => { // View As in Rill Developer sets a mock JWT directly on the client; the provider's // refresher for the "user" context must not replace it with the user's token diff --git a/web-common/src/runtime-client/v2/runtime-client.ts b/web-common/src/runtime-client/v2/runtime-client.ts index 56408003a321..e6f44d9b93e8 100644 --- a/web-common/src/runtime-client/v2/runtime-client.ts +++ b/web-common/src/runtime-client/v2/runtime-client.ts @@ -13,7 +13,6 @@ import { extractErrorMessage, extractErrorStatusCode } from "../../lib/errors"; import { RUNTIME_ACCESS_TOKEN_DEFAULT_TTL, JWT_EXPIRY_WARNING_WINDOW, - CHECK_RUNTIME_STORE_FOR_JWT_INTERVAL, } from "../constants"; import { RequestQueue, createQueueInterceptor } from "./request-queue"; @@ -22,15 +21,20 @@ export type AuthContext = "user" | "mock" | "magic" | "embed"; /** * Fetches a fresh JWT from whoever issued the current one, e.g. by refetching GetProject in Rill Cloud. * Resolves with the new token; rejects when the issuer could not be reached or refused. + * The signal aborts when the client gives up on the call, so the issuer request can be cancelled. */ -export type JwtRefresher = () => Promise; +export type JwtRefresher = (signal: AbortSignal) => Promise; // How long a request waits for a fresh JWT before it fails. const JWT_REFRESH_TIMEOUT = 60_000; +// How long one call to the refresher may take before it is abandoned and retried. +const JWT_REFRESH_ATTEMPT_TIMEOUT = 20_000; // Delay before retrying a failed refresh: doubles from the base up to the max while the issuer keeps failing. const JWT_REFRESH_BACKOFF_BASE = 1_000; const JWT_REFRESH_BACKOFF_MAX = 10_000; +const TimedOut = Symbol("timed out"); + export class RuntimeClient { readonly host: string; readonly instanceId: string; @@ -41,11 +45,14 @@ export class RuntimeClient { private currentJwt: string | undefined; private jwtExpiresAt = 0; private _authContext: AuthContext; - // Refresher registered by RuntimeProvider, bound to the auth context it was registered under - private refreshJwt: JwtRefresher | undefined; - private refreshJwtAuthContext: AuthContext | undefined; + // Resolved and replaced whenever the JWT, the auth context or a refresher changes, or the client + // is disposed, so that requests waiting for a fresh JWT re-check their situation right away + private jwtChanged = createDeferred(); + // Refreshers per auth context: RuntimeProvider registers one for its context and View As one for "mock" + private readonly jwtRefreshers = new Map(); // In-flight refresh shared by every request waiting on it, and the backoff after failures private jwtRefresh: Promise | undefined; + private jwtRefreshAbort: AbortController | undefined; private jwtRefreshBackoff = 0; private jwtRefreshNotBefore = 0; private disposed = false; @@ -126,24 +133,35 @@ export class RuntimeClient { const authContextChanged = !!this._authContext && !!authContext && authContext !== this._authContext; - if (jwt !== this.currentJwt) { + const jwtChanged = jwt !== this.currentJwt; + if (jwtChanged) { this.currentJwt = jwt; this.jwtExpiresAt = jwt ? jwtExpiry(jwt, Date.now()) : 0; } if (authContext) this._authContext = authContext; + if (jwtChanged || authContextChanged) this.wakeWaitingRequests(); return authContextChanged; } /** - * Registers the callback that fetches a fresh JWT when the current one is about to expire. - * It stays bound to the auth context it was registered under: a mock JWT that View As sets directly - * on the client must not be replaced by the provider's token, so the refresher is ignored until the - * client is back in that context. + * Registers the callback that fetches a fresh JWT for the current auth context, replacing any earlier + * one for that context; undefined removes it. Refreshers are kept per context: RuntimeProvider's stays + * registered for its context while View As installs a mock token with its own refresher for "mock", + * and applies again once View As ends and the context switches back. */ setJwtRefresher(refreshJwt: JwtRefresher | undefined): void { - this.refreshJwt = refreshJwt; - this.refreshJwtAuthContext = this._authContext; + const authContext = this._authContext; + if (refreshJwt === this.jwtRefreshers.get(authContext)) return; + if (refreshJwt) { + this.jwtRefreshers.set(authContext, refreshJwt); + } else { + this.jwtRefreshers.delete(authContext); + } + // A new issuer starts without the backoff the previous one earned + this.jwtRefreshBackoff = 0; + this.jwtRefreshNotBefore = 0; + this.wakeWaitingRequests(); } /** @@ -189,10 +207,10 @@ export class RuntimeClient { } /** - * Blocks until the JWT is fresh. With a refresher registered it asks the issuer for a new token - * right away, retrying with backoff while the issuer is unreachable; without one it waits for an - * updateJwt() call. Gives up after JWT_REFRESH_TIMEOUT, or at once when the issuer refuses because - * the session or the project access is gone. + * Blocks until the JWT is fresh. With a refresher registered for the current auth context it asks + * the issuer for a new token right away, retrying with backoff while the issuer is unreachable; + * without one it waits for an updateJwt() call. Gives up after JWT_REFRESH_TIMEOUT, or at once + * when the issuer refuses because the session or the project access is gone. */ private async waitForFreshJwt(): Promise { // Embeds have a 24h backend-issued TTL and no refresher; skip client-side expiry check @@ -204,19 +222,25 @@ export class RuntimeClient { if (this.disposed) { throw new DOMException("Request cancelled", "AbortError"); } - if (Date.now() > deadline) { + const remaining = deadline - Date.now(); + if (remaining <= 0) { throw new Error( "Timed out waiting for a fresh JWT" + (lastError ? `: ${extractErrorMessage(lastError)}` : ""), ); } + // Whatever happens below, a change to the JWT or the client wakes us to re-check + const changed = this.jwtChanged.promise; if (!this.canRefreshJwt()) { - // Wait for whoever owns the token (e.g. View As) to call updateJwt() - await sleep(CHECK_RUNTIME_STORE_FOR_JWT_INTERVAL); + // Wait for whoever owns the token (e.g. View As in Rill Developer) to call updateJwt() + await waitAtMost(changed, remaining); continue; } try { - await this.runJwtRefresh(); + await waitAtMost( + Promise.race([this.runJwtRefresh(), changed]), + remaining, + ); } catch (err) { const status = extractErrorStatusCode(err); if (status === 401 || status === 403) { @@ -237,9 +261,13 @@ export class RuntimeClient { } private canRefreshJwt(): boolean { - return ( - !!this.refreshJwt && this.refreshJwtAuthContext === this._authContext - ); + return this.jwtRefreshers.has(this._authContext); + } + + private wakeWaitingRequests(): void { + const { resolve } = this.jwtChanged; + this.jwtChanged = createDeferred(); + resolve(); } /** Runs one refresh, shared by every caller that joins while it is in flight. */ @@ -251,21 +279,31 @@ export class RuntimeClient { } private async refreshJwtOnce(): Promise { + // Pin the refresher and the context it serves: both can change while this attempt waits or runs + // (e.g. View As switching to a mock token), and a token must only be applied to the context it + // was issued for + const authContext = this._authContext; + const refreshJwt = this.jwtRefreshers.get(authContext); + if (!refreshJwt) return; const delay = this.jwtRefreshNotBefore - Date.now(); if (delay > 0) await sleep(delay); - const refreshJwt = this.refreshJwt; - const authContext = this._authContext; - if (this.disposed || !refreshJwt) return; + if (!this.servesCurrentContext(refreshJwt, authContext)) return; + + // A call that stalls is abandoned and cancelled so the retry is not stuck behind it + const abort = new AbortController(); + this.jwtRefreshAbort = abort; try { - const jwt = await refreshJwt(); - // The provider may have switched auth context (View As) meanwhile; its token wins - if ( - this.disposed || - this.refreshJwt !== refreshJwt || - this._authContext !== authContext - ) { - return; + const jwt = await waitAtMost( + refreshJwt(abort.signal), + JWT_REFRESH_ATTEMPT_TIMEOUT, + ); + if (jwt === TimedOut) { + abort.abort(); + throw new Error( + `the issuer did not respond within ${JWT_REFRESH_ATTEMPT_TIMEOUT / 1000}s`, + ); } + if (!this.servesCurrentContext(refreshJwt, authContext)) return; if (!jwt) throw new Error("the issuer returned no JWT"); this.updateJwt(jwt); if (!this.isJwtFresh()) { @@ -280,11 +318,27 @@ export class RuntimeClient { ); this.jwtRefreshNotBefore = Date.now() + this.jwtRefreshBackoff; throw err; + } finally { + if (this.jwtRefreshAbort === abort) this.jwtRefreshAbort = undefined; } } + /** Whether the refresher an attempt pinned is still the one registered for the client's current context. */ + private servesCurrentContext( + refreshJwt: JwtRefresher, + authContext: AuthContext, + ): boolean { + return ( + !this.disposed && + this._authContext === authContext && + this.jwtRefreshers.get(authContext) === refreshJwt + ); + } + dispose(): void { this.disposed = true; + this.jwtRefreshAbort?.abort(); + this.wakeWaitingRequests(); this.requestQueue.clear(); } } @@ -334,6 +388,26 @@ function decodeJwtClaims(jwt: string): Record | undefined { } } +/** Settles like the promise, or resolves with TimedOut once `ms` elapse first; the timer that lost is cleared. */ +function waitAtMost( + promise: Promise, + ms: number, +): Promise { + let timer: ReturnType | undefined; + const timeout = new Promise((resolve) => { + timer = setTimeout(() => resolve(TimedOut), ms); + }); + return Promise.race([promise, timeout]).finally(() => clearTimeout(timer)); +} + +function createDeferred(): { promise: Promise; resolve: () => void } { + let resolve!: () => void; + const promise = new Promise((r) => { + resolve = r; + }); + return { promise, resolve }; +} + function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } From 49e7a84cb9078abbcb613e1dfa95de736dbfa42b Mon Sep 17 00:00:00 2001 From: Nishant Bangarwa Date: Thu, 5 Nov 2026 01:49:07 +0530 Subject: [PATCH 3/3] fix: make the edit layout's runtime credentials follow the View As selection The edit layout forwarded every GetProject JWT to RuntimeProvider with the "user" context, so any user-token update that landed while View As was active (the 15-minute poll, a focus refetch, or an on-demand refresh that was still in flight when View As started) replaced the mock token and the preview showed the user's data under View As. - The layout derives the provider's jwt, authContext and refreshJwt from the View As stores (editSessionRuntimeAuth): the mock user's JWT and a refresher that reissues it through GetDeployment while a mock user is selected, the user's otherwise - createUpdateEditSessionDevJWT no longer registers the refresher itself; the layout owns it through the provider's props - A regression plays the layout and provider through an in-flight user refresh that completes after View As started and checks the mock credentials survive it --- .../updateEditSessionDevJWT.spec.ts | 148 ++++++++++++++---- .../edit-session/updateEditSessionDevJWT.ts | 64 ++++++-- .../[project]/-/edit/+layout.svelte | 34 +++- 3 files changed, 205 insertions(+), 41 deletions(-) diff --git a/web-admin/src/features/edit-session/updateEditSessionDevJWT.spec.ts b/web-admin/src/features/edit-session/updateEditSessionDevJWT.spec.ts index bc65fa0a9fd9..d4f060a909ea 100644 --- a/web-admin/src/features/edit-session/updateEditSessionDevJWT.spec.ts +++ b/web-admin/src/features/edit-session/updateEditSessionDevJWT.spec.ts @@ -1,8 +1,20 @@ import { adminServiceGetDeployment } from "@rilldata/web-admin/client"; -import { createUpdateEditSessionDevJWT } from "@rilldata/web-admin/features/edit-session/updateEditSessionDevJWT"; +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 } from "@rilldata/web-common/runtime-client/v2"; +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) => ({ @@ -10,9 +22,10 @@ vi.mock("@rilldata/web-admin/client", async (importOriginal) => ({ adminServiceGetDeployment: vi.fn(), })); -describe("createUpdateEditSessionDevJWT", () => { +describe("edit session View As", () => { const fetchMock = vi.fn(); const getDeployment = vi.mocked(adminServiceGetDeployment); + const mockUser = { email: "viewer@example.com" }; beforeEach(() => { vi.useFakeTimers(); @@ -26,6 +39,8 @@ describe("createUpdateEditSessionDevJWT", () => { ); vi.stubGlobal("fetch", fetchMock); getDeployment.mockReset(); + selectedMockUserStore.set(null); + selectedMockUserJWT.set(null); }); afterEach(() => { @@ -41,54 +56,129 @@ describe("createUpdateEditSessionDevJWT", () => { }); } - it("reissues an expired mock JWT instead of waiting for a refetch that never comes", async () => { - getDeployment - .mockResolvedValueOnce({ accessToken: "mock-token-1" }) - .mockResolvedValueOnce({ accessToken: "mock-token-2" }); + /** 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(createQueryClient(), client, { - email: "viewer@example.com", - }); - expect(client.getJwt()).toBe("mock-token-1"); + 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"); - // The mock token's TTL passes without a scheduled refetch (e.g. the laptop was asleep) - vi.advanceTimersByTime(RUNTIME_ACCESS_TOKEN_DEFAULT_TTL); - try { - await client.runtimeService.ping({}); - } catch { - // ConnectRPC may reject the minimal mock response; the request still went out - } + 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(); + }); - expect(getDeployment).toHaveBeenCalledTimes(2); - expect(getDeployment.mock.calls[1][0]).toBe("deployment-1"); - expect(getDeployment.mock.calls[1][1].attributes).toMatchObject({ + 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(client.getJwt()).toBe("mock-token-2"); - const headers = fetchMock.mock.calls[0][1]?.headers as Headers; - expect(headers.get("Authorization")).toBe("Bearer mock-token-2"); + expect(body.accessTokenTtlSeconds).toBe( + RUNTIME_ACCESS_TOKEN_DEFAULT_TTL / 1000, + ); + expect(signal).toBe(abort.signal); + expect(get(selectedMockUserJWT)).toBe("mock-token-2"); }); - it("restores the edit session JWT and the user context when View As ends", async () => { - getDeployment.mockResolvedValueOnce({ accessToken: "mock-token" }); + 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((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"); - await updateDevJWT(queryClient, client, { email: "viewer@example.com" }); - await updateDevJWT(queryClient, client, null); + // 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"); - expect(client.getJwt()).toBe("edit-session-jwt"); + // 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"); }); }); diff --git a/web-admin/src/features/edit-session/updateEditSessionDevJWT.ts b/web-admin/src/features/edit-session/updateEditSessionDevJWT.ts index acfdf60017da..2aea541daef1 100644 --- a/web-admin/src/features/edit-session/updateEditSessionDevJWT.ts +++ b/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,23 +33,16 @@ 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); - // Drop the mock token's refresher before leaving its context; the provider's applies again for the user - if (client.authContext === "mock") client.setJwtRefresher(undefined); client.updateJwt(editSessionJwt, "user"); } else { try { const accessToken = await issueMockUserJwt(deploymentId, mockUser); selectedMockUserJWT.set(accessToken); client.updateJwt(accessToken, "mock"); - // The mock token expires like the user's; reissue it when a request finds it expired - // (e.g. after a laptop sleep) instead of leaving the preview waiting for a token that never comes - client.setJwtRefresher(async (signal) => { - const jwt = await issueMockUserJwt(deploymentId, mockUser, signal); - selectedMockUserJWT.set(jwt); - return jwt; - }); } catch { // Reset the user and make sure we are not in a errored state. selectedMockUserStore.set(prevMockUser); @@ -63,6 +60,53 @@ export function createUpdateEditSessionDevJWT( }; } +/** + * 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, diff --git a/web-admin/src/routes/[organization]/[project]/-/edit/+layout.svelte b/web-admin/src/routes/[organization]/[project]/-/edit/+layout.svelte index d8543185c4f5..35e94d5f7281 100644 --- a/web-admin/src/routes/[organization]/[project]/-/edit/+layout.svelte +++ b/web-admin/src/routes/[organization]/[project]/-/edit/+layout.svelte @@ -13,6 +13,10 @@ import BranchDeploymentStopped from "@rilldata/web-admin/features/branches/BranchDeploymentStopped.svelte"; import EditSessionLoading from "@rilldata/web-admin/features/edit-session/EditSessionLoading.svelte"; import EditSessionTimeoutBanner from "@rilldata/web-admin/features/edit-session/EditSessionTimeoutBanner.svelte"; + import { + createMockUserJwtRefresher, + editSessionRuntimeAuth, + } from "@rilldata/web-admin/features/edit-session/updateEditSessionDevJWT"; import ProjectHeader from "../../../../../features/projects/header/ProjectHeader.svelte"; import { baseGetProjectQueryOptions, @@ -21,6 +25,10 @@ import SlimProjectHeader from "@rilldata/web-admin/features/projects/SlimProjectHeader.svelte"; import { getThemedLogoUrl } from "@rilldata/web-admin/features/themes/organization-logo"; import ErrorPage from "@rilldata/web-common/components/ErrorPage.svelte"; + import { + selectedMockUserJWT, + selectedMockUserStore, + } from "@rilldata/web-common/features/dashboards/granular-access-policies/stores"; import FileAndResourceWatcher from "@rilldata/web-common/features/entity-management/FileAndResourceWatcher.svelte"; import { themeControl } from "@rilldata/web-common/features/themes/theme-control"; import { editorRoutePrefix } from "@rilldata/web-common/layout/navigation/editor-routing"; @@ -85,11 +93,27 @@ $: jwt = $projectQuery.data?.jwt ?? null; // Re-requests the branch's runtime JWT when a request finds it expired instead of waiting for the next poll. - $: refreshJwt = createRuntimeJwtRefresher("user", { + $: userRefreshJwt = createRuntimeJwtRefresher("user", { organization, project, branch, }); + // View As installs a mock user's JWT; it is reissued the same way when it expires. + // Keyed on the deployment id rather than the deployment object so the refresher keeps its identity across polls. + $: deploymentId = deployment?.id; + $: mockRefreshJwt = + $selectedMockUserStore && deploymentId + ? createMockUserJwtRefresher(deploymentId, $selectedMockUserStore) + : undefined; + // The provider follows the View As selection: a GetProject update must not push the user's JWT + // back onto the client while a mock user is selected (see editSessionRuntimeAuth). + $: runtimeAuth = editSessionRuntimeAuth({ + userJwt: jwt ?? undefined, + userRefreshJwt, + mockUser: $selectedMockUserStore, + mockJwt: $selectedMockUserJWT, + mockRefreshJwt, + }); // Flipped when the user clicks "Start deployment" on a stopped deployment; // keeps the UI in loading state while the backend transitions STOPPED → PENDING → RUNNING. @@ -189,7 +213,13 @@ /> {:else if isReady && deployment?.id && instanceId && runtimeHost && jwt} {#key `${runtimeHost}::${instanceId}::${hasPrimaryDeployment}`} - + {#if !inProjectWelcomePage}