@gajae-code/coding-agent 0.17.2 → 0.17.4

package/src/modes/components/user-message-selector.ts

@@ -14,9 +14,10 @@ interface UserMessageItem {
  */
 class UserMessageList implements Component {
 	#selectedIndex: number = 0;
+	#selectionInFlight = false;
 	onSelect?: (entryId: string) => void;
 	onCancel?: () => void;
-	#maxVisible: number = 10; // Max messages visible
+	#availableRows = Number.POSITIVE_INFINITY;
 
 	constructor(private readonly messages: UserMessageItem[]) {
 		// Store messages in chronological order (oldest to newest)
@@ -24,6 +25,10 @@ class UserMessageList implements Component {
 		this.#selectedIndex = Math.max(0, this.messages.length - 1);
 	}
 
+	setAvailableRows(rows: number): void {
+		this.#availableRows = Math.max(3, rows);
+	}
+
 	invalidate(): void {
 		// No cached state to invalidate currently
 	}
@@ -36,12 +41,15 @@ class UserMessageList implements Component {
 			return lines;
 		}
 
-		// Calculate visible range with scrolling
+		// Each entry uses two content rows plus a separator. Leave one row for
+		// the scroll position whenever the full history does not fit.
+		const needsScrollIndicator = this.messages.length * 3 > this.#availableRows;
+		const maxVisible = Math.max(1, Math.floor((this.#availableRows - (needsScrollIndicator ? 1 : 0)) / 3));
 		const startIndex = Math.max(
 			0,
-			Math.min(this.#selectedIndex - Math.floor(this.#maxVisible / 2), this.messages.length - this.#maxVisible),
+			Math.min(this.#selectedIndex - Math.floor(maxVisible / 2), this.messages.length - maxVisible),
 		);
-		const endIndex = Math.min(startIndex + this.#maxVisible, this.messages.length);
+		const endIndex = Math.min(startIndex + maxVisible, this.messages.length);
 
 		// Render visible messages (2 lines per message + blank line)
 		for (let i = startIndex; i < endIndex; i++) {
@@ -77,6 +85,7 @@ class UserMessageList implements Component {
 	}
 
 	handleInput(keyData: string): void {
+		if (this.#selectionInFlight) return;
 		// Up arrow - go to previous (older) message, wrap to bottom when at top
 		if (matchesKey(keyData, "up")) {
 			this.#selectedIndex = this.#selectedIndex === 0 ? this.messages.length - 1 : this.#selectedIndex - 1;
@@ -89,6 +98,7 @@ class UserMessageList implements Component {
 		else if (matchesKey(keyData, "enter") || matchesKey(keyData, "return") || keyData === "\n") {
 			const selected = this.messages[this.#selectedIndex];
 			if (selected && this.onSelect) {
+				this.#selectionInFlight = true;
 				this.onSelect(selected.id);
 			}
 		}
@@ -102,21 +112,43 @@ class UserMessageList implements Component {
 }
 
 /**
- * Component that renders a user message selector for branching
+ * Component that renders a prompt selector for creating a new session
  */
 export class UserMessageSelectorComponent extends Container {
 	#messageList: UserMessageList;
-
-	constructor(messages: UserMessageItem[], onSelect: (entryId: string) => void, onCancel: () => void) {
+	#header: Container;
+	#footer: Container;
+
+	constructor(
+		messages: UserMessageItem[],
+		onSelect: (entryId: string) => void,
+		onCancel: () => void,
+		private readonly getViewportRows: () => number,
+	) {
 		super();
 
 		// 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 Spacer(1));
-		this.addChild(new DynamicBorder());
-		this.addChild(new Spacer(1));
+		this.#header = new Container();
+		this.#header.addChild(new Spacer(1));
+		this.#header.addChild(new Text(theme.bold("Fork from Prompt"), 1, 0));
+		this.#header.addChild(
+			new Text(
+				theme.fg(
+					"muted",
+					"Creates a new session with history before the prompt; shared files stay in the same cwd.",
+				),
+				1,
+				0,
+			),
+		);
+		this.#header.addChild(
+			new Text(theme.fg("muted", "The selected prompt is restored for editing, not submitted."), 1, 0),
+		);
+		this.#header.addChild(new Text(theme.fg("muted", "↑/↓ move · Enter select · Esc cancel"), 1, 0));
+		this.#header.addChild(new Spacer(1));
+		this.#header.addChild(new DynamicBorder());
+		this.#header.addChild(new Spacer(1));
+		this.addChild(this.#header);
 
 		// Create message list
 		this.#messageList = new UserMessageList(messages);
@@ -126,8 +158,10 @@ export class UserMessageSelectorComponent extends Container {
 		this.addChild(this.#messageList);
 
 		// Add bottom border
-		this.addChild(new Spacer(1));
-		this.addChild(new DynamicBorder());
+		this.#footer = new Container();
+		this.#footer.addChild(new Spacer(1));
+		this.#footer.addChild(new DynamicBorder());
+		this.addChild(this.#footer);
 
 		// Auto-cancel if no messages
 		if (messages.length === 0) {
@@ -135,6 +169,14 @@ export class UserMessageSelectorComponent extends Container {
 		}
 	}
 
+	render(width: number): string[] {
+		const chromeRows = this.#header.render(width).length + this.#footer.render(width).length;
+		// Reserve the pinned status/composer rows outside the overlay. Recompute
+		// against current terminal rows and wrapped header height on every render.
+		this.#messageList.setAvailableRows(this.getViewportRows() - chromeRows - 2);
+		return super.render(width);
+	}
+
 	getMessageList(): UserMessageList {
 		return this.#messageList;
 	}

package/src/modes/controllers/command-controller.ts

@@ -38,7 +38,16 @@ import { replaceTabs } from "../../tools/render-utils";
 import { getDisplayChangelogEntries } from "../../utils/changelog";
 import { copyToClipboard } from "../../utils/clipboard";
 import { openPath } from "../../utils/open";
-import { setSessionTerminalTitle } from "../../utils/title-generator";
+import {
+	beginSessionTitleGeneration,
+	invalidateSessionTitleGeneration,
+	isSessionTitleGenerationCurrent,
+} from "../../utils/session-title-generation";
+import {
+	buildConversationTitleInput,
+	generateSessionTitle,
+	setSessionTerminalTitle,
+} from "../../utils/title-generator";
 import { addChatChild, prepareTranscriptRebuild, syncPendingExecutionComponents } from "../utils/ui-helpers";
 
 type HindsightModule = typeof import("../../hindsight");
@@ -84,6 +93,7 @@ export function buildHelpMarkdown(keybindings: Pick<KeybindingsManager, "getAcce
 		"| Resume another session | `/resume` |",
 		"| Show session details | `/session info` |",
 		"| Delete current session transcript/artifacts | `/session delete` |",
+		"| Fork from an earlier prompt into a new session | `/fork` |",
 		`| Select a model | \`/model\` or ${selectModelKey} |`,
 		`| Queue a message for the next turn | ${queueMessageKey} |`,
 		`| Select or edit a queued message | ${dequeueMessageKey} |`,
@@ -1109,31 +1119,7 @@ export class CommandController {
 	}
 
 	async handleForkCommand(): Promise<void> {
-		if (this.ctx.session.isStreaming) {
-			this.ctx.showWarning("Wait for the current response to finish or abort it before forking.");
-			return;
-		}
-
-		const success = await this.ctx.session.fork();
-		if (this.ctx.isStopped?.()) return;
-		if (!success) {
-			this.ctx.showError("Fork failed (session not persisted or cancelled)");
-			return;
-		}
-		clearInteractiveActivityLoaders(this.ctx);
-		stopInteractiveActivityIndicator(this.ctx, { foregroundSettled: true });
-		this.ctx.resetIrcSidebarSession();
-
-		this.ctx.statusLine.invalidate();
-		this.ctx.updateEditorTopBorder();
-
-		const sessionFile = this.ctx.session.sessionFile;
-		const shortPath = sessionFile ? sessionFile.split("/").pop() : "new session";
-		this.ctx.chatContainer.addChild(new Spacer(1));
-		this.ctx.chatContainer.addChild(
-			new Text(`${theme.fg("accent", `${theme.status.success} Session forked to ${shortPath}`)}`, 1, 1),
-		);
-		this.ctx.ui.requestRender();
+		this.ctx.showUserMessageSelector();
 	}
 
 	async handleMoveCommand(targetPath: string): Promise<void> {
@@ -1186,23 +1172,63 @@ export class CommandController {
 		}
 	}
 
-	async handleRenameCommand(title: string): Promise<void> {
+	async handleRenameCommand(title?: string): Promise<void> {
 		try {
+			if (!title?.trim()) {
+				const generation = beginSessionTitleGeneration(this.ctx.sessionManager);
+				const input = buildConversationTitleInput(this.ctx.session.messages);
+				if (!input) {
+					this.ctx.showError("Nothing to summarize yet — pass a title: /rename <title>");
+					return;
+				}
+
+				const generated = await generateSessionTitle(
+					input,
+					this.ctx.session.modelRegistry,
+					this.ctx.settings,
+					this.ctx.session.credentialSessionId,
+					this.ctx.session.model,
+					provider => this.ctx.session.agent.metadataForProvider(provider),
+				);
+				if (
+					!isSessionTitleGenerationCurrent(this.ctx.sessionManager, generation) ||
+					buildConversationTitleInput(this.ctx.session.messages) !== input
+				)
+					return;
+				if (!generated) {
+					this.ctx.showError("Could not generate a session title — pass one: /rename <title>");
+					return;
+				}
+
+				const stored = await this.ctx.sessionManager.setSessionName(generated, "user");
+				if (!stored) {
+					this.ctx.showError("Session name cannot be empty.");
+					return;
+				}
+				this.#finishRename();
+				return;
+			}
+
+			invalidateSessionTitleGeneration(this.ctx.sessionManager);
 			const stored = await this.ctx.sessionManager.setSessionName(title, "user");
 			if (!stored) {
 				this.ctx.showError("Session name cannot be empty.");
 				return;
 			}
-			const name = this.ctx.sessionManager.getSessionName()!;
-			setSessionTerminalTitle(name, this.ctx.sessionManager.getCwd());
-			this.ctx.statusLine.invalidate();
-			this.ctx.updateEditorBorderColor();
-			this.ctx.showStatus(`Session renamed to "${name}".`);
+			this.#finishRename();
 		} catch (err) {
 			this.ctx.showError(`Rename failed: ${err instanceof Error ? err.message : String(err)}`);
 		}
 	}
 
+	#finishRename(): void {
+		const name = this.ctx.sessionManager.getSessionName()!;
+		setSessionTerminalTitle(name, this.ctx.sessionManager.getCwd());
+		this.ctx.statusLine.invalidate();
+		this.ctx.updateEditorBorderColor();
+		this.ctx.showStatus(`Session renamed to "${name}".`);
+	}
+
 	async handleBashCommand(command: string, excludeFromContext = false): Promise<void> {
 		const isDeferred = this.ctx.session.isStreaming;
 		const component = new BashExecutionComponent(command, this.ctx.ui, excludeFromContext);

package/src/modes/controllers/extension-ui-controller.ts

@@ -146,31 +146,55 @@ export class ExtensionUiController {
 		const widgetsBelow = [...this.#hookWidgetsBelow.entries()];
 		const activeHookCustomComponent = this.#activeHookCustomComponent;
 		const activeHookCustomOverlay = this.#activeHookCustomOverlay;
+		const activeHookCustomCancel = this.#activeHookCustomCancel;
 		return () => {
+			const errors: unknown[] = [];
+			const attempt = (cleanup: () => void): void => {
+				try {
+					cleanup();
+				} catch (error) {
+					errors.push(error);
+				}
+			};
 			for (const unsubscribe of terminalInputUnsubscribers) {
-				unsubscribe();
-				this.#extensionTerminalInputUnsubscribers.delete(unsubscribe);
+				if (!this.#extensionTerminalInputUnsubscribers.delete(unsubscribe)) continue;
+				attempt(unsubscribe);
 			}
 			let widgetsChanged = false;
 			for (const [key, widget] of widgetsAbove) {
 				if (this.#hookWidgetsAbove.get(key) !== widget) continue;
 				this.#hookWidgetsAbove.delete(key);
-				widget.dispose?.();
+				attempt(() => widget.dispose?.());
 				widgetsChanged = true;
 			}
 			for (const [key, widget] of widgetsBelow) {
 				if (this.#hookWidgetsBelow.get(key) !== widget) continue;
 				this.#hookWidgetsBelow.delete(key);
-				widget.dispose?.();
+				attempt(() => widget.dispose?.());
 				widgetsChanged = true;
 			}
 			if (
 				this.#activeHookCustomComponent === activeHookCustomComponent &&
-				this.#activeHookCustomOverlay === activeHookCustomOverlay
+				this.#activeHookCustomOverlay === activeHookCustomOverlay &&
+				this.#activeHookCustomCancel === activeHookCustomCancel
 			) {
-				this.#clearActiveHookCustom();
+				this.#activeHookCustomComponent = undefined;
+				this.#activeHookCustomOverlay = undefined;
+				this.#activeHookCustomCancel = undefined;
+				attempt(() => activeHookCustomCancel?.());
+				attempt(() => activeHookCustomComponent?.dispose?.());
+				attempt(() => activeHookCustomOverlay?.hide());
+			}
+			if (widgetsChanged) {
+				attempt(() =>
+					this.#renderHookWidgetContainer(this.ctx.hookWidgetContainerAbove, this.#hookWidgetsAbove, true),
+				);
+				attempt(() =>
+					this.#renderHookWidgetContainer(this.ctx.hookWidgetContainerBelow, this.#hookWidgetsBelow, false),
+				);
+				attempt(() => this.ctx.ui.requestRender());
 			}
-			if (widgetsChanged) this.#rebuildHookWidgets();
+			if (errors.length > 0) throw new AggregateError(errors, "Previous session extension UI cleanup failed");
 		};
 	}
 
@@ -551,6 +575,7 @@ export class ExtensionUiController {
 			getCredentialSessionId: () => this.ctx.session.credentialSessionId,
 			isIdle: () => !this.ctx.session.isStreaming,
 			getActivePromptHandle: () => this.ctx.session.activePromptHandle,
+			getSessionWorkLease: () => this.ctx.session.sessionWorkLease,
 			abort: () => this.ctx.session.abort(),
 			abortPromptAndWait: (handle, options) => this.ctx.session.abortPromptAndWait(handle, options),
 			hasPendingMessages: () => this.ctx.session.queuedMessageCount > 0,
@@ -781,6 +806,7 @@ export class ExtensionUiController {
 			getCredentialSessionId: () => this.ctx.session.credentialSessionId,
 			isIdle: () => !this.ctx.session.isStreaming,
 			getActivePromptHandle: () => this.ctx.session.activePromptHandle,
+			getSessionWorkLease: () => this.ctx.session.sessionWorkLease,
 			abort: () => this.ctx.session.abort(),
 			abortPromptAndWait: (handle, options) => this.ctx.session.abortPromptAndWait(handle, options),
 			hasPendingMessages: () => this.ctx.session.queuedMessageCount > 0,

package/src/modes/controllers/input-controller.ts

@@ -33,6 +33,7 @@ import {
 	locatePastedImageReferenceAroundCursor,
 	parsePastedImagePaths,
 } from "../../utils/pasted-image-path";
+import { invalidateSessionTitleGeneration } from "../../utils/session-title-generation";
 import { generateSessionTitle, setSessionTerminalTitle } from "../../utils/title-generator";
 import { ActionRegistry, APP_ACTION_METADATA } from "../action-registry";
 import { CommandPalette, type CommandPaletteAction, type CommandPaletteEntry } from "../components/command-palette";
@@ -1146,6 +1147,7 @@ export class InputController {
 
 		// Queue input during compaction
 		if (this.ctx.session.isCompacting) {
+			invalidateSessionTitleGeneration(this.ctx.sessionManager);
 			if ((inputImages?.length ?? 0) > 0) {
 				this.ctx.showStatus("Compaction in progress. Retry after it completes to send images.");
 				return;
@@ -1159,6 +1161,7 @@ export class InputController {
 		// prompt to run after the active turn completes (in submission order).
 		// This handles extension commands (execute immediately), prompt template expansion, and queueing
 		if (this.ctx.session.isStreaming) {
+			invalidateSessionTitleGeneration(this.ctx.sessionManager);
 			if (this.#canModifyComposer(composer)) {
 				this.ctx.editor.addToHistory(text);
 				this.ctx.editor.setText("");
@@ -1183,6 +1186,7 @@ export class InputController {
 		}
 
 		// Normal message submission
+		invalidateSessionTitleGeneration(this.ctx.sessionManager);
 		// First, move any pending bash components to chat
 		this.ctx.flushPendingBashComponents();
 

package/src/modes/controllers/selector-controller.ts

@@ -155,7 +155,12 @@ import {
 	MODEL_ONBOARDING_PROVIDER_PRESET_COMMAND,
 	MODEL_ONBOARDING_SETUP_COMMAND,
 } from "../../setup/model-onboarding-guidance";
-import { addApiCompatibleProvider, formatProviderSetupResult } from "../../setup/provider-onboarding";
+import {
+	addApiCompatibleProvider,
+	formatProviderSetupResult,
+	probeOpenAIModelsList,
+	reloadAndRefreshDiscoveryCatalog,
+} from "../../setup/provider-onboarding";
 import { getThinkingLevelMetadata } from "../../thinking";
 import {
 	isConfigurableSearchProviderId,
@@ -1821,15 +1826,48 @@ export class SelectorController {
 		this.showSelector(done => {
 			let wizard: CustomProviderWizardComponent;
 			const submit = async (input: CustomProviderWizardSubmit): Promise<void> => {
+				// Component-owned generation: revision or cancel after submit
+				// start invalidates this attempt. Both completion branches
+				// check it so a stale refresh cannot report success or
+				// rerender over inputs the user has since revised.
+				const generation = wizard.currentSubmitGeneration();
+				const isCurrentSubmission = (): boolean => wizard.isSubmitCurrent(generation);
 				try {
-					const result = await addApiCompatibleProvider(input);
-					await this.ctx.session.modelRegistry.refresh("offline", this.ctx.session.credentialSessionId);
+					const result = await addApiCompatibleProvider({
+						...input,
+						authStorage: this.ctx.session.modelRegistry.authStorage,
+					});
+					let recoveryHint: string | null = null;
+					if (result.discoveryEnabled && !result.preset) {
+						// A discovery-only add has no static models and no cache yet,
+						// so an offline reload alone leaves an empty catalog. Refresh
+						// only the new provider online; a hint is surfaced instead
+						// of unconditional success when it stays unavailable.
+						recoveryHint = await reloadAndRefreshDiscoveryCatalog(
+							this.ctx.session.modelRegistry,
+							result.providerId,
+							this.ctx.session.credentialSessionId,
+						);
+					} else {
+						await this.ctx.session.modelRegistry.refresh("offline", this.ctx.session.credentialSessionId);
+					}
 					await this.ctx.notifyConfigChanged?.();
-					this.ctx.showStatus(formatProviderSetupResult(result));
+					// A late settle after Esc-cancel or a mid-submit revision
+					// must not drive completion UI: done() restores the
+					// composer and could close a selector the user opened
+					// after dismissal, and a stale refresh must not report
+					// success over revised inputs.
+					if (!isCurrentSubmission()) return;
+					this.ctx.showStatus(
+						recoveryHint
+							? `${formatProviderSetupResult(result)}\n${recoveryHint}`
+							: formatProviderSetupResult(result),
+					);
 					wizard.complete();
 					done();
 					this.ctx.ui.requestRender();
 				} catch (err) {
+					if (!isCurrentSubmission()) return;
 					const message = err instanceof Error ? err.message : String(err);
 					wizard.setSubmitError(`Provider setup failed: ${message}`);
 				}
@@ -1843,6 +1881,9 @@ export class SelectorController {
 					this.ctx.ui.requestRender();
 				},
 				() => this.ctx.ui.requestRender(),
+				{
+					discoverModels: async request => probeOpenAIModelsList(request),
+				},
 			);
 			return { component: wizard, focus: wizard };
 		});
@@ -3070,36 +3111,138 @@ export class SelectorController {
 		});
 	}
 
-	showUserMessageSelector(): void {
-		const userMessages = this.ctx.session.getUserMessagesForBranching();
+	#forkBusyReason(): string | undefined {
+		if (this.ctx.session.isStreaming) {
+			return "Wait for the current response to finish or abort it before forking.";
+		}
+		if (this.ctx.session.isCompacting) {
+			return "Wait for compaction to finish or abort it before forking.";
+		}
+		if (this.ctx.hasPendingSubmission()) {
+			return "Wait for the pending message to be submitted or cancel it before forking.";
+		}
+		if (this.ctx.bashComponent || this.ctx.pythonComponent) {
+			return "Wait for the foreground command to finish or cancel it before forking.";
+		}
+		return undefined;
+	}
 
+	showUserMessageSelector(onCommitted?: () => void): void {
+		const openedSessionId = this.ctx.sessionManager.getSessionId();
+		const openedSessionFile = this.ctx.sessionManager.getSessionFile();
+		if (!openedSessionFile) {
+			this.ctx.showError("Fork requires a persisted session");
+			return;
+		}
+		const busyReason = this.#forkBusyReason();
+		if (busyReason) {
+			this.ctx.showWarning(busyReason);
+			return;
+		}
+
+		const userMessages = this.ctx.session.getUserMessagesForBranching();
 		if (userMessages.length === 0) {
 			this.ctx.showStatus("No messages to branch from");
 			return;
 		}
 
+		this.ctx.goalModeController.cancelContinuation();
+		// Close the side-chat composer before the picker takes ownership, including
+		// paths that cancel, fail, or stop while branch() is pending.
+		if (this.ctx.hasActiveBtw()) this.ctx.handleBtwEscape();
+
 		this.showSelector(done => {
+			let selectionInFlight = false;
+			const closeWithoutBranch = (message: string, error = false): void => {
+				done();
+				if (error) {
+					this.ctx.showError(message);
+				} else {
+					this.ctx.showStatus(message);
+				}
+				this.ctx.goalModeController.scheduleContinuation();
+				this.ctx.ui.requestRender();
+			};
 			const selector = new UserMessageSelectorComponent(
 				userMessages.map(m => ({ id: m.entryId, text: m.text })),
 				async entryId => {
-					const result = await this.ctx.session.branch(entryId);
-					if (result.cancelled) {
-						// Hook cancelled the branch
-						done();
-						this.ctx.ui.requestRender();
+					if (selectionInFlight) return;
+					selectionInFlight = true;
+
+					if (
+						this.ctx.sessionManager.getSessionId() !== openedSessionId ||
+						this.ctx.sessionManager.getSessionFile() !== openedSessionFile
+					) {
+						closeWithoutBranch("Fork cancelled because the active session changed", true);
+						return;
+					}
+					const busyReason = this.#forkBusyReason();
+					if (busyReason) {
+						closeWithoutBranch(busyReason);
 						return;
 					}
-					this.ctx.resetIrcSidebarSession();
 
-					this.ctx.rebuildInitialMessages("replace-identity");
-					this.ctx.editor.setText(result.selectedText);
+					let selectedText = userMessages.find(message => message.entryId === entryId)!.text;
+					let branchFailure: unknown;
+					try {
+						const result = await this.ctx.session.branch(entryId);
+						if (this.ctx.isStopped?.()) return;
+						if (result.cancelled) {
+							closeWithoutBranch("Fork cancelled");
+							return;
+						}
+						selectedText = result.selectedText;
+					} catch (error) {
+						if (this.ctx.isStopped?.()) return;
+						const committed =
+							this.ctx.sessionManager.getSessionId() !== openedSessionId ||
+							this.ctx.sessionManager.getSessionFile() !== openedSessionFile;
+						if (!committed) {
+							closeWithoutBranch(`Fork failed: ${error instanceof Error ? error.message : String(error)}`, true);
+							return;
+						}
+						// The successor is already durable. Never describe this as a
+						// cancelled transition or roll it back to the parent's UI.
+						branchFailure = error;
+					}
+
+					try {
+						onCommitted?.();
+					} catch (error) {
+						branchFailure ??= error;
+					}
+					try {
+						this.ctx.resetIrcSidebarSession();
+						this.ctx.resetObserverRegistry();
+						this.#refreshSessionTerminalTitle();
+						this.ctx.statusLine.invalidate();
+						this.ctx.statusLine.setSessionStartTime(Date.now());
+						this.ctx.updateEditorTopBorder();
+						this.ctx.updateEditorBorderColor();
+						this.ctx.rebuildInitialMessages("replace-identity");
+						await this.ctx.reloadTodos();
+					} catch (error) {
+						branchFailure ??= error;
+					}
+					if (this.ctx.isStopped?.()) return;
+					this.ctx.editor.setText(selectedText);
 					done();
-					this.ctx.showStatus("Branched to new session");
+					if (branchFailure !== undefined) {
+						this.ctx.showError(
+							`Fork created, but session setup failed: ${branchFailure instanceof Error ? branchFailure.message : String(branchFailure)}`,
+						);
+					} else {
+						this.ctx.showStatus("Forked to a new session; edit the selected prompt to continue");
+					}
+					// Do not schedule goal continuation over the selected prompt draft.
+					this.ctx.ui.requestRender();
 				},
 				() => {
 					done();
+					this.ctx.goalModeController.scheduleContinuation();
 					this.ctx.ui.requestRender();
 				},
+				() => this.ctx.ui.terminal.rows,
 			);
 			return { component: selector, focus: selector.getMessageList() };
 		});
@@ -3274,6 +3417,9 @@ export class SelectorController {
 						});
 					}
 				},
+				undefined,
+				undefined,
+				(session, starred) => this.ctx.sessionManager.setSessionStarredForPicker(session, starred),
 			);
 			selector.setOnRequestRender(() => this.ctx.ui.requestRender());
 			return { component: selector, focus: selector };

package/src/modes/interactive-mode.ts

@@ -2265,10 +2265,24 @@ export class InteractiveMode implements InteractiveModeContext {
 	}
 
 	#prepareSessionSwitch(cleanupPreviousSessionUi?: () => void): void {
-		this.#btwController.dispose();
-		if (cleanupPreviousSessionUi) cleanupPreviousSessionUi();
-		else this.#extensionUiController.clearExtensionTerminalInputListeners();
-		this.#planModeController.clearReview();
+		const errors: unknown[] = [];
+		try {
+			this.#btwController.dispose();
+		} catch (error) {
+			errors.push(error);
+		}
+		try {
+			if (cleanupPreviousSessionUi) cleanupPreviousSessionUi();
+			else this.#extensionUiController.clearExtensionTerminalInputListeners();
+		} catch (error) {
+			errors.push(error);
+		}
+		try {
+			this.#planModeController.clearReview();
+		} catch (error) {
+			errors.push(error);
+		}
+		if (errors.length > 0) throw new AggregateError(errors, "Previous session UI cleanup failed");
 	}
 
 	async handleClearCommand(): Promise<boolean> {
@@ -2291,7 +2305,6 @@ export class InteractiveMode implements InteractiveModeContext {
 	}
 
 	handleForkCommand(): Promise<void> {
-		this.#btwController.dispose();
 		return this.#commandController.handleForkCommand();
 	}
 
@@ -2299,7 +2312,7 @@ export class InteractiveMode implements InteractiveModeContext {
 		return this.#commandController.handleMoveCommand(targetPath);
 	}
 
-	handleRenameCommand(title: string): Promise<void> {
+	handleRenameCommand(title?: string): Promise<void> {
 		return this.#commandController.handleRenameCommand(title);
 	}
 
@@ -2564,7 +2577,8 @@ export class InteractiveMode implements InteractiveModeContext {
 	}
 
 	showUserMessageSelector(): void {
-		this.#selectorController.showUserMessageSelector();
+		const cleanupPreviousSessionUi = this.#extensionUiController.captureSessionUiCleanup();
+		this.#selectorController.showUserMessageSelector(() => this.#prepareSessionSwitch(cleanupPreviousSessionUi));
 	}
 
 	showTreeSelector(): void {

package/src/modes/runtime-init.ts

@@ -123,6 +123,7 @@ export async function initializeExtensions(session: AgentSession, options: Initi
 			getCredentialSessionId: () => session.credentialSessionId,
 			isIdle: () => !session.isStreaming,
 			getActivePromptHandle: () => session.activePromptHandle,
+			getSessionWorkLease: () => session.sessionWorkLease,
 			abort: () => session.abort(),
 			abortPromptAndWait: (handle, abortOptions) => session.abortPromptAndWait(handle, abortOptions),
 			hasPendingMessages: () => session.queuedMessageCount > 0,