fix(desktop): Branch in new chat drops the question and loses the branched session on restart (#71960)

* fix: Branch in new chat loses the question and the branched session (#issue)

- session.branch on the backend now accepts a count param to truncate
  the parent history to the clicked message, instead of always forking
  the entire transcript. Also returns stored_session_id/messages/info
  so the frontend has parity with session.create's response shape.
- branchCurrentSession (open live chat) now slices history from 0
  instead of from the clicked message index, so the question preceding
  an assistant reply is no longer dropped when branching.
- forkBranch now calls session.branch (not session.create) when
  branching an open live chat, since session.create only persists a DB
  row lazily on first prompt - a branched chat that nobody types into
  never got saved, and vanished as 'session not found' on the next
  app restart. branchStoredSession (branching from the sidebar, no
  live runtime) keeps using session.create as before.

* test: cover session.branch count truncation and open-chat branching

- backend: assert session.branch with a count param only persists the
  first N messages of the live history to the new session.
- frontend: BranchHarness now exposes branchCurrentSession; assert
  branching an open chat from a middle message calls session.branch
  with the parent session id and the correct trimmed count, instead of
  session.create.

apps/desktop/src/app/session/hooks/use-session-actions.test.tsx

@@ -977,19 +977,23 @@ describe('resumeSession failure recovery', () => {
 })
 
 function BranchHarness({
+  activeSessionId = null,
   navigate = vi.fn(),
+  onCurrentReady,
   onReady,
   requestGateway
 }: {
+  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
   requestGateway: <T>(method: string, params?: Record<string, unknown>) => Promise<T>
 }) {
   const ref = <T,>(value: T): MutableRefObject<T> => ({ current: value })
 
   const actions = useSessionActions({
-    activeSessionId: null,
-    activeSessionIdRef: ref<string | null>(null),
+    activeSessionId,
+    activeSessionIdRef: ref<string | null>(activeSessionId),
     busyRef: ref(false),
     creatingSessionRef: ref(false),
     ensureSessionState: () => ({}) as ClientSessionState,
@@ -1008,7 +1012,8 @@ function BranchHarness({
 
   useEffect(() => {
     onReady(actions.branchStoredSession)
-  }, [actions.branchStoredSession, onReady])
+    onCurrentReady?.(actions.branchCurrentSession)
+  }, [actions.branchCurrentSession, actions.branchStoredSession, onCurrentReady, onReady])
 
   return null
 }
@@ -1090,6 +1095,53 @@ describe('branchStoredSession desktop source tagging', () => {
     })
   })
 
+  it('branches an open live chat via session.branch with a trimmed message count (bug #1/#3 fix)', 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: 2,
+          messages: [],
+          info: {}
+        } as never
+      }
+      return {} as never
+    })
+
+    setMessages([
+      { id: 'q1', role: 'user', parts: [{ type: 'text', text: 'question one' }] },
+      { id: 'a1', role: 'assistant', parts: [{ type: 'text', text: 'answer one' }] },
+      { id: 'q2', role: 'user', parts: [{ type: 'text', text: 'question two' }] },
+      { id: 'a2', role: 'assistant', parts: [{ type: 'text', text: 'answer two' }] }
+    ])
+
+    let branchCurrentSession: ((messageId?: string) => Promise<boolean>) | null = null
+    render(
+      <BranchHarness
+        activeSessionId="live-parent"
+        onCurrentReady={branch => (branchCurrentSession = branch)}
+        onReady={() => undefined}
+        requestGateway={requestGateway}
+      />
+    )
+    await waitFor(() => expect(branchCurrentSession).not.toBeNull())
+
+    // Branch from the FIRST assistant reply ("a1"), not the last message —
+    // this is exactly the scenario that used to drop the question (bug #1):
+    // only the clicked message survived instead of everything up to it.
+    await expect(branchCurrentSession!('a1')).resolves.toBe(true)
+
+    expect(requestGateway).toHaveBeenCalledWith('session.branch', {
+      session_id: 'live-parent',
+      count: 2
+    })
+    expect(branchParams).toEqual({ session_id: 'live-parent', count: 2 })
+  })
+
   // #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

apps/desktop/src/app/session/hooks/use-session-actions/index.ts

