Fix branch selector for single message and --no-session mode

- Allow branch selector to open with single user message (changed <= 1 to === 0 check)
- Support in-memory branching for --no-session mode (no files created)
- Add isEnabled() getter to SessionManager
- Update sessionFile getter to return null when sessions disabled
- Update SessionSwitchEvent types to allow null session files
- Add branching tests for single message and --no-session scenarios

fixes #163

packages/ai/test/mistral-debug.test.ts

@@ -69,6 +69,87 @@ describe.skipIf(!process.env.OPENAI_API_KEY)("OpenAI Debug", () => {
 describe.skipIf(!process.env.MISTRAL_API_KEY)("Mistral Debug", () => {
 	const model = getModel("mistral", "devstral-medium-latest");
 
+	it("two subsequent user messages", async () => {
+		const context: Context = {
+			messages: [
+				{ role: "user", content: "Hello", timestamp: Date.now() },
+				{ role: "user", content: "How are you?", timestamp: Date.now() },
+			],
+		};
+		const response = await complete(model, context);
+		console.log("Response:", response.stopReason, response.errorMessage);
+		expect(response.stopReason).not.toBe("error");
+	});
+
+	it("aborted assistant then user message", async () => {
+		const context: Context = {
+			messages: [
+				{ role: "user", content: "Hello", timestamp: Date.now() },
+				{
+					role: "assistant",
+					api: "openai-completions",
+					content: [],
+					provider: "mistral",
+					model: "devstral-medium-latest",
+					usage: {
+						input: 0,
+						output: 0,
+						cacheRead: 0,
+						cacheWrite: 0,
+						totalTokens: 0,
+						cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+					},
+					stopReason: "aborted",
+					timestamp: Date.now(),
+					errorMessage: "Request was aborted.",
+				},
+				{ role: "user", content: "How are you?", timestamp: Date.now() },
+			],
+		};
+		const response = await complete(model, context);
+		console.log("Response:", response.stopReason, response.errorMessage);
+		expect(response.stopReason).not.toBe("error");
+	});
+
+	it("three consecutive user messages (simulating aborted assistant skipped)", async () => {
+		const context: Context = {
+			messages: [
+				{ role: "user", content: "Hello", timestamp: Date.now() },
+				{ role: "user", content: "Ran some command", timestamp: Date.now() },
+				{ role: "user", content: "How are you?", timestamp: Date.now() },
+			],
+		};
+		const response = await complete(model, context);
+		console.log("Response:", response.stopReason, response.errorMessage);
+		expect(response.stopReason).not.toBe("error");
+	});
+
+	it("reproduce 502 from session fixture", async () => {
+		const fs = await import("fs");
+		const path = await import("path");
+		const fixtureData = JSON.parse(fs.readFileSync(path.join(__dirname, "fixtures/mistral.json"), "utf-8"));
+		// Filter out bashExecution and convert to user message like messageTransformer does
+		const messages = fixtureData.map((m: any) => {
+			if (m.role === "bashExecution") {
+				let text = `Ran \`${m.command}\`\n`;
+				if (m.output) {
+					text += "```\n" + m.output + "\n```";
+				} else {
+					text += "(no output)";
+				}
+				return { role: "user", content: [{ type: "text", text }], timestamp: m.timestamp };
+			}
+			return m;
+		});
+		const context: Context = {
+			messages,
+			tools: [weatherTool],
+		};
+		const response = await complete(model, context);
+		console.log("Response:", response.stopReason, response.errorMessage);
+		expect(response.stopReason).not.toBe("error");
+	});
+
 	it("5d. two tool calls + results, no follow-up user", async () => {
 		const context: Context = {
 			messages: [

packages/coding-agent/CHANGELOG.md

@@ -6,6 +6,12 @@
 
 - **HTML export line numbers**: Read tool calls in HTML exports now display line number ranges (e.g., `file.txt:10-20`) when offset/limit parameters are used, matching the TUI display format. Line numbers appear in yellow color for better visibility. ([#166](https://github.com/badlogic/pi-mono/issues/166))
 
+### Fixed
+
+- **Branch selector now works with single message**: Previously the branch selector would not open when there was only one user message. Now it correctly allows branching from any message, including the first one. This is needed for checkpoint hooks to restore state from before the first message. ([#163](https://github.com/badlogic/pi-mono/issues/163))
+
+- **In-memory branching for `--no-session` mode**: Branching now works correctly in `--no-session` mode without creating any session files. The conversation is truncated in memory.
+
 ## [0.18.1] - 2025-12-10
 
 ### Added

packages/coding-agent/docs/hooks.md

@@ -122,8 +122,8 @@ Fired when session changes (`/branch` or session switch).
 
 ```typescript
 pi.on("session_switch", async (event, ctx) => {
-  // event.newSessionFile: string
-  // event.previousSessionFile: string
+  // event.newSessionFile: string | null (null in --no-session mode)
+  // event.previousSessionFile: string | null (null in --no-session mode)
   // event.reason: "branch" | "switch"
 });
 ```

packages/coding-agent/src/core/agent-session.ts

@@ -76,7 +76,7 @@ export interface CompactionResult {
 
 /** Session statistics for /session command */
 export interface SessionStats {
-	sessionFile: string;
+	sessionFile: string | null;
 	sessionId: string;
 	userMessages: number;
 	assistantMessages: number;
@@ -320,9 +320,9 @@ export class AgentSession {
 		return this.agent.getQueueMode();
 	}
 
-	/** Current session file path */
-	get sessionFile(): string {
-		return this.sessionManager.getSessionFile();
+	/** Current session file path, or null if sessions are disabled */
+	get sessionFile(): string | null {
+		return this.sessionManager.isEnabled() ? this.sessionManager.getSessionFile() : null;
 	}
 
 	/** Current session ID */
@@ -966,11 +966,15 @@ export class AgentSession {
 			return { selectedText, skipped: true };
 		}
 
-		// Create branched session
+		// Create branched session (returns null in --no-session mode)
 		const newSessionFile = this.sessionManager.createBranchedSessionFromEntries(entries, entryIndex);
-		this.sessionManager.setSessionFile(newSessionFile);
 
-		// Emit session_switch event
+		// Update session file if we have one (file-based mode)
+		if (newSessionFile !== null) {
+			this.sessionManager.setSessionFile(newSessionFile);
+		}
+
+		// Emit session_switch event (in --no-session mode, both files are null)
 		if (this._hookRunner) {
 			this._hookRunner.setSessionFile(newSessionFile);
 			await this._hookRunner.emit({
@@ -981,7 +985,7 @@ export class AgentSession {
 			});
 		}
 
-		// Reload
+		// Reload messages from entries (works for both file and in-memory mode)
 		const loaded = loadSessionFromEntries(this.sessionManager.loadEntries());
 		this.agent.replaceMessages(loaded.messages);
 

packages/coding-agent/src/core/hooks/types.ts

@@ -86,10 +86,10 @@ export interface SessionStartEvent {
  */
 export interface SessionSwitchEvent {
 	type: "session_switch";
-	/** New session file path */
-	newSessionFile: string;
-	/** Previous session file path */
-	previousSessionFile: string;
+	/** New session file path, or null in --no-session mode */
+	newSessionFile: string | null;
+	/** Previous session file path, or null in --no-session mode */
+	previousSessionFile: string | null;
 	/** Reason for the switch */
 	reason: "branch" | "switch";
 }

packages/coding-agent/src/core/session-manager.ts

@@ -231,6 +231,11 @@ export class SessionManager {
 		this.enabled = false;
 	}
 
+	/** Check if session persistence is enabled */
+	isEnabled(): boolean {
+		return this.enabled;
+	}
+
 	private getSessionDirectory(): string {
 		const cwd = process.cwd();
 		// Replace all path separators and colons (for Windows drive letters) with dashes
@@ -637,32 +642,43 @@ export class SessionManager {
 	/**
 	 * Create a branched session from session entries up to (but not including) a specific entry index.
 	 * This preserves compaction events and all entry types.
-	 * Returns the new session file path.
+	 * Returns the new session file path, or null if in --no-session mode (in-memory only).
 	 */
-	createBranchedSessionFromEntries(entries: SessionEntry[], branchBeforeIndex: number): string {
+	createBranchedSessionFromEntries(entries: SessionEntry[], branchBeforeIndex: number): string | null {
 		const newSessionId = uuidv4();
 		const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
 		const newSessionFile = join(this.sessionDir, `${timestamp}_${newSessionId}.jsonl`);
 
-		// Copy all entries up to (but not including) the branch point
+		// Build new entries list (up to but not including branch point)
+		const newEntries: SessionEntry[] = [];
 		for (let i = 0; i < branchBeforeIndex; i++) {
 			const entry = entries[i];
 
 			if (entry.type === "session") {
 				// Rewrite session header with new ID and branchedFrom
-				const newHeader: SessionHeader = {
+				newEntries.push({
 					...entry,
 					id: newSessionId,
 					timestamp: new Date().toISOString(),
-					branchedFrom: this.sessionFile,
-				};
-				appendFileSync(newSessionFile, JSON.stringify(newHeader) + "\n");
+					branchedFrom: this.enabled ? this.sessionFile : undefined,
+				});
 			} else {
 				// Copy other entries as-is
-				appendFileSync(newSessionFile, JSON.stringify(entry) + "\n");
+				newEntries.push(entry);
 			}
 		}
 
-		return newSessionFile;
+		if (this.enabled) {
+			// Write to file
+			for (const entry of newEntries) {
+				appendFileSync(newSessionFile, JSON.stringify(entry) + "\n");
+			}
+			return newSessionFile;
+		} else {
+			// In-memory mode: replace inMemoryEntries, no file created
+			this.inMemoryEntries = newEntries;
+			this.sessionId = newSessionId;
+			return null;
+		}
 	}
 }

packages/coding-agent/src/modes/interactive/components/user-message-selector.ts

@@ -136,8 +136,8 @@ export class UserMessageSelectorComponent extends Container {
 		this.addChild(new Spacer(1));
 		this.addChild(new DynamicBorder());
 
-		// Auto-cancel if no messages or only one message
-		if (messages.length <= 1) {
+		// Auto-cancel if no messages
+		if (messages.length === 0) {
 			setTimeout(() => onCancel(), 100);
 		}
 	}

packages/coding-agent/src/modes/interactive/interactive-mode.ts

@@ -1214,7 +1214,7 @@ export class InteractiveMode {
 	private showUserMessageSelector(): void {
 		const userMessages = this.session.getUserMessagesForBranching();
 
-		if (userMessages.length <= 1) {
+		if (userMessages.length === 0) {
 			this.showStatus("No messages to branch from");
 			return;
 		}

packages/coding-agent/src/modes/rpc/rpc-types.ts

@@ -65,7 +65,7 @@ export interface RpcSessionState {
 	isStreaming: boolean;
 	isCompacting: boolean;
 	queueMode: "all" | "one-at-a-time";
-	sessionFile: string;
+	sessionFile: string | null;
 	sessionId: string;
 	autoCompactionEnabled: boolean;
 	messageCount: number;

packages/coding-agent/test/agent-session-branching.test.ts

@@ -0,0 +1,158 @@
+/**
+ * Tests for AgentSession branching behavior.
+ *
+ * These tests verify:
+ * - Branching from a single message works
+ * - Branching in --no-session mode (in-memory only)
+ * - getUserMessagesForBranching returns correct entries
+ */
+
+import { existsSync, mkdirSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { Agent, ProviderTransport } from "@mariozechner/pi-agent-core";
+import { getModel } from "@mariozechner/pi-ai";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import { AgentSession } from "../src/core/agent-session.js";
+import { SessionManager } from "../src/core/session-manager.js";
+import { SettingsManager } from "../src/core/settings-manager.js";
+import { codingTools } from "../src/core/tools/index.js";
+
+const API_KEY = process.env.ANTHROPIC_API_KEY || process.env.ANTHROPIC_OAUTH_TOKEN;
+
+describe.skipIf(!API_KEY)("AgentSession branching", () => {
+	let session: AgentSession;
+	let tempDir: string;
+	let sessionManager: SessionManager;
+
+	beforeEach(() => {
+		// Create temp directory for session files
+		tempDir = join(tmpdir(), `pi-branching-test-${Date.now()}`);
+		mkdirSync(tempDir, { recursive: true });
+	});
+
+	afterEach(async () => {
+		if (session) {
+			session.dispose();
+		}
+		if (tempDir && existsSync(tempDir)) {
+			rmSync(tempDir, { recursive: true });
+		}
+	});
+
+	function createSession(noSession: boolean = false) {
+		const model = getModel("anthropic", "claude-sonnet-4-5")!;
+
+		const transport = new ProviderTransport({
+			getApiKey: () => API_KEY,
+		});
+
+		const agent = new Agent({
+			transport,
+			initialState: {
+				model,
+				systemPrompt: "You are a helpful assistant. Be extremely concise, reply with just a few words.",
+				tools: codingTools,
+			},
+		});
+
+		sessionManager = new SessionManager(false);
+		if (noSession) {
+			sessionManager.disable();
+		}
+		const settingsManager = new SettingsManager(tempDir);
+
+		session = new AgentSession({
+			agent,
+			sessionManager,
+			settingsManager,
+		});
+
+		// Must subscribe to enable session persistence
+		session.subscribe(() => {});
+
+		return session;
+	}
+
+	it("should allow branching from single message", async () => {
+		createSession();
+
+		// Send one message
+		await session.prompt("Say hello");
+		await session.agent.waitForIdle();
+
+		// Should have exactly 1 user message available for branching
+		const userMessages = session.getUserMessagesForBranching();
+		expect(userMessages.length).toBe(1);
+		expect(userMessages[0].text).toBe("Say hello");
+
+		// Branch from the first message
+		const result = await session.branch(userMessages[0].entryIndex);
+		expect(result.selectedText).toBe("Say hello");
+		expect(result.skipped).toBe(false);
+
+		// After branching, conversation should be empty (branched before the first message)
+		expect(session.messages.length).toBe(0);
+
+		// Session file should exist (new branch)
+		expect(session.sessionFile).not.toBeNull();
+		expect(existsSync(session.sessionFile!)).toBe(true);
+	});
+
+	it("should support in-memory branching in --no-session mode", async () => {
+		createSession(true);
+
+		// Verify sessions are disabled
+		expect(session.sessionFile).toBeNull();
+
+		// Send one message
+		await session.prompt("Say hi");
+		await session.agent.waitForIdle();
+
+		// Should have 1 user message
+		const userMessages = session.getUserMessagesForBranching();
+		expect(userMessages.length).toBe(1);
+
+		// Verify we have messages before branching
+		expect(session.messages.length).toBeGreaterThan(0);
+
+		// Branch from the first message
+		const result = await session.branch(userMessages[0].entryIndex);
+		expect(result.selectedText).toBe("Say hi");
+		expect(result.skipped).toBe(false);
+
+		// After branching, conversation should be empty
+		expect(session.messages.length).toBe(0);
+
+		// Session file should still be null (no file created)
+		expect(session.sessionFile).toBeNull();
+	});
+
+	it("should branch from middle of conversation", async () => {
+		createSession();
+
+		// Send multiple messages
+		await session.prompt("Say one");
+		await session.agent.waitForIdle();
+
+		await session.prompt("Say two");
+		await session.agent.waitForIdle();
+
+		await session.prompt("Say three");
+		await session.agent.waitForIdle();
+
+		// Should have 3 user messages
+		const userMessages = session.getUserMessagesForBranching();
+		expect(userMessages.length).toBe(3);
+
+		// Branch from second message (keeps first message + response)
+		const secondMessage = userMessages[1];
+		const result = await session.branch(secondMessage.entryIndex);
+		expect(result.selectedText).toBe("Say two");
+
+		// After branching, should have first user message + assistant response
+		expect(session.messages.length).toBe(2);
+		expect(session.messages[0].role).toBe("user");
+		expect(session.messages[1].role).toBe("assistant");
+	});
+});