Studio: restore draft thread (and its RAG docs) across page reloads

assistant-ui mints a fresh __LOCALID_* draft id on every page load, so
RAG docs uploaded under the previous draft id were orphaned after a
refresh — the doc panel queries useThreadDocuments(activeThreadId) and
the new id had nothing.

Persist activeThreadId in localStorage and, on the first settled render
after load, have ActiveThreadSync ask aui to switchToThread(persisted)
when it differs from the freshly-minted draft. Because the draft was
already persisted to the backend by initialize()/ensureThreadRecord
when its first doc was uploaded, the adapter's fetch() resolves it and
aui adopts it as mainThreadId. That keeps aui's mainThreadId and our
activeThreadId unified, so the earlier divergence (uploads under the
persisted id vs chat-completion reading aui's fresh id) can't recur —
unlike the reverted localStorage-only attempt, the chat-adapter's
unstable_threadId now equals the persisted id after the switch.

A one-shot ref ensures we only re-adopt on initial load; user-driven
new-chat / thread switches still flow through normally. If the
persisted draft was never initialized (no doc/message, not in the
backend), switchToThread rejects and we fall back to the fresh draft.

studio/frontend/src/features/chat/runtime-provider.tsx

@@ -1008,17 +1008,51 @@ function ThreadNewChatSwitch({
 function ActiveThreadSync({
   enabled,
 }: { enabled: boolean }): ReactElement | null {
+  const aui = useAui();
+  const isLoading = useAuiState(({ threads }) => threads.isLoading);
   const mainThreadId = useAuiState(({ threads }) => threads.mainThreadId);
   const setActiveThreadId = useChatRuntimeStore(
     (state) => state.setActiveThreadId,
   );
+  // One-shot guard: we only attempt to re-adopt a persisted draft on the
+  // first settled render after a page load, never again for the lifetime
+  // of this provider (so user-driven new-chat / switches aren't fought).
+  const restoreAttemptedRef = useRef(false);
 
   useEffect(() => {
-    if (!enabled) {
+    if (!enabled || isLoading || !mainThreadId) {
       return;
     }
-    setActiveThreadId(mainThreadId ?? null);
-  }, [enabled, mainThreadId, setActiveThreadId]);
+    const persisted = useChatRuntimeStore.getState().activeThreadId;
+    // On a fresh page load aui mints a new `__LOCALID_*` draft. If we have
+    // a different persisted draft id, ask aui to switch to it so any RAG
+    // docs uploaded under that id reattach. The thread was persisted to
+    // the backend when its first doc/message ran initialize(), so the
+    // adapter's fetch() can resolve it. If it can't (a draft that never
+    // got a doc/message), switchToThread rejects and we fall back to the
+    // fresh draft.
+    if (
+      !restoreAttemptedRef.current
+      && persisted
+      && persisted !== mainThreadId
+      && persisted.startsWith("__LOCALID_")
+      && mainThreadId.startsWith("__LOCALID_")
+    ) {
+      restoreAttemptedRef.current = true;
+      const result = aui.threads().switchToThread(persisted) as unknown;
+      if (
+        result
+        && typeof (result as Promise<void>).catch === "function"
+      ) {
+        void (result as Promise<void>).catch(() => {
+          setActiveThreadId(mainThreadId);
+        });
+      }
+      return;
+    }
+    restoreAttemptedRef.current = true;
+    setActiveThreadId(mainThreadId);
+  }, [aui, enabled, isLoading, mainThreadId, setActiveThreadId]);
 
   return null;
 }

studio/frontend/src/features/chat/stores/chat-runtime-store.ts

@@ -36,6 +36,7 @@ export const CHAT_RAG_TOOL_ENABLED_KEY = "unsloth_chat_rag_tool_enabled";
 // localStorage persistence here the user's external pick is silently
 // reset to the default on every page refresh.
 const LAST_EXTERNAL_CHECKPOINT_KEY = "unsloth_chat_last_external_checkpoint";
+const CHAT_ACTIVE_THREAD_KEY = "unsloth_chat_active_thread_id";
 
 function loadLastExternalCheckpoint(): string | null {
   if (typeof window === "undefined") return null;
@@ -610,7 +611,11 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
   defaultChatTemplate: null,
   chatTemplateOverride: null,
   loadedChatTemplateOverride: null,
-  activeThreadId: null,
+  // Persisted so a draft thread's RAG docs reattach after a page reload.
+  // On reload ActiveThreadSync asks assistant-ui to switch to this id
+  // (the thread was persisted to the backend when its first doc/message
+  // initialized it), keeping aui's mainThreadId and this value unified.
+  activeThreadId: loadString(CHAT_ACTIVE_THREAD_KEY, "") || null,
   settingsPanelOpen: false,
   pendingAudioBase64: null,
   pendingAudioName: null,
@@ -755,8 +760,10 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
         activeGgufVariant: ggufVariant ?? null,
       };
     }),
-  setActiveThreadId: (activeThreadId) =>
-    set({ activeThreadId, contextUsage: null }),
+  setActiveThreadId: (activeThreadId) => {
+    saveString(CHAT_ACTIVE_THREAD_KEY, activeThreadId ?? "");
+    set({ activeThreadId, contextUsage: null });
+  },
   setSettingsPanelOpen: (settingsPanelOpen) => set({ settingsPanelOpen }),
   clearCheckpoint: () => {
     // Mirror setCheckpoint's persistence behavior: dropping the