Maximus upgradus

hooks/useClaudeUsage.ts

@@ -4,15 +4,18 @@ import {
   loadClaudeSessionKey,
   saveClaudeSessionKey,
   clearClaudeSessionKey,
+  loadClaudeLastActiveOrg,
+  saveClaudeLastActiveOrg,
+  clearClaudeLastActiveOrg,
 } from "@/lib/storage";
-import type { ClaudeUsageResponse } from "@/types/claude";
+import type { ClaudeAuth, ClaudeUsageResponse } from "@/types/claude";
 
 type Status = "idle" | "loading" | "success" | "error";
 
 async function doFetch(
-  sessionKey: string,
+  auth: ClaudeAuth,
   set: {
-    sessionKey: (v: string | null) => void;
+    auth: (v: ClaudeAuth | null) => void;
     usage: (v: ClaudeUsageResponse | null) => void;
     status: (v: Status) => void;
     error: (v: string | null) => void;
@@ -21,10 +24,13 @@ async function doFetch(
   set.status("loading");
   set.error(null);
   try {
-    const orgs = await fetchClaudeOrgs(sessionKey);
-    if (!orgs.length) throw new Error("NO_ORGS_FOUND");
-    const orgUuid = orgs[0].uuid;
-    const data = await fetchClaudeUsage(sessionKey, orgUuid);
+    let orgUuid = auth.lastActiveOrg;
+    if (!orgUuid) {
+      const orgs = await fetchClaudeOrgs(auth);
+      if (!orgs.length) throw new Error("NO_ORGS_FOUND");
+      orgUuid = orgs[0].uuid;
+    }
+    const data = await fetchClaudeUsage(auth, orgUuid);
     set.usage(data);
     set.status("success");
   } catch (e: unknown) {
@@ -33,29 +39,32 @@ async function doFetch(
     set.status("error");
     if (msg === "TOKEN_EXPIRED") {
       await clearClaudeSessionKey();
-      set.sessionKey(null);
+      set.auth(null);
     }
   }
 }
 
 export function useClaudeUsage() {
-  const [sessionKey, setSessionKey] = useState<string | null>(null);
+  const [auth, setAuth] = useState<ClaudeAuth | null>(null);
   const [pendingKey, setPendingKey] = useState("");
+  const [pendingOrg, setPendingOrg] = useState("");
   const [usage, setUsage] = useState<ClaudeUsageResponse | null>(null);
   const [status, setStatus] = useState<Status>("idle");
   const [error, setError] = useState<string | null>(null);
 
-  const setters = { sessionKey: setSessionKey, usage: setUsage, status: setStatus, error: setError };
+  const setters = { auth: setAuth, usage: setUsage, status: setStatus, error: setError };
 
   useEffect(() => {
-    loadClaudeSessionKey().then((stored) => {
-      if (stored) setSessionKey(stored);
-    });
+    Promise.all([loadClaudeSessionKey(), loadClaudeLastActiveOrg()]).then(
+      ([key, org]) => {
+        if (key) setAuth({ sessionKey: key, lastActiveOrg: org ?? undefined });
+      }
+    );
   }, []);
 
   useEffect(() => {
-    if (sessionKey) void doFetch(sessionKey, setters);
-  }, [sessionKey]);
+    if (auth) void doFetch(auth, setters);
+  }, [auth]);
 
   const saveKey = useCallback(async () => {
     const trimmed = pendingKey.trim();
@@ -63,47 +72,47 @@ export function useClaudeUsage() {
       setError("Session key must start with sk-ant-");
       return;
     }
+    const orgTrimmed = pendingOrg.trim() || undefined;
     await saveClaudeSessionKey(trimmed);
-    setSessionKey(trimmed);
+    if (orgTrimmed) await saveClaudeLastActiveOrg(orgTrimmed);
+    else await clearClaudeLastActiveOrg();
+    setAuth({ sessionKey: trimmed, lastActiveOrg: orgTrimmed });
     setPendingKey("");
+    setPendingOrg("");
     setError(null);
-  }, [pendingKey]);
+  }, [pendingKey, pendingOrg]);
 
   const refresh = useCallback(async () => {
-    if (!sessionKey) return;
-    await doFetch(sessionKey, setters);
-  }, [sessionKey]);
+    if (!auth) return;
+    await doFetch(auth, setters);
+  }, [auth]);
 
-  // Re-reads the session key from storage and fetches fresh data. Call from
-  // useFocusEffect so the dashboard stays current when the key is saved from
-  // another tab. If the string value is unchanged, React bails out on setState
-  // so we trigger the fetch directly.
   const reload = useCallback(async () => {
-    const stored = await loadClaudeSessionKey();
-    if (!stored) {
-      setSessionKey(null);
+    const [key, org] = await Promise.all([loadClaudeSessionKey(), loadClaudeLastActiveOrg()]);
+    if (!key) {
+      setAuth(null);
       setUsage(null);
       setStatus("idle");
       return;
     }
-    // If already set to the same string, setState bails out and [sessionKey]
-    // effect won't fire, so call the fetch directly.
-    setSessionKey(stored);
-    void doFetch(stored, setters);
+    // Creating a new object always triggers the [auth] useEffect even if values are identical
+    setAuth({ sessionKey: key, lastActiveOrg: org ?? undefined });
   }, []);
 
   const clearKey = useCallback(async () => {
     await clearClaudeSessionKey();
-    setSessionKey(null);
+    setAuth(null);
     setUsage(null);
     setStatus("idle");
     setError(null);
   }, []);
 
   return {
-    sessionKey,
+    auth,
     pendingKey,
     setPendingKey,
+    pendingOrg,
+    setPendingOrg,
     usage,
     status,
     error,