From 820a8083d47d371049cb344a33f2ee22a7570a3a Mon Sep 17 00:00:00 2001 From: alelpoan <155192176+alelpoan@users.noreply.github.com> Date: Mon, 27 Jul 2026 17:45:00 +0300 Subject: [PATCH] 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. --- .../hooks/use-session-actions.test.tsx | 58 +++++++++++++++- .../hooks/use-session-actions/index.ts | 34 ++++++--- tests/tui_gateway/test_protocol.py | 69 +++++++++++++++++++ tui_gateway/server.py | 17 ++++- 4 files changed, 163 insertions(+), 15 deletions(-) diff --git a/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx b/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx index 603b45a1a2bd..ddf6571f5511 100644 --- a/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx +++ b/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 + onCurrentReady?: (branchCurrentSession: (messageId?: string) => Promise) => void onReady: (branchStoredSession: (storedSessionId: string, sessionProfile?: string | null) => Promise) => void requestGateway: (method: string, params?: Record) => Promise }) { const ref = (value: T): MutableRefObject => ({ current: value }) const actions = useSessionActions({ - activeSessionId: null, - activeSessionIdRef: ref(null), + activeSessionId, + activeSessionIdRef: ref(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 | undefined + const requestGateway = vi.fn(async (method: string, params?: Record) => { + 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) | null = null + render( + (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 diff --git a/apps/desktop/src/app/session/hooks/use-session-actions/index.ts b/apps/desktop/src/app/session/hooks/use-session-actions/index.ts index 4dbcfa6000ef..fd490666d126 100644 --- a/apps/desktop/src/app/session/hooks/use-session-actions/index.ts +++ b/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('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('session.branch', { + session_id: sourceSessionId, + count: branchMessages.length + }) + : await requestGateway('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) diff --git a/tests/tui_gateway/test_protocol.py b/tests/tui_gateway/test_protocol.py index de2e6c39b014..a32fea2768fe 100644 --- a/tests/tui_gateway/test_protocol.py +++ b/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 diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 95461910856a..863601d68eb3 100644 --- a/tui_gateway/server.py +++ b/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")