feat(chat): add connecting status during organization bootstrap

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
- * 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<string | null>) => {
-  const cachedTokenRef = useRef<string | null>(null);
-  const pendingRef = useRef<Promise<string | null> | null>(null);
-  const [, forceRender] = useState(0);
+  const { data: session } = authClient.useSession();
+  const sessionId = session?.session.id ?? null;
+  const resolverRef = useRef<ReturnType<
+    typeof createSessionTokenResolver
+  > | null>(null);
+  resolverRef.current ??= createSessionTokenResolver();
+  const resolveForSession = resolverRef.current;
 
-  return useCallback(async (): Promise<string | null> => {
-    if (cachedTokenRef.current) {
-      return cachedTokenRef.current;
-    }
-    if (pendingRef.current) {
-      return pendingRef.current;
-    }
-    pendingRef.current = (async (): Promise<string | null> => {
-      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]
+  );
 };

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