From 0a704d469026ad2c3a2bd1b51f5728fe08e6bb3e Mon Sep 17 00:00:00 2001 From: Aidan Date: Tue, 9 Jun 2026 22:28:51 -0700 Subject: [PATCH] [aidan] chat: isolate branch history, fresh session on edit/switch (#62) --- backend/apps/agents/agent_manager.py | 37 ++++++++----------- .../src/app/pages/AgentChat/AgentChat.tsx | 4 ++ .../AgentChat/shell/MessageActionBar.tsx | 6 ++- frontend/src/shared/ws/WebSocketManager.ts | 3 +- 4 files changed, 26 insertions(+), 24 deletions(-) diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index 6db1abcb..955ef29f 100644 --- a/backend/apps/agents/agent_manager.py +++ b/backend/apps/agents/agent_manager.py @@ -1768,27 +1768,20 @@ class AgentManager: except Exception as e: logger.debug(f"thinking_level param injection skipped: {e}") - # MCPActivate fresh-restart path: when the session has prior - # turns AND the user just activated a new MCP, the bundled CLI - # won't re-read mcp_servers from a `resume + fork_session` - # combo (the transport snapshot from the original launch is - # what serves tool schemas). Symptom: model calls hallucinated - # names like `Searchgmail`/`Listemails` instead of the real - # `mcp__google-workspace__query_gmail_emails` because it - # never received the schemas. Soft restart: drop resume + - # sdk_session_id, replay history via the prompt, let the SDK - # build a clean transport with the activated server in its - # mcp_servers dict from the start. Costs one cold-start TTFT - # (~200-400ms) on the auto-continuation turn; that turn is - # already happening anyway because pending_continuation fires - # right after MCPActivate. - if session.needs_fresh_session and session.sdk_session_id: - logger.info( - f"[MCP-DEBUG] Fresh-session restart for {session_id}: dropping " - f"sdk_session_id={session.sdk_session_id} so the new MCP servers " - f"({session.active_mcps}) take effect." - ) - session.sdk_session_id = None + # Fresh-restart path: some session changes must not reuse the + # CLI's resume transcript. MCPActivate needs a new transport so + # tool schemas are reread; branch edits/switches need the model + # to see only _get_branch_messages(session), not facts from the + # old branch's SDK transcript. Soft restart: drop resume + + # sdk_session_id, replay local history via the prompt, let the + # SDK build a clean session from the current app state. + if session.needs_fresh_session: + if session.sdk_session_id: + logger.info( + f"Fresh-session restart for {session_id}: dropping " + f"sdk_session_id={session.sdk_session_id}; active_mcps={session.active_mcps}" + ) + session.sdk_session_id = None session.needs_fresh_session = False session.needs_fork = False # superseded by the fresh restart @@ -3569,6 +3562,7 @@ class AgentManager: ) session.branches[new_branch_id] = new_branch session.active_branch_id = new_branch_id + session.needs_fresh_session = True edited_msg = Message( @@ -3617,6 +3611,7 @@ class AgentManager: if branch_id not in session.branches: raise ValueError(f"Branch {branch_id} not found") session.active_branch_id = branch_id + session.needs_fresh_session = True await ws_manager.send_to_session(session_id, "agent:branch_switched", { "session_id": session_id, "active_branch_id": branch_id, diff --git a/frontend/src/app/pages/AgentChat/AgentChat.tsx b/frontend/src/app/pages/AgentChat/AgentChat.tsx index 94d7dc97..18888f8d 100644 --- a/frontend/src/app/pages/AgentChat/AgentChat.tsx +++ b/frontend/src/app/pages/AgentChat/AgentChat.tsx @@ -927,6 +927,7 @@ const AgentChat: React.FC = ({ sessionId: sessionIdProp, onClose } const isActive = session.status === 'running' || session.status === 'waiting_approval' || session.status === 'draft'; + const branchNavLocked = agentBusy || hasStreaming; const statusStyle = STATUS_STYLES[session.status] || { color: c.text.tertiary, bg: c.bg.secondary }; return ( @@ -1386,11 +1387,14 @@ const AgentChat: React.FC = ({ sessionId: sessionIdProp, onClose ? { currentIndex: Math.max(0, currentBranchIdx), totalBranches: siblings.length, + disabled: branchNavLocked, onPrevious: () => { + if (branchNavLocked) return; const prevBranch = siblings[Math.max(0, currentBranchIdx - 1)]; if (prevBranch && id) dispatch(switchBranch({ sessionId: id, branchId: prevBranch })); }, onNext: () => { + if (branchNavLocked) return; const nextBranch = siblings[Math.min(siblings.length - 1, currentBranchIdx + 1)]; if (nextBranch && id) dispatch(switchBranch({ sessionId: id, branchId: nextBranch })); }, diff --git a/frontend/src/app/pages/AgentChat/shell/MessageActionBar.tsx b/frontend/src/app/pages/AgentChat/shell/MessageActionBar.tsx index 63e3e05d..eb69863e 100644 --- a/frontend/src/app/pages/AgentChat/shell/MessageActionBar.tsx +++ b/frontend/src/app/pages/AgentChat/shell/MessageActionBar.tsx @@ -21,6 +21,7 @@ import FeedbackDialog, { Sentiment } from './FeedbackDialog'; interface BranchNavProps { currentIndex: number; totalBranches: number; + disabled?: boolean; onPrevious: () => void; onNext: () => void; } @@ -108,7 +109,7 @@ const MessageActionBar: React.FC = ({ @@ -127,7 +128,7 @@ const MessageActionBar: React.FC = ({ @@ -199,6 +200,7 @@ export default React.memo(MessageActionBar, (prev, next) => ( && !!prev.onBranch === !!next.onBranch && prev.branchNav?.currentIndex === next.branchNav?.currentIndex && prev.branchNav?.totalBranches === next.branchNav?.totalBranches + && prev.branchNav?.disabled === next.branchNav?.disabled && prev.messageId === next.messageId && prev.sessionId === next.sessionId )); diff --git a/frontend/src/shared/ws/WebSocketManager.ts b/frontend/src/shared/ws/WebSocketManager.ts index 875f09a3..5156d085 100644 --- a/frontend/src/shared/ws/WebSocketManager.ts +++ b/frontend/src/shared/ws/WebSocketManager.ts @@ -22,7 +22,7 @@ import { setTurnLabel, clearTurnLabel, } from '../state/agentsSlice'; -import { streamStart, streamDelta, streamEnd } from '../state/streamingSlice'; +import { streamStart, streamDelta, streamEnd, clearStreamingForSession } from '../state/streamingSlice'; import { addBrowserCardFromBackend, removeBrowserCard, setBrowserCardPosition, setGlowingBrowserCards, GRID_GAP } from '../state/dashboardLayoutSlice'; import { upsertOutput } from '../state/outputsSlice'; import { getAuthToken } from '../config'; @@ -493,6 +493,7 @@ class WebSocketManager { // before its own aux call lands. if (session_id && (data.status === 'completed' || data.status === 'error' || data.status === 'stopped')) { store.dispatch(clearTurnLabel(session_id)); + store.dispatch(clearStreamingForSession(session_id)); } } // Clean spawned browser cards when the PARENT finishes, never per