Add option to fork from any message (user or assistant)

packages/coding-agent/README.md

@@ -175,7 +175,7 @@ Type `/` in the editor to trigger commands. [Extensions](#extensions) can regist
 | `/name <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 <path>       # 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 <path|id>`** - 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.
 

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"}
     ]
   }
 }

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` |

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,19 +3715,15 @@ export class AgentSession {
 			})) as SessionBeforeForkResult | undefined;
 
 			if (result?.cancel) {
-				return { selectedText, cancelled: true };
+				return { cancelled: true };
 			}
 			skipConversationRestore = result?.skipConversationRestore ?? false;
 		}
 
 		// 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

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));