From 5b0bf35175782eb60970b990fd4a16f12fb5ad0a Mon Sep 17 00:00:00 2001 From: -Puter <22245429+puterhimself@users.noreply.github.com> Date: Thu, 23 Jul 2026 16:54:15 +0530 Subject: [PATCH] feat(chat): add connecting status during organization bootstrap --- apps/web/package.json | 1 + .../web/src/components/chat/chat-composer.tsx | 5 +- .../web/src/hooks/chat/use-chat-agent.test.ts | 60 +++++++++ apps/web/src/hooks/chat/use-chat-agent.ts | 34 +++--- .../hooks/use-personal-organization.test.ts | 77 ++++++++++++ .../src/hooks/use-personal-organization.ts | 70 +++++++++++ apps/web/src/hooks/use-project-workspace.ts | 4 +- apps/web/src/lib/flue-transport.test.ts | 115 ++++++++++++++++++ apps/web/src/lib/flue-transport.ts | 82 +++++++++++++ apps/web/src/root.tsx | 114 +++++------------ apps/web/tsconfig.json | 2 +- bun.lock | 2 + packages/auth/package.json | 3 +- packages/auth/src/web/hooks.ts | 51 +++----- .../src/web/session-token-resolver.test.ts | 39 ++++++ .../auth/src/web/session-token-resolver.ts | 59 +++++++++ 16 files changed, 581 insertions(+), 137 deletions(-) create mode 100644 apps/web/src/hooks/chat/use-chat-agent.test.ts create mode 100644 apps/web/src/hooks/use-personal-organization.test.ts create mode 100644 apps/web/src/hooks/use-personal-organization.ts create mode 100644 apps/web/src/lib/flue-transport.test.ts create mode 100644 apps/web/src/lib/flue-transport.ts create mode 100644 packages/auth/src/web/session-token-resolver.test.ts create mode 100644 packages/auth/src/web/session-token-resolver.ts diff --git a/apps/web/package.json b/apps/web/package.json index eb37872..054c429 100644 --- a/apps/web/package.json +++ b/apps/web/package.json @@ -40,6 +40,7 @@ "tailwindcss": "catalog:", "typescript": "catalog:", "vite": "catalog:", + "vitest": "catalog:", "vite-tsconfig-paths": "^6.1.1" } } diff --git a/apps/web/src/components/chat/chat-composer.tsx b/apps/web/src/components/chat/chat-composer.tsx index a68544e..6c0fad8 100644 --- a/apps/web/src/components/chat/chat-composer.tsx +++ b/apps/web/src/components/chat/chat-composer.tsx @@ -11,7 +11,8 @@ import { useChatComposer } from "@/hooks/chat/use-chat-composer"; import type { ChatComposerProps } from "@/lib/chat/types"; export const ChatComposer = ({ error, onSend, status }: ChatComposerProps) => { - const busy = status === "submitted" || status === "streaming"; + const busy = + status === "connecting" || status === "submitted" || status === "streaming"; const composer = useChatComposer({ busy, onSend }); let statusContent: ReactNode = ( @@ -23,7 +24,7 @@ export const ChatComposer = ({ error, onSend, status }: ChatComposerProps) => { statusContent = ( <> - Zopu is responding + {status === "connecting" ? "Preparing Zopu" : "Zopu is responding"} ); } diff --git a/apps/web/src/hooks/chat/use-chat-agent.test.ts b/apps/web/src/hooks/chat/use-chat-agent.test.ts new file mode 100644 index 0000000..487ba70 --- /dev/null +++ b/apps/web/src/hooks/chat/use-chat-agent.test.ts @@ -0,0 +1,60 @@ +import { beforeEach, describe, expect, test, vi } from "vitest"; + +import { useChatAgent } from "./use-chat-agent"; + +const mocks = vi.hoisted(() => ({ + agent: { + error: undefined, + historyReady: true, + messages: [], + sendMessage: vi.fn(() => Promise.resolve()), + status: "idle" as const, + }, + organization: {} as { error?: Error; organizationId?: string }, + useFlueAgent: vi.fn(), +})); + +vi.mock("@flue/react", () => ({ + useFlueAgent: mocks.useFlueAgent, +})); +vi.mock("@/hooks/use-personal-organization", () => ({ + usePersonalOrganization: () => mocks.organization, +})); + +describe("useChatAgent", () => { + beforeEach(() => { + mocks.organization = {}; + mocks.agent.sendMessage.mockClear(); + mocks.useFlueAgent.mockReset(); + mocks.useFlueAgent.mockReturnValue(mocks.agent); + }); + + test("does not use the agent before organization bootstrap completes", async () => { + const chat = useChatAgent(); + + expect(mocks.useFlueAgent).toHaveBeenCalledWith({ + id: undefined, + live: "sse", + name: "zopu", + }); + expect(chat.status).toBe("connecting"); + await expect(chat.sendMessage("too early")).rejects.toThrow( + "Personal organization is still being prepared" + ); + expect(mocks.agent.sendMessage).not.toHaveBeenCalled(); + }); + + test("uses the ensured organization id for the agent session", async () => { + mocks.organization = { organizationId: "org-a" }; + + const chat = useChatAgent(); + expect(mocks.useFlueAgent).toHaveBeenCalledWith({ + id: "org-a", + live: "sse", + name: "zopu", + }); + + await chat.sendMessage("ready"); + expect(mocks.agent.sendMessage).toHaveBeenCalledWith("ready"); + }); +}); diff --git a/apps/web/src/hooks/chat/use-chat-agent.ts b/apps/web/src/hooks/chat/use-chat-agent.ts index b0845fb..e90fb25 100644 --- a/apps/web/src/hooks/chat/use-chat-agent.ts +++ b/apps/web/src/hooks/chat/use-chat-agent.ts @@ -1,36 +1,38 @@ -import { api } from "@code/backend/convex/_generated/api"; import { useFlueAgent } from "@flue/react"; -import { useQuery } from "convex/react"; +import { usePersonalOrganization } from "@/hooks/use-personal-organization"; import { CHAT_AGENT } from "@/lib/chat/constants"; import type { ChatAgentState } from "@/lib/chat/types"; -import { setSendRequestId } from "@/root"; export const useChatAgent = (): ChatAgentState => { - // The global Zopu agent instance id is the user's current personal - // organization id. This is the hard tenancy boundary: the server rejects an - // instance id that is not the caller's own organization. - const organization = useQuery(api.organizations.getCurrent); + // The agent session is not created until the authenticated user's personal + // organization has been ensured and its id is available. + const organization = usePersonalOrganization(); const agent = useFlueAgent({ ...CHAT_AGENT, - id: organization?._id, + id: organization.organizationId, }); const sendMessage = async (message: string): Promise => { - // Generate a stable idempotency request id once per send, immediately - // before the message is submitted. It is reused across fetch/stream - // retries for this send and reset only on the next send, so the server - // can capture normalized evidence keyed by (organization, request id) - // without creating duplicates. - setSendRequestId(crypto.randomUUID()); + if (!organization.organizationId) { + throw ( + organization.error ?? + new Error("Personal organization is still being prepared") + ); + } await agent.sendMessage(message); }; + let { status } = agent; + if (!organization.organizationId) { + status = organization.error ? "error" : "connecting"; + } + return { - error: agent.error, + error: organization.error ?? agent.error, historyReady: agent.historyReady, messages: agent.messages, sendMessage, - status: agent.status, + status, }; }; diff --git a/apps/web/src/hooks/use-personal-organization.test.ts b/apps/web/src/hooks/use-personal-organization.test.ts new file mode 100644 index 0000000..3a1dc23 --- /dev/null +++ b/apps/web/src/hooks/use-personal-organization.test.ts @@ -0,0 +1,77 @@ +import type { EffectCallback } from "react"; +import { beforeEach, describe, expect, test, vi } from "vitest"; + +import { usePersonalOrganization } from "./use-personal-organization"; + +const mocks = vi.hoisted(() => ({ + cleanup: undefined as (() => void) | undefined, + ensurePersonalOrganization: vi.fn(), + session: { + data: { session: { id: "session-a" } }, + } as { data: { session: { id: string } } | null }, + state: null as { + error?: Error; + organizationId?: string; + sessionId: string; + } | null, +})); + +vi.mock("@code/auth/web", () => ({ + authClient: { useSession: () => mocks.session }, +})); +vi.mock("@code/backend/convex/_generated/api", () => ({ + api: { organizations: { ensurePersonalOrganization: "ensure-org" } }, +})); +vi.mock("convex/react", () => ({ + useMutation: () => mocks.ensurePersonalOrganization, +})); +vi.mock("react", () => ({ + useEffect: (effect: EffectCallback) => { + mocks.cleanup?.(); + const cleanup = effect(); + mocks.cleanup = typeof cleanup === "function" ? cleanup : undefined; + }, + useState: () => [ + mocks.state, + (next: { error?: Error; organizationId?: string; sessionId: string }) => { + mocks.state = next; + }, + ], +})); + +describe("usePersonalOrganization", () => { + beforeEach(() => { + mocks.cleanup?.(); + mocks.cleanup = undefined; + mocks.ensurePersonalOrganization.mockReset(); + mocks.session = { data: { session: { id: "session-a" } } }; + mocks.state = null; + }); + + test("an authenticated session ensures its org before exposing the id", async () => { + const ensured = Promise.withResolvers<{ _id: string }>(); + mocks.ensurePersonalOrganization.mockReturnValue(ensured.promise); + + expect(usePersonalOrganization()).toEqual({}); + expect(mocks.ensurePersonalOrganization).toHaveBeenCalledTimes(1); + + ensured.resolve({ _id: "org-a" }); + await ensured.promise; + await Promise.resolve(); + + expect(usePersonalOrganization()).toEqual({ + error: undefined, + organizationId: "org-a", + }); + }); + + test("logout cannot expose the prior session organization", async () => { + mocks.ensurePersonalOrganization.mockResolvedValue({ _id: "org-a" }); + usePersonalOrganization(); + await Promise.resolve(); + await Promise.resolve(); + + mocks.session = { data: null }; + expect(usePersonalOrganization()).toEqual({}); + }); +}); diff --git a/apps/web/src/hooks/use-personal-organization.ts b/apps/web/src/hooks/use-personal-organization.ts new file mode 100644 index 0000000..67bd811 --- /dev/null +++ b/apps/web/src/hooks/use-personal-organization.ts @@ -0,0 +1,70 @@ +import { authClient } from "@code/auth/web"; +import { api } from "@code/backend/convex/_generated/api"; +import type { Id } from "@code/backend/convex/_generated/dataModel"; +import { useMutation } from "convex/react"; +import { useEffect, useState } from "react"; + +interface OrganizationBootstrapState { + readonly error?: Error; + readonly organizationId?: Id<"organizations">; + readonly sessionId: string; +} + +export interface PersonalOrganizationState { + readonly error?: Error; + readonly organizationId?: Id<"organizations">; +} + +/** Ensure the authenticated session has its personal tenancy boundary. */ +export const usePersonalOrganization = (): PersonalOrganizationState => { + const { data: session } = authClient.useSession(); + const sessionId = session?.session.id; + const ensurePersonalOrganization = useMutation( + api.organizations.ensurePersonalOrganization + ); + const [bootstrap, setBootstrap] = useState( + null + ); + + useEffect(() => { + if (!sessionId) { + return; + } + + let active = true; + const bootstrapOrganization = async (): Promise => { + try { + const organization = await ensurePersonalOrganization(); + if (active) { + setBootstrap({ + organizationId: organization._id, + sessionId, + }); + } + } catch (caughtError: unknown) { + if (active) { + setBootstrap({ + error: + caughtError instanceof Error + ? caughtError + : new Error(String(caughtError)), + sessionId, + }); + } + } + }; + void bootstrapOrganization(); + + return () => { + active = false; + }; + }, [ensurePersonalOrganization, sessionId]); + + if (!sessionId || bootstrap?.sessionId !== sessionId) { + return {}; + } + return { + error: bootstrap.error, + organizationId: bootstrap.organizationId, + }; +}; diff --git a/apps/web/src/hooks/use-project-workspace.ts b/apps/web/src/hooks/use-project-workspace.ts index 1a47016..15ad87c 100644 --- a/apps/web/src/hooks/use-project-workspace.ts +++ b/apps/web/src/hooks/use-project-workspace.ts @@ -1,14 +1,14 @@ import { api } from "@code/backend/convex/_generated/api"; import type { Id } from "@code/backend/convex/_generated/dataModel"; +import { useFlueClient } from "@flue/react"; import { useAction, useMutation, useQuery } from "convex/react"; import { useState } from "react"; -import { flueClient } from "@/root"; - const errorMessage = (error: unknown) => error instanceof Error ? error.message : String(error); export const useProjectWorkspace = () => { + const flueClient = useFlueClient(); const projects = useQuery(api.projects.list); const [selectedProjectId, setSelectedProjectId] = useState | null>(null); diff --git a/apps/web/src/lib/flue-transport.test.ts b/apps/web/src/lib/flue-transport.test.ts new file mode 100644 index 0000000..93b758f --- /dev/null +++ b/apps/web/src/lib/flue-transport.test.ts @@ -0,0 +1,115 @@ +import { createFlueClient } from "@flue/sdk"; +import { describe, expect, test } from "vitest"; + +import { createFlueFetch } from "./flue-transport"; + +const BASE_URL = new URL("https://flue.example/api/"); + +describe("createFlueFetch", () => { + test("overlapping SDK agent sends keep distinct request IDs", async () => { + const firstResponse = Promise.withResolvers(); + const secondResponse = Promise.withResolvers(); + const bothStarted = Promise.withResolvers(); + const capturedHeaders: Headers[] = []; + let requestCount = 0; + let requestId = 0; + + const fetchImpl: typeof fetch = (input, init) => { + capturedHeaders.push(new Headers(init?.headers)); + requestCount += 1; + if (requestCount === 2) { + bothStarted.resolve(true); + } + if (requestCount === 1) { + return firstResponse.promise; + } + if (requestCount === 2) { + return secondResponse.promise; + } + throw new Error(`Unexpected request: ${String(input)}`); + }; + const client = createFlueClient({ + baseUrl: BASE_URL.toString(), + fetch: createFlueFetch({ + baseUrl: BASE_URL, + fetchImpl, + generateRequestId: () => `request-${(requestId += 1)}`, + }), + headers: { authorization: "Bearer current-jwt" }, + }); + + const firstSend = client.agents.send("zopu", "org-a", { + message: "first", + }); + const secondSend = client.agents.send("zopu", "org-a", { + message: "second", + }); + await bothStarted.promise; + + expect(capturedHeaders).toHaveLength(2); + expect(capturedHeaders[0]?.get("x-zopu-request-id")).toBe("request-1"); + expect(capturedHeaders[1]?.get("x-zopu-request-id")).toBe("request-2"); + expect(capturedHeaders[0]?.get("authorization")).toBe("Bearer current-jwt"); + expect(capturedHeaders[1]?.get("authorization")).toBe("Bearer current-jwt"); + + firstResponse.resolve( + Response.json( + { + offset: "1", + streamUrl: "http://internal/streams/first", + submissionId: "submission-1", + }, + { status: 202 } + ) + ); + secondResponse.resolve( + Response.json( + { + offset: "2", + streamUrl: "http://internal/streams/second", + submissionId: "submission-2", + }, + { status: 202 } + ) + ); + + await expect(Promise.all([firstSend, secondSend])).resolves.toEqual([ + { + offset: "1", + streamUrl: "https://flue.example/streams/first", + submissionId: "submission-1", + }, + { + offset: "2", + streamUrl: "https://flue.example/streams/second", + submissionId: "submission-2", + }, + ]); + }); + + test("history and stream observation requests remain untagged", async () => { + const capturedHeaders: Headers[] = []; + const fetchImpl: typeof fetch = (_input, init) => { + capturedHeaders.push(new Headers(init?.headers)); + return Promise.resolve(new Response(null, { status: 204 })); + }; + const flueFetch = createFlueFetch({ + baseUrl: BASE_URL, + fetchImpl, + generateRequestId: () => "must-not-be-used", + }); + + await flueFetch("https://flue.example/api/agents/zopu/org-a?view=history", { + headers: { authorization: "Bearer current-jwt" }, + method: "GET", + }); + await flueFetch("https://flue.example/api/agents/zopu/org-a?view=updates", { + headers: { authorization: "Bearer current-jwt" }, + method: "GET", + }); + + expect(capturedHeaders).toHaveLength(2); + expect(capturedHeaders[0]?.has("x-zopu-request-id")).toBe(false); + expect(capturedHeaders[1]?.has("x-zopu-request-id")).toBe(false); + }); +}); diff --git a/apps/web/src/lib/flue-transport.ts b/apps/web/src/lib/flue-transport.ts new file mode 100644 index 0000000..ac838a1 --- /dev/null +++ b/apps/web/src/lib/flue-transport.ts @@ -0,0 +1,82 @@ +interface FlueFetchOptions { + readonly baseUrl: URL; + readonly fetchImpl?: typeof fetch; + readonly generateRequestId?: () => string; +} + +/** + * Build the Flue transport used by the browser client. + * + * The installed Flue SDK does not expose per-call headers on `agents.send`. + * Its custom `fetch` seam does receive the resolved URL and method, so request + * IDs are attached here to each agent admission POST. History and Durable + * Streams observation requests are GETs and intentionally remain untagged. + */ +export const createFlueFetch = ({ + baseUrl, + fetchImpl = fetch, + generateRequestId = crypto.randomUUID, +}: FlueFetchOptions): typeof fetch => { + const basePath = baseUrl.pathname.replace(/\/+$/u, ""); + + return async (input, init) => { + const inputUrl = + typeof input === "string" || input instanceof URL ? input : input.url; + const requestUrl = new URL(inputUrl, baseUrl); + const method = ( + init?.method ?? (input instanceof Request ? input.method : "GET") + ).toUpperCase(); + const relativePath = requestUrl.pathname.startsWith(`${basePath}/`) + ? requestUrl.pathname.slice(basePath.length) + : requestUrl.pathname; + const pathSegments = relativePath.split("/").filter(Boolean); + const isAgentAdmission = + method === "POST" && + pathSegments.length === 3 && + pathSegments[0] === "agents"; + + let requestInit = init; + if (isAgentAdmission) { + const headers = new Headers( + input instanceof Request ? input.headers : undefined + ); + for (const [key, value] of new Headers(init?.headers).entries()) { + headers.set(key, value); + } + headers.set("x-zopu-request-id", generateRequestId()); + requestInit = { ...init, headers }; + } + + const response = await fetchImpl(input, requestInit); + if ( + !response.ok || + !response.headers.get("content-type")?.includes("application/json") + ) { + return response; + } + + const body = (await response.clone().json()) as unknown; + if ( + typeof body !== "object" || + body === null || + !("streamUrl" in body) || + typeof body.streamUrl !== "string" + ) { + return response; + } + + const streamUrl = new URL(body.streamUrl); + streamUrl.protocol = baseUrl.protocol; + streamUrl.host = baseUrl.host; + const headers = new Headers(response.headers); + headers.set("location", streamUrl.toString()); + return Response.json( + { ...body, streamUrl: streamUrl.toString() }, + { + headers, + status: response.status, + statusText: response.statusText, + } + ); + }; +}; diff --git a/apps/web/src/root.tsx b/apps/web/src/root.tsx index 3040c5b..eb98fe3 100644 --- a/apps/web/src/root.tsx +++ b/apps/web/src/root.tsx @@ -1,10 +1,11 @@ -import { authClient, WebAuthProvider } from "@code/auth/web"; +import { useConvexAccessToken, WebAuthProvider } from "@code/auth/web"; import { env } from "@code/env/web"; import { Toaster } from "@code/ui/components/sonner"; import { FlueProvider } from "@flue/react"; import "./index.css"; import { createFlueClient } from "@flue/sdk"; +import { useMemo } from "react"; import { isRouteErrorResponse, Links, @@ -16,6 +17,7 @@ import { import type { Route } from "./+types/root"; import { ThemeProvider } from "./components/theme-provider"; +import { createFlueFetch } from "./lib/flue-transport"; export const links: Route.LinksFunction = () => [ { href: "https://fonts.googleapis.com", rel: "preconnect" }, @@ -31,87 +33,33 @@ export const links: Route.LinksFunction = () => [ ]; const flueBaseUrl = new URL(env.VITE_FLUE_URL); +const flueFetch = createFlueFetch({ baseUrl: flueBaseUrl }); -// --------------------------------------------------------------------------- -// Authenticated Flue request headers. -// -// The access token is the Better Auth/Convex JWT, resolved once per session -// through the installed `convex.token` plugin endpoint and cached. The stable -// per-send request id is generated by the chat hook immediately before a user -// sends a message and is reused across fetch/stream retries for that send, so -// it is a reliable idempotency key for server-side evidence capture. It is -// reset only when a new message is sent. -// --------------------------------------------------------------------------- - -let cachedAccessToken: string | null = null; -let currentSendRequestId: string | null = null; - -/** Set the idempotency request id for the next send. Called by the chat hook. */ -export const setSendRequestId = (id: string): void => { - currentSendRequestId = id; -}; - -const resolveAccessToken = async (): Promise => { - if (cachedAccessToken) { - return cachedAccessToken; - } - const { data } = await authClient.convex.token({ - fetchOptions: { throw: false }, - }); - cachedAccessToken = data?.token ?? null; - return cachedAccessToken; -}; - -const resolveFlueHeaders = async (): Promise> => { - const headers: Record = {}; - const accessToken = await resolveAccessToken(); - if (accessToken) { - headers.authorization = `Bearer ${accessToken}`; - } - if (currentSendRequestId) { - headers["x-zopu-request-id"] = currentSendRequestId; - } - return headers; -}; - -const flueFetch: typeof fetch = async (input, init) => { - const response = await fetch(input, init); - if ( - !response.ok || - !response.headers.get("content-type")?.includes("application/json") - ) { - return response; - } - - const body = (await response.clone().json()) as unknown; - if ( - typeof body !== "object" || - body === null || - !("streamUrl" in body) || - typeof body.streamUrl !== "string" - ) { - return response; - } - - const streamUrl = new URL(body.streamUrl); - streamUrl.protocol = flueBaseUrl.protocol; - streamUrl.host = flueBaseUrl.host; - const headers = new Headers(response.headers); - headers.set("location", streamUrl.toString()); - return Response.json( - { ...body, streamUrl: streamUrl.toString() }, - { - headers, - status: response.status, - statusText: response.statusText, - } +const AuthenticatedFlueProvider = ({ + children, +}: { + children: React.ReactNode; +}) => { + const resolveAccessToken = useConvexAccessToken(); + const client = useMemo( + () => + createFlueClient({ + baseUrl: flueBaseUrl.toString(), + fetch: flueFetch, + headers: async () => { + const accessToken = await resolveAccessToken(); + const headers: Record = {}; + if (accessToken) { + headers.authorization = `Bearer ${accessToken}`; + } + return headers; + }, + }), + [resolveAccessToken] ); + + return {children}; }; -export const flueClient = createFlueClient({ - baseUrl: flueBaseUrl.toString(), - fetch: flueFetch, - headers: resolveFlueHeaders, -}); export const Layout = ({ children }: { children: React.ReactNode }) => ( @@ -129,8 +77,8 @@ export const Layout = ({ children }: { children: React.ReactNode }) => ( ); const App = () => ( - - + + ( - - + + ); export default App; diff --git a/apps/web/tsconfig.json b/apps/web/tsconfig.json index 41cab95..b4317ca 100644 --- a/apps/web/tsconfig.json +++ b/apps/web/tsconfig.json @@ -1,7 +1,7 @@ { "include": ["**/*", "**/.server/**/*", "**/.client/**/*", ".react-router/types/**/*"], "compilerOptions": { - "lib": ["DOM", "DOM.Iterable", "ES2022"], + "lib": ["DOM", "DOM.Iterable", "ES2024"], "types": ["node", "vite/client"], "target": "ES2022", "module": "ES2022", diff --git a/bun.lock b/bun.lock index 50e06ab..e142821 100644 --- a/bun.lock +++ b/bun.lock @@ -140,6 +140,7 @@ "typescript": "catalog:", "vite": "catalog:", "vite-tsconfig-paths": "^6.1.1", + "vitest": "catalog:", }, }, "packages/agents": { @@ -186,6 +187,7 @@ "@types/react": "~19.2.17", "typescript": "catalog:", "vite": "catalog:", + "vitest": "catalog:", }, }, "packages/backend": { diff --git a/packages/auth/package.json b/packages/auth/package.json index 1401c6c..ce61425 100644 --- a/packages/auth/package.json +++ b/packages/auth/package.json @@ -32,6 +32,7 @@ "@code/config": "workspace:*", "@types/react": "~19.2.17", "vite": "catalog:", - "typescript": "catalog:" + "typescript": "catalog:", + "vitest": "catalog:" } } diff --git a/packages/auth/src/web/hooks.ts b/packages/auth/src/web/hooks.ts index 2ee230b..692fac7 100644 --- a/packages/auth/src/web/hooks.ts +++ b/packages/auth/src/web/hooks.ts @@ -1,9 +1,10 @@ import { useForm } from "@tanstack/react-form"; -import { useCallback, useRef, useState } from "react"; +import { useCallback, useRef } from "react"; import { toast } from "sonner"; import { signInSchema, signUpSchema } from "../shared/forms"; import { authClient } from "./auth-client"; +import { createSessionTokenResolver } from "./session-token-resolver"; interface AuthFormOptions { readonly onSuccess?: () => void; @@ -56,41 +57,27 @@ export const useSignUpForm = ({ onSuccess }: AuthFormOptions = {}) => }); /** - * Resolve the current Better Auth/Convex JWT access token for the signed-in - * session. Uses the installed `convex.token` plugin endpoint (the supported - * API) and caches the result so repeated fetches within a session reuse it. - * Returns null when the user is not authenticated. + * Resolve the Better Auth/Convex JWT for the current signed-in session. The + * installed `convex.token` plugin endpoint remains the token source; changing + * or clearing the Better Auth session invalidates the cached value first. */ export const useConvexAccessToken = (): (() => Promise) => { - const cachedTokenRef = useRef(null); - const pendingRef = useRef | null>(null); - const [, forceRender] = useState(0); + const { data: session } = authClient.useSession(); + const sessionId = session?.session.id ?? null; + const resolverRef = useRef | null>(null); + resolverRef.current ??= createSessionTokenResolver(); + const resolveForSession = resolverRef.current; - return useCallback(async (): Promise => { - if (cachedTokenRef.current) { - return cachedTokenRef.current; - } - if (pendingRef.current) { - return pendingRef.current; - } - pendingRef.current = (async (): Promise => { - try { + return useCallback( + () => + resolveForSession(sessionId, async () => { const { data } = await authClient.convex.token({ fetchOptions: { throw: false }, }); - const token = data?.token ?? null; - cachedTokenRef.current = token; - if (!token) { - forceRender((n) => n + 1); - } - return token; - } catch { - cachedTokenRef.current = null; - return null; - } finally { - pendingRef.current = null; - } - })(); - return await pendingRef.current; - }, []); + return data?.token ?? null; + }), + [resolveForSession, sessionId] + ); }; diff --git a/packages/auth/src/web/session-token-resolver.test.ts b/packages/auth/src/web/session-token-resolver.test.ts new file mode 100644 index 0000000..b41dd6a --- /dev/null +++ b/packages/auth/src/web/session-token-resolver.test.ts @@ -0,0 +1,39 @@ +import { describe, expect, test, vi } from "vitest"; + +import { createSessionTokenResolver } from "./session-token-resolver"; + +describe("createSessionTokenResolver", () => { + test("logout clears the cached JWT before another session resolves", async () => { + const resolveToken = createSessionTokenResolver(); + const loadFirstToken = vi.fn(() => Promise.resolve("jwt-for-user-a")); + + await expect(resolveToken("session-a", loadFirstToken)).resolves.toBe( + "jwt-for-user-a" + ); + await expect( + resolveToken("session-a", () => Promise.resolve("unexpected-token")) + ).resolves.toBe("jwt-for-user-a"); + expect(loadFirstToken).toHaveBeenCalledTimes(1); + + const loadWhileLoggedOut = vi.fn(() => Promise.resolve("stale-jwt")); + await expect(resolveToken(null, loadWhileLoggedOut)).resolves.toBeNull(); + expect(loadWhileLoggedOut).not.toHaveBeenCalled(); + + const loadSecondToken = vi.fn(() => Promise.resolve("jwt-for-user-b")); + await expect(resolveToken("session-b", loadSecondToken)).resolves.toBe( + "jwt-for-user-b" + ); + expect(loadSecondToken).toHaveBeenCalledTimes(1); + }); + + test("a direct user switch never returns the previous session token", async () => { + const resolveToken = createSessionTokenResolver(); + + await expect( + resolveToken("session-a", () => Promise.resolve("jwt-for-user-a")) + ).resolves.toBe("jwt-for-user-a"); + await expect( + resolveToken("session-b", () => Promise.resolve("jwt-for-user-b")) + ).resolves.toBe("jwt-for-user-b"); + }); +}); diff --git a/packages/auth/src/web/session-token-resolver.ts b/packages/auth/src/web/session-token-resolver.ts new file mode 100644 index 0000000..1fc1122 --- /dev/null +++ b/packages/auth/src/web/session-token-resolver.ts @@ -0,0 +1,59 @@ +type TokenLoader = () => Promise; + +/** Cache one Better Auth/Convex JWT per concrete session. */ +export const createSessionTokenResolver = () => { + let cachedSessionId: string | null = null; + let cachedToken: string | null = null; + let pending: { + readonly sessionId: string; + readonly promise: Promise; + } | null = null; + + return async ( + sessionId: string | null, + loadToken: TokenLoader + ): Promise => { + if (!sessionId) { + cachedSessionId = null; + cachedToken = null; + pending = null; + return null; + } + + if (cachedSessionId !== sessionId) { + cachedSessionId = sessionId; + cachedToken = null; + pending = null; + } + if (cachedToken) { + return cachedToken; + } + if (pending?.sessionId === sessionId) { + return pending.promise; + } + + const requestedSessionId = sessionId; + const promise = (async (): Promise => { + try { + const token = await loadToken(); + if (cachedSessionId === requestedSessionId) { + cachedToken = token; + } + return token; + } catch { + if (cachedSessionId === requestedSessionId) { + cachedToken = null; + } + return null; + } + })(); + pending = { promise, sessionId: requestedSessionId }; + try { + return await promise; + } finally { + if (pending?.promise === promise) { + pending = null; + } + } + }; +};