@@ -1102,6 +1102,7 @@ export function useSessionActions({
   const forkBranch = useCallback(
     async (
       branchMessages: BranchMessage[],
+      sourceSessionId: null | string,
       parentStoredId: null | string,
       cwd?: string,
       profile?: null | string
@@ -1119,14 +1120,19 @@ export function useSessionActions({
         await ensureGatewayProfile(profile)
 
         // No title: the backend auto-names the branch from its parent's lineage.
-        const branched = await requestGateway<SessionCreateResponse>('session.create', {
-          cols: 96,
-          source: 'desktop',
-          ...(cwd && { cwd }),
-          ...(profile ? { profile } : {}),
-          messages: branchMessages.map(({ content, role }) => ({ content, role })),
-          ...(parentStoredId && { parent_session_id: parentStoredId })
-        })
+        const branched = sourceSessionId
+          ? await requestGateway<SessionCreateResponse>('session.branch', {
+              session_id: sourceSessionId,
+              count: branchMessages.length
+            })
+          : await requestGateway<SessionCreateResponse>('session.create', {
+              cols: 96,
+              source: 'desktop',
+              ...(cwd && { cwd }),
+              ...(profile ? { profile } : {}),
+              messages: branchMessages.map(({ content, role }) => ({ content, role })),
+              ...(parentStoredId && { parent_session_id: parentStoredId })
+            })
 
         const routedSessionId = branched.stored_session_id ?? branched.session_id
         const preview = branchMessages.map(({ content }) => content).find(Boolean) ?? null
@@ -1212,7 +1218,7 @@ export function useSessionActions({
         ? messages.findIndex(message => message.id === messageId)
         : messages.findLastIndex(message => message.role === 'assistant' || message.role === 'user')
 
-      const start = at >= 0 ? at : Math.max(messages.length - 1, 0)
+      const start = 0
       const end = at >= 0 ? at + 1 : messages.length
       const branchMessages = toBranchMessages(messages.slice(start, end))
 
@@ -1229,7 +1235,13 @@ export function useSessionActions({
       // must stay on that thread's backend (cache hit for an open session).
       const profile = await resolveSessionProfile(selectedStoredSessionIdRef.current)
 
-      return forkBranch(branchMessages, selectedStoredSessionIdRef.current, $currentCwd.get().trim(), profile)
+      return forkBranch(
+        branchMessages,
+        activeSessionIdRef.current,
+        selectedStoredSessionIdRef.current,
+        $currentCwd.get().trim(),
+        profile
+      )
     },
     [activeSessionIdRef, busyRef, copy, forkBranch, selectedStoredSessionIdRef]
   )
@@ -1261,7 +1273,7 @@ export function useSessionActions({
           return false
         }
 
-        return await forkBranch(branchMessages, stored?.id ?? storedSessionId, stored?.cwd?.trim(), profile)
+        return await forkBranch(branchMessages, null, stored?.id ?? storedSessionId, stored?.cwd?.trim(), profile)
       } catch (err) {
         notifyError(err, copy.branchFailed)
 

tests/tui_gateway/test_protocol.py

@@ -1289,6 +1289,75 @@ def test_session_branch_persists_branched_from_marker(server, monkeypatch):
     assert kwargs["model_config"] == {"_branched_from": parent_key}
 
 
+def test_session_branch_with_count_truncates_history(server, monkeypatch):
+    """Branch-from-a-specific-message support (issue: Branch in new chat
+    loses the question): the desktop client passes ``count`` to keep only
+    the first N messages of the parent's live history - everything after
+    the clicked message must NOT be copied into the branch.
+    """
+    append_calls = []
+
+    class _DB:
+        def get_session_title(self, _key):
+            return "parent-title"
+
+        def get_next_title_in_lineage(self, base):
+            return f"{base} 2"
+
+        def create_session(self, new_key, **kwargs):
+            return new_key
+
+        def append_message(self, **kwargs):
+            append_calls.append(kwargs)
+            return None
+
+        def set_session_title(self, _key, _title):
+            return None
+
+    monkeypatch.setattr(server, "_get_db", lambda: _DB())
+    monkeypatch.setattr(server, "_resolve_model", lambda: "test/model")
+    monkeypatch.setattr(server, "_new_session_key", lambda: "20260101_000001_child0")
+    monkeypatch.setattr(
+        server,
+        "_make_agent",
+        lambda _sid, key, session_id=None, session_db=None, **_kwargs: types.SimpleNamespace(
+            model="test/model", session_id=session_id or key
+        ),
+    )
+    monkeypatch.setattr(server, "_init_session", lambda *_a, **_k: None)
+    monkeypatch.setattr(server, "_set_session_context", lambda *_a, **_k: [])
+    monkeypatch.setattr(server, "_clear_session_context", lambda *_a, **_k: None)
+    monkeypatch.setattr(server, "_session_cwd", lambda _s: "/tmp/branch-cwd")
+
+    parent_sid = "parent01"
+    parent_key = "20260101_000000_parent"
+    server._sessions[parent_sid] = {
+        "session_key": parent_key,
+        "history": [
+            {"role": "user", "content": "question one"},
+            {"role": "assistant", "content": "answer one"},
+            {"role": "user", "content": "question two"},
+            {"role": "assistant", "content": "answer two"},
+        ],
+        "history_lock": threading.Lock(),
+        "cols": 80,
+    }
+
+    resp = server.handle_request(
+        {
+            "id": "b1",
+            "method": "session.branch",
+            "params": {"session_id": parent_sid, "count": 2},
+        }
+    )
+
+    assert "error" not in resp, resp
+    assert len(append_calls) == 2
+    assert append_calls[0]["content"] == "question one"
+    assert append_calls[1]["content"] == "answer one"
+    assert resp["result"]["message_count"] == 2
+
+
 def test_session_branch_forwards_original_timestamps(server, monkeypatch):
     """TUI /branch must copy the parent's messages WITH their original
     timestamps — append_message otherwise stamps time.time() at INSERT and

tui_gateway/server.py

@@ -10312,6 +10312,9 @@ def _(rid, params: dict) -> dict:
             history = [dict(msg) for msg in session.get("history", [])]
         if not history:
             return _err(rid, 4008, "nothing to branch — send a message first")
+        count = params.get("count")
+        if isinstance(count, int) and count > 0:
+            history = history[:count]
         new_key = _new_session_key()
         new_sid = uuid.uuid4().hex[:8]
         source = _session_source(session)
@@ -10415,7 +10418,19 @@ def _(rid, params: dict) -> dict:
         if lease is not None:
             lease.release()
         return _err(rid, 5000, f"agent init failed on branch: {e}")
-    return _ok(rid, {"session_id": new_sid, "title": title, "parent": old_key})
+    branched_session = _sessions.get(new_sid)
+    return _ok(
+        rid,
+        {
+            "session_id": new_sid,
+            "stored_session_id": new_key,
+            "title": title,
+            "parent": old_key,
+            "message_count": len(history),
+            "messages": _history_to_messages(history),
+            "info": _session_info(agent, branched_session),
+        },
+    )
 
 
 @method("session.interrupt")