From 07b9090259b889ce95068c419534ca1fae62a85e Mon Sep 17 00:00:00 2001 From: DannyFengTianYu Date: Sun, 2 Aug 2026 22:09:16 +0800 Subject: [PATCH] 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) --- .../hooks/use-session-actions.test.tsx | 108 ++++++++++++++- .../hooks/use-session-actions/index.ts | 76 ++++++++--- .../hooks/use-session-actions/utils.test.ts | 42 ++++++ .../hooks/use-session-actions/utils.ts | 71 ++++++++++ tests/test_tui_gateway_server.py | 124 ++++++++++++++++++ tui_gateway/methods_session.py | 36 ++++- 6 files changed, 435 insertions(+), 22 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 f2f2ba9d6e..8c408e5a87 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 @@ -1112,19 +1112,30 @@ function BranchHarness({ navigate = vi.fn(), onCurrentReady, onReady, - requestGateway + onRefs, + requestGateway, + selectedStoredSessionId = null }: { activeSessionId?: string | null navigate?: ReturnType onCurrentReady?: (branchCurrentSession: (messageId?: string) => Promise) => void onReady: (branchStoredSession: (storedSessionId: string, sessionProfile?: string | null) => Promise) => void + onRefs?: (refs: { + activeSessionIdRef: MutableRefObject + selectedStoredSessionIdRef: MutableRefObject + }) => void requestGateway: (method: string, params?: Record) => Promise + selectedStoredSessionId?: string | null }) { const ref = (value: T): MutableRefObject => ({ current: value }) + const activeSessionIdRef = ref(activeSessionId) + const selectedStoredSessionIdRef = ref(selectedStoredSessionId) + + onRefs?.({ activeSessionIdRef, selectedStoredSessionIdRef }) const actions = useSessionActions({ activeSessionId, - activeSessionIdRef: ref(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()), - selectedStoredSessionId: null, - selectedStoredSessionIdRef: ref(null), + selectedStoredSessionId, + selectedStoredSessionIdRef, sessionStateByRuntimeIdRef: ref(new Map()), 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 | 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: 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) | null = null + render( + (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 + selectedStoredSessionIdRef: MutableRefObject + } | 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) | null = null + render( + (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 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 1971ac1ce8..20bdf2352b 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 @@ -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 => { creatingSessionRef.current = true @@ -1353,7 +1355,7 @@ export function useSessionActions({ const branched = sourceSessionId ? await requestGateway('session.branch', { session_id: sourceSessionId, - count: branchMessages.length + ...(branchCount !== undefined ? { count: branchCount } : {}) }) : await requestGateway('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 diff --git a/apps/desktop/src/app/session/hooks/use-session-actions/utils.test.ts b/apps/desktop/src/app/session/hooks/use-session-actions/utils.test.ts index bfe95ae015..b010a6f878 100644 --- a/apps/desktop/src/app/session/hooks/use-session-actions/utils.test.ts +++ b/apps/desktop/src/app/session/hooks/use-session-actions/utils.test.ts @@ -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' } diff --git a/apps/desktop/src/app/session/hooks/use-session-actions/utils.ts b/apps/desktop/src/app/session/hooks/use-session-actions/utils.ts index df7049b61b..b6cd683b59 100644 --- a/apps/desktop/src/app/session/hooks/use-session-actions/utils.ts +++ b/apps/desktop/src/app/session/hooks/use-session-actions/utils.ts @@ -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, diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index b6d17ae421..bfd1bfd66f 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -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" diff --git a/tui_gateway/methods_session.py b/tui_gateway/methods_session.py index ffd5cda771..fab90809f5 100644 --- a/tui_gateway/methods_session.py +++ b/tui_gateway/methods_session.py @@ -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")