From 32e7f883c687c8b5ff3a39d99e4e501b25fb94a4 Mon Sep 17 00:00:00 2001 From: Nimrod Gutman Date: Sat, 14 Mar 2026 15:49:47 +0200 Subject: [PATCH] fix(btw): keep side questions isolated in chat --- src/agents/btw.test.ts | 53 ++++++++++++++++++++++++++++ src/agents/btw.ts | 9 ++++- src/tui/tui-command-handlers.test.ts | 14 +++++--- src/tui/tui-command-handlers.ts | 9 +++++ src/tui/tui-event-handlers.test.ts | 43 ++++++++++++++++++++++ src/tui/tui-event-handlers.ts | 18 +++++++++- src/tui/tui.ts | 29 +++++++++++++++ 7 files changed, 169 insertions(+), 6 deletions(-) diff --git a/src/agents/btw.test.ts b/src/agents/btw.test.ts index f4a1e061daa7..478d0dcc74a2 100644 --- a/src/agents/btw.test.ts +++ b/src/agents/btw.test.ts @@ -586,6 +586,59 @@ describe("runBtwSideQuestion", () => { expect(result).toEqual({ text: "323" }); }); + it("falls back when the active run snapshot leaf no longer exists", async () => { + getActiveEmbeddedRunSnapshotMock.mockReturnValue({ + transcriptLeafId: "assistant-gone", + }); + branchMock.mockImplementationOnce(() => { + throw new Error("Entry 3235c7c4 not found"); + }); + streamSimpleMock.mockReturnValue( + makeAsyncEvents([ + { + type: "done", + reason: "stop", + message: { + role: "assistant", + content: [{ type: "text", text: "323" }], + provider: "anthropic", + api: "anthropic-messages", + model: "claude-sonnet-4-5", + stopReason: "stop", + usage: { + input: 1, + output: 2, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 3, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + timestamp: Date.now(), + }, + }, + ]), + ); + + const result = await runBtwSideQuestion({ + cfg: {} as never, + agentDir: "/tmp/agent", + provider: "anthropic", + model: "claude-sonnet-4-5", + question: "What is 17 * 19?", + sessionEntry: createSessionEntry(), + resolvedReasoningLevel: "off", + opts: {}, + isNewSession: false, + }); + + expect(branchMock).toHaveBeenCalledWith("assistant-gone"); + expect(resetLeafMock).toHaveBeenCalled(); + expect(result).toEqual({ text: "323" }); + expect(diagDebugMock).toHaveBeenCalledWith( + expect.stringContaining("btw snapshot leaf unavailable: sessionId=session-1"), + ); + }); + it("returns the BTW answer and retries transcript persistence after a session lock", async () => { acquireSessionWriteLockMock .mockRejectedValueOnce( diff --git a/src/agents/btw.ts b/src/agents/btw.ts index 08844578749d..3a2b2aaab4ce 100644 --- a/src/agents/btw.ts +++ b/src/agents/btw.ts @@ -305,7 +305,14 @@ export async function runBtwSideQuestion( } else if (activeRunSnapshot) { inFlightPrompt = activeRunSnapshot.inFlightPrompt; if (activeRunSnapshot.transcriptLeafId && sessionManager.branch) { - sessionManager.branch(activeRunSnapshot.transcriptLeafId); + try { + sessionManager.branch(activeRunSnapshot.transcriptLeafId); + } catch (error) { + diag.debug( + `btw snapshot leaf unavailable: sessionId=${sessionId} leaf=${activeRunSnapshot.transcriptLeafId} err=${String(error)}`, + ); + sessionManager.resetLeaf?.(); + } } else { sessionManager.resetLeaf?.(); } diff --git a/src/tui/tui-command-handlers.test.ts b/src/tui/tui-command-handlers.test.ts index 87f68217e3a0..026b63350bee 100644 --- a/src/tui/tui-command-handlers.test.ts +++ b/src/tui/tui-command-handlers.test.ts @@ -21,6 +21,7 @@ function createHarness(params?: { const addSystem = vi.fn(); const requestRender = vi.fn(); const noteLocalRunId = vi.fn(); + const noteLocalBtwRunId = vi.fn(); const loadHistory = params?.loadHistory ?? (vi.fn().mockResolvedValue(undefined) as LoadHistoryMock); const setActivityStatus = params?.setActivityStatus ?? (vi.fn() as SetActivityStatusMock); @@ -49,7 +50,9 @@ function createHarness(params?: { formatSessionKey: vi.fn(), applySessionInfoFromPatch: vi.fn(), noteLocalRunId, + noteLocalBtwRunId, forgetLocalRunId: vi.fn(), + forgetLocalBtwRunId: vi.fn(), requestExit: vi.fn(), }); @@ -64,6 +67,7 @@ function createHarness(params?: { loadHistory, setActivityStatus, noteLocalRunId, + noteLocalBtwRunId, state, }; } @@ -115,15 +119,17 @@ describe("tui command handlers", () => { it("sends /btw without hijacking the active main run", async () => { const setActivityStatus = vi.fn(); - const { handleCommand, sendChat, addUser, noteLocalRunId, state } = createHarness({ - activeChatRunId: "run-main", - setActivityStatus, - }); + const { handleCommand, sendChat, addUser, noteLocalRunId, noteLocalBtwRunId, state } = + createHarness({ + activeChatRunId: "run-main", + setActivityStatus, + }); await handleCommand("/btw what changed?"); expect(addUser).not.toHaveBeenCalled(); expect(noteLocalRunId).not.toHaveBeenCalled(); + expect(noteLocalBtwRunId).toHaveBeenCalledTimes(1); expect(state.activeChatRunId).toBe("run-main"); expect(setActivityStatus).not.toHaveBeenCalledWith("sending"); expect(setActivityStatus).not.toHaveBeenCalledWith("waiting"); diff --git a/src/tui/tui-command-handlers.ts b/src/tui/tui-command-handlers.ts index 9a6d63f53d83..491674af45ff 100644 --- a/src/tui/tui-command-handlers.ts +++ b/src/tui/tui-command-handlers.ts @@ -43,7 +43,9 @@ type CommandHandlerContext = { formatSessionKey: (key: string) => string; applySessionInfoFromPatch: (result: SessionsPatchResult) => void; noteLocalRunId: (runId: string) => void; + noteLocalBtwRunId?: (runId: string) => void; forgetLocalRunId?: (runId: string) => void; + forgetLocalBtwRunId?: (runId: string) => void; requestExit: () => void; }; @@ -70,7 +72,9 @@ export function createCommandHandlers(context: CommandHandlerContext) { formatSessionKey, applySessionInfoFromPatch, noteLocalRunId, + noteLocalBtwRunId, forgetLocalRunId, + forgetLocalBtwRunId, requestExit, } = context; @@ -513,6 +517,8 @@ export function createCommandHandlers(context: CommandHandlerContext) { noteLocalRunId(runId); state.activeChatRunId = runId; setActivityStatus("sending"); + } else { + noteLocalBtwRunId?.(runId); } tui.requestRender(); await client.sendChat({ @@ -528,6 +534,9 @@ export function createCommandHandlers(context: CommandHandlerContext) { tui.requestRender(); } } catch (err) { + if (isBtw) { + forgetLocalBtwRunId?.(runId); + } if (!isBtw && state.activeChatRunId) { forgetLocalRunId?.(state.activeChatRunId); } diff --git a/src/tui/tui-event-handlers.test.ts b/src/tui/tui-event-handlers.test.ts index 2e1046fc650e..2073afe308d8 100644 --- a/src/tui/tui-event-handlers.test.ts +++ b/src/tui/tui-event-handlers.test.ts @@ -79,12 +79,19 @@ describe("tui-event-handlers: handleAgentEvent", () => { const setActivityStatus = vi.fn(); const loadHistory = vi.fn(); const localRunIds = new Set(); + const localBtwRunIds = new Set(); const noteLocalRunId = (runId: string) => { localRunIds.add(runId); }; const forgetLocalRunId = localRunIds.delete.bind(localRunIds); const isLocalRunId = localRunIds.has.bind(localRunIds); const clearLocalRunIds = localRunIds.clear.bind(localRunIds); + const noteLocalBtwRunId = (runId: string) => { + localBtwRunIds.add(runId); + }; + const forgetLocalBtwRunId = localBtwRunIds.delete.bind(localBtwRunIds); + const isLocalBtwRunId = localBtwRunIds.has.bind(localBtwRunIds); + const clearLocalBtwRunIds = localBtwRunIds.clear.bind(localBtwRunIds); return { chatLog, @@ -94,9 +101,13 @@ describe("tui-event-handlers: handleAgentEvent", () => { setActivityStatus, loadHistory, noteLocalRunId, + noteLocalBtwRunId, forgetLocalRunId, isLocalRunId, clearLocalRunIds, + forgetLocalBtwRunId, + isLocalBtwRunId, + clearLocalBtwRunIds, }; }; @@ -117,6 +128,9 @@ describe("tui-event-handlers: handleAgentEvent", () => { loadHistory: context.loadHistory, isLocalRunId: context.isLocalRunId, forgetLocalRunId: context.forgetLocalRunId, + isLocalBtwRunId: context.isLocalBtwRunId, + forgetLocalBtwRunId: context.forgetLocalBtwRunId, + clearLocalBtwRunIds: context.clearLocalBtwRunIds, }); return { ...context, @@ -259,6 +273,35 @@ describe("tui-event-handlers: handleAgentEvent", () => { expect(tui.requestRender).toHaveBeenCalledTimes(1); }); + it("keeps a local BTW result visible when its empty final chat event arrives", () => { + const { state, btw, loadHistory, noteLocalBtwRunId, handleBtwEvent, handleChatEvent } = + createHandlersHarness({ + state: { activeChatRunId: null }, + }); + + noteLocalBtwRunId("run-btw"); + handleBtwEvent({ + kind: "btw", + runId: "run-btw", + sessionKey: state.currentSessionKey, + question: "what changed?", + text: "nothing important", + } satisfies BtwEvent); + + handleChatEvent({ + runId: "run-btw", + sessionKey: state.currentSessionKey, + state: "final", + } satisfies ChatEvent); + + expect(loadHistory).not.toHaveBeenCalled(); + expect(btw.showResult).toHaveBeenCalledWith({ + question: "what changed?", + text: "nothing important", + isError: undefined, + }); + }); + it("does not cross-match canonical session keys from different agents", () => { const { chatLog, handleChatEvent } = createHandlersHarness({ state: { diff --git a/src/tui/tui-event-handlers.ts b/src/tui/tui-event-handlers.ts index 2549bcbba384..6fda2d851638 100644 --- a/src/tui/tui-event-handlers.ts +++ b/src/tui/tui-event-handlers.ts @@ -36,6 +36,9 @@ type EventHandlerContext = { isLocalRunId?: (runId: string) => boolean; forgetLocalRunId?: (runId: string) => void; clearLocalRunIds?: () => void; + isLocalBtwRunId?: (runId: string) => boolean; + forgetLocalBtwRunId?: (runId: string) => void; + clearLocalBtwRunIds?: () => void; }; export function createEventHandlers(context: EventHandlerContext) { @@ -50,6 +53,9 @@ export function createEventHandlers(context: EventHandlerContext) { isLocalRunId, forgetLocalRunId, clearLocalRunIds, + isLocalBtwRunId, + forgetLocalBtwRunId, + clearLocalBtwRunIds, } = context; const finalizedRuns = new Map(); const sessionRuns = new Map(); @@ -88,6 +94,7 @@ export function createEventHandlers(context: EventHandlerContext) { sessionRuns.clear(); streamAssembler = new TuiStreamAssembler(); clearLocalRunIds?.(); + clearLocalBtwRunIds?.(); btw.clear(); }; @@ -202,7 +209,7 @@ export function createEventHandlers(context: EventHandlerContext) { } } noteSessionRun(evt.runId); - if (!state.activeChatRunId) { + if (!state.activeChatRunId && !isLocalBtwRunId?.(evt.runId)) { state.activeChatRunId = evt.runId; } if (evt.state === "delta") { @@ -214,7 +221,14 @@ export function createEventHandlers(context: EventHandlerContext) { setActivityStatus("streaming"); } if (evt.state === "final") { + const isLocalBtwRun = isLocalBtwRunId?.(evt.runId) ?? false; const wasActiveRun = state.activeChatRunId === evt.runId; + if (!evt.message && isLocalBtwRun) { + forgetLocalBtwRunId?.(evt.runId); + noteFinalizedRun(evt.runId); + tui.requestRender(); + return; + } if (!evt.message) { maybeRefreshHistoryForRun(evt.runId, { allowLocalWithoutDisplayableFinal: true, @@ -262,12 +276,14 @@ export function createEventHandlers(context: EventHandlerContext) { }); } if (evt.state === "aborted") { + forgetLocalBtwRunId?.(evt.runId); const wasActiveRun = state.activeChatRunId === evt.runId; chatLog.addSystem("run aborted"); terminateRun({ runId: evt.runId, wasActiveRun, status: "aborted" }); maybeRefreshHistoryForRun(evt.runId); } if (evt.state === "error") { + forgetLocalBtwRunId?.(evt.runId); const wasActiveRun = state.activeChatRunId === evt.runId; chatLog.addSystem(`run error: ${evt.errorMessage ?? "unknown"}`); terminateRun({ runId: evt.runId, wasActiveRun, status: "error" }); diff --git a/src/tui/tui.ts b/src/tui/tui.ts index 143dfb7d62fb..b9c67e76a292 100644 --- a/src/tui/tui.ts +++ b/src/tui/tui.ts @@ -344,6 +344,7 @@ export async function runTui(opts: TuiOptions) { let showThinking = false; let pairingHintShown = false; const localRunIds = new Set(); + const localBtwRunIds = new Set(); const deliverDefault = opts.deliver ?? false; const autoMessage = opts.message?.trim(); @@ -498,6 +499,29 @@ export async function runTui(opts: TuiOptions) { localRunIds.clear(); }; + const noteLocalBtwRunId = (runId: string) => { + if (!runId) { + return; + } + localBtwRunIds.add(runId); + if (localBtwRunIds.size > 200) { + const [first] = localBtwRunIds; + if (first) { + localBtwRunIds.delete(first); + } + } + }; + + const forgetLocalBtwRunId = (runId: string) => { + localBtwRunIds.delete(runId); + }; + + const isLocalBtwRunId = (runId: string) => localBtwRunIds.has(runId); + + const clearLocalBtwRunIds = () => { + localBtwRunIds.clear(); + }; + const client = await GatewayChatClient.connect({ url: opts.url, token: opts.token, @@ -825,6 +849,9 @@ export async function runTui(opts: TuiOptions) { isLocalRunId, forgetLocalRunId, clearLocalRunIds, + isLocalBtwRunId, + forgetLocalBtwRunId, + clearLocalBtwRunIds, }); const requestExit = () => { @@ -856,7 +883,9 @@ export async function runTui(opts: TuiOptions) { setActivityStatus, formatSessionKey, noteLocalRunId, + noteLocalBtwRunId, forgetLocalRunId, + forgetLocalBtwRunId, requestExit, });