From 41e43b6c7bd83f788e7cc44a81f030d932cf7d0d Mon Sep 17 00:00:00 2001
From: Roland Tannous <rolandtannous@gravityq.ai>
Date: Wed, 27 May 2026 21:33:51 +0400
Subject: [PATCH] Studio: persist activeThreadId across reloads so RAG docs
 survive

When the user uploads a doc to a brand-new chat (a draft thread with
an assistant-ui __LOCALID_* id), the backend stores rag_documents
rows scoped to that id. On page reload assistant-ui mints a fresh
__LOCALID_* for the new mainThreadId, so useThreadDocuments asks the
backend for docs under the NEW id and gets nothing, even though the
original rows are still on disk under the OLD id.

  - chat-runtime-store: persist activeThreadId in localStorage via
    a new CHAT_ACTIVE_THREAD_KEY, restore on init, save on every
    setActiveThreadId call (including clears, which write '').
  - ActiveThreadSync: when aui's freshly-minted mainThreadId is a
    __LOCALID_* and we already have a persisted __LOCALID_* draft,
    keep ours instead of overwriting. This only affects RAG/doc
    lookup; aui's chat history for the new draft starts empty
    either way, so there's no regression for users who don't have
    attached docs.

User-initiated thread switches (new chat, switching to a saved
thread, deleting the current thread) all go through setActiveThreadId
with the new id (or null), so they correctly replace/clear the
persisted value.
---
 .../src/features/chat/runtime-provider.tsx        | 15 +++++++++++++++
 .../features/chat/stores/chat-runtime-store.ts    | 11 ++++++++---
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx
index d01383b309..bdf4e528d8 100644
--- a/studio/frontend/src/features/chat/runtime-provider.tsx
+++ b/studio/frontend/src/features/chat/runtime-provider.tsx
@@ -1017,6 +1017,21 @@ function ActiveThreadSync({
     if (!enabled) {
       return;
     }
+    const persisted = useChatRuntimeStore.getState().activeThreadId;
+    // On page reload aui mints a fresh `__LOCALID_*` for its draft thread.
+    // If we already have a persisted draft id, keep it so any RAG docs
+    // attached to that draft id stay visible — assistant-ui's internal
+    // chat history for the new mainThreadId starts empty either way, so
+    // this only affects what we treat as the active id for doc/RAG lookup.
+    if (
+      mainThreadId
+      && persisted
+      && persisted !== mainThreadId
+      && persisted.startsWith("__LOCALID_")
+      && mainThreadId.startsWith("__LOCALID_")
+    ) {
+      return;
+    }
     setActiveThreadId(mainThreadId ?? null);
   }, [enabled, mainThreadId, setActiveThreadId]);
 
diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
index 1d1df695eb..05c77de896 100644
--- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
+++ b/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,9 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
   defaultChatTemplate: null,
   chatTemplateOverride: null,
   loadedChatTemplateOverride: null,
-  activeThreadId: null,
+  // Persisted across reloads so RAG documents uploaded under a draft
+  // thread id stay attached after the page refreshes.
+  activeThreadId: loadString(CHAT_ACTIVE_THREAD_KEY, "") || null,
   settingsPanelOpen: false,
   pendingAudioBase64: null,
   pendingAudioName: null,
@@ -755,8 +758,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