fix(desktop): preserve complete history when branching

Read the durable display transcript when creating a branch instead of copying the compacted model projection. Hydrate the Desktop branch boundary from persisted history, avoid stale whole-chat counts, and seed the new tile from the backend snapshot. Add regression coverage for compacted histories, visible-message counts, selected prefixes, and hydration races. (cherry picked from commit c3d2d759ae104395adde215b2e293ccf8e895684)
2026-08-02 22:09:16 +08:00
parent b066f2b373
commit 07b9090259
6 changed files with 435 additions and 22 deletions
@@ -1112,19 +1112,30 @@ function BranchHarness({
  navigate = vi.fn(),
  onCurrentReady,
  onReady,
-  requestGateway
+  onRefs,
+  requestGateway,
+  selectedStoredSessionId = null
 }: {
  activeSessionId?: string | null
  navigate?: ReturnType<typeof vi.fn>
  onCurrentReady?: (branchCurrentSession: (messageId?: string) => Promise<boolean>) => void
  onReady: (branchStoredSession: (storedSessionId: string, sessionProfile?: string | null) => Promise<boolean>) => void
+  onRefs?: (refs: {
+    activeSessionIdRef: MutableRefObject<string | null>
+    selectedStoredSessionIdRef: MutableRefObject<string | null>
+  }) => void
  requestGateway: <T>(method: string, params?: Record<string, unknown>) => Promise<T>
+  selectedStoredSessionId?: string | null
 }) {
  const ref = <T,>(value: T): MutableRefObject<T> => ({ current: value })
+  const activeSessionIdRef = ref<string | null>(activeSessionId)
+  const selectedStoredSessionIdRef = ref<string | null>(selectedStoredSessionId)
+
+  onRefs?.({ activeSessionIdRef, selectedStoredSessionIdRef })

  const actions = useSessionActions({
    activeSessionId,
-    activeSessionIdRef: ref<string | null>(activeSessionId),
+    activeSessionIdRef,
    busyRef: ref(false),
    creatingSessionRef: ref(false),
    ensureSessionState: () => ({}) as ClientSessionState,
@@ -1134,8 +1145,8 @@ function BranchHarness({
    requestGateway,
    resetViewSync: vi.fn(),
    runtimeIdByStoredSessionIdRef: ref(new Map<string, string>()),
-    selectedStoredSessionId: null,
-    selectedStoredSessionIdRef: ref<string | null>(null),
+    selectedStoredSessionId,
+    selectedStoredSessionIdRef,
    sessionStateByRuntimeIdRef: ref(new Map<string, ClientSessionState>()),
    syncSessionStateToView: vi.fn(),
    updateSessionState: () => ({}) as ClientSessionState
@@ -1276,6 +1287,95 @@ describe('branchStoredSession desktop source tagging', () => {
    expect(branchParams).toEqual({ session_id: 'live-parent', count: 2 })
  })

+  it('hydrates the complete persisted display transcript before branching a compacted live chat', async () => {
+    let branchParams: Record<string, unknown> | undefined
+
+    const requestGateway = vi.fn(async (method: string, params?: Record<string, unknown>) => {
+      if (method === 'session.branch') {
+        branchParams = params
+
+        return {
+          session_id: 'branch-runtime',
+          stored_session_id: 'branch-stored',
+          title: 'Branch',
+          message_count: 4,
+          messages: [],
+          info: {}
+        } as never
+      }
+
+      return {} as never
+    })
+
+    setSessions([storedSession({ id: 'stored-parent', message_count: 4 })])
+    setMessages([
+      { id: 'summary', role: 'assistant', parts: [{ type: 'text', text: 'compact summary' }] },
+      { id: 'tail-user', role: 'user', parts: [{ type: 'text', text: 'second question' }] },
+      { id: 'tail-assistant', role: 'assistant', parts: [{ type: 'text', text: 'second answer' }] }
+    ])
+    vi.mocked(getAllSessionMessages).mockResolvedValue({
+      messages: [
+        { content: 'first question', role: 'user', timestamp: 1 },
+        { content: 'first answer', role: 'assistant', timestamp: 2 },
+        { content: 'second question', role: 'user', timestamp: 3 },
+        { content: 'second answer', role: 'assistant', timestamp: 4 }
+      ],
+      session_id: 'stored-parent'
+    } as never)
+
+    let branchCurrentSession: ((messageId?: string) => Promise<boolean>) | null = null
+    render(
+      <BranchHarness
+        activeSessionId="live-parent"
+        onCurrentReady={branch => (branchCurrentSession = branch)}
+        onReady={() => undefined}
+        requestGateway={requestGateway}
+        selectedStoredSessionId="stored-parent"
+      />
+    )
+    await waitFor(() => expect(branchCurrentSession).not.toBeNull())
+
+    await expect(branchCurrentSession!()).resolves.toBe(true)
+
+    expect(getAllSessionMessages).toHaveBeenCalledWith('stored-parent', undefined)
+    expect(branchParams).toEqual({ session_id: 'live-parent' })
+  })
+
+  it('aborts if the active runtime changes while the branch transcript is hydrating', async () => {
+    let refs: {
+      activeSessionIdRef: MutableRefObject<string | null>
+      selectedStoredSessionIdRef: MutableRefObject<string | null>
+    } | null = null
+
+    const requestGateway = vi.fn(async () => ({}) as never)
+
+    setMessages([{ id: 'q1', role: 'user', parts: [{ type: 'text', text: 'question' }] }])
+    vi.mocked(getAllSessionMessages).mockImplementation(async () => {
+      refs!.activeSessionIdRef.current = 'live-other'
+
+      return {
+        messages: [{ content: 'question', role: 'user', timestamp: 1 }],
+        session_id: 'stored-parent'
+      } as never
+    })
+
+    let branchCurrentSession: ((messageId?: string) => Promise<boolean>) | null = null
+    render(
+      <BranchHarness
+        activeSessionId="live-parent"
+        onCurrentReady={branch => (branchCurrentSession = branch)}
+        onReady={() => undefined}
+        onRefs={value => (refs = value)}
+        requestGateway={requestGateway}
+        selectedStoredSessionId="stored-parent"
+      />
+    )
+    await waitFor(() => expect(branchCurrentSession).not.toBeNull())
+
+    await expect(branchCurrentSession!()).resolves.toBe(false)
+    expect(requestGateway).not.toHaveBeenCalledWith('session.branch', expect.anything())
+  })
+
  // #67603: right-clicking a session outside the paginated sidebar window is a
  // cache miss. Resolve its owning profile (cache → active → cross-profile) and
  // swap to it before reading the transcript / creating the branch, so the fork
@@ -92,6 +92,7 @@ import {
  resolveResumedBusy,
  resolveSessionProfile,
  resolveStoredSession,
+  selectBranchMessages,
  sessionMatchesStoredId,
  sessionShouldHaveTranscript,
  toBranchMessages,
@@ -1335,7 +1336,8 @@ export function useSessionActions({
      sourceSessionId: null | string,
      parentStoredId: null | string,
      cwd?: string,
-      profile?: null | string
+      profile?: null | string,
+      branchCount?: number
    ): Promise<boolean> => {
      creatingSessionRef.current = true

@@ -1353,7 +1355,7 @@ export function useSessionActions({
        const branched = sourceSessionId
          ? await requestGateway<SessionCreateResponse>('session.branch', {
              session_id: sourceSessionId,
-              count: branchMessages.length
+              ...(branchCount !== undefined ? { count: branchCount } : {})
            })
          : await requestGateway<SessionCreateResponse>('session.create', {
              cols: 96,
@@ -1364,8 +1366,12 @@ export function useSessionActions({
              ...(parentStoredId && { parent_session_id: parentStoredId })
            })

+        const responseBranchMessages =
+          sourceSessionId && branched.messages?.length ? toBranchMessages(toChatMessages(branched.messages)) : []
+
+        const effectiveBranchMessages = responseBranchMessages.length ? responseBranchMessages : branchMessages
        const routedSessionId = branched.stored_session_id ?? branched.session_id
-        const preview = branchMessages.map(({ content }) => content).find(Boolean) ?? null
+        const preview = effectiveBranchMessages.map(({ content }) => content).find(Boolean) ?? null
        // Draft until submit: nest under the parent at the parent's recency so it
        // doesn't bubble to the top until a real message lands (backend persists
        // + auto-names it then). The selected row survives refreshes (sessionsToKeep).
@@ -1390,7 +1396,7 @@ export function useSessionActions({
          branched.session_id,
          state => ({
            ...state,
-            messages: branchMessages.map(({ source }) => source),
+            messages: effectiveBranchMessages.map(({ source }) => source),
            busy: false,
            awaitingResponse: false
          }),
@@ -1444,15 +1450,52 @@ export function useSessionActions({
        return false
      }

+      const startingActiveSessionId = activeSessionIdRef.current
      const messages = $messages.get()
+      const storedSessionId = selectedStoredSessionIdRef.current
+      const startingRouteToken = getRouteToken()
+      const startingCwd = $currentCwd.get().trim()

-      const at = messageId
-        ? messages.findIndex(message => message.id === messageId)
-        : messages.findLastIndex(message => message.role === 'assistant' || message.role === 'user')
+      // The live atom may be a compacted model projection. Read the durable
+      // display projection before choosing the branch prefix so a whole-chat
+      // branch does not inherit only the summary/tail. If the backend is
+      // temporarily unavailable, retain the local snapshot and let the branch
+      // RPC make its own authoritative read.
+      let authoritativeMessages: ChatMessage[] | null = null
+      const profile = await resolveSessionProfile(storedSessionId)

-      const start = 0
-      const end = at >= 0 ? at + 1 : messages.length
-      const branchMessages = toBranchMessages(messages.slice(start, end))
+      if (storedSessionId) {
+        try {
+          const persisted = await getAllSessionMessages(storedSessionId, profile)
+          const hydrated = toChatMessages(persisted.messages)
+
+          if (hydrated.length) {
+            authoritativeMessages = hydrated
+          }
+        } catch {
+          // The branch RPC has a backend-side display projection fallback.
+        }
+      }
+
+      const drift = sessionContextDrift({
+        startRouteToken: startingRouteToken,
+        nowRouteToken: getRouteToken(),
+        startSelectedStoredId: storedSessionId,
+        nowSelectedStoredId: selectedStoredSessionIdRef.current
+      })
+
+      const runtimeChanged = activeSessionIdRef.current !== startingActiveSessionId
+      const selectionChanged = selectedStoredSessionIdRef.current !== storedSessionId
+
+      if (drift || runtimeChanged || selectionChanged) {
+        console.warn('[branch-drift-abort]', drift ?? 'runtime-or-selection-changed', {
+          phase: 'transcript-hydration'
+        })
+
+        return false
+      }
+
+      const branchMessages = selectBranchMessages(messages, authoritativeMessages, messageId)

      if (!branchMessages.length) {
        notify({ kind: 'warning', title: copy.nothingToBranch, message: copy.branchNoText })
@@ -1465,17 +1508,16 @@ export function useSessionActions({
      // The open chat's owning profile, NOT the picker's / launch profile —
      // /profile only retargets new chats, so a branch of an existing thread
      // must stay on that thread's backend (cache hit for an open session).
-      const profile = await resolveSessionProfile(selectedStoredSessionIdRef.current)
-
      return forkBranch(
        branchMessages,
-        activeSessionIdRef.current,
-        selectedStoredSessionIdRef.current,
-        $currentCwd.get().trim(),
-        profile
+        startingActiveSessionId,
+        storedSessionId,
+        startingCwd,
+        profile,
+        messageId ? branchMessages.length : undefined
      )
    },
-    [activeSessionIdRef, busyRef, copy, forkBranch, selectedStoredSessionIdRef]
+    [activeSessionIdRef, busyRef, copy, forkBranch, getRouteToken, selectedStoredSessionIdRef]
  )

  // Branch any listed session, not just the open one. Reads the target's stored
@@ -29,6 +29,7 @@ import {
  reconcileResumeMessages,
  removeRepresentedLocalLiveProjection,
  resolveResumedBusy,
+  selectBranchMessages,
  sessionMatchesStoredId,
  sessionShouldHaveTranscript,
  toBranchMessages
@@ -250,6 +251,47 @@ describe('toBranchMessages', () => {
  })
 })

+describe('selectBranchMessages', () => {
+  it('uses the complete authoritative transcript for a whole-chat branch', () => {
+    const local = [msg('summary', 'assistant', 'compact summary'), msg('tail', 'assistant', 'latest answer')]
+
+    const authoritative = [
+      msg('old-user', 'user', 'first question', { rowId: 11 }),
+      msg('old-assistant', 'assistant', 'first answer', { rowId: 12 }),
+      msg('tail-user', 'user', 'latest question', { rowId: 13 }),
+      msg('tail-assistant', 'assistant', 'latest answer', { rowId: 14 })
+    ]
+
+    expect(selectBranchMessages(local, authoritative).map(message => message.content)).toEqual([
+      'first question',
+      'first answer',
+      'latest question',
+      'latest answer'
+    ])
+  })
+
+  it('maps a clicked local bubble to the authoritative row before slicing', () => {
+    const local = [
+      msg('tail-user', 'user', 'latest question', { rowId: 13 }),
+      msg('tail-assistant', 'assistant', 'latest answer', { rowId: 14 })
+    ]
+
+    const authoritative = [
+      msg('old-user', 'user', 'first question', { rowId: 11 }),
+      msg('old-assistant', 'assistant', 'first answer', { rowId: 12 }),
+      msg('tail-user', 'user', 'latest question', { rowId: 13 }),
+      msg('tail-assistant', 'assistant', 'latest answer', { rowId: 14 })
+    ]
+
+    expect(selectBranchMessages(local, authoritative, 'tail-assistant').map(message => message.content)).toEqual([
+      'first question',
+      'first answer',
+      'latest question',
+      'latest answer'
+    ])
+  })
+})
+
 describe('chatPartsEquivalent', () => {
  it('returns true for identical text parts', () => {
    const partA = { type: 'text' as const, text: 'Hello world' }
@@ -1132,6 +1132,77 @@ export const toBranchMessages = (messages: ChatMessage[]): BranchMessage[] =>
    .map(message => ({ content: chatMessageText(message), role: message.role, source: message }))
    .filter(({ content, role }) => content.trim() && (role === 'assistant' || role === 'user'))

+/**
+ * Choose the transcript used to seed an open-chat branch.
+ *
+ * The local renderer can hold a compacted model projection, while the REST
+ * transcript contains the complete display projection. Use the latter for a
+ * whole-chat branch. When branching from a clicked bubble, map that bubble by
+ * durable row id first and by same-role/text ordinal as a legacy fallback; if
+ * it cannot be mapped, keep the local prefix rather than silently choosing a
+ * different point in the conversation.
+ */
+export function selectBranchMessages(
+  localMessages: ChatMessage[],
+  authoritativeMessages: ChatMessage[] | null,
+  messageId?: string
+): BranchMessage[] {
+  const localIndex = messageId ? localMessages.findIndex(message => message.id === messageId) : -1
+
+  if (!authoritativeMessages?.length) {
+    return toBranchMessages(localMessages.slice(0, localIndex >= 0 ? localIndex + 1 : localMessages.length))
+  }
+
+  if (!messageId) {
+    return toBranchMessages(authoritativeMessages)
+  }
+
+  if (localIndex < 0) {
+    return toBranchMessages(localMessages)
+  }
+
+  const target = localMessages[localIndex]
+
+  let authoritativeIndex =
+    target.rowId === undefined
+      ? -1
+      : authoritativeMessages.findIndex(message => message.rowId !== undefined && message.rowId === target.rowId)
+
+  // Strip `@image:` directive lines the same way the persisted→ChatMessage
+  // conversion does (extractImageRefs lifts them into attachmentRefs), so a
+  // local optimistic bubble and its authoritative twin compare equal.
+  const comparableText = (message: ChatMessage) =>
+    textWithoutEmbeddedImages(chatMessageText(message))
+      .replace(/^@image:[^\n]*\n?/gm, '')
+      .trim()
+
+  if (authoritativeIndex < 0) {
+    const targetText = comparableText(target)
+
+    const targetOrdinal = localMessages
+      .slice(0, localIndex + 1)
+      .filter(message => message.role === target.role && comparableText(message) === targetText).length
+
+    let ordinal = 0
+
+    authoritativeIndex = authoritativeMessages.findIndex(message => {
+      if (message.role !== target.role || comparableText(message) !== targetText) {
+        return false
+      }
+
+      ordinal += 1
+
+      return ordinal === targetOrdinal
+    })
+  }
+
+  if (authoritativeIndex < 0) {
+    return toBranchMessages(localMessages.slice(0, localIndex + 1))
+  }
+
+  return toBranchMessages(authoritativeMessages.slice(0, authoritativeIndex + 1))
+}
+
 export function upsertOptimisticSession(
  created: SessionCreateResponse,
  id: string,
@@ -13183,6 +13183,130 @@ def test_session_branch_installs_parent_profile_secret_scope(monkeypatch, tmp_pa
            server._sessions.pop(k, None)


+def test_session_branch_uses_persisted_display_history_after_compaction(monkeypatch, tmp_path):
+    """A live branch must copy the complete visible transcript, not the compacted model tail."""
+    profile_home = tmp_path / "profiles" / "mlperf"
+    profile_home.mkdir(parents=True)
+    seen: dict = {"msgs": []}
+
+    display_history = [
+        {"role": "user", "content": "first question", "timestamp": 1.0},
+        {"role": "assistant", "content": "first answer", "timestamp": 2.0},
+        {"role": "assistant", "content": "", "tool_calls": [{"id": "call-1"}]},
+        {"role": "tool", "content": "tool output", "tool_call_id": "call-1"},
+        {"role": "user", "content": "second question", "timestamp": 3.0},
+        {"role": "assistant", "content": "second answer", "timestamp": 4.0},
+    ]
+
+    class LaunchDB:
+        def get_session_title(self, _key):
+            return "launch"
+
+    class ProfileDB:
+        def __init__(self, db_path=None):
+            seen.setdefault("inits", 0)
+            seen["inits"] += 1
+
+        def get_session_title(self, _key):
+            return "parent"
+
+        def get_next_title_in_lineage(self, current):
+            return f"{current} (branch)"
+
+        def get_resume_conversations(self, key):
+            assert key == "parent-key"
+            # The model projection has already been compacted to a summary + tail;
+            # the display projection still contains every visible turn.
+            return (
+                [{"role": "assistant", "content": "compact summary"}],
+                display_history,
+            )
+
+        def create_session(self, _new_key, **_kwargs):
+            return None
+
+        def append_message(self, **kwargs):
+            seen["msgs"].append(kwargs)
+
+        def append_messages_batch(self, session_id, messages, **kwargs):
+            for message in messages:
+                seen["msgs"].append(dict(message, session_id=session_id))
+            return list(range(1, len(messages) + 1))
+
+        def set_session_title(self, _key, _title):
+            return True
+
+        def get_session(self, key):
+            return {"id": key, "cwd": str(tmp_path)}
+
+        def update_session_cwd(self, *args, **kwargs):
+            return None
+
+        def close(self):
+            return None
+
+    class FakeAgent:
+        model = "test-model"
+        session_id = None
+
+    parent = {
+        "session_key": "parent-key",
+        # This is the model-fed projection after compaction: the old turns are
+        # absent here even though the display projection above retains them.
+        "history": [
+            {"role": "assistant", "content": "compact summary"},
+            {"role": "user", "content": "second question"},
+            {"role": "assistant", "content": "second answer"},
+        ],
+        "history_lock": threading.Lock(),
+        "running": False,
+        "cols": 80,
+        "profile_home": str(profile_home),
+        "source": "tui",
+        "agent": FakeAgent(),
+        "created_at": 1.0,
+        "last_active": 1.0,
+        "cwd": str(tmp_path),
+    }
+    server._sessions["parent"] = parent
+    monkeypatch.setattr(server, "_get_db", lambda: LaunchDB())
+    monkeypatch.setattr("hermes_state.SessionDB", ProfileDB)
+    monkeypatch.setattr(server, "_claim_active_session_slot", lambda *args, **kwargs: (None, None))
+    monkeypatch.setattr(server, "_make_agent", lambda *args, **kwargs: FakeAgent())
+    monkeypatch.setattr(server, "_set_session_context", lambda *args, **kwargs: {})
+    monkeypatch.setattr(server, "_clear_session_context", lambda *args, **kwargs: None)
+    monkeypatch.setattr(server, "_resolve_model", lambda: "test-model")
+    monkeypatch.setattr(server, "_session_cwd", lambda _session: str(tmp_path))
+    monkeypatch.setattr(server, "_register_session_cwd", lambda *args, **kwargs: None)
+    monkeypatch.setattr(server, "_attach_worker", lambda *args, **kwargs: None)
+
+    try:
+        response = server.handle_request(
+            {
+                "id": "1",
+                "method": "session.branch",
+                "params": {"session_id": "parent", "count": 4},
+            }
+        )
+
+        assert "result" in response, response
+        assert [message["content"] for message in seen["msgs"]] == [
+            "first question",
+            "first answer",
+            "second question",
+            "second answer",
+        ]
+        assert [message["text"] for message in response["result"]["messages"]] == [
+            "first question",
+            "first answer",
+            "second question",
+            "second answer",
+        ]
+    finally:
+        for key in list(server._sessions):
+            server._sessions.pop(key, None)
+
+
 def test_pending_title_finalizer_uses_session_profile_db(monkeypatch, tmp_path):
    """Post-turn pending_title must land in the session profile store."""
    profile_home = tmp_path / "profiles" / "mlperf"
@@ -2813,7 +2813,41 @@ def _(rid, params: dict) -> dict:
            return _db_unavailable_error(rid, code=5008)
        old_key = session["session_key"]
        with session["history_lock"]:
-            history = [dict(msg) for msg in session.get("history", [])]
+            in_memory_history = [
+                dict(msg)
+                for msg in list(session.get("display_history_prefix") or []) + list(session.get("history", []))
+                if isinstance(msg, dict)
+            ]
+
+        def _visible_branch_history(messages):
+            visible = []
+            for message in messages or []:
+                if not isinstance(message, dict) or message.get("role") not in {"user", "assistant"}:
+                    continue
+                if not _coerce_message_text(message.get("content")).strip():
+                    continue
+                # Keep the FULL row — the copy loop below preserves reasoning
+                # fields and timeline-marker tags (display_kind/display_metadata,
+                # #82756); a minimal role/content copy would silently drop them.
+                visible.append(dict(message))
+            return visible
+
+        # The live session history is the model projection. After compaction it
+        # may contain only a summary and the protected tail, while the persisted
+        # display projection still contains the complete visible transcript. A
+        # branch must snapshot the latter; otherwise the child permanently loses
+        # every turn archived before the fork.
+        history = None
+        get_resume_conversations = getattr(db, "get_resume_conversations", None)
+        if callable(get_resume_conversations):
+            try:
+                _, display_history = get_resume_conversations(old_key)
+                display_history = _reconcile_display_with_live(display_history, in_memory_history)
+                history = _visible_branch_history(display_history)
+            except Exception:
+                logger.debug("branch display projection read failed", exc_info=True)
+        if not history:
+            history = _visible_branch_history(in_memory_history)
        if not history:
            return _err(rid, 4008, "nothing to branch — send a message first")
        count = params.get("count")