feat(chat): add connecting status during organization bootstrap

apps/web/package.json

@@ -40,6 +40,7 @@
     "tailwindcss": "catalog:",
     "typescript": "catalog:",
     "vite": "catalog:",
+    "vitest": "catalog:",
     "vite-tsconfig-paths": "^6.1.1"
   }
 }

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 = (
       <>
         <LoaderCircle className="size-3.5 animate-spin" />
-        Zopu is responding
+        {status === "connecting" ? "Preparing Zopu" : "Zopu is responding"}
       </>
     );
   }

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");
+  });
+});

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
+  // The agent session is not created until the authenticated user's personal
-  // organization id. This is the hard tenancy boundary: the server rejects an
+  // organization has been ensured and its id is available.
-  // instance id that is not the caller's own organization.
+  const organization = usePersonalOrganization();
-  const organization = useQuery(api.organizations.getCurrent);
   const agent = useFlueAgent({
     ...CHAT_AGENT,
-    id: organization?._id,
+    id: organization.organizationId,
   });
 
   const sendMessage = async (message: string): Promise<void> => {
-    // Generate a stable idempotency request id once per send, immediately
+    if (!organization.organizationId) {
-    // before the message is submitted. It is reused across fetch/stream
+      throw (
-    // retries for this send and reset only on the next send, so the server
+        organization.error ??
-    // can capture normalized evidence keyed by (organization, request id)
+        new Error("Personal organization is still being prepared")
-    // without creating duplicates.
+      );
-    setSendRequestId(crypto.randomUUID());
+    }
     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,
   };
 };

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({});
+  });
+});

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<OrganizationBootstrapState | null>(
+    null
+  );
+
+  useEffect(() => {
+    if (!sessionId) {
+      return;
+    }
+
+    let active = true;
+    const bootstrapOrganization = async (): Promise<void> => {
+      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,
+  };
+};

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<Id<"projects"> | null>(null);

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<Response>();
+    const secondResponse = Promise.withResolvers<Response>();
+    const bothStarted = Promise.withResolvers<boolean>();
+    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);
+  });
+});

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,
+      }
+    );
+  };
+};

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 });
 
-// ---------------------------------------------------------------------------
+const AuthenticatedFlueProvider = ({
-// Authenticated Flue request headers.
+  children,
-//
+}: {
-// The access token is the Better Auth/Convex JWT, resolved once per session
+  children: React.ReactNode;
-// 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
+  const resolveAccessToken = useConvexAccessToken();
-// sends a message and is reused across fetch/stream retries for that send, so
+  const client = useMemo(
-// it is a reliable idempotency key for server-side evidence capture. It is
+    () =>
-// reset only when a new message is sent.
+      createFlueClient({
-// ---------------------------------------------------------------------------
+        baseUrl: flueBaseUrl.toString(),
-
+        fetch: flueFetch,
-let cachedAccessToken: string | null = null;
+        headers: async () => {
-let currentSendRequestId: string | null = null;
+          const accessToken = await resolveAccessToken();
-
+          const headers: Record<string, string> = {};
-/** Set the idempotency request id for the next send. Called by the chat hook. */
+          if (accessToken) {
-export const setSendRequestId = (id: string): void => {
+            headers.authorization = `Bearer ${accessToken}`;
-  currentSendRequestId = id;
+          }
-};
+          return headers;
-
+        },
-const resolveAccessToken = async (): Promise<string | null> => {
+      }),
-  if (cachedAccessToken) {
+    [resolveAccessToken]
-    return cachedAccessToken;
-  }
-  const { data } = await authClient.convex.token({
-    fetchOptions: { throw: false },
-  });
-  cachedAccessToken = data?.token ?? null;
-  return cachedAccessToken;
-};
-
-const resolveFlueHeaders = async (): Promise<Record<string, string>> => {
-  const headers: Record<string, string> = {};
-  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,
-    }
   );
+
+  return <FlueProvider client={client}>{children}</FlueProvider>;
 };
-export const flueClient = createFlueClient({
-  baseUrl: flueBaseUrl.toString(),
-  fetch: flueFetch,
-  headers: resolveFlueHeaders,
-});
 export const Layout = ({ children }: { children: React.ReactNode }) => (
   <html lang="en">
     <head>
@@ -129,8 +77,8 @@ export const Layout = ({ children }: { children: React.ReactNode }) => (
 );
 
 const App = () => (
-  <FlueProvider client={flueClient}>
+  <WebAuthProvider>
-    <WebAuthProvider>
+    <AuthenticatedFlueProvider>
       <ThemeProvider
         attribute="class"
         defaultTheme="dark"
@@ -142,8 +90,8 @@ const App = () => (
         </div>
         <Toaster richColors />
       </ThemeProvider>
-    </WebAuthProvider>
+    </AuthenticatedFlueProvider>
-  </FlueProvider>
+  </WebAuthProvider>
 );
 
 export default App;

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",

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": {

packages/auth/package.json

@@ -32,6 +32,7 @@
     "@code/config": "workspace:*",
     "@types/react": "~19.2.17",
     "vite": "catalog:",
-    "typescript": "catalog:"
+    "typescript": "catalog:",
+    "vitest": "catalog:"
   }
 }

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
+ * Resolve the Better Auth/Convex JWT for the current signed-in session. The
- * session. Uses the installed `convex.token` plugin endpoint (the supported
+ * installed `convex.token` plugin endpoint remains the token source; changing
- * API) and caches the result so repeated fetches within a session reuse it.
+ * or clearing the Better Auth session invalidates the cached value first.
- * Returns null when the user is not authenticated.
  */
 export const useConvexAccessToken = (): (() => Promise<string | null>) => {
-  const cachedTokenRef = useRef<string | null>(null);
+  const { data: session } = authClient.useSession();
-  const pendingRef = useRef<Promise<string | null> | null>(null);
+  const sessionId = session?.session.id ?? null;
-  const [, forceRender] = useState(0);
+  const resolverRef = useRef<ReturnType<
+    typeof createSessionTokenResolver
+  > | null>(null);
+  resolverRef.current ??= createSessionTokenResolver();
+  const resolveForSession = resolverRef.current;
 
-  return useCallback(async (): Promise<string | null> => {
+  return useCallback(
-    if (cachedTokenRef.current) {
+    () =>
-      return cachedTokenRef.current;
+      resolveForSession(sessionId, async () => {
-    }
-    if (pendingRef.current) {
-      return pendingRef.current;
-    }
-    pendingRef.current = (async (): Promise<string | null> => {
-      try {
         const { data } = await authClient.convex.token({
           fetchOptions: { throw: false },
         });
-        const token = data?.token ?? null;
+        return data?.token ?? null;
-        cachedTokenRef.current = token;
+      }),
-        if (!token) {
+    [resolveForSession, sessionId]
-          forceRender((n) => n + 1);
+  );
-        }
-        return token;
-      } catch {
-        cachedTokenRef.current = null;
-        return null;
-      } finally {
-        pendingRef.current = null;
-      }
-    })();
-    return await pendingRef.current;
-  }, []);
 };

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");
+  });
+});

packages/auth/src/web/session-token-resolver.ts

@@ -0,0 +1,59 @@
+type TokenLoader = () => Promise<string | null>;
+
+/** 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<string | null>;
+  } | null = null;
+
+  return async (
+    sessionId: string | null,
+    loadToken: TokenLoader
+  ): Promise<string | null> => {
+    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<string | null> => {
+      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;
+      }
+    }
+  };
+};