From c027ec192ef5e4e69f4c02973081f5b3da67aa20 Mon Sep 17 00:00:00 2001 From: Neodon <82944+neodon@users.noreply.github.com> Date: Fri, 3 Apr 2026 13:44:22 -0500 Subject: [PATCH] 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 | 5 ++-- .../src/features/chat/api/chat-adapter.ts | 16 ++++++---- .../frontend/src/features/chat/chat-page.tsx | 15 ++++++++-- .../src/features/chat/runtime-provider.tsx | 30 +++++++++++++++++-- 4 files changed, 54 insertions(+), 12 deletions(-) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index b23372b766..86c22ce25a 100644 --- a/studio/backend/core/inference/tools.py +++ b/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] diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index e287daf33a..3d1bce2905 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/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 }; diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 08450c7ec7..1dbff145ee 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/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"); diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index 02f792b509..7e6cd8e1dd 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/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 ( - + {initialThreadId && } {!initialThreadId && newThreadNonce && (