diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 4acc122cd..1f0ec07f9 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -69,6 +69,10 @@ - Added `app.retry` as an `Alt+R` keybinding for retrying the last failed or aborted assistant turn ([#2790](https://github.com/can1357/oh-my-pi/issues/2790)). +### Added + +- Added `b branch` promotion for completed `/btw` answers, creating a branch that preserves the side-question input and full assistant response including thinking blocks. + ## [16.0.5] - 2026-06-17 ### Added diff --git a/packages/coding-agent/src/modes/components/btw-panel.ts b/packages/coding-agent/src/modes/components/btw-panel.ts index 91c461779..a1cd3ed42 100644 --- a/packages/coding-agent/src/modes/components/btw-panel.ts +++ b/packages/coding-agent/src/modes/components/btw-panel.ts @@ -58,6 +58,10 @@ export class BtwPanelComponent extends Container { this.#rebuild(); } + isBranchable(): boolean { + return this.#state === "complete" && this.#answer.trim().length > 0; + } + close(): void { this.#closed = true; } @@ -85,7 +89,7 @@ export class BtwPanelComponent extends Container { case "running": return theme.fg("muted", "Esc cancel /btw"); case "complete": - return theme.fg("muted", "Esc dismiss"); + return theme.fg("muted", this.isBranchable() ? "b branch · Esc dismiss" : "Esc dismiss"); case "aborted": return theme.fg("warning", `${theme.status.warning} Cancelled · Esc dismiss`); case "error": diff --git a/packages/coding-agent/src/modes/controllers/btw-controller.ts b/packages/coding-agent/src/modes/controllers/btw-controller.ts index 73a44cce1..11ac6e615 100644 --- a/packages/coding-agent/src/modes/controllers/btw-controller.ts +++ b/packages/coding-agent/src/modes/controllers/btw-controller.ts @@ -1,3 +1,4 @@ +import type { AssistantMessage } from "@oh-my-pi/pi-ai"; import { prompt } from "@oh-my-pi/pi-utils"; import btwUserPrompt from "../../prompts/system/btw-user.md" with { type: "text" }; import { BtwPanelComponent } from "../components/btw-panel"; @@ -7,10 +8,37 @@ interface BtwRequest { component: BtwPanelComponent; abortController: AbortController; question: string; + leafId: string | null; +} + +function assistantMessageWithReplyText(assistantMessage: AssistantMessage, replyText: string): AssistantMessage { + const content: AssistantMessage["content"] = []; + let replacedText = false; + for (const part of assistantMessage.content) { + if (part.type === "thinking") { + content.push({ type: "thinking", thinking: part.thinking }); + continue; + } + if (part.type === "redactedThinking") continue; + if (part.type !== "text") { + content.push(part); + continue; + } + if (replacedText) continue; + content.push({ type: "text", text: replyText }); + replacedText = true; + } + if (!replacedText) content.push({ type: "text", text: replyText }); + return { ...assistantMessage, content, providerPayload: undefined }; } export class BtwController { #activeRequest: BtwRequest | undefined; + #lastQuestion: string | undefined; + #lastReplyText: string | undefined; + #lastAssistantMessage: AssistantMessage | undefined; + #lastLeafId: string | null | undefined; + #branchInFlight = false; constructor(private readonly ctx: InteractiveModeContext) {} @@ -18,6 +46,29 @@ export class BtwController { return this.#activeRequest !== undefined; } + canBranch(): boolean { + return ( + !this.#branchInFlight && + this.#activeRequest?.component.isBranchable() === true && + this.#lastQuestion !== undefined && + this.#lastReplyText !== undefined && + this.#lastAssistantMessage !== undefined && + this.#lastLeafId !== null && + this.#lastLeafId === this.ctx.sessionManager.getLeafId() + ); + } + + async handleBranch(): Promise { + if (!this.canBranch() || !this.#lastQuestion || !this.#lastAssistantMessage) return false; + this.#branchInFlight = true; + try { + await this.ctx.handleBtwBranch(this.#lastQuestion, this.#lastAssistantMessage); + return true; + } finally { + this.#branchInFlight = false; + } + } + handleEscape(): boolean { if (!this.#activeRequest) return false; this.#closeActiveRequest({ abort: this.#activeRequest.abortController.signal.aborted === false }); @@ -47,6 +98,7 @@ export class BtwController { component: new BtwPanelComponent({ question: trimmedQuestion, tui: this.ctx.ui }), abortController: new AbortController(), question: trimmedQuestion, + leafId: this.ctx.sessionManager.getLeafId(), }; this.ctx.btwContainer.clear(); this.ctx.btwContainer.addChild(request.component); @@ -58,7 +110,7 @@ export class BtwController { async #runRequest(request: BtwRequest): Promise { try { const promptText = prompt.render(btwUserPrompt, { question: request.question }); - const { replyText } = await this.ctx.session.runEphemeralTurn({ + const { replyText, assistantMessage } = await this.ctx.session.runEphemeralTurn({ promptText, onTextDelta: delta => { if (this.#isActiveRequest(request)) { @@ -75,6 +127,14 @@ export class BtwController { request.component.setAnswer(replyText); } request.component.markComplete(); + if (request.component.isBranchable()) { + this.#lastQuestion = request.question; + this.#lastReplyText = replyText; + this.#lastAssistantMessage = assistantMessageWithReplyText(assistantMessage, replyText); + this.#lastLeafId = request.leafId; + } else { + this.#clearBranchState(); + } } catch (error) { if (!this.#isActiveRequest(request)) { return; @@ -91,6 +151,7 @@ export class BtwController { const request = this.#activeRequest; if (!request) return; this.#activeRequest = undefined; + this.#clearBranchState(); if (options.abort) { request.abortController.abort(); } @@ -99,6 +160,13 @@ export class BtwController { this.ctx.ui.requestRender(); } + #clearBranchState(): void { + this.#lastQuestion = undefined; + this.#lastReplyText = undefined; + this.#lastAssistantMessage = undefined; + this.#lastLeafId = undefined; + } + #isActiveRequest(request: BtwRequest): boolean { return this.#activeRequest === request; } diff --git a/packages/coding-agent/src/modes/controllers/input-controller.ts b/packages/coding-agent/src/modes/controllers/input-controller.ts index ff09e3e9c..921d5516b 100644 --- a/packages/coding-agent/src/modes/controllers/input-controller.ts +++ b/packages/coding-agent/src/modes/controllers/input-controller.ts @@ -78,6 +78,7 @@ export class InputController { #enhancedPaste?: EnhancedPasteController; #focusedLeftTapListenerInstalled = false; + #btwBranchListenerInstalled = false; // Tap counter for the double-← gesture; reset whenever a quiet gap // (>= LEFT_DOUBLE_TAP_MAX_GAP_MS) starts a fresh sequence. See // #detectLeftDoubleTap. @@ -143,6 +144,16 @@ export class InputController { return { consume: true }; }); } + if (!this.#btwBranchListenerInstalled) { + this.#btwBranchListenerInstalled = true; + this.ctx.ui.addInputListener(data => { + if (!matchesKey(data, "b")) return undefined; + if (!this.ctx.canBranchBtw()) return undefined; + if (this.ctx.editor.getText().trim()) return undefined; + void this.ctx.handleBtwBranchKey(); + return { consume: true }; + }); + } this.ctx.editor.onEscape = () => { // Active context maintenance owns Esc: auto/manual compaction, // handoff generation, and auto-retry backoff all advertise diff --git a/packages/coding-agent/src/modes/interactive-mode.ts b/packages/coding-agent/src/modes/interactive-mode.ts index cc3fef1a2..3a43ba35a 100644 --- a/packages/coding-agent/src/modes/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive-mode.ts @@ -3724,6 +3724,34 @@ export class InteractiveMode implements InteractiveModeContext { return this.#btwController.handleEscape(); } + canBranchBtw(): boolean { + return this.#btwController.canBranch(); + } + + handleBtwBranchKey(): Promise { + return this.#btwController.handleBranch(); + } + + async handleBtwBranch(question: string, assistantMessage: AssistantMessage): Promise { + try { + const result = await this.session.branchFromBtw(question, assistantMessage); + if (result.cancelled) { + this.showStatus("/btw branch cancelled", { dim: true }); + return; + } + this.#btwController.dispose(); + this.#omfgController.dispose(); + this.chatContainer.clear(); + this.renderInitialMessages({ clearTerminalHistory: true }); + this.updateEditorBorderColor(); + this.showStatus( + result.sessionFile ? `Branched /btw to ${path.basename(result.sessionFile)}` : "Branched /btw", + ); + } catch (error) { + this.showError(`Cannot branch /btw: ${error instanceof Error ? error.message : String(error)}`); + } + } + handleOmfgCommand(complaint: string): Promise { return this.#omfgController.start(complaint); } diff --git a/packages/coding-agent/src/modes/types.ts b/packages/coding-agent/src/modes/types.ts index f3adf48a1..aa4ee1695 100644 --- a/packages/coding-agent/src/modes/types.ts +++ b/packages/coding-agent/src/modes/types.ts @@ -339,6 +339,9 @@ export interface InteractiveModeContext { handleTanCommand(work: string): Promise; hasActiveBtw(): boolean; handleBtwEscape(): boolean; + handleBtwBranchKey(): Promise; + canBranchBtw(): boolean; + handleBtwBranch(question: string, assistantMessage: AssistantMessage): Promise; handleOmfgCommand(complaint: string): Promise; hasActiveOmfg(): boolean; handleOmfgEscape(): boolean; diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 1fcc8818d..774b6dcf6 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -11201,6 +11201,94 @@ export class AgentSession { return { selectedText, cancelled: false }; } + async branchFromBtw( + question: string, + assistantMessage: AssistantMessage, + ): Promise<{ cancelled: boolean; sessionFile: string | undefined }> { + const previousSessionFile = this.sessionFile; + if (!this.sessionManager.getSessionFile()) { + throw new Error("Cannot branch /btw: session is not persisted"); + } + + const leafId = this.sessionManager.getLeafId(); + if (!leafId) { + throw new Error("Cannot branch /btw: current session has no leaf"); + } + + if ( + this.isBashRunning || + this.isEvalRunning || + this.isCompacting || + this.isGeneratingHandoff || + this.isRetrying + ) { + throw new Error("Cannot branch /btw while session maintenance or user work is still running"); + } + + if (this.#extensionRunner?.hasHandlers("session_before_branch")) { + const result = (await this.#extensionRunner.emit({ + type: "session_before_branch", + entryId: leafId, + })) as SessionBeforeBranchResult | undefined; + + if (result?.cancel) { + return { cancelled: true, sessionFile: previousSessionFile }; + } + } + + await this.#cancelPostPromptTasks(); + if ( + this.isBashRunning || + this.isEvalRunning || + this.isCompacting || + this.isGeneratingHandoff || + this.isRetrying + ) { + throw new Error("Cannot branch /btw while session maintenance or user work is still running"); + } + + this.#pendingNextTurnMessages = []; + this.#scheduledHiddenNextTurnGeneration = undefined; + this.agent.replaceQueues([], []); + if (this.isStreaming) { + await this.abort({ goalReason: "internal", reason: "branching /btw" }); + this.agent.replaceQueues([], []); + } + await this.sessionManager.flush(); + this.#cancelOwnAsyncJobs(); + + this.sessionManager.createBranchedSession(leafId); + this.sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: question }], + timestamp: Date.now(), + }); + this.sessionManager.appendMessage(assistantMessage); + this.#syncTodoPhasesFromBranch(); + this.#freshProviderSessionId = undefined; + this.#syncAgentSessionId(); + this.#rekeyHindsightMemoryForCurrentSessionId(); + this.#rekeyMnemopiMemoryForCurrentSessionId(); + this.#resetHindsightConversationTrackingIfHindsight(); + this.#resetMnemopiConversationTrackingIfMnemopi(); + + const sessionContext = this.buildDisplaySessionContext(); + await this.#restoreMCPSelectionsForSessionContext(sessionContext); + + if (this.#extensionRunner) { + await this.#extensionRunner.emit({ + type: "session_branch", + previousSessionFile, + }); + } + + this.agent.replaceMessages(sessionContext.messages); + this.#advisorRuntime?.reset(); + this.#closeCodexProviderSessionsForHistoryRewrite(); + + return { cancelled: false, sessionFile: this.sessionFile }; + } + // ========================================================================= // Tree Navigation // ========================================================================= diff --git a/packages/coding-agent/test/agent-session-btw-branch.test.ts b/packages/coding-agent/test/agent-session-btw-branch.test.ts new file mode 100644 index 000000000..62fdb04fb --- /dev/null +++ b/packages/coding-agent/test/agent-session-btw-branch.test.ts @@ -0,0 +1,295 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; +import * as fs from "node:fs"; +import * as os from "node:os"; +import * as path from "node:path"; +import { Agent } from "@oh-my-pi/pi-agent-core"; +import type { AssistantMessage } from "@oh-my-pi/pi-ai"; +import { createMockModel, type MockHandler } from "@oh-my-pi/pi-ai/providers/mock"; +import { getBundledModel } from "@oh-my-pi/pi-catalog/models"; +import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; +import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import type { ExtensionRunner } from "@oh-my-pi/pi-coding-agent/extensibility/extensions"; +import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; +import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage"; +import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; +import { Snowflake } from "@oh-my-pi/pi-utils"; + +function createBtwAssistant(): AssistantMessage { + return { + role: "assistant", + content: [ + { type: "thinking", thinking: "Check the failure mode first.", thinkingSignature: "sig" }, + { type: "text", text: "The fix is to branch the side answer." }, + ], + api: "anthropic-messages", + provider: "anthropic", + model: "claude-sonnet-4-5", + usage: { + input: 1, + output: 2, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 3, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: Date.now(), + }; +} + +describe("AgentSession.branchFromBtw", () => { + let tempDir: string; + let session: AgentSession | undefined; + let authStorage: AuthStorage | undefined; + + beforeEach(() => { + tempDir = path.join(os.tmpdir(), `pi-btw-branch-test-${Snowflake.next()}`); + fs.mkdirSync(tempDir, { recursive: true }); + }); + + afterEach(async () => { + await session?.dispose(); + authStorage?.close(); + await fs.promises + .rm(tempDir, { recursive: true, force: true, maxRetries: 5, retryDelay: 50 }) + .catch(() => undefined); + vi.restoreAllMocks(); + }); + + async function createSession(options?: { + persisted?: boolean; + extensionRunner?: ExtensionRunner; + handler?: MockHandler; + }) { + const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; + const mock = createMockModel({ handler: options?.handler ?? (() => ({ content: ["unused"] })) }); + const agent = new Agent({ + getApiKey: () => "test-key", + initialState: { model, systemPrompt: ["Test"], tools: [] }, + streamFn: mock.stream, + }); + const sessionManager = + options?.persisted === false ? SessionManager.inMemory() : SessionManager.create(tempDir, tempDir); + const settings = Settings.isolated({ "compaction.enabled": false }); + authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); + const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); + authStorage.setRuntimeApiKey("anthropic", "test-key"); + session = new AgentSession({ + agent, + sessionManager, + settings, + modelRegistry, + extensionRunner: options?.extensionRunner, + }); + return session; + } + + it("creates a persisted branch with the /btw user input and complete assistant message", async () => { + const activeSession = await createSession(); + activeSession.sessionManager.appendMessage({ role: "user", content: "seed", timestamp: Date.now() - 2 }); + activeSession.sessionManager.appendMessage({ + role: "assistant", + content: [{ type: "text", text: "seed response" }], + api: "anthropic-messages", + provider: "anthropic", + model: "claude-sonnet-4-5", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: Date.now() - 1, + }); + activeSession.agent.replaceMessages(activeSession.sessionManager.buildSessionContext().messages); + await activeSession.sessionManager.flush(); + const originalFile = activeSession.sessionFile; + expect(originalFile).toBeDefined(); + const originalRaw = fs.readFileSync(originalFile!, "utf8"); + const assistantMessage = createBtwAssistant(); + + const result = await activeSession.branchFromBtw("why did this fail?", assistantMessage); + + expect(result.cancelled).toBe(false); + expect(result.sessionFile).toBe(activeSession.sessionFile); + expect(result.sessionFile).toBeDefined(); + expect(result.sessionFile).not.toBe(originalFile); + expect(fs.readFileSync(originalFile!, "utf8")).toBe(originalRaw); + const messages = activeSession.messages; + expect(messages.at(-2)).toMatchObject({ role: "user", content: [{ type: "text", text: "why did this fail?" }] }); + expect(messages.at(-1)).toEqual(assistantMessage); + }); + + it("honors session_before_branch cancellation without creating a branch", async () => { + const emit = vi.fn(async () => ({ cancel: true })); + const extensionRunner = { + hasHandlers: vi.fn((eventType: string) => eventType === "session_before_branch"), + emit, + } as unknown as ExtensionRunner; + const activeSession = await createSession({ extensionRunner }); + activeSession.sessionManager.appendMessage({ role: "user", content: "seed", timestamp: Date.now() }); + await activeSession.sessionManager.flush(); + const originalFile = activeSession.sessionFile; + + const result = await activeSession.branchFromBtw("question", createBtwAssistant()); + + expect(result).toEqual({ cancelled: true, sessionFile: originalFile }); + expect(activeSession.sessionFile).toBe(originalFile); + expect(emit).toHaveBeenCalledWith({ + type: "session_before_branch", + entryId: activeSession.sessionManager.getLeafId(), + }); + }); + + it("syncs promoted /btw messages into live context even when hooks skip conversation restore", async () => { + const extensionRunner = { + hasHandlers: vi.fn((eventType: string) => eventType === "session_before_branch"), + emit: vi.fn(async () => ({ skipConversationRestore: true })), + } as unknown as ExtensionRunner; + const activeSession = await createSession({ extensionRunner }); + activeSession.sessionManager.appendMessage({ role: "user", content: "seed", timestamp: Date.now() }); + activeSession.agent.replaceMessages(activeSession.sessionManager.buildSessionContext().messages); + await activeSession.sessionManager.flush(); + const assistantMessage = createBtwAssistant(); + + const result = await activeSession.branchFromBtw("question", assistantMessage); + + expect(result.cancelled).toBe(false); + const messages = activeSession.messages; + expect(messages.at(-2)).toMatchObject({ role: "user", content: [{ type: "text", text: "question" }] }); + expect(messages.at(-1)).toEqual(assistantMessage); + }); + + it("aborts an in-flight main stream before switching to the /btw branch", async () => { + const providerStarted = Promise.withResolvers(); + const activeSession = await createSession({ + handler: () => { + providerStarted.resolve(); + return { content: ["main response should not move"], delayMs: 60_000 }; + }, + }); + activeSession.sessionManager.appendMessage({ role: "user", content: "seed", timestamp: Date.now() }); + await activeSession.sessionManager.flush(); + + const promptPromise = activeSession.prompt("main prompt"); + await providerStarted.promise; + expect(activeSession.isStreaming).toBe(true); + await activeSession.followUp("queued follow-up should not move"); + expect(activeSession.queuedMessageCount).toBe(1); + + const assistantMessage = createBtwAssistant(); + const result = await activeSession.branchFromBtw("question", assistantMessage); + await promptPromise; + + expect(result.cancelled).toBe(false); + const messages = activeSession.messages; + expect(messages.at(-2)).toMatchObject({ role: "user", content: [{ type: "text", text: "question" }] }); + expect(messages.at(-1)).toEqual(assistantMessage); + expect(messages).not.toContainEqual( + expect.objectContaining({ + role: "assistant", + content: [{ type: "text", text: "main response should not move" }], + }), + ); + expect(activeSession.queuedMessageCount).toBe(0); + expect(messages).not.toContainEqual( + expect.objectContaining({ + role: "user", + content: [{ type: "text", text: "queued follow-up should not move" }], + }), + ); + }); + + it("refuses to branch /btw while user bash work is still running", async () => { + const activeSession = await createSession(); + activeSession.sessionManager.appendMessage({ role: "user", content: "seed", timestamp: Date.now() }); + await activeSession.sessionManager.flush(); + + const bashPromise = activeSession.executeBash('bun -e "await Bun.sleep(60_000)"', () => undefined, { + useUserShell: false, + }); + while (!activeSession.isBashRunning) await Bun.sleep(1); + + await expect(activeSession.branchFromBtw("question", createBtwAssistant())).rejects.toThrow( + "Cannot branch /btw while session maintenance or user work is still running", + ); + + activeSession.abortBash(); + await bashPromise.catch(() => undefined); + }); + + it("refuses to branch /btw while user Python work is still running", async () => { + const activeSession = await createSession(); + activeSession.sessionManager.appendMessage({ role: "user", content: "seed", timestamp: Date.now() }); + await activeSession.sessionManager.flush(); + const abortController = new AbortController(); + const execution = Promise.withResolvers().promise; + activeSession.trackEvalExecution(execution, abortController).catch(() => undefined); + expect(activeSession.isEvalRunning).toBe(true); + + await expect(activeSession.branchFromBtw("question", createBtwAssistant())).rejects.toThrow( + "Cannot branch /btw while session maintenance or user work is still running", + ); + + abortController.abort(); + }); + + it("refuses to branch /btw while context maintenance is running", async () => { + const activeSession = await createSession(); + activeSession.sessionManager.appendMessage({ role: "user", content: "seed", timestamp: Date.now() }); + await activeSession.sessionManager.flush(); + const sessionWithMaintenance = activeSession as AgentSession & { _maintenanceForTest?: boolean }; + Object.defineProperty(sessionWithMaintenance, "isCompacting", { + get: () => sessionWithMaintenance._maintenanceForTest === true, + }); + sessionWithMaintenance._maintenanceForTest = true; + + await expect(activeSession.branchFromBtw("question", createBtwAssistant())).rejects.toThrow( + "Cannot branch /btw while session maintenance or user work is still running", + ); + }); + + it("cancels post-prompt work after branch hooks before switching sessions", async () => { + const hookRelease = Promise.withResolvers(); + const extensionRunner = { + hasHandlers: vi.fn((eventType: string) => eventType === "session_before_branch"), + emit: vi.fn(async () => { + await hookRelease.promise; + return undefined; + }), + } as unknown as ExtensionRunner; + const activeSession = await createSession({ extensionRunner }); + activeSession.sessionManager.appendMessage({ role: "user", content: "seed", timestamp: Date.now() }); + await activeSession.sessionManager.flush(); + activeSession.queueDeferredMessage({ + role: "custom", + customType: "test-hidden-message", + content: "hidden", + display: false, + timestamp: Date.now(), + }); + expect(activeSession.hasPostPromptWork).toBe(true); + + const branchPromise = activeSession.branchFromBtw("question", createBtwAssistant()); + await Promise.resolve(); + expect(activeSession.hasPostPromptWork).toBe(true); + + hookRelease.resolve(); + const result = await branchPromise; + + expect(result.cancelled).toBe(false); + expect(activeSession.hasPostPromptWork).toBe(false); + }); + + it("throws for in-memory sessions", async () => { + const activeSession = await createSession({ persisted: false }); + activeSession.sessionManager.appendMessage({ role: "user", content: "seed", timestamp: Date.now() }); + + await expect(activeSession.branchFromBtw("question", createBtwAssistant())).rejects.toThrow( + "Cannot branch /btw: session is not persisted", + ); + }); +}); diff --git a/packages/coding-agent/test/input-controller-keybindings.test.ts b/packages/coding-agent/test/input-controller-keybindings.test.ts index f2f97f172..2f736cbf7 100644 --- a/packages/coding-agent/test/input-controller-keybindings.test.ts +++ b/packages/coding-agent/test/input-controller-keybindings.test.ts @@ -72,6 +72,8 @@ async function createContext() { retry, }; const updatePendingMessagesDisplay = vi.fn(); + const handleBtwBranchKey = vi.fn(async () => true); + const canBranchBtw = vi.fn(() => false); const editor: FakeEditor = { setText(text: string) { editorText = text; @@ -152,6 +154,8 @@ async function createContext() { showModelSelector, updateEditorBorderColor: vi.fn(), hasActiveBtw: vi.fn(() => false), + handleBtwBranchKey, + canBranchBtw, showError: vi.fn(), showStatus: vi.fn(), } as unknown as InteractiveModeContext; @@ -170,6 +174,9 @@ async function createContext() { retry, abort, resetDisplay, + handleBtwBranchKey, + addInputListener, + canBranchBtw, }, }; } @@ -281,6 +288,47 @@ describe("InputController keybinding setup", () => { expect(editor.getText()).toBe(""); }); + it("routes b to branch a branchable /btw panel", async () => { + const { InputController, ctx, spies } = await createContext(); + (ctx.canBranchBtw as unknown as { mockReturnValue(value: boolean): void }).mockReturnValue(true); + const controller = new InputController(ctx); + + controller.setupKeyHandlers(); + const listener = spies.addInputListener.mock.calls[1]?.[0]; + expect(listener).toBeDefined(); + const result = listener?.("b"); + + expect(result).toEqual({ consume: true }); + expect(spies.handleBtwBranchKey).toHaveBeenCalledTimes(1); + }); + + it("lets b fall through while the editor has draft text", async () => { + const { InputController, ctx, editor, spies } = await createContext(); + (ctx.canBranchBtw as unknown as { mockReturnValue(value: boolean): void }).mockReturnValue(true); + editor.setText("build a branch"); + const controller = new InputController(ctx); + + controller.setupKeyHandlers(); + const listener = spies.addInputListener.mock.calls[1]?.[0]; + expect(listener).toBeDefined(); + const result = listener?.("b"); + + expect(result).toBeUndefined(); + expect(spies.handleBtwBranchKey).not.toHaveBeenCalled(); + }); + + it("lets b fall through when /btw is not branchable", async () => { + const { InputController, ctx, spies } = await createContext(); + const controller = new InputController(ctx); + + controller.setupKeyHandlers(); + const listener = spies.addInputListener.mock.calls[1]?.[0]; + expect(listener).toBeDefined(); + const result = listener?.("b"); + + expect(result).toBeUndefined(); + expect(spies.handleBtwBranchKey).not.toHaveBeenCalled(); + }); it("empty Enter aborts the active stream when queued messages are pending", async () => { const { InputController, ctx, editor, spies } = await createContext(); const session = ctx.session as unknown as { isStreaming: boolean; queuedMessageCount: number }; diff --git a/packages/coding-agent/test/modes/controllers/btw-controller.test.ts b/packages/coding-agent/test/modes/controllers/btw-controller.test.ts index fad9d3d8c..5e1194e1e 100644 --- a/packages/coding-agent/test/modes/controllers/btw-controller.test.ts +++ b/packages/coding-agent/test/modes/controllers/btw-controller.test.ts @@ -1,5 +1,6 @@ import { beforeAll, describe, expect, it, vi } from "bun:test"; import type { AssistantMessage, Usage } from "@oh-my-pi/pi-ai"; +import { BtwPanelComponent } from "@oh-my-pi/pi-coding-agent/modes/components/btw-panel"; import { BtwController } from "@oh-my-pi/pi-coding-agent/modes/controllers/btw-controller"; import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; import type { InteractiveModeContext } from "@oh-my-pi/pi-coding-agent/modes/types"; @@ -48,18 +49,42 @@ function makeFakeSession( } function makeCtx(session: InteractiveModeContext["session"], btwContainer = new Container()): InteractiveModeContext { + let leafId: string | null = "leaf-1"; return { ui: { requestRender: vi.fn(), requestComponentRender: vi.fn() } as unknown as TUI, btwContainer, session, + sessionManager: { getLeafId: () => leafId } as unknown as InteractiveModeContext["sessionManager"], showStatus: vi.fn(), showError: vi.fn(), - } as unknown as InteractiveModeContext; + handleBtwBranch: vi.fn(async () => {}), + setTestLeafId(nextLeafId: string | null) { + leafId = nextLeafId; + }, + } as unknown as InteractiveModeContext & { setTestLeafId(nextLeafId: string | null): void }; } beforeAll(async () => { await initTheme(); }); +async function drainBtwRequest(): Promise { + await Promise.resolve(); + await Promise.resolve(); +} + +describe("BtwPanelComponent", () => { + it("is branchable only after a complete non-empty answer", () => { + const ui = { requestRender: vi.fn(), requestComponentRender: vi.fn() } as unknown as TUI; + const panel = new BtwPanelComponent({ question: "Question?", tui: ui }); + + expect(panel.isBranchable()).toBe(false); + panel.setAnswer(" "); + panel.markComplete(); + expect(panel.isBranchable()).toBe(false); + panel.setAnswer("Answer"); + expect(panel.isBranchable()).toBe(true); + }); +}); describe("BtwController", () => { it("dispatches the question to runEphemeralTurn with the btw prompt wrapper and a fresh signal", async () => { @@ -87,10 +112,8 @@ describe("BtwController", () => { it("replaces a previous request by aborting it before issuing the next runEphemeralTurn", async () => { const signals: AbortSignal[] = []; - let firstRelease!: () => void; - const firstPromise = new Promise(resolve => { - firstRelease = () => resolve({ replyText: "first", assistantMessage: createAssistantMessage("first") }); - }); + const first = Promise.withResolvers(); + const firstPromise = first.promise; const runEphemeralTurn = vi .fn<(args: RunEphemeralTurnArgs) => Promise>() .mockImplementationOnce(async args => { @@ -116,11 +139,11 @@ describe("BtwController", () => { expect(signals[1]?.aborted).toBe(false); expect(btwContainer.children).toHaveLength(1); // Allow the orphaned first request to finish to keep the test clean. - firstRelease(); + first.resolve({ replyText: "first", assistantMessage: createAssistantMessage("first") }); }); it("clears the panel when the active request is dismissed via Escape", async () => { - const runEphemeralTurn = vi.fn(async () => new Promise(() => {})); + const runEphemeralTurn = vi.fn(async () => Promise.withResolvers().promise); const btwContainer = new Container(); const ctx = makeCtx(makeFakeSession(runEphemeralTurn), btwContainer); const controller = new BtwController(ctx); @@ -158,4 +181,229 @@ describe("BtwController", () => { expect(runEphemeralTurn).not.toHaveBeenCalled(); expect(ctx.showError).toHaveBeenCalled(); }); + + it("does not allow branch while /btw is still running", async () => { + const runEphemeralTurn = vi.fn(async () => Promise.withResolvers().promise); + const ctx = makeCtx(makeFakeSession(runEphemeralTurn)); + const controller = new BtwController(ctx); + + await controller.start("Question?"); + + expect(controller.canBranch()).toBe(false); + }); + + it("does not allow branch when the completed answer has no originating leaf", async () => { + const assistantMessage = createAssistantMessage("Answer"); + const runEphemeralTurn = vi.fn(async () => ({ replyText: "Answer", assistantMessage })); + const ctx = makeCtx(makeFakeSession(runEphemeralTurn)) as InteractiveModeContext & { + setTestLeafId(nextLeafId: string | null): void; + }; + ctx.setTestLeafId(null); + const controller = new BtwController(ctx); + + await controller.start("Question?"); + await drainBtwRequest(); + + expect(controller.canBranch()).toBe(false); + }); + + it("allows branch after a complete non-empty reply", async () => { + const assistantMessage = createAssistantMessage("Answer"); + const runEphemeralTurn = vi.fn(async () => ({ replyText: "Answer", assistantMessage })); + const ctx = makeCtx(makeFakeSession(runEphemeralTurn)); + const controller = new BtwController(ctx); + + await controller.start("Question?"); + await drainBtwRequest(); + + expect(controller.canBranch()).toBe(true); + }); + + it("does not allow branch after a complete empty reply", async () => { + const runEphemeralTurn = vi.fn(async () => ({ + replyText: " ", + assistantMessage: createAssistantMessage(" "), + })); + const ctx = makeCtx(makeFakeSession(runEphemeralTurn)); + const controller = new BtwController(ctx); + + await controller.start("Question?"); + await drainBtwRequest(); + + expect(controller.canBranch()).toBe(false); + }); + + it("does not allow branch after aborted or errored requests", async () => { + const abortedRun = vi.fn(async () => Promise.withResolvers().promise); + const abortedController = new BtwController(makeCtx(makeFakeSession(abortedRun))); + await abortedController.start("Question?"); + expect(abortedController.handleEscape()).toBe(true); + expect(abortedController.canBranch()).toBe(false); + + const erroredRun = vi.fn(async () => { + throw new Error("boom"); + }); + const erroredController = new BtwController(makeCtx(makeFakeSession(erroredRun))); + await erroredController.start("Question?"); + await drainBtwRequest(); + expect(erroredController.canBranch()).toBe(false); + }); + + it("handleBranch returns false and does not call the context when not branchable", async () => { + const runEphemeralTurn = vi.fn(async () => ({ replyText: "", assistantMessage: createAssistantMessage("") })); + const ctx = makeCtx(makeFakeSession(runEphemeralTurn)); + const controller = new BtwController(ctx); + + await controller.start("Question?"); + await drainBtwRequest(); + + expect(await controller.handleBranch()).toBe(false); + expect(ctx.handleBtwBranch).not.toHaveBeenCalled(); + }); + + it("handleBranch calls the context with the question and full assistant message when branchable", async () => { + const assistantMessage = createAssistantMessage("Answer"); + const runEphemeralTurn = vi.fn(async () => ({ replyText: "Answer", assistantMessage })); + const ctx = makeCtx(makeFakeSession(runEphemeralTurn)); + const controller = new BtwController(ctx); + + await controller.start("Question?"); + await drainBtwRequest(); + + expect(await controller.handleBranch()).toBe(true); + expect(ctx.handleBtwBranch).toHaveBeenCalledWith("Question?", assistantMessage); + }); + + it("branches the sanitized reply text while preserving non-text assistant content", async () => { + const assistantMessage: AssistantMessage = { + ...createAssistantMessage("raw repeated repeated repeated"), + content: [ + { + type: "thinking", + thinking: "Keep this reasoning.", + thinkingSignature: "signed-for-ephemeral-prompt", + itemId: "item-1", + }, + { type: "redactedThinking", data: "encrypted-ephemeral-thinking" }, + { type: "text", text: "raw repeated repeated repeated" }, + { type: "text", text: "raw duplicate tail" }, + ], + }; + const runEphemeralTurn = vi.fn(async () => ({ replyText: "sanitized", assistantMessage })); + const ctx = makeCtx(makeFakeSession(runEphemeralTurn)); + const controller = new BtwController(ctx); + + await controller.start("Question?"); + await drainBtwRequest(); + + expect(await controller.handleBranch()).toBe(true); + expect(ctx.handleBtwBranch).toHaveBeenCalledWith("Question?", { + ...assistantMessage, + content: [ + { type: "thinking", thinking: "Keep this reasoning." }, + { type: "text", text: "sanitized" }, + ], + }); + }); + + it("branches the sanitized reply text without native replay payload metadata", async () => { + const providerPayload = { + type: "openaiResponsesHistory" as const, + provider: "openai-codex", + dt: true, + items: [{ type: "reasoning", encrypted_content: "raw-ephemeral-output" }], + }; + const assistantMessage: AssistantMessage = { + ...createAssistantMessage("raw ephemeral output"), + api: "openai-codex-responses", + provider: "openai-codex", + model: "gpt-5-codex", + content: [ + { type: "thinking", thinking: "reasoning", thinkingSignature: "native-signature", itemId: "rs_1" }, + { type: "text", text: "raw ephemeral output" }, + ], + providerPayload, + }; + const runEphemeralTurn = vi.fn(async () => ({ replyText: "sanitized", assistantMessage })); + const ctx = makeCtx(makeFakeSession(runEphemeralTurn)); + const controller = new BtwController(ctx); + + await controller.start("Question?"); + await drainBtwRequest(); + + expect(await controller.handleBranch()).toBe(true); + expect(ctx.handleBtwBranch).toHaveBeenCalledWith("Question?", { + ...assistantMessage, + content: [ + { type: "thinking", thinking: "reasoning" }, + { type: "text", text: "sanitized" }, + ], + providerPayload: undefined, + }); + }); + + it("ignores duplicate branch requests while branch promotion is in flight", async () => { + const assistantMessage = createAssistantMessage("Answer"); + const runEphemeralTurn = vi.fn(async () => ({ replyText: "Answer", assistantMessage })); + const ctx = makeCtx(makeFakeSession(runEphemeralTurn)); + const branchStarted = Promise.withResolvers(); + const releaseBranch = Promise.withResolvers(); + ctx.handleBtwBranch = vi.fn(async () => { + branchStarted.resolve(); + await releaseBranch.promise; + }); + const controller = new BtwController(ctx); + + await controller.start("Question?"); + await drainBtwRequest(); + + const firstBranch = controller.handleBranch(); + await branchStarted.promise; + + expect(controller.canBranch()).toBe(false); + expect(await controller.handleBranch()).toBe(false); + expect(ctx.handleBtwBranch).toHaveBeenCalledTimes(1); + + releaseBranch.resolve(); + expect(await firstBranch).toBe(true); + }); + + it("does not branch a completed answer after the session leaf changes", async () => { + const assistantMessage = createAssistantMessage("Answer"); + const runEphemeralTurn = vi.fn(async () => ({ replyText: "Answer", assistantMessage })); + const ctx = makeCtx(makeFakeSession(runEphemeralTurn)) as InteractiveModeContext & { + setTestLeafId(nextLeafId: string | null): void; + }; + const controller = new BtwController(ctx); + + await controller.start("Question?"); + await drainBtwRequest(); + expect(controller.canBranch()).toBe(true); + + ctx.setTestLeafId("leaf-2"); + + expect(controller.canBranch()).toBe(false); + expect(await controller.handleBranch()).toBe(false); + expect(ctx.handleBtwBranch).not.toHaveBeenCalled(); + }); + + it("clears stored branch state on escape and dispose", async () => { + const runEphemeralTurn = vi.fn(async () => ({ + replyText: "Answer", + assistantMessage: createAssistantMessage("Answer"), + })); + const escapeController = new BtwController(makeCtx(makeFakeSession(runEphemeralTurn))); + await escapeController.start("Question?"); + await drainBtwRequest(); + expect(escapeController.canBranch()).toBe(true); + expect(escapeController.handleEscape()).toBe(true); + expect(escapeController.canBranch()).toBe(false); + + const disposeController = new BtwController(makeCtx(makeFakeSession(runEphemeralTurn))); + await disposeController.start("Question?"); + await drainBtwRequest(); + expect(disposeController.canBranch()).toBe(true); + disposeController.dispose(); + expect(disposeController.canBranch()).toBe(false); + }); });