fix(studio): ensure first chat tool call starts in session sandbox (#4810)

Fixes #4809

On a new Studio chat, the first tool call could start before the frontend
initializes the thread ID. That meant the first request could go out without
a session_id, so the backend started the tool in the shared sandbox root
instead of the chat's session sandbox.

Frontend:
- Eagerly initialize the thread when switching to a new chat
- Resolve the thread ID once at request time and keep it stable through
  async model-load waits
- Disable ActiveThreadSync during new-chat initialization to prevent
  stale thread IDs from being written back
- Add error handling for thread initialization failures
- Clear activeThreadId on all compare-mode entry paths to prevent
  cross-session leakage
- Fix exitCompare to restore context usage from the saved view
- Coerce falsy thread IDs to undefined for consistent backend/frontend
  fallback behavior
- Use _default as the image sessionId fallback to match the backend

Backend:
- Use ~/studio_sandbox/_default when a request arrives without a session_id

studio/backend/core/inference/tools.py

@@ -34,7 +34,8 @@ _MAX_OUTPUT_CHARS = 8000  # truncate long output
 _BASH_BLOCKED_WORDS = {"rm", "sudo", "dd", "chmod", "mkfs", "shutdown", "reboot"}
 
 # Per-session working directories so each chat thread gets its own sandbox.
-# Falls back to a shared ~/studio_sandbox/ for API callers without a session_id.
+# Falls back to a shared ~/studio_sandbox/_default for API callers without a
+# session_id.
 _workdirs: dict[str, str] = {}
 
 
@@ -55,7 +56,7 @@ def _get_workdir(session_id: str | None = None) -> str:
             if not os.path.realpath(workdir).startswith(os.path.realpath(sandbox_root)):
                 workdir = os.path.join(sandbox_root, "_invalid")
         else:
-            workdir = sandbox_root
+            workdir = os.path.join(sandbox_root, "_default")
         os.makedirs(workdir, exist_ok = True)
         _workdirs[key] = workdir
     return _workdirs[key]

studio/frontend/src/features/chat/api/chat-adapter.ts

@@ -421,6 +421,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
   return {
     async *run({ messages, abortSignal, unstable_threadId }) {
       let runtime = useChatRuntimeStore.getState();
+      // Capture the thread ID once at the start so it stays stable even if
+      // the user switches chats while waiting for model load / auto-load.
+      const resolvedThreadId =
+        (unstable_threadId ?? runtime.activeThreadId) || undefined;
 
       // Wait for in-progress model load to finish before inferring
       if (runtime.modelLoading) {
@@ -473,14 +477,14 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
         }
         runtime.clearPendingAudio();
       }
-      const useAdapter = await resolveUseAdapter(unstable_threadId);
+      const useAdapter = await resolveUseAdapter(resolvedThreadId);
 
       // ── Audio model path (non-streaming) ─────────────────────
       const activeModel = runtime.models.find(
         (m) => m.id === params.checkpoint,
       );
       if (activeModel?.isAudio && !activeModel?.hasAudioInput) {
-        const threadKey = unstable_threadId || "__default";
+        const threadKey = resolvedThreadId || "__default";
         runtime.setThreadRunning(threadKey, true);
         try {
           yield {
@@ -527,7 +531,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
         return;
       }
 
-      const threadKey = unstable_threadId || "__default";
+      const threadKey = resolvedThreadId || "__default";
       let waitingFirstChunk = true;
       let firstTokenSettled = false;
       const streamStartTime = Date.now();
@@ -600,7 +604,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                     const mins = useChatRuntimeStore.getState().toolCallTimeout;
                     return mins >= 9999 ? 9999 : mins * 60;
                   })(),
-                  session_id: unstable_threadId || undefined,
+                  session_id: resolvedThreadId,
                 }
               : {}),
           },
