diff --git a/packages/coding-agent/README.md b/packages/coding-agent/README.md index e37ecae1..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 | @@ -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..bd195834 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,7 +922,17 @@ Response: } ``` -If an extension cancelled the fork: +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, `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", @@ -934,7 +944,7 @@ If an extension cancelled the fork: #### 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 170ae765..54f3f8a3 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -3627,23 +3627,83 @@ 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 previousSessionFile = this.sessionFile; 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. + // + // 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); + }); + return { selectedText: "", cancelled }; + } + + const selectedText = this._extractMessageText(selectedEntry.message.content); + + // 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 }; + } + + /** + * 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; @@ -3655,7 +3715,7 @@ export class AgentSession { })) as SessionBeforeForkResult | undefined; if (result?.cancel) { - return { selectedText, cancelled: true }; + return { cancelled: true }; } skipConversationRestore = result?.skipConversationRestore ?? false; } @@ -3663,11 +3723,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 +3743,7 @@ export class AgentSession { this.agent.replaceMessages(sessionContext.messages); } - return { selectedText, cancelled: false }; + return { cancelled: false }; } // ========================================================================= @@ -3852,7 +3908,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; @@ -3922,26 +3978,66 @@ 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). + * 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 containing a tool call whose result lives in a descendant entry) are + * excluded — see _isForkableAssistant. */ - 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 { + // 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 }); } } return result; } - private _extractUserMessageText(content: string | Array<{ type: string; text?: string }>): string { + /** + * 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)) { 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 e5c9c547..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,12 +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 } /** - * 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[] = []; @@ -33,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; } @@ -48,23 +52,26 @@ class UserMessageList implements Component { 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) - 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) + // Second line: position + what forking here does const position = i + 1; - const metadata = ` Message ${position} of ${this.messages.length}`; - const metadataLine = theme.fg("muted", metadata); - lines.push(metadataLine); + 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 } @@ -104,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; @@ -115,7 +123,16 @@ 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 any message: an assistant reply continues from that answer, a question rewinds to re-ask it", + ), + 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..02658f7c 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -4296,29 +4296,44 @@ export class InteractiveMode { } private showUserMessageSelector(): void { - const userMessages = this.session.getUserMessagesForForking(); + const messages = this.session.getForkableMessages(); - if (userMessages.length === 0) { + if (messages.length === 0) { this.showStatus("No messages to fork from"); return; } + // 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( - userMessages.map((m) => ({ id: m.entryId, text: m.text })), + items, async (entryId) => { - const result = await this.session.fork(entryId); - if (result.cancelled) { - // Extension cancelled the fork + 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; + } + + 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.ui.requestRender(); - 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(); - this.editor.setText(result.selectedText); - done(); - this.showStatus("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 629b1448..f02ad292 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-client.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-client.ts @@ -629,9 +629,10 @@ export class RpcClient { /** * 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 f2025f1e..9a8e9d9c 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -2127,7 +2127,7 @@ export async function runRpcMode(session: AgentSession, modelFallbackMessage?: s } 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 b81a20bf..b8952a5e 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-types.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-types.ts @@ -378,7 +378,7 @@ export type RpcResponse = 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.test.ts b/packages/coding-agent/test/agent-session-fork.test.ts new file mode 100644 index 00000000..4e285112 --- /dev/null +++ b/packages/coding-agent/test/agent-session-fork.test.ts @@ -0,0 +1,338 @@ +/** + * 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); + }); +}); + +// 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, + }; +} +// 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() { + 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("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; + + 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); + }); + + 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/interactive-mode-fork.test.ts b/packages/coding-agent/test/interactive-mode-fork.test.ts new file mode 100644 index 00000000..67682fed --- /dev/null +++ b/packages/coding-agent/test/interactive-mode-fork.test.ts @@ -0,0 +1,163 @@ +/** + * Unit coverage for the interactive `/fork` selector wiring + * (InteractiveMode.showUserMessageSelector). + * + * 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 (role-labeled, all messages). +const captured: Array<{ + items: Array<{ id: string; text: string; role: "user" | "assistant" }>; + 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; role: "user" | "assistant" }>, + public onSelect: (entryId: string) => void | Promise, + public onCancel: () => void, + ) { + captured.push({ items, onSelect, onCancel }); + } + getMessageList() { + return {}; + } + } + return { ...actual, UserMessageSelectorComponent: MockUserMessageSelectorComponent }; +}); + +import { InteractiveMode } from "../src/modes/interactive/interactive-mode.js"; + +interface FakeOverrides { + 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: { + getForkableMessages: vi.fn(() => overrides.messages ?? []), + fork: vi.fn(async () => { + if (overrides.forkThrows) throw overrides.forkThrows; + return overrides.forkResult ?? { cancelled: false, selectedText: "prefill" }; + }), + }, + 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 at any message", () => { + beforeEach(() => { + captured.length = 0; + }); + + it("lists all user and assistant messages with their roles (no action row)", () => { + const fake = makeFakeThis({ + messages: [ + { entryId: "u1", text: "first", role: "user" }, + { entryId: "a1", text: "answer", role: "assistant" }, + { entryId: "u2", text: "second", role: "user" }, + ], + }); + invoke(fake); + + expect(fake.showSelector).toHaveBeenCalledOnce(); + expect(captured).toHaveLength(1); + const { items } = captured[0]; + expect(items.map((i) => i.id)).toEqual(["u1", "a1", "u2"]); + expect(items.map((i) => i.role)).toEqual(["user", "assistant", "user"]); + }); + + 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("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); + + 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("forking at a user message pre-fills the editor with the selected text", async () => { + const fake = makeFakeThis({ + messages: [{ entryId: "u1", text: "first", role: "user" }], + forkResult: { cancelled: false, selectedText: "first" }, + }); + invoke(fake); + + await captured[0].onSelect("u1"); + + 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 the fork is cancelled", async () => { + const fake = makeFakeThis({ + messages: [{ entryId: "u1", text: "first", role: "user" }], + forkResult: { cancelled: true, selectedText: "" }, + }); + invoke(fake); + + await captured[0].onSelect("u1"); + + 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("surfaces an error instead of crashing when fork() throws", async () => { + const fake = makeFakeThis({ + messages: [{ entryId: "a1", text: "answer", role: "assistant" }], + forkThrows: new Error("Entry not found"), + }); + invoke(fake); + + await captured[0].onSelect("a1"); + + 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.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/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/coding-agent/test/user-message-selector.test.ts b/packages/coding-agent/test/user-message-selector.test.ts new file mode 100644 index 00000000..b1b21605 --- /dev/null +++ b/packages/coding-agent/test/user-message-selector.test.ts @@ -0,0 +1,81 @@ +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 tied to the correct message line", () => { + const component = new UserMessageSelectorComponent(makeItems(), vi.fn(), vi.fn()); + const lines = component.getMessageList().render(80); + + // 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", () => { + 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/src/client/api.ts b/packages/dashboard/src/client/api.ts index 7137644b..b0318ac1 100644 --- a/packages/dashboard/src/client/api.ts +++ b/packages/dashboard/src/client/api.ts @@ -155,7 +155,9 @@ 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 })), tree: (key: string) => request<{ roots: SessionTreeNodeDto[]; leafId: string | null }>(`/api/runtimes/${key}/tree`), diff --git a/packages/dashboard/src/client/screens/session.tsx b/packages/dashboard/src/client/screens/session.tsx index 9d3e2f6c..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([]); @@ -1269,11 +1271,20 @@ 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; + } + // 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); @@ -1282,6 +1293,9 @@ 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."); + async function openStatsPopover() { setShowStatsPopover(true); setStatsPopoverError(undefined); @@ -2336,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 d40635cb..e5a50aad 100644 --- a/packages/dashboard/src/client/styles/app.css +++ b/packages/dashboard/src/client/styles/app.css @@ -2761,7 +2761,8 @@ details.thinking .thinking-body { border-color: 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/test/client/screens.test.tsx b/packages/dashboard/test/client/screens.test.tsx index d1a21dd4..d62157f3 100644 --- a/packages/dashboard/test/client/screens.test.tsx +++ b/packages/dashboard/test/client/screens.test.tsx @@ -5791,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 () => {}); @@ -5819,6 +5821,81 @@ describe("dashboard client regressions", () => { expect((el.querySelector("textarea") as HTMLTextAreaElement).value).toBe("original prompt"); }); + 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: { forkasst: createSessionViewState("forkasst") }, + fleet: () => ({ runtimes: [], diskSessions: [] }), + hydrateSession, + 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(); + await new Promise((resolve) => setTimeout(resolve, 10)); + // 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.fork).toHaveBeenCalledWith("forkasst", "a1"); + expect(hydrateSession).toHaveBeenCalledWith("forkasst"); + expect(refreshDiskSessions).toHaveBeenCalledOnce(); + // 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 () => { + 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; + 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 = { diff --git a/packages/dashboard/test/runtime-pool.test.ts b/packages/dashboard/test/runtime-pool.test.ts index 20bb4b90..0d557f46 100644 --- a/packages/dashboard/test/runtime-pool.test.ts +++ b/packages/dashboard/test/runtime-pool.test.ts @@ -185,6 +185,8 @@ 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 })), listSessions: vi.fn(async () => []), switchSession: vi.fn(async () => ({ cancelled: false })), prompt: vi.fn(async () => {}), 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