Studio: reopen a chat on the branch it was left on (#13109)

* Studio: reopen a chat on the branch it was left on

* Keep the saved branch when a chat delete fails

* Reopen on the newest leaf below the saved head

* Trim comments in PR scope

---------

Co-authored-by: Maheswar Kumar <110882203+mahiatlinux@users.noreply.github.com>
Co-authored-by: mahiatlinux <mahiatlinux@users.noreply.github.com>

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

@@ -284,7 +284,9 @@ import {
   createAnnotationsFile,
 } from "./utils/document-annotations";
 import { requestTemporaryPromptQueueStop } from "./utils/prompt-queue-boundary";
+import { savedBranchHead } from "./utils/branch-head";
 import { estimateContextUsage } from "./utils/estimate-chat-tokens";
+import { orderBySelectedBranch } from "./utils/message-order";
 import { isAssistantLocalThreadId } from "./utils/thread-ids";
 import {
   consumeProjectSourcesPending,
@@ -4236,10 +4238,9 @@ export function ChatPage({
   }, [currentProjectId, navigate, search]);
 
   const exitCompare = useCallback(() => {
-    // Prefer the explicit save; fall back to the last non-compare view so the composer + menu path
-    // also returns where the user started.
+    // the composer and menu exit paths rely on the last non-compare view.
     const saved = viewBeforeCompareRef.current ?? lastNonCompareViewRef.current;
-    // No saved view (compare opened by direct URL); fall back to a fresh chat.
+    // direct compare URLs have no saved view, so return to a fresh chat.
     if (!saved) {
       navigate({ to: "/chat" });
       return;
@@ -4252,10 +4253,13 @@ export function ChatPage({
       void listStoredChatMessages(threadId)
         .then((messages) => {
           const store = useChatRuntimeStore.getState();
-          const usage = savedUsageFor(messages, store) ?? estimateContextUsage(messages);
+          const branch = orderBySelectedBranch(
+            messages,
+            savedBranchHead(threadId, messages),
+          );
+          const usage = savedUsageFor(branch, store) ?? estimateContextUsage(branch);
           if (!usage) return;
-          // Key by the thread this restore read, like the history loader: the await above can outlast a
-          // switch away, and an unkeyed write would file this usage under the incoming thread.
+          // key usage by the restored thread because this read can outlast a thread switch.
           store.setThreadContextUsage(threadId, usage);
           if (store.activeThreadId === threadId) {
             store.setContextUsage(usage);

studio/frontend/src/features/chat/hooks/use-branch-head-recorder.ts

@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { useAuiState } from "@assistant-ui/react";
+import { useEffect } from "react";
+import { useChatRuntimeStore } from "../stores/chat-runtime-store";
+import { writeBranchHead } from "../utils/branch-head";
+import { isThreadIncognito } from "../utils/chat-history-storage";
+
+// empty threads are loading; `incognito` changes when a temporary chat gets its persistent id.
+export function useBranchHeadRecorder(): void {
+  const remoteId = useAuiState(({ threadListItem }) => threadListItem.remoteId);
+  const headId = useAuiState(({ thread }) => thread.messages.at(-1)?.id);
+  const incognito = useChatRuntimeStore((s) => s.incognito);
+
+  useEffect(() => {
+    if (!remoteId || !headId || isThreadIncognito(remoteId)) return;
+    writeBranchHead(remoteId, headId);
+  }, [headId, incognito, remoteId]);
+}

studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts

@@ -15,6 +15,7 @@ import {
   listStoredChatThreadsWithMessages,
   updateStoredChatThread,
 } from "../utils/chat-history-storage";
+import { clearBranchHead } from "../utils/branch-head";
 import { clearComposerDraft } from "../utils/composer-draft";
 import { offerToDeleteKeptSandboxes } from "../utils/offer-kept-sandbox-files";
 import { stopChatThread } from "../utils/stop-chat-thread";
@@ -344,10 +345,9 @@ export async function deleteChatItems(
     cancelIfRunning(id);
   }
 
-  // Drop saved composer drafts so deleted threads leave no orphan keys.
+  // clear composer drafts so deleted threads leave no orphan keys.
   for (const id of threadIds) clearComposerDraft(id);
 
-  // Optimistic tombstone: hide immediately; roll back on backend error.
   markChatThreadsDeleted(threadIds);
   notifyChatHistoryUpdated();
 
@@ -358,9 +358,9 @@ export async function deleteChatItems(
 
   try {
     const kept = await deleteStoredChatThreads(threadIds, args);
-    // Whether or not deletion was asked for: a sandbox that could not be removed leaves files with
-    // no card to reach them from, and the chat is already gone, so this offer is the only notice
-    // and the only retry.
+    // retain the head until deletion succeeds so a restored chat reopens on the same branch.
+    for (const id of threadIds) clearBranchHead(id);
+    // offer recovery for sandbox files that outlive a deleted chat and lose their sidebar entry.
     offerToDeleteKeptSandboxes(kept);
   } catch (error) {
     removeChatThreadTombstones(threadIds);

studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx

@@ -99,6 +99,7 @@ import {
   ndjsonBody,
   type ConversationJsonlLayout,
 } from "../utils/ndjson";
+import { savedBranchHead } from "../utils/branch-head";
 import { orderByParentChain } from "../utils/message-order";
 import { liveThreadBranch } from "../utils/live-thread-head";
 import { unwrapPastedTextContent } from "../utils/pasted-text.ts";
@@ -243,7 +244,7 @@ async function loadConversationMessages(
     includeSiblings = true,
     includeInstructions = true,
   } = options;
-  // Read before the storage await: switching chats meanwhile would point the lookup at another thread.
+  // read before awaiting storage so a chat switch cannot redirect the lookup
   const liveBranch = liveThreadBranch(threadId);
   const [raw, instructions] = await Promise.all([
     listStoredChatMessages(threadId),
@@ -253,10 +254,10 @@ async function loadConversationMessages(
     toast.info(emptyMessage);
     return null;
   }
-  // No parentId = legacy flat thread (already DB createdAt-sorted); walking the chain would invert order.
+  // preserve DB order for legacy flat threads because parent-chain traversal would reverse it
   const hasParentIds = raw.some((m) => (m as { parentId?: unknown }).parentId != null);
   if (!hasParentIds) return [...instructions, ...raw];
-  const headId = liveBranchHeadId(liveBranch, raw);
+  const headId = branchHeadId(threadId, liveBranch, raw);
   return [
     ...instructions,
     ...orderByParentChain(raw, { includeSiblings, headId }),
@@ -264,7 +265,7 @@ async function loadConversationMessages(
 }
 
 async function chatInstructionsTurn(threadId: string): Promise<MessageRecord[]> {
-  // A debounced or held edit is not on the row yet, and the next reply already runs with it.
+  // pending instruction edits are absent from storage even though the next reply already uses them
   await settleThreadScopedSettingsForCopy(threadId);
   const thread = await getStoredChatThread(threadId);
   if (!thread) return [];
@@ -286,13 +287,15 @@ async function chatInstructionsTurn(threadId: string): Promise<MessageRecord[]>
   ];
 }
 
-// Newest saved turn of the branch on screen: a reply still generating is not stored yet, and falling back to the newest leaf would export the reply it replaces.
-function liveBranchHeadId(
+// use the visible stored turn because a generating reply is absent and the newest leaf may replace it
+function branchHeadId(
+  threadId: string,
   liveBranch: string[] | null,
   raw: Array<{ id: string }>,
 ): string | null | undefined {
-  // An empty list is no opinion, not an empty branch: switching chats sets remoteId before the history load refills the view.
-  if (!liveBranch?.length) return undefined;
+  // an empty live branch means history has not refilled after a chat switch
+  // off-screen chats export the branch that reopening them would show
+  if (!liveBranch?.length) return savedBranchHead(threadId, raw);
   const storedIds = new Set(raw.map((m) => m.id));
   return [...liveBranch].reverse().find((id) => storedIds.has(id)) ?? null;
 }
@@ -301,7 +304,7 @@ function exportTs(): string {
   return new Date().toISOString().slice(0, 19).replace(/:/g, "-");
 }
 
-// Attachments live in msg.attachments[].content, not msg.content, so flatten both here or they'd be dropped on export.
+// attachment content is separate from message content and must be flattened for export
 function messageToText(msg: { content: unknown; attachments?: unknown }): string {
   const parts: string[] = [];
   const main = contentBlocksToText(msg.content);
@@ -1006,11 +1009,11 @@ export async function buildFineTuneJsonl(
     const hasParentIds = raw.some(
       (m) => (m as { parentId?: unknown }).parentId != null,
     );
-    // Chain only: retries/regenerations leave sibling branches, and mixing alternate replies into one conversation corrupts the training targets.
+    // exclude sibling retries because alternate replies would corrupt training targets
     const ordered = hasParentIds
       ? (orderByParentChain(raw, {
           includeSiblings: false,
-          headId: liveBranchHeadId(liveBranch, raw),
+          headId: branchHeadId(id, liveBranch, raw),
         }) as typeof raw)
       : raw;
     const turns = messagesToFineTuneTurns([...instructions, ...ordered]);

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

@@ -186,7 +186,12 @@ import {
   onChatAttachmentDeleted,
 } from "./utils/chat-attachment-events";
 import { chatHistoryClearBoundary } from "./utils/chat-history-clear-boundary";
-import { createParentResolver } from "./utils/message-order";
+import { useBranchHeadRecorder } from "./hooks/use-branch-head-recorder";
+import { savedBranchHead } from "./utils/branch-head";
+import {
+  createParentResolver,
+  orderBySelectedBranch,
+} from "./utils/message-order";
 import { estimateContextUsage } from "./utils/estimate-chat-tokens";
 import {
   awaitStoredChatThreadWrites,
@@ -2666,8 +2671,12 @@ function useStudioRuntimeAdapters(
           }
         }
 
-        // Restore context usage from last assistant message if model matches.
-        const lastAssistant = [...msgs]
+        // select the persisted branch for import and context-usage restoration
+        const headId = savedBranchHead(remoteId, msgs);
+        const branch = orderBySelectedBranch(msgs, headId);
+
+        // restore usage from the selected branch's last assistant message when the model matches
+        const lastAssistant = [...branch]
           .reverse()
           .find((m) => m.role === "assistant");
         const savedUsage = (lastAssistant?.metadata as Record<string, unknown>)
@@ -2683,47 +2692,38 @@ function useStudioRuntimeAdapters(
             }
           | undefined;
         const store = useChatRuntimeStore.getState();
-        // Window check applies only when a local GGUF window is known; external providers have
-        // loadedContextLength === null. llama.cpp stops at the window, so a saved count past it is stale;
-        // MLX runs past it by design, and a thread whose recount is unsupported would never get another.
+        // reject usage beyond a local GGUF window; MLX and external providers have no local bound
         const localLimit = store.loadedIsGguf ? store.loadedContextLength : null;
         const withinLocalLimit =
           !localLimit || (savedUsage?.contextTokens ?? savedUsage?.totalTokens ?? 0) <= localLimit;
-        // Legacy unscoped usage (no modelId) is trusted only when a known local
-        // window bounds the totals, so an old local turn can't be misattributed
-        // to a newly-selected external provider.
+        // trust legacy usage only when a known local window prevents cross-provider attribution
         const modelMatches = savedUsage?.modelId
           ? savedUsage.modelId === store.params.checkpoint
           : typeof store.loadedContextLength === "number" &&
             store.loadedContextLength > 0;
-        // The value, not a boolean: the writes below need the narrowing.
+        // retain the usage object so TypeScript preserves narrowing for later writes
         const restoredUsage =
           savedUsage && withinLocalLimit && modelMatches ? savedUsage : null;
-        const shownUsage = restoredUsage ?? estimateContextUsage(msgs);
+        const shownUsage = restoredUsage ?? estimateContextUsage(branch);
         if (shownUsage) {
-          // Key by the thread this loader read, not whichever is active when the await resolves: a switch
-          // inside it would file this thread's usage under the incoming one.
+          // key usage by the loaded thread because the active thread may change during awaits
           store.setThreadContextUsage(remoteId, shownUsage);
           if (store.activeThreadId === remoteId) {
             store.setContextUsage(shownUsage);
           }
         }
-        // Only when nothing was restored: saved usage is the last completion's exact totals, and
-        // refreshContextUsage does NOT stand down for usage already there, so it would overwrite them with
-        // an estimate whose completionTokens is 0. A thread opened after a model switch fails modelMatches
-        // and still gets priced (#7450). Primary pane only: a compare pane never owns the global bar.
+        // recount only without exact usage; model switches qualify, and compare panes never own the bar (#7450)
         if (!restoredUsage && modelType === "base" && !pairId) {
           void refreshContextUsage({ threadId: remoteId });
         }
 
-        // If any message has a stored parentId, reconstruct the tree so retries load as branches rather
-        // than a flat list, inferring sequential parents for old messages in mixed threads. Fall
-        // back to fromArray for fully legacy threads.
+        // rebuild branches when parentIds exist; infer sequential parents for mixed legacy threads
         const hasParentIds = msgs.some((m) => m.parentId != null);
         if (hasParentIds) {
           const resolveParent = createParentResolver();
           return completeLoad(
             {
+              headId,
               messages: msgs.map((m) => ({
                 parentId: resolveParent(m),
                 message: toThreadMessage(m),
@@ -3604,7 +3604,7 @@ function ActiveBranchRegistrar({
       try {
         return aui.thread().getState().messages;
       } catch {
-        // No thread mounted yet; the recount falls back to the stored records.
+        // without a mounted thread, recount from stored records
         return null;
       }
     });
@@ -3614,9 +3614,13 @@ function ActiveBranchRegistrar({
   return null;
 }
 
-// Price whichever thread the bar points at whenever it has nothing to show. Only two paths reach
-// it: a model change empties contextUsageByThreadId while a mounted thread does not rerun
-// its history loader, and on a deep link the loader and status can each land before the other.
+// include hidden panes because background replies also move the branch head
+function BranchHeadRecorder(): ReactElement | null {
+  useBranchHeadRecorder();
+  return null;
+}
+
+// recount an empty bar after a model change or when deep-link loading races model status
 function ThreadContextUsageRecount({
   enabled,
 }: { enabled: boolean }): ReactElement | null {
@@ -3624,8 +3628,7 @@ function ThreadContextUsageRecount({
   const checkpoint = useChatRuntimeStore((s) => s.params.checkpoint);
   const loadedContextLength = useChatRuntimeStore((s) => s.loadedContextLength);
   const modelLoading = useChatRuntimeStore((s) => s.modelLoading);
-  // A DEPENDENCY, not just a guard: nothing else here changes when a run ends, so a count skipped
-  // for being busy would never be retried. Every run, since that is what the endpoint refuses on.
+  // subscribe to run state so a recount skipped while busy is retried when the run ends
   const runActive = useChatRuntimeStore((s) =>
     Object.values(s.runningByThreadId).some(Boolean),
   );
@@ -4045,6 +4048,7 @@ export function ChatRuntimeProvider({
           newThreadSwitchStateRef={newThreadSwitchStateRef}
         />
         <CancelRegistrar />
+        <BranchHeadRecorder />
         {initialThreadId && (
           <ThreadAutoSwitch
             threadId={initialThreadId}

studio/frontend/src/features/chat/utils/branch-head.ts

@@ -0,0 +1,44 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { type ParentLinkedMessage, resolveSavedBranchHead } from "./message-order";
+
+// storage keeps every branch but not which was open; this prefix joins account-switch purging.
+const BRANCH_HEAD_PREFIX = "unsloth_chat_branch_head:";
+
+function branchHeadKey(threadId: string): string {
+  return `${BRANCH_HEAD_PREFIX}${threadId}`;
+}
+
+// best effort like composer drafts because storage can be blocked or full.
+function readBranchHead(threadId: string): string | null {
+  try {
+    return window.localStorage.getItem(branchHeadKey(threadId));
+  } catch {
+    return null;
+  }
+}
+
+export function writeBranchHead(threadId: string, messageId: string): void {
+  try {
+    window.localStorage.setItem(branchHeadKey(threadId), messageId);
+  } catch {
+    // storage failures are non-fatal
+  }
+}
+
+export function clearBranchHead(threadId: string): void {
+  try {
+    window.localStorage.removeItem(branchHeadKey(threadId));
+  } catch {
+    // storage failures are non-fatal
+  }
+}
+
+/** saved branch head for `threadId`; undefined selects the newest stored row. */
+export function savedBranchHead(
+  threadId: string,
+  messages: ParentLinkedMessage[],
+): string | undefined {
+  return resolveSavedBranchHead(messages, readBranchHead(threadId));
+}

studio/frontend/src/features/chat/utils/message-order.ts

@@ -76,6 +76,44 @@ export function orderBySelectedBranch<T extends ParentLinkedMessage>(
   return chain.reverse();
 }
 
+// remote turns can extend a saved head; use the newest leaf because assistant-ui drops descendants.
+export function resolveSavedBranchHead<T extends ParentLinkedMessage>(
+  messages: T[],
+  savedHeadId: string | null | undefined,
+): string | undefined {
+  if (!savedHeadId) return undefined;
+  const sorted = messages.slice().sort(compareStoredMessages);
+  const resolveParent = createParentResolver();
+  const children = new Map<string, string[]>();
+  const order = new Map<string, number>();
+  sorted.forEach((message, index) => {
+    order.set(message.id, index);
+    const parentId = resolveParent(message);
+    if (parentId == null) return;
+    const siblings = children.get(parentId) ?? [];
+    siblings.push(message.id);
+    children.set(parentId, siblings);
+  });
+  if (!order.has(savedHeadId)) return undefined;
+  let headId = savedHeadId;
+  let headOrder = -1;
+  const seen = new Set<string>();
+  const pending = [savedHeadId];
+  while (pending.length > 0) {
+    const id = pending.pop()!;
+    if (seen.has(id)) continue;
+    seen.add(id);
+    const below = children.get(id);
+    if (below?.length) {
+      pending.push(...below);
+    } else if ((order.get(id) ?? -1) > headOrder) {
+      headId = id;
+      headOrder = order.get(id) ?? -1;
+    }
+  }
+  return headId;
+}
+
 // follow the newest parent chain because response slots can predate the next user message.
 export function orderByParentChain<T extends ParentLinkedMessage>(
   messages: T[],

studio/frontend/src/features/chat/utils/refresh-context-usage.ts

@@ -14,11 +14,11 @@ import { countChatInputTokens } from "../api/chat-api";
 import { isExternalModelId } from "../external-providers";
 import { useChatRuntimeStore } from "../stores/chat-runtime-store";
 import type { MessageRecord } from "../types";
+import { savedBranchHead } from "./branch-head";
 import { listStoredChatMessages } from "./chat-history-storage";
 import { orderBySelectedBranch } from "./message-order";
 
-// Per thread, not per module: in compare mode a hidden pane's history load would otherwise
-// invalidate the visible thread's count, blanking the bar.
+// key generations by thread so a hidden compare pane cannot invalidate the visible count.
 const refreshGenerations = new Map<string | null, number>();
 
 function nextGeneration(threadKey: string | null): number {
@@ -206,10 +206,9 @@ export async function refreshContextUsage(
     const liveBranch = readOwnBranch();
 
     let runMessages: readonly ThreadMessage[];
-    // Re-read before publishing, so a turn sent while this count was in flight drops it.
+    // re-read before publishing so a turn sent during the count invalidates it.
     let countedBranch: string | null = null;
-    // The stored fallback's witness: storage records and ThreadMessages hash differently, so only
-    // ids survive both, and the last one moves as soon as a turn is sent.
+    // only the last id witnesses the fallback because stored and runtime records hash differently.
     let countedLastId: string | null = null;
     const fromLiveBranch = Boolean(liveBranch && liveBranch.length > 0);
     if (fromLiveBranch) {
@@ -217,24 +216,19 @@ export async function refreshContextUsage(
     } else {
       const records = threadId ? await listStoredChatMessages(threadId) : [];
       if (stale()) return;
-      runMessages = orderBySelectedBranch(records).map(
-        storedMessageToRunMessage,
-      );
+      runMessages = orderBySelectedBranch(
+        records,
+        threadId ? savedBranchHead(threadId, records) : undefined,
+      ).map(storedMessageToRunMessage);
     }
 
-    // /chat/count_tokens always 503s on images and /apply-template swaps each for a marker.
-    // Declining before the hash keeps the base64 out of it and out of a request body that can
-    // reach megabytes, both synchronous on the UI thread.
+    // /chat/count_tokens rejects images; hashing or sending their base64 can block the UI thread.
     if (messagesContainImage(runMessages)) return;
 
-    // The real request replays the newest user audio as audio_base64 but toOpenAIMessages has no
-    // audio branch, so counting would price a text-only prompt. Decline as images do.
+    // audio is omitted from toOpenAIMessages, so counting would price only its text.
     if (findLatestUserAudioBase64(runMessages)) return;
 
-    // Same for video: the request replays the clip as video_base64 and llama-server expands it
-    // into frames, while toOpenAIMessages has no video branch, so the bar would show room the
-    // window does not have. Declining also keeps up to 85 MB of base64 out of
-    // branchSignature's JSON.stringify, the same main-thread cost as the image bail.
+    // video counting misses server-side frames; declining also avoids hashing up to 85 MB.
     if (findLatestUserVideoBase64(runMessages)) return;
 
     if (fromLiveBranch) {

studio/frontend/tests/chat-export-system-prompt.test.ts

@@ -159,6 +159,7 @@ function loadExporters(
     threadScopedDefault: (key: string) => INSTALLATION_DEFAULTS[key],
     composerProjectByPendingThread: new Map(),
     ...liveThreadHead,
+    savedBranchHead: () => undefined,
     orderByParentChain,
     unwrapPastedTextContent,
     toolResultModelText: (result: unknown) => result,

studio/frontend/tests/current-date-context-usage.test.ts

@@ -38,6 +38,7 @@ test("a prompt setting change clears local usage even during a run", async () =>
         useChatRuntimeStore: { getState: () => state },
       },
       "./message-order": messageOrder,
+      "./branch-head": { savedBranchHead: () => undefined },
       "./chat-history-storage": {},
     },
   );
@@ -74,6 +75,7 @@ test("a prompt setting change leaves external usage intact", async () => {
         useChatRuntimeStore: { getState: () => state },
       },
       "./message-order": messageOrder,
+      "./branch-head": { savedBranchHead: () => undefined },
       "./chat-history-storage": {},
     },
   );

studio/frontend/tests/finetune-export-branches.test.ts

@@ -7,7 +7,10 @@ import vm from "node:vm";
 import ts from "typescript";
 
 import * as liveThreadHead from "../src/features/chat/utils/live-thread-head.ts";
-import { orderByParentChain } from "../src/features/chat/utils/message-order.ts";
+import {
+  orderByParentChain,
+  resolveSavedBranchHead,
+} from "../src/features/chat/utils/message-order.ts";
 import { unwrapPastedTextContent } from "../src/features/chat/utils/pasted-text.ts";
 import { readSrc } from "./helpers/kit.ts";
 
@@ -49,7 +52,10 @@ function storedMessages(
   }));
 }
 
-function loadBuilder(chats: Record<string, StoredMessage[]>) {
+function loadBuilder(
+  chats: Record<string, StoredMessage[]>,
+  savedHeads: Record<string, string> = {},
+) {
   const javascript = ts.transpileModule(
     [
       sliceSource(
@@ -78,6 +84,8 @@ function loadBuilder(chats: Record<string, StoredMessage[]>) {
     getStoredChatThread: async () => undefined,
     settleThreadScopedSettingsForCopy: async () => {},
     ...liveThreadHead,
+    savedBranchHead: (id: string, raw: StoredMessage[]) =>
+      resolveSavedBranchHead(raw, savedHeads[id]),
     orderByParentChain,
     unwrapPastedTextContent,
   } as Record<string, unknown>;
@@ -89,8 +97,9 @@ async function fineTuneMessages(
   chats: Record<string, StoredMessage[]>,
   liveBranch: string[] | null = null,
   openedEarlier: string[] = [],
+  savedHeads: Record<string, string> = {},
 ) {
-  const build = loadBuilder(chats);
+  const build = loadBuilder(chats, savedHeads);
   const unregisters = liveBranch
     ? [...openedEarlier, "open"].map((remoteId) =>
         liveThreadHead.registerLiveThreadView({
@@ -164,3 +173,15 @@ test("chat fine-tune data keeps chats opened earlier on their newest reply", asy
     ],
   );
 });
+
+test("chat fine-tune data uses the reply a closed chat was left on", async () => {
+  assert.deepEqual(
+    await fineTuneMessages({ open: regenerated() }, null, [], { open: "a1" }),
+    [
+      [
+        { role: "user", content: "Name one fruit." },
+        { role: "assistant", content: "Apples." },
+      ],
+    ],
+  );
+});

studio/frontend/tests/markdown-export-branches.test.ts

@@ -118,6 +118,7 @@ function loadExporters(
     canMergeConversationExport,
     planChatItemSources,
     ...liveThreadHead,
+    savedBranchHead: () => undefined,
     orderByParentChain,
     createConversationMarkdownBuilder,
     createConversationMarkdownExporter,

studio/frontend/tests/reopened-chat-branch.test.ts

@@ -0,0 +1,85 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import assert from "node:assert/strict";
+import test from "node:test";
+
+import {
+  orderBySelectedBranch,
+  resolveSavedBranchHead,
+} from "../src/features/chat/utils/message-order.ts";
+
+const RETRIED_FIRST_REPLY = [
+  { id: "u1", parentId: null, createdAt: 1, role: "user" },
+  { id: "a1", parentId: "u1", createdAt: 2, role: "assistant" },
+  { id: "u2", parentId: "a1", createdAt: 3, role: "user" },
+  { id: "a2", parentId: "u2", createdAt: 4, role: "assistant" },
+  { id: "u3", parentId: "a2", createdAt: 5, role: "user" },
+  { id: "a3", parentId: "u3", createdAt: 6, role: "assistant" },
+  { id: "a1-retry", parentId: "u1", createdAt: 7, role: "assistant" },
+];
+
+const ids = (rows: { id: string }[]) => rows.map((row) => row.id);
+
+test("the branch left open is reopened, not the newest row's", () => {
+  const headId = resolveSavedBranchHead(RETRIED_FIRST_REPLY, "a3");
+  assert.equal(headId, "a3");
+  assert.deepEqual(ids(orderBySelectedBranch(RETRIED_FIRST_REPLY, headId)), [
+    "u1",
+    "a1",
+    "u2",
+    "a2",
+    "u3",
+    "a3",
+  ]);
+});
+
+test("with nothing saved the newest row's branch is shown, as before", () => {
+  const headId = resolveSavedBranchHead(RETRIED_FIRST_REPLY, null);
+  assert.equal(headId, undefined);
+  assert.deepEqual(ids(orderBySelectedBranch(RETRIED_FIRST_REPLY, headId)), [
+    "u1",
+    "a1-retry",
+  ]);
+});
+
+test("a saved head whose row is gone falls back to the newest row", () => {
+  assert.equal(resolveSavedBranchHead(RETRIED_FIRST_REPLY, "deleted"), undefined);
+});
+
+test("turns added under the saved head since are part of its branch", () => {
+  const extended = [
+    ...RETRIED_FIRST_REPLY,
+    { id: "u4", parentId: "a3", createdAt: 8, role: "user" },
+    { id: "a4", parentId: "u4", createdAt: 9, role: "assistant" },
+  ];
+  assert.equal(resolveSavedBranchHead(extended, "a3"), "a4");
+});
+
+test("a saved head with retried children follows the newest child", () => {
+  const retriedLast = [
+    ...RETRIED_FIRST_REPLY,
+    { id: "a3-retry", parentId: "u3", createdAt: 8, role: "assistant" },
+  ];
+  assert.equal(resolveSavedBranchHead(retriedLast, "u3"), "a3-retry");
+});
+
+test("the newest turn under the saved head wins over a newer sibling", () => {
+  // a remote retry under u3 may precede continuing from a3
+  const continuedOlder = [
+    ...RETRIED_FIRST_REPLY,
+    { id: "a3-retry", parentId: "u3", createdAt: 8, role: "assistant" },
+    { id: "u4", parentId: "a3", createdAt: 9, role: "user" },
+    { id: "a4", parentId: "u4", createdAt: 10, role: "assistant" },
+  ];
+  assert.equal(resolveSavedBranchHead(continuedOlder, "u3"), "a4");
+});
+
+test("legacy rows without parents chain through to the newest", () => {
+  const legacy = [
+    { id: "u1", createdAt: 1, role: "user" },
+    { id: "a1", createdAt: 2, role: "assistant" },
+    { id: "u2", createdAt: 3, role: "user" },
+  ];
+  assert.equal(resolveSavedBranchHead(legacy, "a1"), "u2");
+});

studio/frontend/tests/saved-branch-head.test.ts

@@ -0,0 +1,137 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import assert from "node:assert/strict";
+import test from "node:test";
+
+import { loadWithStubs } from "./helpers/module-stubs.ts";
+
+type View = { remoteId?: string; messages: { id: string }[]; incognito: boolean };
+
+// rerun the mocked effect only when dependencies change, matching React
+function recorder(temporary: Set<string>) {
+  const writes: [string, string][] = [];
+  let view: View = { messages: [], incognito: false };
+  let previousDeps: unknown[] | null = null;
+  const runRecorder = loadWithStubs<{ useBranchHeadRecorder: () => void }>(
+    new URL("../src/features/chat/hooks/use-branch-head-recorder.ts", import.meta.url),
+    {
+      react: {
+        useEffect: (effect: () => void, deps: unknown[]) => {
+          if (previousDeps && deps.every((dep, i) => dep === previousDeps?.[i])) return;
+          previousDeps = deps;
+          effect();
+        },
+      },
+      "@assistant-ui/react": {
+        useAuiState: (select: (state: unknown) => unknown) =>
+          select({
+            threadListItem: { remoteId: view.remoteId },
+            thread: { messages: view.messages },
+          }),
+      },
+      "../stores/chat-runtime-store": {
+        useChatRuntimeStore: (select: (state: unknown) => unknown) =>
+          select({ incognito: view.incognito }),
+      },
+      "../utils/branch-head": {
+        writeBranchHead: (threadId: string, messageId: string) =>
+          writes.push([threadId, messageId]),
+      },
+      "../utils/chat-history-storage": {
+        isThreadIncognito: (threadId: string) => temporary.has(threadId),
+      },
+    },
+  ).useBranchHeadRecorder;
+  return {
+    writes,
+    render(next: View) {
+      view = next;
+      runRecorder();
+    },
+  };
+}
+
+const branch = (...ids: string[]) => ids.map((id) => ({ id }));
+
+test("the recorder saves the head of the branch on screen as it moves", () => {
+  const { writes, render } = recorder(new Set());
+  render({ remoteId: "t", messages: [], incognito: false });
+  render({ remoteId: "t", messages: branch("u1", "a1", "u2", "a2"), incognito: false });
+  render({ remoteId: "t", messages: branch("u1", "a1-retry"), incognito: false });
+  render({ remoteId: "t", messages: branch("u1", "a1", "u2", "a2"), incognito: false });
+  assert.deepEqual(writes, [
+    ["t", "a2"],
+    ["t", "a1-retry"],
+    ["t", "a2"],
+  ]);
+});
+
+test("a temporary chat saved after a branch switch records the branch on screen", () => {
+  const temporary = new Set(["t"]);
+  const { writes, render } = recorder(temporary);
+  render({ remoteId: "t", messages: branch("u1", "a1", "u2", "a2"), incognito: true });
+  render({ remoteId: "t", messages: branch("u1", "a1-retry"), incognito: true });
+  render({ remoteId: "t", messages: branch("u1", "a1", "u2", "a2"), incognito: true });
+  assert.deepEqual(writes, []);
+  // saving unmarks the thread and disables incognito without moving the head
+  temporary.delete("t");
+  render({ remoteId: "t", messages: branch("u1", "a1", "u2", "a2"), incognito: false });
+  assert.deepEqual(writes, [["t", "a2"]]);
+});
+
+type DeleteModule = {
+  deleteChatItems: (
+    items: { type: "single"; id: string }[],
+    activeId: string | undefined,
+    onSelect: () => void,
+  ) => Promise<void>;
+};
+
+function sidebar(deleteResult: () => Promise<string[]>) {
+  const cleared: string[] = [];
+  const restored: string[][] = [];
+  const module = loadWithStubs<DeleteModule>(
+    new URL("../src/features/chat/hooks/use-chat-sidebar-items.ts", import.meta.url),
+    {
+      react: { useEffect: () => {}, useMemo: () => {}, useState: () => [] },
+      "../api/chat-api": {
+        CHAT_HISTORY_UPDATED_EVENT: "chat-history-updated",
+        notifyChatHistoryUpdated: () => {},
+      },
+      "../stores/chat-runtime-store": {
+        useChatRuntimeStore: { getState: () => ({ setActiveThreadId: () => {} }) },
+      },
+      "../utils/chat-history-storage": { deleteStoredChatThreads: deleteResult },
+      "../utils/branch-head": { clearBranchHead: (id: string) => cleared.push(id) },
+      "../utils/composer-draft": { clearComposerDraft: () => {} },
+      "../utils/offer-kept-sandbox-files": { offerToDeleteKeptSandboxes: () => {} },
+      "../utils/stop-chat-thread": { stopChatThread: () => {} },
+      "../utils/chat-thread-tombstones": {
+        markChatThreadsDeleted: () => {},
+        removeChatThreadTombstones: (ids: string[]) => restored.push(ids),
+      },
+      "../utils/prompt-queue-boundary": { requestPromptQueueStop: () => {} },
+      "../utils/repair-legacy-chat-titles": { repairLegacyChatTitles: () => {} },
+    },
+  );
+  return { ...module, cleared, restored };
+}
+
+test("deleting a chat clears its saved branch once the delete holds", async () => {
+  const { deleteChatItems, cleared } = sidebar(async () => []);
+  await deleteChatItems([{ type: "single", id: "t" }], undefined, () => {});
+  assert.deepEqual(cleared, ["t"]);
+});
+
+test("a rejected delete keeps the saved branch of the chat it brings back", async () => {
+  const { deleteChatItems, cleared, restored } = sidebar(async () => {
+    throw new Error("delete failed");
+  });
+  await assert.rejects(
+    deleteChatItems([{ type: "single", id: "t" }], undefined, () => {}),
+    /delete failed/,
+  );
+  assert.deepEqual(restored, [["t"]]);
+  assert.deepEqual(cleared, []);
+});

studio/frontend/tests/sharegpt-export-branches.test.ts

@@ -86,6 +86,7 @@ function loadExporters(
     getStoredChatThread: async () => undefined,
     settleThreadScopedSettingsForCopy: async () => {},
     ...liveThreadHead,
+    savedBranchHead: () => undefined,
     orderByParentChain,
     exportFormatIncludesSiblings,
     ndjsonBody,