From 6e8cb46cba27a6bb44f53d401711a7fde75e1ae4 Mon Sep 17 00:00:00 2001 From: Test Date: Sun, 9 Aug 2026 08:50:01 +0200 Subject: [PATCH 1/7] chore: open PR for issue 439 From cde4a60277767a97ff50d9ea505f2e02c354853c Mon Sep 17 00:00:00 2001 From: Test Date: Sun, 9 Aug 2026 09:26:07 +0200 Subject: [PATCH 2/7] Add fork from current state (including last response) (#439) --- .../coding-agent/src/core/agent-session.ts | 66 +++++++++++++-- .../components/user-message-selector.ts | 30 ++++++- .../src/modes/interactive/interactive-mode.ts | 36 ++++++++- .../coding-agent/src/modes/rpc/rpc-client.ts | 10 +++ .../coding-agent/src/modes/rpc/rpc-mode.ts | 5 ++ .../coding-agent/src/modes/rpc/rpc-types.ts | 2 + .../test/agent-session-fork-current.test.ts | 80 +++++++++++++++++++ .../test/rpc-fork-current.test.ts | 31 +++++++ .../session-manager/tree-traversal.test.ts | 24 ++++++ packages/dashboard/src/client/api.ts | 1 + .../dashboard/src/client/screens/session.tsx | 17 ++++ packages/dashboard/src/client/styles/app.css | 22 ++++- packages/dashboard/src/server/server.ts | 4 + .../dashboard/test/client/screens.test.tsx | 31 +++++++ packages/dashboard/test/runtime-pool.test.ts | 3 + packages/dashboard/test/server.test.ts | 8 ++ test.sh | 11 +++ 17 files changed, 365 insertions(+), 16 deletions(-) create mode 100644 packages/coding-agent/test/agent-session-fork-current.test.ts create mode 100644 packages/coding-agent/test/rpc-fork-current.test.ts diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 170ae765..a5683390 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -3636,7 +3636,6 @@ export class AgentSession { * - cancelled: True if an extension cancelled the fork */ async fork(entryId: string): Promise<{ selectedText: string; cancelled: boolean }> { - const previousSessionFile = this.sessionFile; const selectedEntry = this.sessionManager.getEntry(entryId); if (!selectedEntry || selectedEntry.type !== "message" || selectedEntry.message.role !== "user") { @@ -3645,6 +3644,61 @@ export class AgentSession { const selectedText = this._extractUserMessageText(selectedEntry.message.content); + // Existing fork semantics: rewind to *before* the selected user message by + // branching from its parent, so the selected message (and everything after + // it) is dropped and its text is offered as editor pre-fill. + const { cancelled } = await this._performFork(entryId, (previousSessionFile) => { + if (!selectedEntry.parentId) { + this.sessionManager.newSession({ parentSession: previousSessionFile }); + } else { + this.sessionManager.createBranchedSession(selectedEntry.parentId); + } + }); + + return { selectedText, cancelled }; + } + + /** + * Create a fork from the current state, including the last model response. + * + * Unlike fork(), which rewinds to *before* a selected user message, this + * branches from the current leaf entry — so the new branch's tail is the + * latest entry (typically the last assistant response). There is no editor + * pre-fill; the new session is ready for a fresh turn on top of the captured + * history. + * + * @returns Object with `cancelled` (true if an extension cancelled the fork, + * or if there is no current leaf to fork from — i.e. an empty session). + */ + async forkFromCurrent(): Promise<{ cancelled: boolean }> { + const leafId = this.sessionManager.getLeafId(); + + // Nothing to fork from (fresh/empty session). + if (!leafId) { + return { cancelled: true }; + } + + return this._performFork(leafId, () => { + // Branch from the leaf itself so the last entry is included. + this.sessionManager.createBranchedSession(leafId); + }); + } + + /** + * Shared fork machinery: emit the cancellable session_before_fork event, + * clear pending state, create the branch via the supplied strategy, reload + * the conversation, and emit session_fork. + * + * @param entryId Entry the fork is anchored to (reported to extensions). + * @param branch Strategy that creates the branched/new session. Receives the + * previous session file so callers can set it as the parent when needed. + */ + private async _performFork( + entryId: string, + branch: (previousSessionFile: string | undefined) => void, + ): Promise<{ cancelled: boolean }> { + const previousSessionFile = this.sessionFile; + let skipConversationRestore = false; // Emit session_before_fork event (can be cancelled) @@ -3655,7 +3709,7 @@ export class AgentSession { })) as SessionBeforeForkResult | undefined; if (result?.cancel) { - return { selectedText, cancelled: true }; + return { cancelled: true }; } skipConversationRestore = result?.skipConversationRestore ?? false; } @@ -3663,11 +3717,7 @@ export class AgentSession { // Clear pending messages (bound to old session state) this._pendingNextTurnMessages = []; - if (!selectedEntry.parentId) { - this.sessionManager.newSession({ parentSession: previousSessionFile }); - } else { - this.sessionManager.createBranchedSession(selectedEntry.parentId); - } + branch(previousSessionFile); this.agent.sessionId = this.sessionManager.getSessionId(); // Reload messages from entries (works for both file and in-memory mode) @@ -3687,7 +3737,7 @@ export class AgentSession { this.agent.replaceMessages(sessionContext.messages); } - return { selectedText, cancelled: false }; + return { cancelled: false }; } // ========================================================================= diff --git a/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts b/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts index e5c9c547..b3d42cee 100644 --- a/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts +++ b/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts @@ -6,8 +6,12 @@ interface UserMessageItem { id: string; // Entry ID in the session text: string; // The message text timestamp?: string; // Optional timestamp if available + isAction?: boolean; // Special action row (e.g. "fork from current state") rather than a history message } +/** Sentinel entry id for the "fork from current state (include last response)" action row. */ +export const FORK_FROM_CURRENT_ID = "__fork_from_current__"; + /** * Custom user message list component with selection */ @@ -45,6 +49,7 @@ class UserMessageList implements Component { const endIndex = Math.min(startIndex + this.maxVisible, this.messages.length); // Render visible messages (2 lines per message + blank line) + const totalRealMessages = this.messages.filter((m) => !m.isAction).length; for (let i = startIndex; i < endIndex; i++) { const message = this.messages[i]; const isSelected = i === this.selectedIndex; @@ -55,14 +60,25 @@ class UserMessageList implements Component { // First line: cursor + message const cursor = isSelected ? theme.fg("accent", "› ") : " "; const maxMsgWidth = width - 2; // Account for cursor (2 chars) + + if (message.isAction) { + // Distinct styling for the action row (e.g. "fork from current state"). + const label = truncateToWidth(`⎇ ${normalizedMessage}`, maxMsgWidth); + const styled = isSelected ? theme.fg("accent", theme.bold(label)) : theme.fg("accent", label); + lines.push(cursor + styled); + lines.push(theme.fg("muted", " Branch here, including the last response")); + lines.push(""); // Blank line between entries + continue; + } + const truncatedMsg = truncateToWidth(normalizedMessage, maxMsgWidth); const messageLine = cursor + (isSelected ? theme.bold(truncatedMsg) : truncatedMsg); lines.push(messageLine); - // Second line: metadata (position in history) - const position = i + 1; - const metadata = ` Message ${position} of ${this.messages.length}`; + // Second line: metadata (position in history, counting only real messages) + const position = this.messages.slice(0, i + 1).filter((m) => !m.isAction).length; + const metadata = ` Message ${position} of ${totalRealMessages}`; const metadataLine = theme.fg("muted", metadata); lines.push(metadataLine); lines.push(""); // Blank line between messages @@ -115,7 +131,13 @@ export class UserMessageSelectorComponent extends Container { // Add header this.addChild(new Spacer(1)); this.addChild(new Text(theme.bold("Branch from Message"), 1, 0)); - this.addChild(new Text(theme.fg("muted", "Select a message to create a new branch from that point"), 1, 0)); + this.addChild( + new Text( + theme.fg("muted", "Pick a message to rewind and re-run, or fork from current state (keeps last response)"), + 1, + 0, + ), + ); this.addChild(new Spacer(1)); this.addChild(new DynamicBorder()); this.addChild(new Spacer(1)); diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index d4884d89..d149878f 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -121,7 +121,7 @@ import { TasksPanelComponent } from "./components/tasks-panel.js"; import { ToolExecutionComponent } from "./components/tool-execution.js"; import { TreeSelectorComponent } from "./components/tree-selector.js"; import { UserMessageComponent } from "./components/user-message.js"; -import { UserMessageSelectorComponent } from "./components/user-message-selector.js"; +import { FORK_FROM_CURRENT_ID, UserMessageSelectorComponent } from "./components/user-message-selector.js"; import { getAvailableThemes, getAvailableThemesWithPaths, @@ -4297,16 +4297,46 @@ export class InteractiveMode { private showUserMessageSelector(): void { const userMessages = this.session.getUserMessagesForForking(); + const hasCurrentState = this.sessionManager.getLeafId() !== null; - if (userMessages.length === 0) { + if (userMessages.length === 0 && !hasCurrentState) { this.showStatus("No messages to fork from"); return; } + // Build the selector list: history messages (rewind + re-run), plus a + // trailing "fork from current state" action that keeps the last response. + const items: Array<{ id: string; text: string; isAction?: boolean }> = userMessages.map((m) => ({ + id: m.entryId, + text: m.text, + })); + if (hasCurrentState) { + items.push({ + id: FORK_FROM_CURRENT_ID, + text: "Fork from current state (include last response)", + isAction: true, + }); + } + this.showSelector((done) => { const selector = new UserMessageSelectorComponent( - userMessages.map((m) => ({ id: m.entryId, text: m.text })), + items, async (entryId) => { + if (entryId === FORK_FROM_CURRENT_ID) { + const result = await this.session.forkFromCurrent(); + if (result.cancelled) { + done(); + this.ui.requestRender(); + return; + } + + this.resetChatDisplay(); + this.editor.setText(""); + done(); + this.showStatus("Branched to new session (including last response)"); + return; + } + const result = await this.session.fork(entryId); if (result.cancelled) { // Extension cancelled the fork diff --git a/packages/coding-agent/src/modes/rpc/rpc-client.ts b/packages/coding-agent/src/modes/rpc/rpc-client.ts index 629b1448..6dc1e297 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-client.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-client.ts @@ -626,6 +626,16 @@ export class RpcClient { return this.getData(response); } + /** + * Fork from the current state, including the last model response. + * Unlike fork(), this branches from the current leaf (no editor pre-fill). + * @returns Object with `cancelled` (if an extension cancelled, or the session is empty) + */ + async forkCurrent(): Promise<{ cancelled: boolean }> { + const response = await this.send({ type: "fork_current" }); + return this.getData(response); + } + /** * Get messages available for forking. */ diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index f2025f1e..5bdd97c8 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -2126,6 +2126,11 @@ export async function runRpcMode(session: AgentSession, modelFallbackMessage?: s return success(id, "fork", { text: result.selectedText, cancelled: result.cancelled }); } + case "fork_current": { + const result = await session.forkFromCurrent(); + return success(id, "fork_current", { cancelled: result.cancelled }); + } + case "get_fork_messages": { const messages = session.getUserMessagesForForking(); return success(id, "get_fork_messages", { messages }); diff --git a/packages/coding-agent/src/modes/rpc/rpc-types.ts b/packages/coding-agent/src/modes/rpc/rpc-types.ts index b81a20bf..579310ef 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-types.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-types.ts @@ -85,6 +85,7 @@ export type RpcCommand = | { id?: string; type: "switch_session"; sessionPath: string } | { id?: string; type: "delete_session"; sessionPath: string } | { id?: string; type: "fork"; entryId: string } + | { id?: string; type: "fork_current" } | { id?: string; type: "get_fork_messages" } | { id?: string; type: "get_tree" } | { @@ -373,6 +374,7 @@ export type RpcResponse = | { id?: string; type: "response"; command: "switch_session"; success: true; data: { cancelled: boolean } } | { id?: string; type: "response"; command: "delete_session"; success: true; data: { method: "trash" | "unlink" } } | { id?: string; type: "response"; command: "fork"; success: true; data: { text: string; cancelled: boolean } } + | { id?: string; type: "response"; command: "fork_current"; success: true; data: { cancelled: boolean } } | { id?: string; type: "response"; diff --git a/packages/coding-agent/test/agent-session-fork-current.test.ts b/packages/coding-agent/test/agent-session-fork-current.test.ts new file mode 100644 index 00000000..d06ed035 --- /dev/null +++ b/packages/coding-agent/test/agent-session-fork-current.test.ts @@ -0,0 +1,80 @@ +/** + * Tests for AgentSession.forkFromCurrent(). + * + * Unlike fork(), which rewinds to *before* a selected user message, forkFromCurrent() + * branches from the current leaf — so the new branch includes the last model response. + * + * These run offline (no live API): the conversation is constructed by appending + * messages directly to the in-memory SessionManager, then forkFromCurrent() reloads + * the branch via buildSessionContext + replaceMessages. + */ + +import { afterEach, describe, expect, it } from "vitest"; +import { createHarnessWithExtensions, type Harness } from "./test-harness.js"; +import { assistantMsg, userMsg } from "./utilities.js"; + +describe("AgentSession.forkFromCurrent", () => { + let harness: Harness; + + afterEach(() => { + harness?.cleanup(); + }); + + it("includes the last model response in the forked branch", async () => { + harness = await createHarnessWithExtensions(); + const { session, sessionManager } = harness; + + // Conversation: q1 -> a1 -> q2 -> a2 (a2 is the last model response) + sessionManager.appendMessage(userMsg("q1")); + sessionManager.appendMessage(assistantMsg("a1")); + sessionManager.appendMessage(userMsg("q2")); + const lastAssistantId = sessionManager.appendMessage(assistantMsg("a2")); + expect(sessionManager.getLeafId()).toBe(lastAssistantId); + + const result = await session.forkFromCurrent(); + expect(result.cancelled).toBe(false); + + // The full conversation survives and the tail is the last assistant response. + expect(session.messages).toHaveLength(4); + const tail = session.messages.at(-1)!; + expect(tail.role).toBe("assistant"); + expect(JSON.stringify(tail)).toContain("a2"); + + // A new session was branched (fresh id), still tracking the conversation. + expect(sessionManager.getLeafId()).not.toBeNull(); + }); + + it("is a no-op that returns cancelled for an empty session", async () => { + harness = await createHarnessWithExtensions(); + const { session, sessionManager } = harness; + + expect(sessionManager.getLeafId()).toBeNull(); + + const result = await session.forkFromCurrent(); + expect(result.cancelled).toBe(true); + expect(session.messages).toHaveLength(0); + }); + + it("can be cancelled by a session_before_fork extension handler", async () => { + harness = await createHarnessWithExtensions({ + extensionFactories: [ + (dreb) => { + dreb.on("session_before_fork", async () => ({ cancel: true })); + }, + ], + }); + const { session, sessionManager } = harness; + + sessionManager.appendMessage(userMsg("q1")); + sessionManager.appendMessage(assistantMsg("a1")); + const leafBefore = sessionManager.getLeafId(); + const entriesBefore = sessionManager.getEntries(); + + const result = await session.forkFromCurrent(); + expect(result.cancelled).toBe(true); + + // Cancellation must not branch or mutate the session. + expect(sessionManager.getLeafId()).toBe(leafBefore); + expect(sessionManager.getEntries()).toEqual(entriesBefore); + }); +}); diff --git a/packages/coding-agent/test/rpc-fork-current.test.ts b/packages/coding-agent/test/rpc-fork-current.test.ts new file mode 100644 index 00000000..02708ae7 --- /dev/null +++ b/packages/coding-agent/test/rpc-fork-current.test.ts @@ -0,0 +1,31 @@ +import { describe, expect, it, vi } from "vitest"; +import { RpcClient } from "../src/modes/rpc/rpc-client.js"; + +describe("RpcClient.forkCurrent", () => { + it("sends fork_current with no arguments and unwraps the cancelled flag", async () => { + const client = new RpcClient() as any; + const data = { cancelled: false }; + client.send = vi.fn().mockResolvedValue({ + type: "response", + command: "fork_current", + success: true, + data, + }); + + await expect(client.forkCurrent()).resolves.toEqual(data); + expect(client.send).toHaveBeenCalledWith({ type: "fork_current" }); + expect(client.send.mock.calls[0]).toHaveLength(1); + }); + + it("propagates a cancelled fork (e.g. empty session or extension cancel)", async () => { + const client = new RpcClient() as any; + client.send = vi.fn().mockResolvedValue({ + type: "response", + command: "fork_current", + success: true, + data: { cancelled: true }, + }); + + await expect(client.forkCurrent()).resolves.toEqual({ cancelled: true }); + }); +}); diff --git a/packages/coding-agent/test/session-manager/tree-traversal.test.ts b/packages/coding-agent/test/session-manager/tree-traversal.test.ts index bf86868e..91778bc7 100644 --- a/packages/coding-agent/test/session-manager/tree-traversal.test.ts +++ b/packages/coding-agent/test/session-manager/tree-traversal.test.ts @@ -440,6 +440,30 @@ describe("createBranchedSession", () => { expect(entries[1].id).toBe(id2); }); + it("branching from the current leaf keeps the last assistant response (fork from current state)", () => { + const session = SessionManager.inMemory(); + + // Conversation: u1 -> a1 -> u2 -> a2 (a2 is the last model response) + session.appendMessage(userMsg("q1")); + session.appendMessage(assistantMsg("a1")); + session.appendMessage(userMsg("q2")); + const lastAssistantId = session.appendMessage(assistantMsg("a2")); + + // "Fork from current state" branches from the current leaf itself. + expect(session.getLeafId()).toBe(lastAssistantId); + session.createBranchedSession(session.getLeafId()!); + + // The whole conversation is preserved, and the tail is the last assistant + // response — unlike the user-message rewind fork, which drops it. + const messages = session + .getEntries() + .filter((e): e is Extract => e.type === "message"); + expect(messages).toHaveLength(4); + const tail = messages.at(-1)!; + expect(tail.id).toBe(lastAssistantId); + expect(tail.message.role).toBe("assistant"); + }); + it("preserves the current session name when branching before its metadata entry", () => { const session = SessionManager.inMemory(); session.appendMessage(userMsg("question")); diff --git a/packages/dashboard/src/client/api.ts b/packages/dashboard/src/client/api.ts index 7137644b..56167f1e 100644 --- a/packages/dashboard/src/client/api.ts +++ b/packages/dashboard/src/client/api.ts @@ -158,6 +158,7 @@ export const api = { request<{ messages: Array<{ entryId: string; text: string }> }>(`/api/runtimes/${key}/fork-messages`), fork: (key: string, entryId: string) => request<{ text: string; cancelled: boolean }>(`/api/runtimes/${key}/fork`, json({ entryId })), + forkCurrent: (key: string) => request<{ cancelled: boolean }>(`/api/runtimes/${key}/fork-current`, json({})), tree: (key: string) => request<{ roots: SessionTreeNodeDto[]; leafId: string | null }>(`/api/runtimes/${key}/tree`), navigateTree: (key: string, targetId: string) => request<{ cancelled: boolean; editorText?: string }>(`/api/runtimes/${key}/tree`, json({ targetId })), diff --git a/packages/dashboard/src/client/screens/session.tsx b/packages/dashboard/src/client/screens/session.tsx index 9d3e2f6c..d11b21df 100644 --- a/packages/dashboard/src/client/screens/session.tsx +++ b/packages/dashboard/src/client/screens/session.tsx @@ -1282,6 +1282,19 @@ export function SessionScreen(props: { store: AppStore; sessionKey: string }): J } } + async function forkFromCurrentState() { + setForkError(undefined); + try { + await api.forkCurrent(props.sessionKey); + // No composer pre-fill: the branch already includes the last response. + await props.store.hydrateSession(props.sessionKey); + await props.store.refreshDiskSessions(); + setShowForkModal(false); + } catch (err) { + setForkError(err instanceof Error ? err.message : String(err)); + } + } + async function openStatsPopover() { setShowStatsPopover(true); setStatsPopoverError(undefined); @@ -2327,6 +2340,10 @@ export function SessionScreen(props: { store: AppStore; sessionKey: string }): J

{forkError()}

 + 0} fallback={

loading forkable messages…

}>
diff --git a/packages/dashboard/src/client/styles/app.css b/packages/dashboard/src/client/styles/app.css index d40635cb..bea4906b 100644 --- a/packages/dashboard/src/client/styles/app.css +++ b/packages/dashboard/src/client/styles/app.css @@ -2757,10 +2757,30 @@ details.thinking .thinking-body { gap: var(--space-2); } -.fork-message:hover { +.fork-current-btn { + border: var(--hairline); + border-radius: var(--radius); + background: var(--bg); + color: var(--text); + font: inherit; + text-align: left; + padding: var(--space-2); + cursor: pointer; + display: flex; + gap: var(--space-2); + margin-bottom: var(--space-2); + border-color: var(--accent, var(--text)); +} + +.fork-message:hover, +.fork-current-btn:hover { border-color: var(--text); } +.fork-current-btn .fork-entry-id { + color: var(--accent, var(--text)); +} + .fork-entry-id { color: var(--muted); font-size: var(--fs-small); diff --git a/packages/dashboard/src/server/server.ts b/packages/dashboard/src/server/server.ts index 179b2ced..40cd63a1 100644 --- a/packages/dashboard/src/server/server.ts +++ b/packages/dashboard/src/server/server.ts @@ -954,6 +954,10 @@ export function createDashboardServer(options: DashboardServerOptions): Dashboar withRuntime(req, res, (h) => h.client.fork(entryId)); }); + app.post("/api/runtimes/:key/fork-current", (req, res) => { + withRuntime(req, res, (h) => h.client.forkCurrent()); + }); + app.get("/api/runtimes/:key/tree", (req, res) => { withRuntime(req, res, (h) => h.client.getTree()); }); diff --git a/packages/dashboard/test/client/screens.test.tsx b/packages/dashboard/test/client/screens.test.tsx index d1a21dd4..f30ce613 100644 --- a/packages/dashboard/test/client/screens.test.tsx +++ b/packages/dashboard/test/client/screens.test.tsx @@ -81,6 +81,7 @@ vi.mock("../../src/client/api.js", () => ({ branch: vi.fn(async () => ({ branch: null })), forkMessages: vi.fn(async () => ({ messages: [] })), fork: vi.fn(async () => ({ text: "", cancelled: false })), + forkCurrent: vi.fn(async () => ({ cancelled: false })), dailyCost: vi.fn(async () => ({ cost: 0.42 })), settings: vi.fn(async () => ({ defaultProvider: "anthropic", defaultModel: "m1" })), devices: vi.fn(async () => ({ devices: [] })), @@ -5819,6 +5820,36 @@ describe("dashboard client regressions", () => { expect((el.querySelector("textarea") as HTMLTextAreaElement).value).toBe("original prompt"); }); + it("fork modal can fork from current state without prefilling the composer", async () => { + vi.mocked(api.forkMessages).mockResolvedValue({ messages: [{ entryId: "u1", text: "original prompt" }] }); + vi.mocked(api.forkCurrent).mockResolvedValue({ cancelled: false }); + const store = makeStore() as any; + const hydrateSession = vi.fn(async () => {}); + const refreshDiskSessions = vi.fn(async () => {}); + const fakeStore = { + ...store, + sessions: { forkcur: createSessionViewState("forkcur") }, + fleet: () => ({ runtimes: [], diskSessions: [] }), + hydrateSession, + refreshDiskSessions, + }; + const el = mount(() => ); + (el.querySelector(".session-bar .right .switcher:last-child") as HTMLButtonElement).click(); + await new Promise((resolve) => setTimeout(resolve, 0)); + [...el.querySelectorAll("button")].find((button) => button.textContent?.includes("fork"))?.click(); + await new Promise((resolve) => setTimeout(resolve, 10)); + // Mocks are not auto-cleared between tests; ignore any api.fork calls from prior tests. + vi.mocked(api.fork).mockClear(); + (el.querySelector(".fork-current-btn") as HTMLButtonElement).click(); + await new Promise((resolve) => setTimeout(resolve, 10)); + expect(api.forkCurrent).toHaveBeenCalledWith("forkcur"); + expect(api.fork).not.toHaveBeenCalled(); + expect(hydrateSession).toHaveBeenCalledWith("forkcur"); + expect(refreshDiskSessions).toHaveBeenCalledOnce(); + // No composer pre-fill: the branch already includes the last response. + expect((el.querySelector("textarea") as HTMLTextAreaElement).value).toBe(""); + }); + it("session stats popover shows the detailed stats breakdown", async () => { const store = makeStore() as any; const fakeStore = { diff --git a/packages/dashboard/test/runtime-pool.test.ts b/packages/dashboard/test/runtime-pool.test.ts index 20bb4b90..1adfe4d2 100644 --- a/packages/dashboard/test/runtime-pool.test.ts +++ b/packages/dashboard/test/runtime-pool.test.ts @@ -185,6 +185,9 @@ export function makeFakeClient() { importJsonl: vi.fn(async () => ({ cancelled: false })), getTree: vi.fn(async () => ({ roots: [], leafId: null })), navigateTree: vi.fn(async () => ({ cancelled: false })), + getForkMessages: vi.fn(async () => []), + fork: vi.fn(async () => ({ text: "", cancelled: false })), + forkCurrent: vi.fn(async () => ({ cancelled: false })), listSessions: vi.fn(async () => []), switchSession: vi.fn(async () => ({ cancelled: false })), prompt: vi.fn(async () => {}), diff --git a/packages/dashboard/test/server.test.ts b/packages/dashboard/test/server.test.ts index 2434f7f0..fe16a3c9 100644 --- a/packages/dashboard/test/server.test.ts +++ b/packages/dashboard/test/server.test.ts @@ -1083,6 +1083,13 @@ describe("dashboard server — fleet and runtimes", () => { body: JSON.stringify({ targetId: "entry-1" }), }).then((r) => r.json()), ).resolves.toEqual({ cancelled: false }); + await expect( + fetch(`${base}/api/runtimes/${key}/fork-current`, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({}), + }).then((r) => r.json()), + ).resolves.toEqual({ cancelled: false }); await expect(fetch(`${base}/api/runtimes/${key}/sessions`).then((r) => r.json())).resolves.toEqual({ sessions: [], }); @@ -1107,6 +1114,7 @@ describe("dashboard server — fleet and runtimes", () => { expect(clients[0].importJsonl).toHaveBeenCalledWith("/tmp/session.jsonl"); expect(clients[0].getTree).toHaveBeenCalled(); expect(clients[0].navigateTree).toHaveBeenCalledWith("entry-1"); + expect(clients[0].forkCurrent).toHaveBeenCalled(); expect(clients[0].listSessions).toHaveBeenCalled(); expect(clients[0].switchSession).toHaveBeenCalledWith("/tmp/session.jsonl"); expect(clients[1].getDailyCost).toHaveBeenCalled(); diff --git a/test.sh b/test.sh index 88597a94..939585f4 100755 --- a/test.sh +++ b/test.sh @@ -1,6 +1,17 @@ #!/usr/bin/env bash set -e +# When invoked from a git hook (e.g. husky pre-commit), git exports repo-location +# variables (GIT_DIR, GIT_INDEX_FILE, GIT_WORK_TREE, …) into the environment. +# These leak into tests that shell out to `git` in throwaway temp repos +# (git-update.test.ts, tools.test.ts's .gitignore cases, …), redirecting their +# `git init`/`clone`/`commit` at the parent repo and making them fail — even +# though the same tests pass when run standalone. Unset the location-pinning +# vars so subprocess git operations resolve against their own cwd. Identity +# vars (GIT_AUTHOR_*/GIT_COMMITTER_*) are intentionally left intact. +unset GIT_DIR GIT_WORK_TREE GIT_INDEX_FILE GIT_PREFIX GIT_COMMON_DIR \ + GIT_NAMESPACE GIT_OBJECT_DIRECTORY GIT_ALTERNATE_OBJECT_DIRECTORIES + # Skip local LLM tests (ollama, lmstudio) — no local server expected in CI/hooks export DREB_NO_LOCAL_LLM=1 From a198676fc7ac2a0b35a84054389eac1c59494589 Mon Sep 17 00:00:00 2001 From: Test Date: Sun, 9 Aug 2026 10:59:58 +0200 Subject: [PATCH 3/7] Inform user when fork is cancelled; test + dedupe fork paths (#439) --- .../src/modes/interactive/interactive-mode.ts | 32 ++-- .../test/agent-session-fork-current.test.ts | 47 +++++ .../test/interactive-mode-fork.test.ts | 165 ++++++++++++++++++ .../dashboard/src/client/screens/session.tsx | 33 ++-- .../dashboard/test/client/screens.test.tsx | 33 ++++ 5 files changed, 276 insertions(+), 34 deletions(-) create mode 100644 packages/coding-agent/test/interactive-mode-fork.test.ts diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index d149878f..1f2eb6c8 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -4322,33 +4322,27 @@ export class InteractiveMode { const selector = new UserMessageSelectorComponent( items, async (entryId) => { - if (entryId === FORK_FROM_CURRENT_ID) { - const result = await this.session.forkFromCurrent(); - if (result.cancelled) { - done(); - this.ui.requestRender(); - return; - } + const isCurrent = entryId === FORK_FROM_CURRENT_ID; + const result = isCurrent ? await this.session.forkFromCurrent() : await this.session.fork(entryId); - this.resetChatDisplay(); - this.editor.setText(""); - done(); - this.showStatus("Branched to new session (including last response)"); - return; - } - - const result = await this.session.fork(entryId); if (result.cancelled) { - // Extension cancelled the fork + // Empty session (nothing to branch from) or an extension vetoed + // the fork — tell the user rather than silently dismissing the + // selector. done(); - this.ui.requestRender(); + this.showStatus("Fork cancelled — no new branch was created"); return; } this.resetChatDisplay(); - this.editor.setText(result.selectedText); + // The current-state branch already includes the last response, so + // there is nothing to re-ask; the message branch pre-fills the editor + // with the selected question for re-running. + this.editor.setText((result as { selectedText?: string }).selectedText ?? ""); done(); - this.showStatus("Branched to new session"); + this.showStatus( + isCurrent ? "Branched to new session (including last response)" : "Branched to new session", + ); }, () => { done(); diff --git a/packages/coding-agent/test/agent-session-fork-current.test.ts b/packages/coding-agent/test/agent-session-fork-current.test.ts index d06ed035..3455005e 100644 --- a/packages/coding-agent/test/agent-session-fork-current.test.ts +++ b/packages/coding-agent/test/agent-session-fork-current.test.ts @@ -77,4 +77,51 @@ describe("AgentSession.forkFromCurrent", () => { expect(sessionManager.getLeafId()).toBe(leafBefore); expect(sessionManager.getEntries()).toEqual(entriesBefore); }); + + it("emits session_fork after the branch is created", async () => { + const forkEvents: Array<{ type: string }> = []; + harness = await createHarnessWithExtensions({ + extensionFactories: [ + (dreb) => { + dreb.on("session_fork", async (event) => { + forkEvents.push(event); + }); + }, + ], + }); + const { session, sessionManager } = harness; + + sessionManager.appendMessage(userMsg("q1")); + sessionManager.appendMessage(assistantMsg("a1")); + + const result = await session.forkFromCurrent(); + expect(result.cancelled).toBe(false); + + // AC5: the after-fork event must fire exactly once on the new leaf-based path. + expect(forkEvents).toHaveLength(1); + expect(forkEvents[0].type).toBe("session_fork"); + }); + + it("does not emit session_fork when a session_before_fork handler cancels", async () => { + const forkEvents: Array<{ type: string }> = []; + harness = await createHarnessWithExtensions({ + extensionFactories: [ + (dreb) => { + dreb.on("session_before_fork", async () => ({ cancel: true })); + dreb.on("session_fork", async (event) => { + forkEvents.push(event); + }); + }, + ], + }); + const { session, sessionManager } = harness; + + sessionManager.appendMessage(userMsg("q1")); + sessionManager.appendMessage(assistantMsg("a1")); + + const result = await session.forkFromCurrent(); + expect(result.cancelled).toBe(true); + // A vetoed fork must not fire the after-fork event. + expect(forkEvents).toHaveLength(0); + }); }); diff --git a/packages/coding-agent/test/interactive-mode-fork.test.ts b/packages/coding-agent/test/interactive-mode-fork.test.ts new file mode 100644 index 00000000..728185a1 --- /dev/null +++ b/packages/coding-agent/test/interactive-mode-fork.test.ts @@ -0,0 +1,165 @@ +/** + * Unit coverage for the interactive `/fork` selector wiring + * (InteractiveMode.showUserMessageSelector). + * + * Exercised without a full TUI: the selector component is mocked to capture the + * onSelect callback, and the method is invoked via prototype.call with a + * hand-built `this` (the same pattern as interactive-mode-status.test.ts). + */ + +import { beforeEach, describe, expect, it, vi } from "vitest"; + +// Capture every UserMessageSelectorComponent construction so tests can drive the +// onSelect callback and assert the items list (incl. the fork-from-current row). +const captured: Array<{ + items: Array<{ id: string; text: string; isAction?: boolean }>; + onSelect: (entryId: string) => void | Promise; + onCancel: () => void; +}> = []; + +vi.mock("../src/modes/interactive/components/user-message-selector.js", async (importOriginal) => { + const actual = await importOriginal>(); + class MockUserMessageSelectorComponent { + constructor( + public items: Array<{ id: string; text: string; isAction?: boolean }>, + public onSelect: (entryId: string) => void | Promise, + public onCancel: () => void, + ) { + captured.push({ items, onSelect, onCancel }); + } + getMessageList() { + return {}; + } + } + return { ...actual, UserMessageSelectorComponent: MockUserMessageSelectorComponent }; +}); + +import { FORK_FROM_CURRENT_ID } from "../src/modes/interactive/components/user-message-selector.js"; +import { InteractiveMode } from "../src/modes/interactive/interactive-mode.js"; + +interface FakeOverrides { + userMessages?: Array<{ entryId: string; text: string }>; + leafId?: string | null; + forkFromCurrentResult?: { cancelled: boolean }; + forkResult?: { cancelled: boolean; selectedText: string }; +} + +function makeFakeThis(overrides: FakeOverrides = {}) { + const done = vi.fn(); + const fake = { + session: { + getUserMessagesForForking: vi.fn(() => overrides.userMessages ?? []), + forkFromCurrent: vi.fn(async () => overrides.forkFromCurrentResult ?? { cancelled: false }), + fork: vi.fn(async () => overrides.forkResult ?? { cancelled: false, selectedText: "prefill" }), + }, + sessionManager: { getLeafId: vi.fn(() => overrides.leafId ?? null) }, + showStatus: vi.fn(), + resetChatDisplay: vi.fn(), + editor: { setText: vi.fn() }, + ui: { requestRender: vi.fn() }, + // showSelector(builder) invokes the builder with a `done` callback and keeps + // whatever component it returns; here we just run the builder synchronously. + showSelector: vi.fn((builder: (done: () => void) => unknown) => builder(done)), + _done: done, + }; + return fake; +} + +function invoke(fake: ReturnType) { + ( + InteractiveMode as unknown as { prototype: { showUserMessageSelector: () => void } } + ).prototype.showUserMessageSelector.call(fake); +} + +describe("InteractiveMode.showUserMessageSelector — fork-from-current wiring", () => { + beforeEach(() => { + captured.length = 0; + }); + + it("appends the fork-from-current action row when there is a current leaf", () => { + const fake = makeFakeThis({ + userMessages: [ + { entryId: "u1", text: "first" }, + { entryId: "u2", text: "second" }, + ], + leafId: "a2", + }); + invoke(fake); + + expect(fake.showSelector).toHaveBeenCalledOnce(); + expect(captured).toHaveLength(1); + const { items } = captured[0]; + // History messages first, then the trailing action row. + expect(items.map((i) => i.id)).toEqual(["u1", "u2", FORK_FROM_CURRENT_ID]); + const action = items.at(-1)!; + expect(action.isAction).toBe(true); + expect(action.text.toLowerCase()).toContain("current state"); + }); + + it("omits the action row and shows nothing to fork when there is no leaf and no messages", () => { + const fake = makeFakeThis({ userMessages: [], leafId: null }); + invoke(fake); + + expect(fake.showSelector).not.toHaveBeenCalled(); + expect(fake.showStatus).toHaveBeenCalledWith("No messages to fork from"); + }); + + it("still offers the action row when there is a leaf but no user messages", () => { + const fake = makeFakeThis({ userMessages: [], leafId: "a1" }); + invoke(fake); + + expect(fake.showSelector).toHaveBeenCalledOnce(); + expect(captured[0].items.map((i) => i.id)).toEqual([FORK_FROM_CURRENT_ID]); + }); + + it("routes the action row to forkFromCurrent() and resets the editor on success", async () => { + const fake = makeFakeThis({ + userMessages: [{ entryId: "u1", text: "first" }], + leafId: "a1", + forkFromCurrentResult: { cancelled: false }, + }); + invoke(fake); + + await captured[0].onSelect(FORK_FROM_CURRENT_ID); + + expect(fake.session.forkFromCurrent).toHaveBeenCalledOnce(); + expect(fake.session.fork).not.toHaveBeenCalled(); + expect(fake.resetChatDisplay).toHaveBeenCalledOnce(); + expect(fake.editor.setText).toHaveBeenCalledWith(""); // no re-ask pre-fill + expect(fake._done).toHaveBeenCalledOnce(); + expect(fake.showStatus).toHaveBeenCalledWith("Branched to new session (including last response)"); + }); + + it("informs the user and does not reset the editor when forkFromCurrent is cancelled", async () => { + const fake = makeFakeThis({ + userMessages: [{ entryId: "u1", text: "first" }], + leafId: "a1", + forkFromCurrentResult: { cancelled: true }, + }); + invoke(fake); + + await captured[0].onSelect(FORK_FROM_CURRENT_ID); + + expect(fake.session.forkFromCurrent).toHaveBeenCalledOnce(); + expect(fake.resetChatDisplay).not.toHaveBeenCalled(); + expect(fake.editor.setText).not.toHaveBeenCalled(); + expect(fake._done).toHaveBeenCalledOnce(); + expect(fake.showStatus).toHaveBeenCalledWith("Fork cancelled — no new branch was created"); + }); + + it("routes a history message to fork() and pre-fills the editor with the selected text", async () => { + const fake = makeFakeThis({ + userMessages: [{ entryId: "u1", text: "first" }], + leafId: "a1", + forkResult: { cancelled: false, selectedText: "first" }, + }); + invoke(fake); + + await captured[0].onSelect("u1"); + + expect(fake.session.fork).toHaveBeenCalledWith("u1"); + expect(fake.session.forkFromCurrent).not.toHaveBeenCalled(); + expect(fake.editor.setText).toHaveBeenCalledWith("first"); + expect(fake.showStatus).toHaveBeenCalledWith("Branched to new session"); + }); +}); diff --git a/packages/dashboard/src/client/screens/session.tsx b/packages/dashboard/src/client/screens/session.tsx index d11b21df..a82d7187 100644 --- a/packages/dashboard/src/client/screens/session.tsx +++ b/packages/dashboard/src/client/screens/session.tsx @@ -1269,11 +1269,18 @@ export function SessionScreen(props: { store: AppStore; sessionKey: string }): J } } - async function selectForkMessage(entryId: string) { + // Shared completion for both fork flows: run the fork action, inform the user + // (keeping the modal open) if no branch was created, otherwise pre-fill the + // composer when the action returns re-ask text, refresh, and close. + async function finishFork(action: () => Promise<{ cancelled: boolean; text?: string }>, cancelMessage: string) { setForkError(undefined); try { - const result = await api.fork(props.sessionKey, entryId); - if (!result.cancelled) setComposerText(result.text); + const result = await action(); + if (result.cancelled) { + setForkError(cancelMessage); + return; + } + if (result.text !== undefined) setComposerText(result.text); await props.store.hydrateSession(props.sessionKey); await props.store.refreshDiskSessions(); setShowForkModal(false); @@ -1282,18 +1289,14 @@ export function SessionScreen(props: { store: AppStore; sessionKey: string }): J } } - async function forkFromCurrentState() { - setForkError(undefined); - try { - await api.forkCurrent(props.sessionKey); - // No composer pre-fill: the branch already includes the last response. - await props.store.hydrateSession(props.sessionKey); - await props.store.refreshDiskSessions(); - setShowForkModal(false); - } catch (err) { - setForkError(err instanceof Error ? err.message : String(err)); - } - } + const selectForkMessage = (entryId: string) => + finishFork(() => api.fork(props.sessionKey, entryId), "Fork cancelled — no new branch was created."); + + const forkFromCurrentState = () => + finishFork( + () => api.forkCurrent(props.sessionKey), + "Can't fork from current state — the session is empty or the fork was cancelled.", + ); async function openStatsPopover() { setShowStatsPopover(true); diff --git a/packages/dashboard/test/client/screens.test.tsx b/packages/dashboard/test/client/screens.test.tsx index f30ce613..65b12d61 100644 --- a/packages/dashboard/test/client/screens.test.tsx +++ b/packages/dashboard/test/client/screens.test.tsx @@ -5850,6 +5850,39 @@ describe("dashboard client regressions", () => { expect((el.querySelector("textarea") as HTMLTextAreaElement).value).toBe(""); }); + it("fork modal informs the user and stays open when fork-from-current is cancelled", async () => { + vi.mocked(api.forkMessages).mockResolvedValue({ messages: [{ entryId: "u1", text: "original prompt" }] }); + // Empty session / extension veto → backend returns cancelled with no branch created. + vi.mocked(api.forkCurrent).mockResolvedValue({ cancelled: true }); + const store = makeStore() as any; + const hydrateSession = vi.fn(async () => {}); + const refreshDiskSessions = vi.fn(async () => {}); + const fakeStore = { + ...store, + sessions: { forkcancel: createSessionViewState("forkcancel") }, + fleet: () => ({ runtimes: [], diskSessions: [] }), + hydrateSession, + refreshDiskSessions, + }; + const el = mount(() => ); + (el.querySelector(".session-bar .right .switcher:last-child") as HTMLButtonElement).click(); + await new Promise((resolve) => setTimeout(resolve, 0)); + [...el.querySelectorAll("button")].find((button) => button.textContent?.includes("fork"))?.click(); + await new Promise((resolve) => setTimeout(resolve, 10)); + // Ignore the mount-time hydration; assert only what the fork handler does. + hydrateSession.mockClear(); + refreshDiskSessions.mockClear(); + (el.querySelector(".fork-current-btn") as HTMLButtonElement).click(); + await new Promise((resolve) => setTimeout(resolve, 10)); + expect(api.forkCurrent).toHaveBeenCalledWith("forkcancel"); + // The user is informed: the modal stays open with a message, and no session + // churn happens as if a branch had been created. + expect(el.querySelector(".fork-current-btn")).not.toBeNull(); + expect(el.querySelector(".pair-error")?.textContent ?? "").toMatch(/can't fork|cancelled/i); + expect(hydrateSession).not.toHaveBeenCalled(); + expect(refreshDiskSessions).not.toHaveBeenCalled(); + }); + it("session stats popover shows the detailed stats breakdown", async () => { const store = makeStore() as any; const fakeStore = { From 50e792be92e1d99adf10e1dafc60f7e4d1d559db Mon Sep 17 00:00:00 2001 From: Test Date: Sun, 9 Aug 2026 12:03:08 +0200 Subject: [PATCH 4/7] Test message-fork cancelled path keeps modal open (#439) --- .../dashboard/test/client/screens.test.tsx | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/packages/dashboard/test/client/screens.test.tsx b/packages/dashboard/test/client/screens.test.tsx index 65b12d61..4882f58e 100644 --- a/packages/dashboard/test/client/screens.test.tsx +++ b/packages/dashboard/test/client/screens.test.tsx @@ -5883,6 +5883,41 @@ describe("dashboard client regressions", () => { expect(refreshDiskSessions).not.toHaveBeenCalled(); }); + it("fork modal informs the user and stays open when a message fork is cancelled", async () => { + vi.mocked(api.forkMessages).mockResolvedValue({ messages: [{ entryId: "u1", text: "original prompt" }] }); + // Extension veto → api.fork returns cancelled with no branch created. + vi.mocked(api.fork).mockResolvedValue({ text: "", cancelled: true }); + const store = makeStore() as any; + const hydrateSession = vi.fn(async () => {}); + const refreshDiskSessions = vi.fn(async () => {}); + const fakeStore = { + ...store, + sessions: { forkmsgcancel: createSessionViewState("forkmsgcancel") }, + fleet: () => ({ runtimes: [], diskSessions: [] }), + hydrateSession, + refreshDiskSessions, + }; + const el = mount(() => ); + (el.querySelector(".session-bar .right .switcher:last-child") as HTMLButtonElement).click(); + await new Promise((resolve) => setTimeout(resolve, 0)); + [...el.querySelectorAll("button")].find((button) => button.textContent?.includes("fork"))?.click(); + await new Promise((resolve) => setTimeout(resolve, 10)); + // Ignore the mount-time hydration; assert only what the fork handler does. + hydrateSession.mockClear(); + refreshDiskSessions.mockClear(); + (el.querySelector(".fork-message") as HTMLButtonElement).click(); + await new Promise((resolve) => setTimeout(resolve, 10)); + expect(api.fork).toHaveBeenCalledWith("forkmsgcancel", "u1"); + // The shared finishFork helper must inform the user for the message-fork flow too: + // the modal stays open with a message, the composer is not pre-filled, and no + // session churn happens as if a branch had been created. + expect(el.querySelector(".fork-message")).not.toBeNull(); + expect(el.querySelector(".pair-error")?.textContent ?? "").toMatch(/no new branch|cancelled/i); + expect((el.querySelector("textarea") as HTMLTextAreaElement).value).toBe(""); + expect(hydrateSession).not.toHaveBeenCalled(); + expect(refreshDiskSessions).not.toHaveBeenCalled(); + }); + it("session stats popover shows the detailed stats breakdown", async () => { const store = makeStore() as any; const fakeStore = { From 21707341aa8e4b8702c619867b332ba8395ab171 Mon Sep 17 00:00:00 2001 From: Test Date: Tue, 11 Aug 2026 20:22:28 +0200 Subject: [PATCH 5/7] Fork from any message: role-aware fork, drop fork-from-current (#439) --- .../coding-agent/src/core/agent-session.ts | 92 +++++---- .../components/user-message-selector.ts | 55 +++-- .../src/modes/interactive/interactive-mode.ts | 61 +++--- .../coding-agent/src/modes/rpc/rpc-client.ts | 15 +- .../coding-agent/src/modes/rpc/rpc-mode.ts | 7 +- .../coding-agent/src/modes/rpc/rpc-types.ts | 4 +- .../test/agent-session-branching.test.ts | 8 +- .../test/agent-session-fork-current.test.ts | 127 ------------ .../test/agent-session-fork.test.ts | 188 ++++++++++++++++++ .../test/interactive-mode-fork.test.ts | 112 +++++------ .../test/rpc-fork-current.test.ts | 31 --- packages/coding-agent/test/rpc-fork.test.ts | 42 ++++ packages/dashboard/src/client/api.ts | 5 +- .../dashboard/src/client/screens/session.tsx | 20 +- packages/dashboard/src/client/styles/app.css | 25 +-- packages/dashboard/src/server/server.ts | 4 - .../dashboard/test/client/screens.test.tsx | 66 ++---- packages/dashboard/test/runtime-pool.test.ts | 1 - packages/dashboard/test/server.test.ts | 8 - 19 files changed, 426 insertions(+), 445 deletions(-) delete mode 100644 packages/coding-agent/test/agent-session-fork-current.test.ts create mode 100644 packages/coding-agent/test/agent-session-fork.test.ts delete mode 100644 packages/coding-agent/test/rpc-fork-current.test.ts create mode 100644 packages/coding-agent/test/rpc-fork.test.ts diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index a5683390..cdf17be5 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -3627,26 +3627,49 @@ export class AgentSession { } /** - * Create a fork from a specific entry. + * Create a fork from a specific entry. The fork point may be any user or + * assistant message in the transcript; branch semantics depend on the role: + * + * - **Assistant message** -> the new branch *includes* the selected response + * (and everything before it); no editor pre-fill. "Continue from this answer." + * Forking at the last assistant message keeps the entire current state. + * - **User message** -> rewind to *before* the selected message (branch from its + * parent, dropping the message and everything after it) and offer its text as + * editor pre-fill. "Edit / re-ask this question." + * * Emits before_fork/fork session events to extensions. * - * @param entryId ID of the entry to fork from + * @param entryId ID of the message entry to fork from * @returns Object with: - * - selectedText: The text of the selected user message (for editor pre-fill) + * - selectedText: The selected user message text for editor pre-fill (empty + * when forking at an assistant message). * - cancelled: True if an extension cancelled the fork */ async fork(entryId: string): Promise<{ selectedText: string; cancelled: boolean }> { const selectedEntry = this.sessionManager.getEntry(entryId); - if (!selectedEntry || selectedEntry.type !== "message" || selectedEntry.message.role !== "user") { + if ( + !selectedEntry || + selectedEntry.type !== "message" || + (selectedEntry.message.role !== "user" && selectedEntry.message.role !== "assistant") + ) { throw new Error("Invalid entry ID for forking"); } - const selectedText = this._extractUserMessageText(selectedEntry.message.content); + if (selectedEntry.message.role === "assistant") { + // Continue-from-answer: branch from the assistant entry itself so it (and + // everything before it) is retained. No editor pre-fill. + const { cancelled } = await this._performFork(entryId, () => { + this.sessionManager.createBranchedSession(entryId); + }); + return { selectedText: "", cancelled }; + } + + const selectedText = this._extractMessageText(selectedEntry.message.content); - // Existing fork semantics: rewind to *before* the selected user message by - // branching from its parent, so the selected message (and everything after - // it) is dropped and its text is offered as editor pre-fill. + // Rewind to *before* the selected user message by branching from its parent, + // so the selected message (and everything after it) is dropped and its text is + // offered as editor pre-fill. const { cancelled } = await this._performFork(entryId, (previousSessionFile) => { if (!selectedEntry.parentId) { this.sessionManager.newSession({ parentSession: previousSessionFile }); @@ -3658,32 +3681,6 @@ export class AgentSession { return { selectedText, cancelled }; } - /** - * Create a fork from the current state, including the last model response. - * - * Unlike fork(), which rewinds to *before* a selected user message, this - * branches from the current leaf entry — so the new branch's tail is the - * latest entry (typically the last assistant response). There is no editor - * pre-fill; the new session is ready for a fresh turn on top of the captured - * history. - * - * @returns Object with `cancelled` (true if an extension cancelled the fork, - * or if there is no current leaf to fork from — i.e. an empty session). - */ - async forkFromCurrent(): Promise<{ cancelled: boolean }> { - const leafId = this.sessionManager.getLeafId(); - - // Nothing to fork from (fresh/empty session). - if (!leafId) { - return { cancelled: true }; - } - - return this._performFork(leafId, () => { - // Branch from the leaf itself so the last entry is included. - this.sessionManager.createBranchedSession(leafId); - }); - } - /** * Shared fork machinery: emit the cancellable session_before_fork event, * clear pending state, create the branch via the supplied strategy, reload @@ -3902,7 +3899,7 @@ export class AgentSession { if (targetEntry.type === "message" && targetEntry.message.role === "user") { // User message: leaf = parent (null if root), text goes to editor newLeafId = targetEntry.parentId; - editorText = this._extractUserMessageText(targetEntry.message.content); + editorText = this._extractMessageText(targetEntry.message.content); } else if (targetEntry.type === "custom_message") { // Custom message: leaf = parent (null if root), text goes to editor newLeafId = targetEntry.parentId; @@ -3972,26 +3969,35 @@ export class AgentSession { } /** - * Get all user messages from session for fork selector. + * Get all forkable messages (user *and* assistant) for the fork selector. + * + * Each entry carries its role so callers can label it and choose the right + * fork semantics (assistant = continue-from-answer, user = rewind + re-ask). + * Assistant turns with no renderable text (pure tool-call turns) still appear + * as fork points, with a generic label. */ - getUserMessagesForForking(): Array<{ entryId: string; text: string }> { + getForkableMessages(): Array<{ entryId: string; text: string; role: "user" | "assistant" }> { const entries = this.sessionManager.getEntries(); - const result: Array<{ entryId: string; text: string }> = []; + const result: Array<{ entryId: string; text: string; role: "user" | "assistant" }> = []; for (const entry of entries) { if (entry.type !== "message") continue; - if (entry.message.role !== "user") continue; + const role = entry.message.role; + if (role !== "user" && role !== "assistant") continue; - const text = this._extractUserMessageText(entry.message.content); - if (text) { - result.push({ entryId: entry.id, text }); + const text = this._extractMessageText(entry.message.content); + if (role === "user") { + // Preserve existing behavior: skip empty user messages. + if (text) result.push({ entryId: entry.id, text, role }); + } else { + result.push({ entryId: entry.id, text: text || "(assistant response)", role }); } } return result; } - private _extractUserMessageText(content: string | Array<{ type: string; text?: string }>): string { + private _extractMessageText(content: string | Array<{ type: string; text?: string }>): string { if (typeof content === "string") return content; if (Array.isArray(content)) { return content diff --git a/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts b/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts index b3d42cee..80059cb1 100644 --- a/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts +++ b/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts @@ -4,16 +4,16 @@ import { DynamicBorder } from "./dynamic-border.js"; interface UserMessageItem { id: string; // Entry ID in the session - text: string; // The message text + text: string; // The message text (preview) + role: "user" | "assistant"; // Whose message this is — drives fork semantics timestamp?: string; // Optional timestamp if available - isAction?: boolean; // Special action row (e.g. "fork from current state") rather than a history message } -/** Sentinel entry id for the "fork from current state (include last response)" action row. */ -export const FORK_FROM_CURRENT_ID = "__fork_from_current__"; - /** - * Custom user message list component with selection + * Custom message list component with selection. Lists both user and assistant + * messages as fork points; the role determines the branch semantics: + * - assistant → continue from that answer (branch includes it) + * - user → rewind to before it and re-ask (editor pre-filled) */ class UserMessageList implements Component { private messages: UserMessageItem[] = []; @@ -37,7 +37,7 @@ class UserMessageList implements Component { const lines: string[] = []; if (this.messages.length === 0) { - lines.push(theme.fg("muted", " No user messages found")); + lines.push(theme.fg("muted", " No messages found")); return lines; } @@ -49,38 +49,29 @@ class UserMessageList implements Component { const endIndex = Math.min(startIndex + this.maxVisible, this.messages.length); // Render visible messages (2 lines per message + blank line) - const totalRealMessages = this.messages.filter((m) => !m.isAction).length; for (let i = startIndex; i < endIndex; i++) { const message = this.messages[i]; const isSelected = i === this.selectedIndex; + const isAssistant = message.role === "assistant"; // Normalize message to single line const normalizedMessage = message.text.replace(/\n/g, " ").trim(); - // First line: cursor + message + // First line: cursor + role badge + message preview const cursor = isSelected ? theme.fg("accent", "› ") : " "; - const maxMsgWidth = width - 2; // Account for cursor (2 chars) - - if (message.isAction) { - // Distinct styling for the action row (e.g. "fork from current state"). - const label = truncateToWidth(`⎇ ${normalizedMessage}`, maxMsgWidth); - const styled = isSelected ? theme.fg("accent", theme.bold(label)) : theme.fg("accent", label); - lines.push(cursor + styled); - lines.push(theme.fg("muted", " Branch here, including the last response")); - lines.push(""); // Blank line between entries - continue; - } - - const truncatedMsg = truncateToWidth(normalizedMessage, maxMsgWidth); - const messageLine = cursor + (isSelected ? theme.bold(truncatedMsg) : truncatedMsg); + const badgeText = isAssistant ? "[Assistant] " : "[You] "; + const badge = theme.fg(isAssistant ? "accent" : "muted", badgeText); + const maxMsgWidth = width - 2 - badgeText.length; // cursor (2) + badge + const truncatedMsg = truncateToWidth(normalizedMessage, Math.max(0, maxMsgWidth)); + const messageLine = cursor + badge + (isSelected ? theme.bold(truncatedMsg) : truncatedMsg); lines.push(messageLine); - // Second line: metadata (position in history, counting only real messages) - const position = this.messages.slice(0, i + 1).filter((m) => !m.isAction).length; - const metadata = ` Message ${position} of ${totalRealMessages}`; - const metadataLine = theme.fg("muted", metadata); - lines.push(metadataLine); + // Second line: position + what forking here does + const position = i + 1; + const hint = isAssistant ? "continue from here" : "rewind & re-ask"; + const metadata = ` Message ${position} of ${this.messages.length} · ${hint}`; + lines.push(theme.fg("muted", metadata)); lines.push(""); // Blank line between messages } @@ -120,7 +111,8 @@ class UserMessageList implements Component { } /** - * Component that renders a user message selector for branching + * Component that renders a message selector for branching. Any user or assistant + * message is a valid fork point. */ export class UserMessageSelectorComponent extends Container { private messageList: UserMessageList; @@ -133,7 +125,10 @@ export class UserMessageSelectorComponent extends Container { this.addChild(new Text(theme.bold("Branch from Message"), 1, 0)); this.addChild( new Text( - theme.fg("muted", "Pick a message to rewind and re-run, or fork from current state (keeps last response)"), + theme.fg( + "muted", + "Pick any message: an assistant reply continues from that answer, a question rewinds to re-ask it", + ), 1, 0, ), diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 1f2eb6c8..02658f7c 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -121,7 +121,7 @@ import { TasksPanelComponent } from "./components/tasks-panel.js"; import { ToolExecutionComponent } from "./components/tool-execution.js"; import { TreeSelectorComponent } from "./components/tree-selector.js"; import { UserMessageComponent } from "./components/user-message.js"; -import { FORK_FROM_CURRENT_ID, UserMessageSelectorComponent } from "./components/user-message-selector.js"; +import { UserMessageSelectorComponent } from "./components/user-message-selector.js"; import { getAvailableThemes, getAvailableThemesWithPaths, @@ -4296,53 +4296,44 @@ export class InteractiveMode { } private showUserMessageSelector(): void { - const userMessages = this.session.getUserMessagesForForking(); - const hasCurrentState = this.sessionManager.getLeafId() !== null; + const messages = this.session.getForkableMessages(); - if (userMessages.length === 0 && !hasCurrentState) { + if (messages.length === 0) { this.showStatus("No messages to fork from"); return; } - // Build the selector list: history messages (rewind + re-run), plus a - // trailing "fork from current state" action that keeps the last response. - const items: Array<{ id: string; text: string; isAction?: boolean }> = userMessages.map((m) => ({ - id: m.entryId, - text: m.text, - })); - if (hasCurrentState) { - items.push({ - id: FORK_FROM_CURRENT_ID, - text: "Fork from current state (include last response)", - isAction: true, - }); - } + // Every user or assistant message is a fork point. Assistant replies branch + // so the answer is kept (continue from here); user messages rewind to before + // the question and pre-fill the editor (re-ask). + const items = messages.map((m) => ({ id: m.entryId, text: m.text, role: m.role })); this.showSelector((done) => { const selector = new UserMessageSelectorComponent( items, async (entryId) => { - const isCurrent = entryId === FORK_FROM_CURRENT_ID; - const result = isCurrent ? await this.session.forkFromCurrent() : await this.session.fork(entryId); + try { + const result = await this.session.fork(entryId); + if (result.cancelled) { + // An extension vetoed the fork — tell the user rather than + // silently dismissing the selector. + done(); + this.showStatus("Fork cancelled — no new branch was created"); + return; + } - if (result.cancelled) { - // Empty session (nothing to branch from) or an extension vetoed - // the fork — tell the user rather than silently dismissing the - // selector. + this.resetChatDisplay(); + // Assistant forks return empty text (nothing to re-ask); user forks + // pre-fill the editor with the selected question. + this.editor.setText(result.selectedText); done(); - this.showStatus("Fork cancelled — no new branch was created"); - return; + this.showStatus("Branched to new session"); + } catch (err) { + // Forking can throw (e.g. a stale entry or a filesystem error while + // writing the branch). Surface it instead of crashing the TUI. + done(); + this.showStatus(`Fork failed: ${err instanceof Error ? err.message : String(err)}`); } - - this.resetChatDisplay(); - // The current-state branch already includes the last response, so - // there is nothing to re-ask; the message branch pre-fills the editor - // with the selected question for re-running. - this.editor.setText((result as { selectedText?: string }).selectedText ?? ""); - done(); - this.showStatus( - isCurrent ? "Branched to new session (including last response)" : "Branched to new session", - ); }, () => { done(); diff --git a/packages/coding-agent/src/modes/rpc/rpc-client.ts b/packages/coding-agent/src/modes/rpc/rpc-client.ts index 6dc1e297..f02ad292 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-client.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-client.ts @@ -626,22 +626,13 @@ export class RpcClient { return this.getData(response); } - /** - * Fork from the current state, including the last model response. - * Unlike fork(), this branches from the current leaf (no editor pre-fill). - * @returns Object with `cancelled` (if an extension cancelled, or the session is empty) - */ - async forkCurrent(): Promise<{ cancelled: boolean }> { - const response = await this.send({ type: "fork_current" }); - return this.getData(response); - } - /** * Get messages available for forking. */ - async getForkMessages(): Promise> { + async getForkMessages(): Promise> { const response = await this.send({ type: "get_fork_messages" }); - return this.getData<{ messages: Array<{ entryId: string; text: string }> }>(response).messages; + return this.getData<{ messages: Array<{ entryId: string; text: string; role: "user" | "assistant" }> }>(response) + .messages; } /** diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index 5bdd97c8..9a8e9d9c 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -2126,13 +2126,8 @@ export async function runRpcMode(session: AgentSession, modelFallbackMessage?: s return success(id, "fork", { text: result.selectedText, cancelled: result.cancelled }); } - case "fork_current": { - const result = await session.forkFromCurrent(); - return success(id, "fork_current", { cancelled: result.cancelled }); - } - case "get_fork_messages": { - const messages = session.getUserMessagesForForking(); + const messages = session.getForkableMessages(); return success(id, "get_fork_messages", { messages }); } diff --git a/packages/coding-agent/src/modes/rpc/rpc-types.ts b/packages/coding-agent/src/modes/rpc/rpc-types.ts index 579310ef..b8952a5e 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-types.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-types.ts @@ -85,7 +85,6 @@ export type RpcCommand = | { id?: string; type: "switch_session"; sessionPath: string } | { id?: string; type: "delete_session"; sessionPath: string } | { id?: string; type: "fork"; entryId: string } - | { id?: string; type: "fork_current" } | { id?: string; type: "get_fork_messages" } | { id?: string; type: "get_tree" } | { @@ -374,13 +373,12 @@ export type RpcResponse = | { id?: string; type: "response"; command: "switch_session"; success: true; data: { cancelled: boolean } } | { id?: string; type: "response"; command: "delete_session"; success: true; data: { method: "trash" | "unlink" } } | { id?: string; type: "response"; command: "fork"; success: true; data: { text: string; cancelled: boolean } } - | { id?: string; type: "response"; command: "fork_current"; success: true; data: { cancelled: boolean } } | { id?: string; type: "response"; command: "get_fork_messages"; success: true; - data: { messages: Array<{ entryId: string; text: string }> }; + data: { messages: Array<{ entryId: string; text: string; role: "user" | "assistant" }> }; } | { id?: string; diff --git a/packages/coding-agent/test/agent-session-branching.test.ts b/packages/coding-agent/test/agent-session-branching.test.ts index a3fa35de..0eb941af 100644 --- a/packages/coding-agent/test/agent-session-branching.test.ts +++ b/packages/coding-agent/test/agent-session-branching.test.ts @@ -4,7 +4,7 @@ * These tests verify: * - Forking from a single message works * - Forking in --no-session mode (in-memory only) - * - getUserMessagesForForking returns correct entries + * - getForkableMessages returns correct entries */ import { existsSync, mkdirSync, rmSync } from "node:fs"; @@ -80,7 +80,7 @@ describe.skipIf(process.env.DREB_SKIP_LIVE_API === "1" || !API_KEY)("AgentSessio await session.agent.waitForIdle(); // Should have exactly 1 user message available for forking - const userMessages = session.getUserMessagesForForking(); + const userMessages = session.getForkableMessages().filter((m) => m.role === "user"); expect(userMessages.length).toBe(1); expect(userMessages[0].text).toBe("Say hello"); @@ -108,7 +108,7 @@ describe.skipIf(process.env.DREB_SKIP_LIVE_API === "1" || !API_KEY)("AgentSessio await session.agent.waitForIdle(); // Should have 1 user message - const userMessages = session.getUserMessagesForForking(); + const userMessages = session.getForkableMessages().filter((m) => m.role === "user"); expect(userMessages.length).toBe(1); // Verify we have messages before forking @@ -140,7 +140,7 @@ describe.skipIf(process.env.DREB_SKIP_LIVE_API === "1" || !API_KEY)("AgentSessio await session.agent.waitForIdle(); // Should have 3 user messages - const userMessages = session.getUserMessagesForForking(); + const userMessages = session.getForkableMessages().filter((m) => m.role === "user"); expect(userMessages.length).toBe(3); // Fork from second message (keeps first message + response) diff --git a/packages/coding-agent/test/agent-session-fork-current.test.ts b/packages/coding-agent/test/agent-session-fork-current.test.ts deleted file mode 100644 index 3455005e..00000000 --- a/packages/coding-agent/test/agent-session-fork-current.test.ts +++ /dev/null @@ -1,127 +0,0 @@ -/** - * Tests for AgentSession.forkFromCurrent(). - * - * Unlike fork(), which rewinds to *before* a selected user message, forkFromCurrent() - * branches from the current leaf — so the new branch includes the last model response. - * - * These run offline (no live API): the conversation is constructed by appending - * messages directly to the in-memory SessionManager, then forkFromCurrent() reloads - * the branch via buildSessionContext + replaceMessages. - */ - -import { afterEach, describe, expect, it } from "vitest"; -import { createHarnessWithExtensions, type Harness } from "./test-harness.js"; -import { assistantMsg, userMsg } from "./utilities.js"; - -describe("AgentSession.forkFromCurrent", () => { - let harness: Harness; - - afterEach(() => { - harness?.cleanup(); - }); - - it("includes the last model response in the forked branch", async () => { - harness = await createHarnessWithExtensions(); - const { session, sessionManager } = harness; - - // Conversation: q1 -> a1 -> q2 -> a2 (a2 is the last model response) - sessionManager.appendMessage(userMsg("q1")); - sessionManager.appendMessage(assistantMsg("a1")); - sessionManager.appendMessage(userMsg("q2")); - const lastAssistantId = sessionManager.appendMessage(assistantMsg("a2")); - expect(sessionManager.getLeafId()).toBe(lastAssistantId); - - const result = await session.forkFromCurrent(); - expect(result.cancelled).toBe(false); - - // The full conversation survives and the tail is the last assistant response. - expect(session.messages).toHaveLength(4); - const tail = session.messages.at(-1)!; - expect(tail.role).toBe("assistant"); - expect(JSON.stringify(tail)).toContain("a2"); - - // A new session was branched (fresh id), still tracking the conversation. - expect(sessionManager.getLeafId()).not.toBeNull(); - }); - - it("is a no-op that returns cancelled for an empty session", async () => { - harness = await createHarnessWithExtensions(); - const { session, sessionManager } = harness; - - expect(sessionManager.getLeafId()).toBeNull(); - - const result = await session.forkFromCurrent(); - expect(result.cancelled).toBe(true); - expect(session.messages).toHaveLength(0); - }); - - it("can be cancelled by a session_before_fork extension handler", async () => { - harness = await createHarnessWithExtensions({ - extensionFactories: [ - (dreb) => { - dreb.on("session_before_fork", async () => ({ cancel: true })); - }, - ], - }); - const { session, sessionManager } = harness; - - sessionManager.appendMessage(userMsg("q1")); - sessionManager.appendMessage(assistantMsg("a1")); - const leafBefore = sessionManager.getLeafId(); - const entriesBefore = sessionManager.getEntries(); - - const result = await session.forkFromCurrent(); - expect(result.cancelled).toBe(true); - - // Cancellation must not branch or mutate the session. - expect(sessionManager.getLeafId()).toBe(leafBefore); - expect(sessionManager.getEntries()).toEqual(entriesBefore); - }); - - it("emits session_fork after the branch is created", async () => { - const forkEvents: Array<{ type: string }> = []; - harness = await createHarnessWithExtensions({ - extensionFactories: [ - (dreb) => { - dreb.on("session_fork", async (event) => { - forkEvents.push(event); - }); - }, - ], - }); - const { session, sessionManager } = harness; - - sessionManager.appendMessage(userMsg("q1")); - sessionManager.appendMessage(assistantMsg("a1")); - - const result = await session.forkFromCurrent(); - expect(result.cancelled).toBe(false); - - // AC5: the after-fork event must fire exactly once on the new leaf-based path. - expect(forkEvents).toHaveLength(1); - expect(forkEvents[0].type).toBe("session_fork"); - }); - - it("does not emit session_fork when a session_before_fork handler cancels", async () => { - const forkEvents: Array<{ type: string }> = []; - harness = await createHarnessWithExtensions({ - extensionFactories: [ - (dreb) => { - dreb.on("session_before_fork", async () => ({ cancel: true })); - dreb.on("session_fork", async (event) => { - forkEvents.push(event); - }); - }, - ], - }); - const { session, sessionManager } = harness; - - sessionManager.appendMessage(userMsg("q1")); - sessionManager.appendMessage(assistantMsg("a1")); - - const result = await session.forkFromCurrent(); - expect(result.cancelled).toBe(true); - // A vetoed fork must not fire the after-fork event. - expect(forkEvents).toHaveLength(0); - }); -}); diff --git a/packages/coding-agent/test/agent-session-fork.test.ts b/packages/coding-agent/test/agent-session-fork.test.ts new file mode 100644 index 00000000..351d391a --- /dev/null +++ b/packages/coding-agent/test/agent-session-fork.test.ts @@ -0,0 +1,188 @@ +/** + * Tests for AgentSession.fork() — forking at any transcript position. + * + * The fork point may be any user or assistant message. Role determines semantics: + * - assistant → branch *includes* the selected response (continue from that answer), + * with no editor pre-fill. Forking at the last assistant keeps the whole state. + * - user → rewind to *before* the selected question (drop it and everything after) + * and offer its text as editor pre-fill. + * + * These run offline (no live API): the conversation is constructed by appending + * messages directly to the in-memory SessionManager, then fork() reloads the branch + * via buildSessionContext + replaceMessages. + */ + +import { afterEach, describe, expect, it } from "vitest"; +import { createHarnessWithExtensions, type Harness } from "./test-harness.js"; +import { assistantMsg, userMsg } from "./utilities.js"; + +describe("AgentSession.fork — any message", () => { + let harness: Harness; + + afterEach(() => { + harness?.cleanup(); + }); + + it("forking at the last assistant message keeps the full state (incl. last response)", async () => { + harness = await createHarnessWithExtensions(); + const { session, sessionManager } = harness; + + // q1 -> a1 -> q2 -> a2 (a2 is the last model response) + sessionManager.appendMessage(userMsg("q1")); + sessionManager.appendMessage(assistantMsg("a1")); + sessionManager.appendMessage(userMsg("q2")); + const lastAssistantId = sessionManager.appendMessage(assistantMsg("a2")); + expect(sessionManager.getLeafId()).toBe(lastAssistantId); + + const result = await session.fork(lastAssistantId); + expect(result.cancelled).toBe(false); + // No re-ask pre-fill when forking at an assistant message. + expect(result.selectedText).toBe(""); + + // The full conversation survives and the tail is the last assistant response. + expect(session.messages).toHaveLength(4); + const tail = session.messages.at(-1)!; + expect(tail.role).toBe("assistant"); + expect(JSON.stringify(tail)).toContain("a2"); + expect(sessionManager.getLeafId()).not.toBeNull(); + }); + + it("forking at an earlier assistant message includes up to and including it", async () => { + harness = await createHarnessWithExtensions(); + const { session, sessionManager } = harness; + + sessionManager.appendMessage(userMsg("q1")); + const a1Id = sessionManager.appendMessage(assistantMsg("a1")); + sessionManager.appendMessage(userMsg("q2")); + sessionManager.appendMessage(assistantMsg("a2")); + + const result = await session.fork(a1Id); + expect(result.cancelled).toBe(false); + expect(result.selectedText).toBe(""); + + // Branch is q1 -> a1; q2/a2 are dropped and the tail is a1. + expect(session.messages).toHaveLength(2); + const tail = session.messages.at(-1)!; + expect(tail.role).toBe("assistant"); + expect(JSON.stringify(tail)).toContain("a1"); + }); + + it("forking at a user message rewinds to before it and offers its text as pre-fill", async () => { + harness = await createHarnessWithExtensions(); + const { session, sessionManager } = harness; + + sessionManager.appendMessage(userMsg("q1")); + sessionManager.appendMessage(assistantMsg("a1")); + const q2Id = sessionManager.appendMessage(userMsg("q2")); + sessionManager.appendMessage(assistantMsg("a2")); + + const result = await session.fork(q2Id); + expect(result.cancelled).toBe(false); + // The selected question is offered for editing/re-asking. + expect(result.selectedText).toBe("q2"); + + // Branch is q1 -> a1 (before q2); q2 and everything after are dropped. + expect(session.messages).toHaveLength(2); + const tail = session.messages.at(-1)!; + expect(tail.role).toBe("assistant"); + expect(JSON.stringify(tail)).toContain("a1"); + }); + + it("throws for a non-message / invalid entry id", async () => { + harness = await createHarnessWithExtensions(); + const { session, sessionManager } = harness; + sessionManager.appendMessage(userMsg("q1")); + + await expect(session.fork("does-not-exist")).rejects.toThrow(/Invalid entry ID/); + }); + + it("can be cancelled by a session_before_fork extension handler", async () => { + harness = await createHarnessWithExtensions({ + extensionFactories: [ + (dreb) => { + dreb.on("session_before_fork", async () => ({ cancel: true })); + }, + ], + }); + const { session, sessionManager } = harness; + + sessionManager.appendMessage(userMsg("q1")); + const a1Id = sessionManager.appendMessage(assistantMsg("a1")); + const leafBefore = sessionManager.getLeafId(); + const entriesBefore = sessionManager.getEntries(); + + const result = await session.fork(a1Id); + expect(result.cancelled).toBe(true); + + // Cancellation must not branch or mutate the session. + expect(sessionManager.getLeafId()).toBe(leafBefore); + expect(sessionManager.getEntries()).toEqual(entriesBefore); + }); + + it("skips the conversation restore when a session_before_fork handler requests it", async () => { + harness = await createHarnessWithExtensions({ + extensionFactories: [ + (dreb) => { + dreb.on("session_before_fork", async () => ({ skipConversationRestore: true })); + }, + ], + }); + const { session, sessionManager } = harness; + + sessionManager.appendMessage(userMsg("q1")); + const a1Id = sessionManager.appendMessage(assistantMsg("a1")); + + const result = await session.fork(a1Id); + expect(result.cancelled).toBe(false); + + // The branch was created (new leaf tracked)... + expect(sessionManager.getLeafId()).not.toBeNull(); + // ...but agent.replaceMessages was skipped, so the in-memory conversation is + // NOT reloaded from the branch (stays empty here, since nothing was streamed). + expect(session.messages).toHaveLength(0); + }); + + it("emits session_fork exactly once after the branch is created", async () => { + const forkEvents: Array<{ type: string }> = []; + harness = await createHarnessWithExtensions({ + extensionFactories: [ + (dreb) => { + dreb.on("session_fork", async (event) => { + forkEvents.push(event); + }); + }, + ], + }); + const { session, sessionManager } = harness; + + sessionManager.appendMessage(userMsg("q1")); + const a1Id = sessionManager.appendMessage(assistantMsg("a1")); + + const result = await session.fork(a1Id); + expect(result.cancelled).toBe(false); + expect(forkEvents).toHaveLength(1); + expect(forkEvents[0].type).toBe("session_fork"); + }); + + it("does not emit session_fork when a session_before_fork handler cancels", async () => { + const forkEvents: Array<{ type: string }> = []; + harness = await createHarnessWithExtensions({ + extensionFactories: [ + (dreb) => { + dreb.on("session_before_fork", async () => ({ cancel: true })); + dreb.on("session_fork", async (event) => { + forkEvents.push(event); + }); + }, + ], + }); + const { session, sessionManager } = harness; + + sessionManager.appendMessage(userMsg("q1")); + const a1Id = sessionManager.appendMessage(assistantMsg("a1")); + + const result = await session.fork(a1Id); + expect(result.cancelled).toBe(true); + expect(forkEvents).toHaveLength(0); + }); +}); diff --git a/packages/coding-agent/test/interactive-mode-fork.test.ts b/packages/coding-agent/test/interactive-mode-fork.test.ts index 728185a1..67682fed 100644 --- a/packages/coding-agent/test/interactive-mode-fork.test.ts +++ b/packages/coding-agent/test/interactive-mode-fork.test.ts @@ -2,17 +2,19 @@ * Unit coverage for the interactive `/fork` selector wiring * (InteractiveMode.showUserMessageSelector). * - * Exercised without a full TUI: the selector component is mocked to capture the - * onSelect callback, and the method is invoked via prototype.call with a - * hand-built `this` (the same pattern as interactive-mode-status.test.ts). + * Any user or assistant message is a fork point — there is no separate + * "fork from current state" action row. Exercised without a full TUI: the + * selector component is mocked to capture the onSelect callback, and the method + * is invoked via prototype.call with a hand-built `this` (the same pattern as + * interactive-mode-status.test.ts). */ import { beforeEach, describe, expect, it, vi } from "vitest"; // Capture every UserMessageSelectorComponent construction so tests can drive the -// onSelect callback and assert the items list (incl. the fork-from-current row). +// onSelect callback and assert the items list (role-labeled, all messages). const captured: Array<{ - items: Array<{ id: string; text: string; isAction?: boolean }>; + items: Array<{ id: string; text: string; role: "user" | "assistant" }>; onSelect: (entryId: string) => void | Promise; onCancel: () => void; }> = []; @@ -21,7 +23,7 @@ vi.mock("../src/modes/interactive/components/user-message-selector.js", async (i const actual = await importOriginal>(); class MockUserMessageSelectorComponent { constructor( - public items: Array<{ id: string; text: string; isAction?: boolean }>, + public items: Array<{ id: string; text: string; role: "user" | "assistant" }>, public onSelect: (entryId: string) => void | Promise, public onCancel: () => void, ) { @@ -34,25 +36,24 @@ vi.mock("../src/modes/interactive/components/user-message-selector.js", async (i return { ...actual, UserMessageSelectorComponent: MockUserMessageSelectorComponent }; }); -import { FORK_FROM_CURRENT_ID } from "../src/modes/interactive/components/user-message-selector.js"; import { InteractiveMode } from "../src/modes/interactive/interactive-mode.js"; interface FakeOverrides { - userMessages?: Array<{ entryId: string; text: string }>; - leafId?: string | null; - forkFromCurrentResult?: { cancelled: boolean }; + messages?: Array<{ entryId: string; text: string; role: "user" | "assistant" }>; forkResult?: { cancelled: boolean; selectedText: string }; + forkThrows?: Error; } function makeFakeThis(overrides: FakeOverrides = {}) { const done = vi.fn(); const fake = { session: { - getUserMessagesForForking: vi.fn(() => overrides.userMessages ?? []), - forkFromCurrent: vi.fn(async () => overrides.forkFromCurrentResult ?? { cancelled: false }), - fork: vi.fn(async () => overrides.forkResult ?? { cancelled: false, selectedText: "prefill" }), + getForkableMessages: vi.fn(() => overrides.messages ?? []), + fork: vi.fn(async () => { + if (overrides.forkThrows) throw overrides.forkThrows; + return overrides.forkResult ?? { cancelled: false, selectedText: "prefill" }; + }), }, - sessionManager: { getLeafId: vi.fn(() => overrides.leafId ?? null) }, showStatus: vi.fn(), resetChatDisplay: vi.fn(), editor: { setText: vi.fn() }, @@ -71,95 +72,92 @@ function invoke(fake: ReturnType) { ).prototype.showUserMessageSelector.call(fake); } -describe("InteractiveMode.showUserMessageSelector — fork-from-current wiring", () => { +describe("InteractiveMode.showUserMessageSelector — fork at any message", () => { beforeEach(() => { captured.length = 0; }); - it("appends the fork-from-current action row when there is a current leaf", () => { + it("lists all user and assistant messages with their roles (no action row)", () => { const fake = makeFakeThis({ - userMessages: [ - { entryId: "u1", text: "first" }, - { entryId: "u2", text: "second" }, + messages: [ + { entryId: "u1", text: "first", role: "user" }, + { entryId: "a1", text: "answer", role: "assistant" }, + { entryId: "u2", text: "second", role: "user" }, ], - leafId: "a2", }); invoke(fake); expect(fake.showSelector).toHaveBeenCalledOnce(); expect(captured).toHaveLength(1); const { items } = captured[0]; - // History messages first, then the trailing action row. - expect(items.map((i) => i.id)).toEqual(["u1", "u2", FORK_FROM_CURRENT_ID]); - const action = items.at(-1)!; - expect(action.isAction).toBe(true); - expect(action.text.toLowerCase()).toContain("current state"); + expect(items.map((i) => i.id)).toEqual(["u1", "a1", "u2"]); + expect(items.map((i) => i.role)).toEqual(["user", "assistant", "user"]); }); - it("omits the action row and shows nothing to fork when there is no leaf and no messages", () => { - const fake = makeFakeThis({ userMessages: [], leafId: null }); + it("shows nothing to fork and does not open the selector when there are no messages", () => { + const fake = makeFakeThis({ messages: [] }); invoke(fake); expect(fake.showSelector).not.toHaveBeenCalled(); expect(fake.showStatus).toHaveBeenCalledWith("No messages to fork from"); }); - it("still offers the action row when there is a leaf but no user messages", () => { - const fake = makeFakeThis({ userMessages: [], leafId: "a1" }); + it("forking at an assistant message resets the display with no editor pre-fill", async () => { + const fake = makeFakeThis({ + messages: [{ entryId: "a1", text: "answer", role: "assistant" }], + forkResult: { cancelled: false, selectedText: "" }, + }); invoke(fake); - expect(fake.showSelector).toHaveBeenCalledOnce(); - expect(captured[0].items.map((i) => i.id)).toEqual([FORK_FROM_CURRENT_ID]); + await captured[0].onSelect("a1"); + + expect(fake.session.fork).toHaveBeenCalledWith("a1"); + expect(fake.resetChatDisplay).toHaveBeenCalledOnce(); + expect(fake.editor.setText).toHaveBeenCalledWith(""); + expect(fake._done).toHaveBeenCalledOnce(); + expect(fake.showStatus).toHaveBeenCalledWith("Branched to new session"); }); - it("routes the action row to forkFromCurrent() and resets the editor on success", async () => { + it("forking at a user message pre-fills the editor with the selected text", async () => { const fake = makeFakeThis({ - userMessages: [{ entryId: "u1", text: "first" }], - leafId: "a1", - forkFromCurrentResult: { cancelled: false }, + messages: [{ entryId: "u1", text: "first", role: "user" }], + forkResult: { cancelled: false, selectedText: "first" }, }); invoke(fake); - await captured[0].onSelect(FORK_FROM_CURRENT_ID); + await captured[0].onSelect("u1"); - expect(fake.session.forkFromCurrent).toHaveBeenCalledOnce(); - expect(fake.session.fork).not.toHaveBeenCalled(); - expect(fake.resetChatDisplay).toHaveBeenCalledOnce(); - expect(fake.editor.setText).toHaveBeenCalledWith(""); // no re-ask pre-fill - expect(fake._done).toHaveBeenCalledOnce(); - expect(fake.showStatus).toHaveBeenCalledWith("Branched to new session (including last response)"); + expect(fake.session.fork).toHaveBeenCalledWith("u1"); + expect(fake.editor.setText).toHaveBeenCalledWith("first"); + expect(fake.showStatus).toHaveBeenCalledWith("Branched to new session"); }); - it("informs the user and does not reset the editor when forkFromCurrent is cancelled", async () => { + it("informs the user and does not reset the editor when the fork is cancelled", async () => { const fake = makeFakeThis({ - userMessages: [{ entryId: "u1", text: "first" }], - leafId: "a1", - forkFromCurrentResult: { cancelled: true }, + messages: [{ entryId: "u1", text: "first", role: "user" }], + forkResult: { cancelled: true, selectedText: "" }, }); invoke(fake); - await captured[0].onSelect(FORK_FROM_CURRENT_ID); + await captured[0].onSelect("u1"); - expect(fake.session.forkFromCurrent).toHaveBeenCalledOnce(); expect(fake.resetChatDisplay).not.toHaveBeenCalled(); expect(fake.editor.setText).not.toHaveBeenCalled(); expect(fake._done).toHaveBeenCalledOnce(); expect(fake.showStatus).toHaveBeenCalledWith("Fork cancelled — no new branch was created"); }); - it("routes a history message to fork() and pre-fills the editor with the selected text", async () => { + it("surfaces an error instead of crashing when fork() throws", async () => { const fake = makeFakeThis({ - userMessages: [{ entryId: "u1", text: "first" }], - leafId: "a1", - forkResult: { cancelled: false, selectedText: "first" }, + messages: [{ entryId: "a1", text: "answer", role: "assistant" }], + forkThrows: new Error("Entry not found"), }); invoke(fake); - await captured[0].onSelect("u1"); + await captured[0].onSelect("a1"); - expect(fake.session.fork).toHaveBeenCalledWith("u1"); - expect(fake.session.forkFromCurrent).not.toHaveBeenCalled(); - expect(fake.editor.setText).toHaveBeenCalledWith("first"); - expect(fake.showStatus).toHaveBeenCalledWith("Branched to new session"); + expect(fake.resetChatDisplay).not.toHaveBeenCalled(); + expect(fake._done).toHaveBeenCalledOnce(); + expect(fake.showStatus).toHaveBeenCalledWith("Fork failed: Entry not found"); }); }); diff --git a/packages/coding-agent/test/rpc-fork-current.test.ts b/packages/coding-agent/test/rpc-fork-current.test.ts deleted file mode 100644 index 02708ae7..00000000 --- a/packages/coding-agent/test/rpc-fork-current.test.ts +++ /dev/null @@ -1,31 +0,0 @@ -import { describe, expect, it, vi } from "vitest"; -import { RpcClient } from "../src/modes/rpc/rpc-client.js"; - -describe("RpcClient.forkCurrent", () => { - it("sends fork_current with no arguments and unwraps the cancelled flag", async () => { - const client = new RpcClient() as any; - const data = { cancelled: false }; - client.send = vi.fn().mockResolvedValue({ - type: "response", - command: "fork_current", - success: true, - data, - }); - - await expect(client.forkCurrent()).resolves.toEqual(data); - expect(client.send).toHaveBeenCalledWith({ type: "fork_current" }); - expect(client.send.mock.calls[0]).toHaveLength(1); - }); - - it("propagates a cancelled fork (e.g. empty session or extension cancel)", async () => { - const client = new RpcClient() as any; - client.send = vi.fn().mockResolvedValue({ - type: "response", - command: "fork_current", - success: true, - data: { cancelled: true }, - }); - - await expect(client.forkCurrent()).resolves.toEqual({ cancelled: true }); - }); -}); diff --git a/packages/coding-agent/test/rpc-fork.test.ts b/packages/coding-agent/test/rpc-fork.test.ts new file mode 100644 index 00000000..e34ec68d --- /dev/null +++ b/packages/coding-agent/test/rpc-fork.test.ts @@ -0,0 +1,42 @@ +import { describe, expect, it, vi } from "vitest"; +import { RpcClient } from "../src/modes/rpc/rpc-client.js"; + +describe("RpcClient fork surface", () => { + it("fork() sends the entry id and unwraps { text, cancelled }", async () => { + const client = new RpcClient() as any; + const data = { text: "re-ask me", cancelled: false }; + client.send = vi.fn().mockResolvedValue({ type: "response", command: "fork", success: true, data }); + + await expect(client.fork("e1")).resolves.toEqual(data); + expect(client.send).toHaveBeenCalledWith({ type: "fork", entryId: "e1" }); + }); + + it("fork() propagates an assistant fork (empty re-ask text)", async () => { + const client = new RpcClient() as any; + client.send = vi.fn().mockResolvedValue({ + type: "response", + command: "fork", + success: true, + data: { text: "", cancelled: false }, + }); + + await expect(client.fork("a3")).resolves.toEqual({ text: "", cancelled: false }); + }); + + it("getForkMessages() unwraps messages that carry a role", async () => { + const client = new RpcClient() as any; + const messages = [ + { entryId: "u1", text: "hi", role: "user" as const }, + { entryId: "a1", text: "hello", role: "assistant" as const }, + ]; + client.send = vi.fn().mockResolvedValue({ + type: "response", + command: "get_fork_messages", + success: true, + data: { messages }, + }); + + await expect(client.getForkMessages()).resolves.toEqual(messages); + expect(client.send).toHaveBeenCalledWith({ type: "get_fork_messages" }); + }); +}); diff --git a/packages/dashboard/src/client/api.ts b/packages/dashboard/src/client/api.ts index 56167f1e..b0318ac1 100644 --- a/packages/dashboard/src/client/api.ts +++ b/packages/dashboard/src/client/api.ts @@ -155,10 +155,11 @@ export const api = { commands: (key: string) => request<{ commands: CommandDto[] }>(`/api/runtimes/${key}/commands`), branch: (key: string) => request<{ branch: string | null }>(`/api/runtimes/${key}/branch`), forkMessages: (key: string) => - request<{ messages: Array<{ entryId: string; text: string }> }>(`/api/runtimes/${key}/fork-messages`), + request<{ messages: Array<{ entryId: string; text: string; role: "user" | "assistant" }> }>( + `/api/runtimes/${key}/fork-messages`, + ), fork: (key: string, entryId: string) => request<{ text: string; cancelled: boolean }>(`/api/runtimes/${key}/fork`, json({ entryId })), - forkCurrent: (key: string) => request<{ cancelled: boolean }>(`/api/runtimes/${key}/fork-current`, json({})), tree: (key: string) => request<{ roots: SessionTreeNodeDto[]; leafId: string | null }>(`/api/runtimes/${key}/tree`), navigateTree: (key: string, targetId: string) => request<{ cancelled: boolean; editorText?: string }>(`/api/runtimes/${key}/tree`, json({ targetId })), diff --git a/packages/dashboard/src/client/screens/session.tsx b/packages/dashboard/src/client/screens/session.tsx index a82d7187..838eda59 100644 --- a/packages/dashboard/src/client/screens/session.tsx +++ b/packages/dashboard/src/client/screens/session.tsx @@ -891,7 +891,9 @@ export function SessionScreen(props: { store: AppStore; sessionKey: string }): J const [fileAttachments, setFileAttachments] = createSignal([]); const [historyIndex, setHistoryIndex] = createSignal(); const [showForkModal, setShowForkModal] = createSignal(false); - const [forkMessages, setForkMessages] = createSignal>([]); + const [forkMessages, setForkMessages] = createSignal< + Array<{ entryId: string; text: string; role: "user" | "assistant" }> + >([]); const [forkError, setForkError] = createSignal(); const [showTreeModal, setShowTreeModal] = createSignal(false); const [treeRoots, setTreeRoots] = createSignal([]); @@ -1280,7 +1282,9 @@ export function SessionScreen(props: { store: AppStore; sessionKey: string }): J setForkError(cancelMessage); return; } - if (result.text !== undefined) setComposerText(result.text); + // Only user (re-ask) forks return text; assistant forks return "" and must + // not clobber whatever the user has already typed into the composer. + if (result.text) setComposerText(result.text); await props.store.hydrateSession(props.sessionKey); await props.store.refreshDiskSessions(); setShowForkModal(false); @@ -1292,12 +1296,6 @@ export function SessionScreen(props: { store: AppStore; sessionKey: string }): J const selectForkMessage = (entryId: string) => finishFork(() => api.fork(props.sessionKey, entryId), "Fork cancelled — no new branch was created."); - const forkFromCurrentState = () => - finishFork( - () => api.forkCurrent(props.sessionKey), - "Can't fork from current state — the session is empty or the fork was cancelled.", - ); - async function openStatsPopover() { setShowStatsPopover(true); setStatsPopoverError(undefined); @@ -2343,10 +2341,6 @@ export function SessionScreen(props: { store: AppStore; sessionKey: string }): J

{forkError()}

 - 0} fallback={

loading forkable messages…

}>
@@ -2356,7 +2350,7 @@ export function SessionScreen(props: { store: AppStore; sessionKey: string }): J class="fork-message" onClick={() => selectForkMessage(message.entryId)} > - {message.entryId} + {message.role === "assistant" ? "assistant" : "you"} {message.text} )} diff --git a/packages/dashboard/src/client/styles/app.css b/packages/dashboard/src/client/styles/app.css index bea4906b..e5a50aad 100644 --- a/packages/dashboard/src/client/styles/app.css +++ b/packages/dashboard/src/client/styles/app.css @@ -2757,31 +2757,12 @@ details.thinking .thinking-body { gap: var(--space-2); } -.fork-current-btn { - border: var(--hairline); - border-radius: var(--radius); - background: var(--bg); - color: var(--text); - font: inherit; - text-align: left; - padding: var(--space-2); - cursor: pointer; - display: flex; - gap: var(--space-2); - margin-bottom: var(--space-2); - border-color: var(--accent, var(--text)); -} - -.fork-message:hover, -.fork-current-btn:hover { +.fork-message:hover { border-color: var(--text); } -.fork-current-btn .fork-entry-id { - color: var(--accent, var(--text)); -} - -.fork-entry-id { +.fork-entry-id, +.fork-role { color: var(--muted); font-size: var(--fs-small); flex: 0 0 auto; diff --git a/packages/dashboard/src/server/server.ts b/packages/dashboard/src/server/server.ts index 40cd63a1..179b2ced 100644 --- a/packages/dashboard/src/server/server.ts +++ b/packages/dashboard/src/server/server.ts @@ -954,10 +954,6 @@ export function createDashboardServer(options: DashboardServerOptions): Dashboar withRuntime(req, res, (h) => h.client.fork(entryId)); }); - app.post("/api/runtimes/:key/fork-current", (req, res) => { - withRuntime(req, res, (h) => h.client.forkCurrent()); - }); - app.get("/api/runtimes/:key/tree", (req, res) => { withRuntime(req, res, (h) => h.client.getTree()); }); diff --git a/packages/dashboard/test/client/screens.test.tsx b/packages/dashboard/test/client/screens.test.tsx index 4882f58e..eecef680 100644 --- a/packages/dashboard/test/client/screens.test.tsx +++ b/packages/dashboard/test/client/screens.test.tsx @@ -81,7 +81,6 @@ vi.mock("../../src/client/api.js", () => ({ branch: vi.fn(async () => ({ branch: null })), forkMessages: vi.fn(async () => ({ messages: [] })), fork: vi.fn(async () => ({ text: "", cancelled: false })), - forkCurrent: vi.fn(async () => ({ cancelled: false })), dailyCost: vi.fn(async () => ({ cost: 0.42 })), settings: vi.fn(async () => ({ defaultProvider: "anthropic", defaultModel: "m1" })), devices: vi.fn(async () => ({ devices: [] })), @@ -5792,7 +5791,9 @@ describe("dashboard client regressions", () => { }); it("fork modal rewinds to a selected user message and prefills the composer", async () => { - vi.mocked(api.forkMessages).mockResolvedValue({ messages: [{ entryId: "u1", text: "original prompt" }] }); + vi.mocked(api.forkMessages).mockResolvedValue({ + messages: [{ entryId: "u1", text: "original prompt", role: "user" }], + }); vi.mocked(api.fork).mockResolvedValue({ text: "original prompt", cancelled: false }); const store = makeStore() as any; const hydrateSession = vi.fn(async () => {}); @@ -5820,71 +5821,42 @@ describe("dashboard client regressions", () => { expect((el.querySelector("textarea") as HTMLTextAreaElement).value).toBe("original prompt"); }); - it("fork modal can fork from current state without prefilling the composer", async () => { - vi.mocked(api.forkMessages).mockResolvedValue({ messages: [{ entryId: "u1", text: "original prompt" }] }); - vi.mocked(api.forkCurrent).mockResolvedValue({ cancelled: false }); + it("fork modal forks at an assistant message without prefilling the composer", async () => { + vi.mocked(api.forkMessages).mockResolvedValue({ + messages: [{ entryId: "a1", text: "the answer", role: "assistant" }], + }); + // Assistant forks return empty re-ask text (branch already includes the answer). + vi.mocked(api.fork).mockResolvedValue({ text: "", cancelled: false }); const store = makeStore() as any; const hydrateSession = vi.fn(async () => {}); const refreshDiskSessions = vi.fn(async () => {}); const fakeStore = { ...store, - sessions: { forkcur: createSessionViewState("forkcur") }, + sessions: { forkasst: createSessionViewState("forkasst") }, fleet: () => ({ runtimes: [], diskSessions: [] }), hydrateSession, refreshDiskSessions, }; - const el = mount(() => ); + const el = mount(() => ); (el.querySelector(".session-bar .right .switcher:last-child") as HTMLButtonElement).click(); await new Promise((resolve) => setTimeout(resolve, 0)); [...el.querySelectorAll("button")].find((button) => button.textContent?.includes("fork"))?.click(); await new Promise((resolve) => setTimeout(resolve, 10)); - // Mocks are not auto-cleared between tests; ignore any api.fork calls from prior tests. - vi.mocked(api.fork).mockClear(); - (el.querySelector(".fork-current-btn") as HTMLButtonElement).click(); + // The row is labeled by role. + expect(el.querySelector(".fork-role")?.textContent).toBe("assistant"); + (el.querySelector(".fork-message") as HTMLButtonElement).click(); await new Promise((resolve) => setTimeout(resolve, 10)); - expect(api.forkCurrent).toHaveBeenCalledWith("forkcur"); - expect(api.fork).not.toHaveBeenCalled(); - expect(hydrateSession).toHaveBeenCalledWith("forkcur"); + expect(api.fork).toHaveBeenCalledWith("forkasst", "a1"); + expect(hydrateSession).toHaveBeenCalledWith("forkasst"); expect(refreshDiskSessions).toHaveBeenCalledOnce(); // No composer pre-fill: the branch already includes the last response. expect((el.querySelector("textarea") as HTMLTextAreaElement).value).toBe(""); }); - it("fork modal informs the user and stays open when fork-from-current is cancelled", async () => { - vi.mocked(api.forkMessages).mockResolvedValue({ messages: [{ entryId: "u1", text: "original prompt" }] }); - // Empty session / extension veto → backend returns cancelled with no branch created. - vi.mocked(api.forkCurrent).mockResolvedValue({ cancelled: true }); - const store = makeStore() as any; - const hydrateSession = vi.fn(async () => {}); - const refreshDiskSessions = vi.fn(async () => {}); - const fakeStore = { - ...store, - sessions: { forkcancel: createSessionViewState("forkcancel") }, - fleet: () => ({ runtimes: [], diskSessions: [] }), - hydrateSession, - refreshDiskSessions, - }; - const el = mount(() => ); - (el.querySelector(".session-bar .right .switcher:last-child") as HTMLButtonElement).click(); - await new Promise((resolve) => setTimeout(resolve, 0)); - [...el.querySelectorAll("button")].find((button) => button.textContent?.includes("fork"))?.click(); - await new Promise((resolve) => setTimeout(resolve, 10)); - // Ignore the mount-time hydration; assert only what the fork handler does. - hydrateSession.mockClear(); - refreshDiskSessions.mockClear(); - (el.querySelector(".fork-current-btn") as HTMLButtonElement).click(); - await new Promise((resolve) => setTimeout(resolve, 10)); - expect(api.forkCurrent).toHaveBeenCalledWith("forkcancel"); - // The user is informed: the modal stays open with a message, and no session - // churn happens as if a branch had been created. - expect(el.querySelector(".fork-current-btn")).not.toBeNull(); - expect(el.querySelector(".pair-error")?.textContent ?? "").toMatch(/can't fork|cancelled/i); - expect(hydrateSession).not.toHaveBeenCalled(); - expect(refreshDiskSessions).not.toHaveBeenCalled(); - }); - it("fork modal informs the user and stays open when a message fork is cancelled", async () => { - vi.mocked(api.forkMessages).mockResolvedValue({ messages: [{ entryId: "u1", text: "original prompt" }] }); + vi.mocked(api.forkMessages).mockResolvedValue({ + messages: [{ entryId: "u1", text: "original prompt", role: "user" }], + }); // Extension veto → api.fork returns cancelled with no branch created. vi.mocked(api.fork).mockResolvedValue({ text: "", cancelled: true }); const store = makeStore() as any; diff --git a/packages/dashboard/test/runtime-pool.test.ts b/packages/dashboard/test/runtime-pool.test.ts index 1adfe4d2..0d557f46 100644 --- a/packages/dashboard/test/runtime-pool.test.ts +++ b/packages/dashboard/test/runtime-pool.test.ts @@ -187,7 +187,6 @@ export function makeFakeClient() { navigateTree: vi.fn(async () => ({ cancelled: false })), getForkMessages: vi.fn(async () => []), fork: vi.fn(async () => ({ text: "", cancelled: false })), - forkCurrent: vi.fn(async () => ({ cancelled: false })), listSessions: vi.fn(async () => []), switchSession: vi.fn(async () => ({ cancelled: false })), prompt: vi.fn(async () => {}), diff --git a/packages/dashboard/test/server.test.ts b/packages/dashboard/test/server.test.ts index fe16a3c9..2434f7f0 100644 --- a/packages/dashboard/test/server.test.ts +++ b/packages/dashboard/test/server.test.ts @@ -1083,13 +1083,6 @@ describe("dashboard server — fleet and runtimes", () => { body: JSON.stringify({ targetId: "entry-1" }), }).then((r) => r.json()), ).resolves.toEqual({ cancelled: false }); - await expect( - fetch(`${base}/api/runtimes/${key}/fork-current`, { - method: "POST", - headers: { "content-type": "application/json" }, - body: JSON.stringify({}), - }).then((r) => r.json()), - ).resolves.toEqual({ cancelled: false }); await expect(fetch(`${base}/api/runtimes/${key}/sessions`).then((r) => r.json())).resolves.toEqual({ sessions: [], }); @@ -1114,7 +1107,6 @@ describe("dashboard server — fleet and runtimes", () => { expect(clients[0].importJsonl).toHaveBeenCalledWith("/tmp/session.jsonl"); expect(clients[0].getTree).toHaveBeenCalled(); expect(clients[0].navigateTree).toHaveBeenCalledWith("entry-1"); - expect(clients[0].forkCurrent).toHaveBeenCalled(); expect(clients[0].listSessions).toHaveBeenCalled(); expect(clients[0].switchSession).toHaveBeenCalledWith("/tmp/session.jsonl"); expect(clients[1].getDailyCost).toHaveBeenCalled(); From 87c8e7514aef19eeb1c98d56a577856282b11bd9 Mon Sep 17 00:00:00 2001 From: Test Date: Tue, 11 Aug 2026 21:51:58 +0200 Subject: [PATCH 6/7] Reject unsafe assistant fork targets; add coverage + docs (#439) --- packages/coding-agent/README.md | 2 +- packages/coding-agent/docs/rpc.md | 22 +++- packages/coding-agent/docs/tree.md | 2 +- .../coding-agent/src/core/agent-session.ts | 39 +++++++ .../test/agent-session-fork.test.ts | 109 ++++++++++++++++++ .../test/user-message-selector.test.ts | 75 ++++++++++++ .../dashboard/test/client/screens.test.tsx | 10 +- 7 files changed, 249 insertions(+), 10 deletions(-) create mode 100644 packages/coding-agent/test/user-message-selector.test.ts diff --git a/packages/coding-agent/README.md b/packages/coding-agent/README.md index e37ecae1..7b2952fe 100644 --- a/packages/coding-agent/README.md +++ b/packages/coding-agent/README.md @@ -243,7 +243,7 @@ dreb --fork # Fork specific session file or ID into a new session - Filter modes (Ctrl+O): default → no-tools → user-only → labeled-only → all - Press `L` (Shift+L) to label entries as bookmarks -**`/fork`** - Create a new session file from the current branch. Opens a selector, copies history up to the selected point, and places that message in the editor for modification. +**`/fork`** - Create a new session file by branching from any point in the current conversation. Opens a selector listing every user and assistant message: picking an **assistant** message keeps that response and everything before it (continue from that answer) with an empty editor; picking a **user** message rewinds to before it (dropping it and everything after) and places its text in the editor for re-asking. **`--fork `** - Fork an existing session file or partial session UUID directly from the CLI. This copies the full source session into a new session file in the current project. diff --git a/packages/coding-agent/docs/rpc.md b/packages/coding-agent/docs/rpc.md index 3cc8aeb1..de980cb6 100644 --- a/packages/coding-agent/docs/rpc.md +++ b/packages/coding-agent/docs/rpc.md @@ -906,13 +906,13 @@ The path uses the same unrestricted, cross-project addressing as [`switch_sessio #### fork -Create a new fork from a previous user message. Can be cancelled by a `session_before_fork` extension event handler. Returns the text of the message being forked from. +Create a new fork from any user or assistant message in the transcript. Can be cancelled by a `session_before_fork` extension event handler. For a **user** message the response `text` is that message's text (offered as editor pre-fill for re-asking) and the branch rewinds to before it; for an **assistant** message the branch *includes* that response (continue from that answer) and `text` is empty. Assistant turns that were interrupted (`error`/`aborted`) or are still waiting on tool results are not valid fork points and are rejected. ```json {"type": "fork", "entryId": "abc123"} ``` -Response: +Response (forking at a user message — text is offered as editor pre-fill): ```json { "type": "response", @@ -922,19 +922,29 @@ Response: } ``` +Response (forking at an assistant message — no pre-fill): +```json +{ + "type": "response", + "command": "fork", + "success": true, + "data": {"text": "", "cancelled": false} +} +``` + If an extension cancelled the fork: ```json { "type": "response", "command": "fork", "success": true, - "data": {"text": "The original prompt text...", "cancelled": true} + "data": {"text": "", "cancelled": true} } ``` #### get_fork_messages -Get user messages available for forking. +Get the messages available for forking (both user and assistant). Each entry carries its `role` so callers can label it and choose the right fork semantics. ```json {"type": "get_fork_messages"} @@ -948,8 +958,8 @@ Response: "success": true, "data": { "messages": [ - {"entryId": "abc123", "text": "First prompt..."}, - {"entryId": "def456", "text": "Second prompt..."} + {"entryId": "abc123", "text": "First prompt...", "role": "user"}, + {"entryId": "def456", "text": "The answer...", "role": "assistant"} ] } } diff --git a/packages/coding-agent/docs/tree.md b/packages/coding-agent/docs/tree.md index 81d4a89e..feeb9aa2 100644 --- a/packages/coding-agent/docs/tree.md +++ b/packages/coding-agent/docs/tree.md @@ -10,7 +10,7 @@ Sessions are stored as trees where each entry has an `id` and `parentId`. The "l | Feature | `/fork` | `/tree` | |---------|---------|---------| -| View | Flat list of user messages | Full tree structure | +| View | Flat list of user and assistant messages | Full tree structure | | Action | Extracts path to **new session file** | Changes leaf in **same session** | | Summary | Never | Optional (user prompted) | | Events | `session_before_fork` / `session_fork` | `session_before_tree` / `session_tree` | diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index cdf17be5..a134fc83 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -3659,6 +3659,15 @@ export class AgentSession { if (selectedEntry.message.role === "assistant") { // Continue-from-answer: branch from the assistant entry itself so it (and // everything before it) is retained. No editor pre-fill. + // + // Reject turns that can't be safely branched from (interrupted, or waiting + // on tool results) — branching there would silently produce a branch that + // doesn't match the selected turn. See _isForkableAssistant. + if (!this._isForkableAssistant(selectedEntry.message)) { + throw new Error( + "Cannot fork at this assistant turn: it was interrupted or is still waiting on tool results", + ); + } const { cancelled } = await this._performFork(entryId, () => { this.sessionManager.createBranchedSession(entryId); }); @@ -3975,6 +3984,9 @@ export class AgentSession { * fork semantics (assistant = continue-from-answer, user = rewind + re-ask). * Assistant turns with no renderable text (pure tool-call turns) still appear * as fork points, with a generic label. + * + * Assistant turns that cannot be safely branched from (interrupted turns, or + * turns still waiting on tool results) are excluded — see _isForkableAssistant. */ getForkableMessages(): Array<{ entryId: string; text: string; role: "user" | "assistant" }> { const entries = this.sessionManager.getEntries(); @@ -3990,6 +4002,8 @@ export class AgentSession { // Preserve existing behavior: skip empty user messages. if (text) result.push({ entryId: entry.id, text, role }); } else { + // Only offer assistant turns that can be safely branched from. + if (!this._isForkableAssistant(entry.message as AssistantMessage)) continue; result.push({ entryId: entry.id, text: text || "(assistant response)", role }); } } @@ -3997,6 +4011,31 @@ export class AgentSession { return result; } + /** + * Whether an assistant turn can be safely used as a fork point. + * + * Forking anchors on the entry's ancestors only (SessionManager.getBranch + * walks parentId upward), and errored/aborted turns are dropped by + * transformMessages() before every request. Two kinds of assistant turn + * therefore produce a branch that silently does NOT match what was selected: + * + * - stopReason "error"/"aborted": transformMessages() skips the turn, so the + * reply vanishes from context on the next request (defeating "continue from + * this answer", and risking back-to-back user messages on strict providers). + * - turns containing tool calls: their tool results are *descendant* entries a + * branch cannot include, so transformMessages() substitutes a fabricated + * "No result provided" (isError) result — telling the model a successful + * tool call failed. + * + * A completed answer (the intended "continue from here" target) has a terminal + * stopReason and no unresolved tool calls, so it passes. + */ + private _isForkableAssistant(message: AssistantMessage): boolean { + if (message.stopReason === "error" || message.stopReason === "aborted") return false; + if (Array.isArray(message.content) && message.content.some((c) => c.type === "toolCall")) return false; + return true; + } + private _extractMessageText(content: string | Array<{ type: string; text?: string }>): string { if (typeof content === "string") return content; if (Array.isArray(content)) { diff --git a/packages/coding-agent/test/agent-session-fork.test.ts b/packages/coding-agent/test/agent-session-fork.test.ts index 351d391a..1c0f8b0c 100644 --- a/packages/coding-agent/test/agent-session-fork.test.ts +++ b/packages/coding-agent/test/agent-session-fork.test.ts @@ -186,3 +186,112 @@ describe("AgentSession.fork — any message", () => { expect(forkEvents).toHaveLength(0); }); }); + +// Assistant-turn fixtures for the forkable-listing rules. assistantMsg() builds a +// completed (stopReason "stop") text turn; these override that to exercise the +// turns that must be excluded as fork points. +function erroredAssistant(text: string) { + return { ...assistantMsg(text), stopReason: "error" as const, errorMessage: "boom" }; +} +function abortedAssistant(text: string) { + return { ...assistantMsg(text), stopReason: "aborted" as const }; +} +function toolCallAssistant() { + return { + ...assistantMsg(""), + content: [{ type: "toolCall" as const, id: "tc1", name: "bash", arguments: { cmd: "ls" } }], + stopReason: "toolUse" as const, + }; +} +// A completed turn that renders no text (e.g. thinking-only): still a valid fork +// point, exercised for the "(assistant response)" fallback label. +function emptyTextAssistant() { + return { ...assistantMsg(""), content: [] }; +} + +describe("AgentSession.getForkableMessages", () => { + let harness: Harness; + + afterEach(() => { + harness?.cleanup(); + }); + + it("lists user and assistant messages with their roles and text", async () => { + harness = await createHarnessWithExtensions(); + const { session, sessionManager } = harness; + + sessionManager.appendMessage(userMsg("q1")); + sessionManager.appendMessage(assistantMsg("a1")); + sessionManager.appendMessage(userMsg("q2")); + + const list = session.getForkableMessages(); + expect(list).toEqual([ + { entryId: expect.any(String), text: "q1", role: "user" }, + { entryId: expect.any(String), text: "a1", role: "assistant" }, + { entryId: expect.any(String), text: "q2", role: "user" }, + ]); + }); + + it("skips empty user messages but keeps a renderless assistant turn with a fallback label", async () => { + harness = await createHarnessWithExtensions(); + const { session, sessionManager } = harness; + + sessionManager.appendMessage(userMsg("")); // empty user → dropped + sessionManager.appendMessage(emptyTextAssistant()); // no text but forkable → labeled + + const list = session.getForkableMessages(); + expect(list).toHaveLength(1); + expect(list[0].role).toBe("assistant"); + expect(list[0].text).toBe("(assistant response)"); + }); + + it("excludes errored and aborted assistant turns (they vanish from later requests)", async () => { + harness = await createHarnessWithExtensions(); + const { session, sessionManager } = harness; + + sessionManager.appendMessage(userMsg("q1")); + sessionManager.appendMessage(erroredAssistant("partial-error")); + sessionManager.appendMessage(abortedAssistant("partial-abort")); + const goodId = sessionManager.appendMessage(assistantMsg("real answer")); + + const list = session.getForkableMessages(); + const assistants = list.filter((m) => m.role === "assistant"); + expect(assistants).toHaveLength(1); + expect(assistants[0].entryId).toBe(goodId); + expect(assistants[0].text).toBe("real answer"); + }); + + it("excludes assistant turns with unresolved tool calls (their tool results are descendants)", async () => { + harness = await createHarnessWithExtensions(); + const { session, sessionManager } = harness; + + sessionManager.appendMessage(userMsg("list files")); + sessionManager.appendMessage(toolCallAssistant()); + const finalId = sessionManager.appendMessage(assistantMsg("here are the files")); + + const list = session.getForkableMessages(); + const assistants = list.filter((m) => m.role === "assistant"); + expect(assistants).toHaveLength(1); + expect(assistants[0].entryId).toBe(finalId); + }); + + it("fork() rejects a direct call on an errored assistant entry", async () => { + harness = await createHarnessWithExtensions(); + const { session, sessionManager } = harness; + + sessionManager.appendMessage(userMsg("q1")); + const errId = sessionManager.appendMessage(erroredAssistant("partial")); + + await expect(session.fork(errId)).rejects.toThrow(/interrupted|waiting on tool results/i); + }); + + it("fork() rejects a direct call on a tool-call assistant entry", async () => { + harness = await createHarnessWithExtensions(); + const { session, sessionManager } = harness; + + sessionManager.appendMessage(userMsg("list files")); + const tcId = sessionManager.appendMessage(toolCallAssistant()); + + await expect(session.fork(tcId)).rejects.toThrow(/interrupted|waiting on tool results/i); + }); +}); diff --git a/packages/coding-agent/test/user-message-selector.test.ts b/packages/coding-agent/test/user-message-selector.test.ts new file mode 100644 index 00000000..79b7b356 --- /dev/null +++ b/packages/coding-agent/test/user-message-selector.test.ts @@ -0,0 +1,75 @@ +import { setKeybindings } from "@dreb/tui"; +import { beforeAll, beforeEach, describe, expect, test, vi } from "vitest"; +import { KeybindingsManager } from "../src/core/keybindings.js"; +import { UserMessageSelectorComponent } from "../src/modes/interactive/components/user-message-selector.js"; +import { initTheme } from "../src/modes/interactive/theme/theme.js"; + +beforeAll(() => { + initTheme("dark"); +}); + +beforeEach(() => { + setKeybindings(new KeybindingsManager()); +}); + +const ENTER = "\r"; +const ESC = "\x1b"; + +type Item = { id: string; text: string; role: "user" | "assistant" }; + +function makeItems(): Item[] { + return [ + { id: "u1", text: "first question", role: "user" }, + { id: "a1", text: "the assistant answer", role: "assistant" }, + { id: "u2", text: "second question", role: "user" }, + ]; +} + +describe("UserMessageSelectorComponent render", () => { + test("renders role badges and role-specific fork hints for both roles", () => { + const component = new UserMessageSelectorComponent(makeItems(), vi.fn(), vi.fn()); + const lines = component.getMessageList().render(80); + const blob = lines.join("\n"); + + // Role badges distinguish the two message kinds. + expect(blob).toContain("[Assistant]"); + expect(blob).toContain("[You]"); + // Role-specific hints describe the opposite-consequence actions. + expect(blob).toContain("continue from here"); // assistant + expect(blob).toContain("rewind & re-ask"); // user + // Message previews appear. + expect(blob).toContain("the assistant answer"); + expect(blob).toContain("first question"); + }); + + test("bottom-anchors selection on the most recent message", () => { + const onSelect = vi.fn(); + const component = new UserMessageSelectorComponent(makeItems(), onSelect, vi.fn()); + component.getMessageList().handleInput(ENTER); + // Last item (u2) is selected by default. + expect(onSelect).toHaveBeenCalledWith("u2"); + }); + + test("Escape cancels", () => { + const onCancel = vi.fn(); + const component = new UserMessageSelectorComponent(makeItems(), vi.fn(), onCancel); + component.getMessageList().handleInput(ESC); + expect(onCancel).toHaveBeenCalledOnce(); + }); + + test("does not throw or produce negative-width truncation on a very narrow terminal", () => { + // The [Assistant] badge eats into the width budget; at width 10 the message + // budget goes negative and must be clamped (Math.max(0, …)). + const component = new UserMessageSelectorComponent(makeItems(), vi.fn(), vi.fn()); + expect(() => component.getMessageList().render(10)).not.toThrow(); + const lines = component.getMessageList().render(10); + expect(Array.isArray(lines)).toBe(true); + expect(lines.length).toBeGreaterThan(0); + }); + + test("renders an empty-state message when there are no fork points", () => { + const component = new UserMessageSelectorComponent([], vi.fn(), vi.fn()); + const lines = component.getMessageList().render(80); + expect(lines.join("\n")).toContain("No messages found"); + }); +}); diff --git a/packages/dashboard/test/client/screens.test.tsx b/packages/dashboard/test/client/screens.test.tsx index eecef680..d62157f3 100644 --- a/packages/dashboard/test/client/screens.test.tsx +++ b/packages/dashboard/test/client/screens.test.tsx @@ -5838,6 +5838,11 @@ describe("dashboard client regressions", () => { refreshDiskSessions, }; const el = mount(() => ); + // Pre-type a draft into the composer. The no-clobber guard in finishFork must + // preserve it: an assistant fork returns text "" and must NOT wipe the draft. + const composer = el.querySelector("textarea") as HTMLTextAreaElement; + composer.value = "draft in progress"; + composer.dispatchEvent(new InputEvent("input", { bubbles: true })); (el.querySelector(".session-bar .right .switcher:last-child") as HTMLButtonElement).click(); await new Promise((resolve) => setTimeout(resolve, 0)); [...el.querySelectorAll("button")].find((button) => button.textContent?.includes("fork"))?.click(); @@ -5849,8 +5854,9 @@ describe("dashboard client regressions", () => { expect(api.fork).toHaveBeenCalledWith("forkasst", "a1"); expect(hydrateSession).toHaveBeenCalledWith("forkasst"); expect(refreshDiskSessions).toHaveBeenCalledOnce(); - // No composer pre-fill: the branch already includes the last response. - expect((el.querySelector("textarea") as HTMLTextAreaElement).value).toBe(""); + // No composer pre-fill AND no clobber: the user's in-progress draft survives + // (assistant forks return "" and must not overwrite the composer). + expect((el.querySelector("textarea") as HTMLTextAreaElement).value).toBe("draft in progress"); }); it("fork modal informs the user and stays open when a message fork is cancelled", async () => { From 92bd2ba647076f5655f7f5e7f2f4ae932a635b09 Mon Sep 17 00:00:00 2001 From: Test Date: Wed, 12 Aug 2026 06:57:59 +0200 Subject: [PATCH 7/7] Harden fork tests; fix stale fork docs (#439) --- packages/coding-agent/README.md | 2 +- packages/coding-agent/docs/rpc.md | 4 +- .../coding-agent/src/core/agent-session.ts | 7 ++-- .../test/agent-session-fork.test.ts | 41 +++++++++++++++++++ .../test/user-message-selector.test.ts | 28 ++++++++----- 5 files changed, 65 insertions(+), 17 deletions(-) diff --git a/packages/coding-agent/README.md b/packages/coding-agent/README.md index 7b2952fe..bd977746 100644 --- a/packages/coding-agent/README.md +++ b/packages/coding-agent/README.md @@ -175,7 +175,7 @@ Type `/` in the editor to trigger commands. [Extensions](#extensions) can regist | `/name ` | Set session display name | | `/session` | Show session info (path, tokens, cost) | | `/tree` | Jump to any point in the session and continue from there | -| `/fork` | Create a new session from the current branch | +| `/fork` | Branch a new session from any user or assistant message (assistant = continue from that answer, user = rewind and re-ask) | | `/compact [prompt]` | Manually compact context, optional custom instructions | | `/copy` | Open multi-select message picker to copy any messages to clipboard. Assistant reasoning is excluded by default and offered as a separate, selectable `Thinking` row. | | `/dream` | Consolidate and prune memories — backs up, merges duplicates, scans sessions for patterns | diff --git a/packages/coding-agent/docs/rpc.md b/packages/coding-agent/docs/rpc.md index de980cb6..bd195834 100644 --- a/packages/coding-agent/docs/rpc.md +++ b/packages/coding-agent/docs/rpc.md @@ -932,13 +932,13 @@ Response (forking at an assistant message — no pre-fill): } ``` -If an extension cancelled the fork: +If an extension cancelled the fork, `text` still mirrors what the corresponding successful fork would have returned — the user message's text for a user-message fork, or `""` for an assistant-message fork: ```json { "type": "response", "command": "fork", "success": true, - "data": {"text": "", "cancelled": true} + "data": {"text": "The original prompt text...", "cancelled": true} } ``` diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index a134fc83..54f3f8a3 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -3982,11 +3982,12 @@ export class AgentSession { * * Each entry carries its role so callers can label it and choose the right * fork semantics (assistant = continue-from-answer, user = rewind + re-ask). - * Assistant turns with no renderable text (pure tool-call turns) still appear - * as fork points, with a generic label. + * A forkable assistant turn with no renderable text (e.g. a thinking-only + * turn) still appears as a fork point, with a generic label. * * Assistant turns that cannot be safely branched from (interrupted turns, or - * turns still waiting on tool results) are excluded — see _isForkableAssistant. + * turns containing a tool call whose result lives in a descendant entry) are + * excluded — see _isForkableAssistant. */ getForkableMessages(): Array<{ entryId: string; text: string; role: "user" | "assistant" }> { const entries = this.sessionManager.getEntries(); diff --git a/packages/coding-agent/test/agent-session-fork.test.ts b/packages/coding-agent/test/agent-session-fork.test.ts index 1c0f8b0c..4e285112 100644 --- a/packages/coding-agent/test/agent-session-fork.test.ts +++ b/packages/coding-agent/test/agent-session-fork.test.ts @@ -203,6 +203,19 @@ function toolCallAssistant() { stopReason: "toolUse" as const, }; } +// The dominant real tool-call shape: narration text AND a tool call in the same +// content array. Distinguishes `.some()` (correct — any toolCall block excludes) +// from `.every()`, which the single-element toolCallAssistant() fixture cannot. +function mixedToolCallAssistant() { + return { + ...assistantMsg(""), + content: [ + { type: "text" as const, text: "checking the files" }, + { type: "toolCall" as const, id: "tc2", name: "bash", arguments: { cmd: "ls" } }, + ], + stopReason: "toolUse" as const, + }; +} // A completed turn that renders no text (e.g. thinking-only): still a valid fork // point, exercised for the "(assistant response)" fallback label. function emptyTextAssistant() { @@ -275,6 +288,24 @@ describe("AgentSession.getForkableMessages", () => { expect(assistants[0].entryId).toBe(finalId); }); + it("excludes assistant turns that mix narration text with a tool call (pins .some, not .every)", async () => { + harness = await createHarnessWithExtensions(); + const { session, sessionManager } = harness; + + sessionManager.appendMessage(userMsg("list files")); + // A turn with BOTH text and a toolCall: `.every(isToolCall)` is false here, + // so only `.some(isToolCall)` correctly excludes it. + sessionManager.appendMessage(mixedToolCallAssistant()); + const finalId = sessionManager.appendMessage(assistantMsg("here are the files")); + + const list = session.getForkableMessages(); + const assistants = list.filter((m) => m.role === "assistant"); + expect(assistants).toHaveLength(1); + expect(assistants[0].entryId).toBe(finalId); + // The mixed turn's narration text must not leak in as a fork point. + expect(list.some((m) => m.text === "checking the files")).toBe(false); + }); + it("fork() rejects a direct call on an errored assistant entry", async () => { harness = await createHarnessWithExtensions(); const { session, sessionManager } = harness; @@ -294,4 +325,14 @@ describe("AgentSession.getForkableMessages", () => { await expect(session.fork(tcId)).rejects.toThrow(/interrupted|waiting on tool results/i); }); + + it("fork() rejects a direct call on a mixed text+tool-call assistant entry", async () => { + harness = await createHarnessWithExtensions(); + const { session, sessionManager } = harness; + + sessionManager.appendMessage(userMsg("list files")); + const mixedId = sessionManager.appendMessage(mixedToolCallAssistant()); + + await expect(session.fork(mixedId)).rejects.toThrow(/interrupted|waiting on tool results/i); + }); }); diff --git a/packages/coding-agent/test/user-message-selector.test.ts b/packages/coding-agent/test/user-message-selector.test.ts index 79b7b356..b1b21605 100644 --- a/packages/coding-agent/test/user-message-selector.test.ts +++ b/packages/coding-agent/test/user-message-selector.test.ts @@ -26,20 +26,26 @@ function makeItems(): Item[] { } describe("UserMessageSelectorComponent render", () => { - test("renders role badges and role-specific fork hints for both roles", () => { + test("renders role badges and role-specific fork hints tied to the correct message line", () => { const component = new UserMessageSelectorComponent(makeItems(), vi.fn(), vi.fn()); const lines = component.getMessageList().render(80); - const blob = lines.join("\n"); - // Role badges distinguish the two message kinds. - expect(blob).toContain("[Assistant]"); - expect(blob).toContain("[You]"); - // Role-specific hints describe the opposite-consequence actions. - expect(blob).toContain("continue from here"); // assistant - expect(blob).toContain("rewind & re-ask"); // user - // Message previews appear. - expect(blob).toContain("the assistant answer"); - expect(blob).toContain("first question"); + // Each message renders as: line N = cursor + role badge + preview, + // line N+1 = the role-specific hint. Assert the badge and hint are attached + // to the CORRECT message line, so a swapped badge/hint mapping fails. + const assistantIdx = lines.findIndex((l) => l.includes("the assistant answer")); + expect(assistantIdx).toBeGreaterThanOrEqual(0); + expect(lines[assistantIdx]).toContain("[Assistant]"); + expect(lines[assistantIdx]).not.toContain("[You]"); + expect(lines[assistantIdx + 1]).toContain("continue from here"); + expect(lines[assistantIdx + 1]).not.toContain("rewind & re-ask"); + + const userIdx = lines.findIndex((l) => l.includes("first question")); + expect(userIdx).toBeGreaterThanOrEqual(0); + expect(lines[userIdx]).toContain("[You]"); + expect(lines[userIdx]).not.toContain("[Assistant]"); + expect(lines[userIdx + 1]).toContain("rewind & re-ask"); + expect(lines[userIdx + 1]).not.toContain("continue from here"); }); test("bottom-anchors selection on the most recent message", () => {