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")