@@ -641,7 +645,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                 let parsedResult: string | { text: string; images: string[]; sessionId: string };
                 if (imgIdx !== -1) {
                   const text = rawResult.slice(0, imgIdx);
-                  const sessionId = unstable_threadId || "";
+                  // Fall back to "_default" to match the backend sandbox directory
+                  // used when no session_id is provided (see tools.py _get_workdir).
+                  const sessionId = resolvedThreadId || "_default";
                   try {
                     const images = JSON.parse(rawResult.slice(imgIdx + imgMarker.length)) as string[];
                     parsedResult = { text, images, sessionId };

studio/frontend/src/features/chat/chat-page.tsx

@@ -592,6 +592,9 @@ export function ChatPage(): ReactElement {
   }, []);
   const handleNewCompare = useCallback(() => {
     setView({ mode: "compare", pairId: crypto.randomUUID() });
+    // Clear activeThreadId so compare panes do not inherit the single-chat
+    // thread ID as a fallback for session_id routing.
+    useChatRuntimeStore.getState().setActiveThreadId(null);
     useChatRuntimeStore.getState().setContextUsage(null);
   }, []);
 
@@ -619,6 +622,9 @@ export function ChatPage(): ReactElement {
   const enterCompare = useCallback(() => {
     setViewBeforeCompare((prev) => prev ?? view);
     setView({ mode: "compare", pairId: crypto.randomUUID() });
+    // Clear activeThreadId so compare panes do not inherit the single-chat
+    // thread ID as a fallback for session_id routing.
+    useChatRuntimeStore.getState().setActiveThreadId(null);
     useChatRuntimeStore.getState().setContextUsage(null);
   }, [view]);
 
@@ -626,9 +632,13 @@ export function ChatPage(): ReactElement {
     if (!viewBeforeCompare) return;
     setView(viewBeforeCompare);
     setViewBeforeCompare(null);
-    // Restore context usage from the active thread's last assistant message
+    // Restore context usage from the active thread's last assistant message.
+    // Use the thread ID from the saved view rather than the store, because
+    // activeThreadId may have been cleared on compare entry.
     const store = useChatRuntimeStore.getState();
-    const threadId = store.activeThreadId;
+    const threadId =
+      ("threadId" in viewBeforeCompare ? viewBeforeCompare.threadId : null) ??
+      store.activeThreadId;
     if (threadId) {
       void db.messages
         .where("threadId")
@@ -735,6 +745,7 @@ export function ChatPage(): ReactElement {
           await selectModelRef.current({ id: targetLora.id, isLora: true });
           if (canceled) return;
           setView({ mode: "compare", pairId: crypto.randomUUID() });
+          useChatRuntimeStore.getState().setActiveThreadId(null);
           useChatRuntimeStore.getState().setContextUsage(null);
           clearHandoff();
           console.info("[chat-handoff] loaded lora + opened compare");

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

@@ -679,9 +679,33 @@ function ThreadNewChatSwitch({
   const isLoading = useAuiState(({ threads }) => threads.isLoading);
 
   useEffect(() => {
-    if (!isLoading) {
-      aui.threads().switchToNewThread();
+    if (isLoading) {
+      return;
     }
+
+    let cancelled = false;
+    // Clear immediately so the adapter never picks up a stale thread ID
+    // from a previous chat while we initialize the new one.
+    useChatRuntimeStore.getState().setActiveThreadId(null);
+
+    void (async () => {
+      try {
+        aui.threads().switchToNewThread();
+        const { remoteId } = await aui.threadListItem().initialize();
+        if (!cancelled) {
+          useChatRuntimeStore.getState().setActiveThreadId(remoteId);
+        }
+      } catch (error) {
+        if (!cancelled) {
+          useChatRuntimeStore.getState().setActiveThreadId(null);
+        }
+        console.error("Failed to initialize new chat thread", error);
+      }
+    })();
+
+    return () => {
+      cancelled = true;
+    };
   }, [aui, isLoading, nonce]);
 
   return null;
@@ -730,7 +754,7 @@ export function ChatRuntimeProvider({
 
   return (
     <AssistantRuntimeProvider runtime={runtime} aui={aui}>
-      <ActiveThreadSync enabled={modelType === "base" && !pairId} />
+      <ActiveThreadSync enabled={modelType === "base" && !pairId && !newThreadNonce} />
       {initialThreadId && <ThreadAutoSwitch threadId={initialThreadId} />}
       {!initialThreadId && newThreadNonce && (
         <ThreadNewChatSwitch nonce={newThreadNonce} />