centaurus-cli 2.9.5 → 2.9.7

package/dist/cli-adapter.js

@@ -22,6 +22,7 @@ import { backgroundCommandTool } from './tools/background-command.js';
 import { subAgentTool } from './tools/sub-agent.js';
 import { enterRemoteSessionTool } from './tools/enter-remote-session.js';
 import { workflowTool } from './tools/workflow-tool.js';
+import { fastContextTool } from './tools/fast-context.js';
 import { SubAgentManager } from './services/sub-agent-manager.js';
 import { ShellInputAgent } from './services/shell-input-agent.js';
 import { apiClient } from './services/api-client.js';
@@ -46,6 +47,7 @@ import { BackgroundTaskManager } from './services/background-task-manager.js';
 import { sessionQuotaManager } from './services/session-quota-manager.js';
 import { ollamaService, OllamaService } from './services/ollama-service.js';
 import { workflowStorage } from './services/workflow-storage.js';
+import { CheckpointManager } from './services/checkpoint-manager.js';
 export class CentaurusCLI {
     configManager;
     toolRegistry;
@@ -116,9 +118,16 @@ export class CentaurusCLI {
     onPromptAnswered; // Callback when AI answers a shell prompt
     onShowWorkflowCreatorCallback; // Callback to show workflow creator screen with optional initial steps
     onWorkflowSaveCallback; // Callback when workflow is saved
+    onAiAutoSuggestChange; // Callback for AI auto-suggest setting changes
+    onRevertToCheckpointCallback; // Callback for revert UI update
+    // Checkpoint manager for revert functionality
+    checkpointManager;
+    currentCheckpointId; // Track current checkpoint being created
     // Workflow learning mode state
     workflowLearningActive = false;
     learnedWorkflowSteps = [];
+    // Callback to set input value (e.g., for revert)
+    onSetInputCallback = null;
     constructor() {
         this.configManager = new ConfigManager();
         this.toolRegistry = new ToolRegistry();
@@ -233,6 +242,24 @@ export class CentaurusCLI {
     setOnWorkflowSave(callback) {
         this.onWorkflowSaveCallback = callback;
     }
+    setOnAiAutoSuggestChange(callback) {
+        this.onAiAutoSuggestChange = callback;
+    }
+    setOnRevertToCheckpoint(callback) {
+        this.onRevertToCheckpointCallback = callback;
+    }
+    /**
+     * Get checkpoints for autocomplete dropdown
+     */
+    getCheckpointsForAutocomplete() {
+        if (!this.checkpointManager)
+            return [];
+        return this.checkpointManager.list().map(cp => ({
+            id: cp.id,
+            prompt: cp.prompt,
+            timestamp: new Date(cp.createdAt)
+        }));
+    }
     /**
      * Save a workflow from the workflow creator UI
      */
@@ -395,7 +422,10 @@ Begin executing now, starting with Step 1.`;
                 // starting 'fresh' in the remote user's home/default dir is safer for 'warpify' semantics regardless of type.
                 // However, to avoid regression for WSL/Docker usage where inheritance is expected, we restricts this fix to SSH.
                 const initialCwd = type === 'ssh' ? undefined : this.cwd;
-                context = await detection.handler.connect(command, initialCwd);
+                // Use factory method to create a new handler instance for this session
+                // This prevents state issues where the singleton handler's client is overwritten by nested sessions
+                const newHandler = detection.handler.createNew();
+                context = await newHandler.connect(command, initialCwd);
             }
             this.connectionCommandStack.push(command);
             this.contextManager.pushContext(context);
@@ -939,6 +969,7 @@ Begin executing now, starting with Step 1.`;
         this.toolRegistry.register(subAgentTool);
         this.toolRegistry.register(enterRemoteSessionTool);
         this.toolRegistry.register(workflowTool);
+        this.toolRegistry.register(fastContextTool);
         // Initialize SubAgentManager with tool registry
         SubAgentManager.initialize(this.toolRegistry);
         SubAgentManager.setOnSubAgentCountChange((count) => {
@@ -1279,12 +1310,104 @@ DO NOT use write_to_file, edit_file, or execute_command until the plan is approv
         // Start logging session and log user message
         conversationLogger.startSession();
         conversationLogger.logUserMessage(message);
+        // Initialize and start checkpoint for revert functionality
+        if (!this.checkpointManager) {
+            this.checkpointManager = new CheckpointManager();
+        }
+        // Ensure we have a chat ID for checkpoints (even for the first message)
+        if (!this.currentChatId) {
+            this.currentChatId = localChatStorage.generateChatId();
+        }
+        if (this.currentChatId) {
+            this.checkpointManager.setCurrentChatId(this.currentChatId);
+            this.currentCheckpointId = undefined;
+            // Start checkpoint snapshot for this user message
+            const currentContext = this.contextManager.getCurrentContext();
+            // Build remote session info and handler for non-local contexts
+            let remoteSessionInfo;
+            let remoteHandler;
+            if (currentContext.type !== 'local' && currentContext.handler) {
+                const metadata = currentContext.metadata;
+                remoteSessionInfo = {
+                    hostname: metadata.hostname,
+                    username: metadata.username,
+                    sessionId: currentContext.sessionId,
+                    connectionString: metadata.hostname
+                        ? `${metadata.username || 'user'}@${metadata.hostname}`
+                        : metadata.distroName || metadata.containerId || undefined,
+                };
+                // Use the handler if it implements the RemoteFileHandler interface
+                if (typeof currentContext.handler.readFile === 'function' &&
+                    typeof currentContext.handler.writeFile === 'function' &&
+                    typeof currentContext.handler.executeCommand === 'function' &&
+                    typeof currentContext.handler.isConnected === 'function') {
+                    remoteHandler = currentContext.handler;
+                }
+            }
+            try {
+                const checkpointStartPromise = this.checkpointManager.startCheckpoint({
+                    prompt: message,
+                    cwd: this.cwd,
+                    contextType: currentContext.type,
+                    conversationIndex: this.conversationHistory.length - 1,
+                    uiMessageIndex: this.uiMessageHistory.length,
+                    remoteSessionInfo,
+                    handler: remoteHandler,
+                });
+                let checkpoint = null;
+                if (currentContext.type !== 'local') {
+                    const remoteCheckpointTimeoutMs = 2000;
+                    const timeoutMarker = '__checkpoint_timeout__';
+                    let timeoutHandle = null;
+                    const timeoutPromise = new Promise((resolve) => {
+                        timeoutHandle = setTimeout(() => resolve(timeoutMarker), remoteCheckpointTimeoutMs);
+                    });
+                    const checkpointOrTimeout = await Promise.race([
+                        checkpointStartPromise,
+                        timeoutPromise,
+                    ]);
+                    if (timeoutHandle) {
+                        clearTimeout(timeoutHandle);
+                    }
+                    if (checkpointOrTimeout === timeoutMarker) {
+                        quickLog(`[${new Date().toISOString()}] [Checkpoint] Remote checkpoint start exceeded ${remoteCheckpointTimeoutMs}ms (${currentContext.type}); continuing without blocking AI turn\n`);
+                        void checkpointStartPromise
+                            .then((lateCheckpoint) => {
+                            if (!lateCheckpoint || !this.checkpointManager) {
+                                return;
+                            }
+                            this.checkpointManager.discardCheckpointById(lateCheckpoint.id);
+                            quickLog(`[${new Date().toISOString()}] [Checkpoint] Discarded late checkpoint ${lateCheckpoint.id} created after timeout\n`);
+                        })
+                            .catch((lateError) => {
+                            quickLog(`[${new Date().toISOString()}] [Checkpoint] Late checkpoint creation failed after timeout: ${lateError?.message || lateError}\n`);
+                        });
+                    }
+                    else {
+                        checkpoint = checkpointOrTimeout;
+                    }
+                }
+                else {
+                    checkpoint = await checkpointStartPromise;
+                }
+                if (checkpoint) {
+                    this.currentCheckpointId = checkpoint.id;
+                    quickLog(`[${new Date().toISOString()}] [Checkpoint] Started checkpoint ${checkpoint.id} (${currentContext.type}) for: "${message.slice(0, 50)}..."\n`);
+                }
+            }
+            catch (error) {
+                quickLog(`[${new Date().toISOString()}] [Checkpoint] Failed to start checkpoint: ${error.message}\n`);
+            }
+        }
         try {
             const tools = this.toolRegistry.getSchemas();
             const context = {
                 cwd: this.cwd,
                 contextManager: this.contextManager,
                 cliAdapter: this, // Pass CLI adapter reference for interactive process management
+                checkpointManager: this.checkpointManager, // For session-aware diff tool
+                currentCheckpointId: this.currentCheckpointId, // Active checkpoint for this request
+                currentChatId: this.currentChatId || undefined, // Current chat session ID (for session-wide diffs)
                 requireApproval: async (message, risky, preview, operationType, operationDetails) => {
                     // Special bypass for shell input to running processes:
                     // If the AI is sending input to an existing shell (via shell_input), we bypass the separate approval step.
@@ -1297,10 +1420,10 @@ DO NOT use write_to_file, edit_file, or execute_command until the plan is approv
                     }
                     return true; // Auto-approve if no callback registered
                 },
-                onStreamingOutput: (chunk, type) => {
+                onStreamingOutput: (chunk, type, toolName) => {
                     // Forward streaming output to UI callback
                     if (this.onToolStreamingOutput) {
-                        this.onToolStreamingOutput({ toolName: 'execute_command', chunk, type });
+                        this.onToolStreamingOutput({ toolName: toolName || 'execute_command', chunk, type });
                     }
                 },
             };
@@ -2391,6 +2514,17 @@ DO NOT use write_to_file, edit_file, or execute_command until the plan is approv
             throw new Error(`AI Error: ${error.message}`);
         }
         finally {
+            // Finalize checkpoint if one was created for this message
+            if (this.currentCheckpointId && this.checkpointManager) {
+                try {
+                    await this.checkpointManager.finalizeCheckpoint(this.currentCheckpointId);
+                    quickLog(`[${new Date().toISOString()}] [Checkpoint] Finalized checkpoint ${this.currentCheckpointId}\n`);
+                }
+                catch (error) {
+                    quickLog(`[${new Date().toISOString()}] [Checkpoint] Failed to finalize checkpoint: ${error.message}\n`);
+                }
+                this.currentCheckpointId = undefined;
+            }
             // Clean up abort controller
             this.currentAbortController = undefined;
         }
@@ -2813,12 +2947,18 @@ Start by listing the directory structure to understand what you're working with.
                     const value = args[1].toLowerCase();
                     if (value === 'on') {
                         this.configManager.set('aiAutoSuggest', true);
+                        if (this.onAiAutoSuggestChange) {
+                            this.onAiAutoSuggestChange(true);
+                        }
                         responseMessage = '✅ **AI Auto-Suggestions Enabled**\n\n' +
                             'From now on, I will suggest commands after 5 seconds of inactivity.\n' +
                             'Suggestions will appear in grey text. Use the **Right Arrow** key to accept them.';
                     }
                     else if (value === 'off') {
                         this.configManager.set('aiAutoSuggest', false);
+                        if (this.onAiAutoSuggestChange) {
+                            this.onAiAutoSuggestChange(false);
+                        }
                         responseMessage = '✅ **AI Auto-Suggestions Disabled**\n\n' +
                             'I will no longer provide AI-powered command suggestions.';
                     }
@@ -3075,6 +3215,189 @@ Then try /models local again.`;
             case 'exit':
                 process.exit(0);
                 break;
+            case 'revert': {
+                const checkpointArg = args[0];
+                if (!checkpointArg) {
+                    responseMessage = '❌ Usage: `/revert <checkpoint-id>`\nType `/revert ` and use autocomplete to select a checkpoint.';
+                    break;
+                }
+                if (!this.checkpointManager) {
+                    responseMessage = '❌ No checkpoints available in this session.';
+                    break;
+                }
+                try {
+                    // Smart context matching: determine if we can revert based on current vs checkpoint context
+                    const currentContext = this.contextManager.getCurrentContext();
+                    const allCheckpoints = this.checkpointManager.list();
+                    const targetCheckpoint = allCheckpoints.find(c => c.id === checkpointArg);
+                    if (!targetCheckpoint) {
+                        responseMessage = `❌ Checkpoint "${checkpointArg}" not found.`;
+                        break;
+                    }
+                    // Context validation for revert
+                    const cpContextType = targetCheckpoint.contextType;
+                    const currentContextType = currentContext.type;
+                    if (cpContextType !== 'local' && currentContextType === 'local') {
+                        // Remote checkpoint but we're local — can't reconnect to SSH to revert files
+                        const sessionType = cpContextType.toUpperCase();
+                        const sessionInfo = targetCheckpoint.remoteSessionInfo;
+                        const target = sessionInfo?.connectionString || sessionInfo?.hostname || 'the remote machine';
+                        responseMessage = `❌ This checkpoint was created during a ${sessionType} session (${target}).\n` +
+                            `You are not currently connected to that session.\n\n` +
+                            `Please reconnect to the ${sessionType} session first, then retry /revert.`;
+                        break;
+                    }
+                    if (cpContextType === 'local' && currentContextType !== 'local') {
+                        // Local checkpoint but we're in a remote session — wrong context
+                        responseMessage = `❌ This checkpoint was created in a local session.\n` +
+                            `You are currently in a ${currentContextType.toUpperCase()} session.\n\n` +
+                            `Please exit your remote session first, then retry /revert.`;
+                        break;
+                    }
+                    if (cpContextType !== 'local' && currentContextType !== 'local' && cpContextType !== currentContextType) {
+                        // Both remote but different types (e.g., SSH checkpoint but in Docker now)
+                        responseMessage = `❌ This checkpoint was created in a ${cpContextType.toUpperCase()} session, ` +
+                            `but you are currently in a ${currentContextType.toUpperCase()} session.\n\n` +
+                            `Please connect to the correct ${cpContextType.toUpperCase()} session first.`;
+                        break;
+                    }
+                    // For remote checkpoints with mismatched hosts, validate the connection target
+                    if (cpContextType !== 'local' && targetCheckpoint.remoteSessionInfo) {
+                        const cpHost = targetCheckpoint.remoteSessionInfo.hostname;
+                        const currentHost = currentContext.metadata.hostname;
+                        if (cpHost && currentHost && cpHost !== currentHost) {
+                            responseMessage = `❌ This checkpoint was created on ${cpHost}, ` +
+                                `but you are currently connected to ${currentHost}.\n\n` +
+                                `Please connect to ${cpHost} first, then retry /revert.`;
+                            break;
+                        }
+                    }
+                    // Build handler for remote revert if needed
+                    let revertHandler;
+                    if (cpContextType !== 'local' && currentContext.handler) {
+                        if (typeof currentContext.handler.readFile === 'function' &&
+                            typeof currentContext.handler.writeFile === 'function' &&
+                            typeof currentContext.handler.executeCommand === 'function' &&
+                            typeof currentContext.handler.isConnected === 'function') {
+                            revertHandler = currentContext.handler;
+                        }
+                    }
+                    const result = await this.checkpointManager.revertToCheckpoint(checkpointArg, revertHandler);
+                    // Find checkpoint index for UI truncation
+                    const checkpoints = this.checkpointManager.list();
+                    const checkpointIndex = checkpoints.findIndex(c => c.id === checkpointArg);
+                    // Truncate conversation history to the checkpoint
+                    if (checkpointIndex >= 0) {
+                        const checkpoint = result.checkpoint;
+                        // Populate the input bar with the reverted prompt
+                        if (this.onSetInputCallback) {
+                            this.onSetInputCallback(checkpoint.prompt);
+                        }
+                        // Use conversationIndex from metadata if available (robust)
+                        // This avoids issues with duplicate prompts getting truncated at the wrong occurrence
+                        if (typeof checkpoint.conversationIndex === 'number' &&
+                            checkpoint.conversationIndex >= 0 &&
+                            checkpoint.conversationIndex < this.conversationHistory.length) {
+                            // Truncate to index (exclusive) to remove the message from history
+                            // This allows the user to "edit" the message in the input bar
+                            this.conversationHistory = this.conversationHistory.slice(0, checkpoint.conversationIndex);
+                            quickLog(`[${new Date().toISOString()}] [Revert] Truncated conversation history to index ${checkpoint.conversationIndex} (exclusive)\n`);
+                        }
+                        else {
+                            // Fallback to string matching (legacy or missing index)
+                            const targetPrompt = checkpoint.prompt;
+                            // Find the user message with this prompt in conversationHistory
+                            let foundIndex = -1;
+                            // Use a broader search to find the matching message
+                            // The prompt might be a substring or exact match
+                            const searchPrompt = targetPrompt.slice(0, 50);
+                            for (let i = 0; i < this.conversationHistory.length; i++) {
+                                const msg = this.conversationHistory[i];
+                                if (msg.role === 'user' && msg.content && msg.content.includes(searchPrompt)) {
+                                    foundIndex = i;
+                                    break;
+                                }
+                            }
+                            if (foundIndex >= 0) {
+                                // Keep messages up to but NOT including this user message
+                                this.conversationHistory = this.conversationHistory.slice(0, foundIndex);
+                                quickLog(`[${new Date().toISOString()}] [Revert] Truncated conversation history to index ${foundIndex} (exclusive, using string match)\n`);
+                            }
+                        }
+                        // Fix: Also truncate UI message history to ensure consistency
+                        // Use uiMessageIndex from metadata if available (robust)
+                        if (typeof checkpoint.uiMessageIndex === 'number' &&
+                            checkpoint.uiMessageIndex >= 0 &&
+                            checkpoint.uiMessageIndex < this.uiMessageHistory.length) {
+                            // Truncate to index (exclusive)
+                            this.uiMessageHistory = this.uiMessageHistory.slice(0, checkpoint.uiMessageIndex);
+                            quickLog(`[${new Date().toISOString()}] [Revert] Truncated UI message history to index ${checkpoint.uiMessageIndex} (exclusive)\n`);
+                        }
+                        else {
+                            // Fallback to string matching (legacy or missing index)
+                            const targetPrompt = checkpoint.prompt;
+                            const searchPrompt = targetPrompt.slice(0, 50);
+                            let foundUiIndex = -1;
+                            for (let i = 0; i < this.uiMessageHistory.length; i++) {
+                                const msg = this.uiMessageHistory[i];
+                                if (msg.role === 'user' && msg.content && msg.content.includes(searchPrompt)) {
+                                    foundUiIndex = i;
+                                    break;
+                                }
+                            }
+                            if (foundUiIndex >= 0) {
+                                // Keep messages up to but NOT including this user message
+                                this.uiMessageHistory = this.uiMessageHistory.slice(0, foundUiIndex);
+                                quickLog(`[${new Date().toISOString()}] [Revert] Truncated UI message history to index ${foundUiIndex} (exclusive, using string match)\n`);
+                            }
+                        }
+                        // Remove the reverted checkpoint AND all checkpoints created after it
+                        // They should not appear in autocomplete anymore
+                        this.checkpointManager.removeCheckpointsFrom(checkpointArg);
+                        // Build the success message BEFORE calling handleChatPickerSelection
+                        // This will be passed as a parameter and included in the same React setState call
+                        quickLog(`[${new Date().toISOString()}] [Revert] Building success message for checkpoint "${checkpointArg}"\n`);
+                        const truncatedPrompt = result.checkpoint.prompt.length > 50
+                            ? result.checkpoint.prompt.slice(0, 50) + '...'
+                            : result.checkpoint.prompt;
+                        let revertSuccessMessage = `✅ Reverted to: "${truncatedPrompt}"\n` +
+                            `Restored ${result.restored} files, removed ${result.removed} files.`;
+                        if (result.errors.length > 0) {
+                            revertSuccessMessage += `\n⚠️ Warnings: ${result.errors.join(', ')}`;
+                        }
+                        // Save the truncated state to disk
+                        if (this.currentChatId) {
+                            quickLog(`[${new Date().toISOString()}] [Revert] Saving truncated chat to disk (chatId: ${this.currentChatId})\n`);
+                            quickLog(`[${new Date().toISOString()}] [Revert] uiMessageHistory.length BEFORE save: ${this.uiMessageHistory.length}\n`);
+                            this.saveCurrentChat();
+                            // Reload the chat to force UI update ensures consistency
+                            // This simulates a /chat resume command which correctly populates the UI
+                            // Pass skipLoadedMessage=true and the revert success message
+                            // The success message is appended to restoredMessages in handleChatPickerSelection
+                            // This ensures it's part of the same React setState call, avoiding race conditions
+                            quickLog(`[${new Date().toISOString()}] [Revert] Calling handleChatPickerSelection with revertSuccessMessage="${revertSuccessMessage.substring(0, 50)}..."\n`);
+                            const currentContext = this.contextManager.getCurrentContext();
+                            const preserveRemote = currentContext.type !== 'local' && typeof currentContext.handler?.isConnected === 'function'
+                                ? currentContext.handler.isConnected()
+                                : currentContext.type !== 'local';
+                            await this.handleChatPickerSelection(this.currentChatId, true, revertSuccessMessage, preserveRemote);
+                            // Clear responseMessage so it doesn't get sent via onDirectMessageCallback
+                            // (the message was already included in the restored messages above)
+                            quickLog(`[${new Date().toISOString()}] [Revert] handleChatPickerSelection completed, clearing responseMessage\n`);
+                            responseMessage = '';
+                        }
+                        else {
+                            // Fallback: if no current chat, show the message via normal callback
+                            quickLog(`[${new Date().toISOString()}] [Revert] No currentChatId, using fallback responseMessage\n`);
+                            responseMessage = revertSuccessMessage;
+                        }
+                    }
+                }
+                catch (error) {
+                    responseMessage = `❌ Failed to revert: ${error.message}`;
+                }
+                break;
+            }
             case 'add-command':
             case 'add-command-auto-detect':
                 // Handle custom command auto-detect management
@@ -3945,7 +4268,7 @@ Create once, run many times across different machines.`;
     /**
      * Load a chat from local storage and restore it
      */
-    loadChat(chatId) {
+    loadChat(chatId, options) {
         const chat = localChatStorage.loadChat(chatId);
         if (!chat) {
             return false;
@@ -3957,11 +4280,32 @@ Create once, run many times across different machines.`;
             tool_calls: msg.tool_calls,
             tool_call_id: msg.tool_call_id,
         }));
+        // IMPORTANT: Also restore UI message history from saved chat
+        // This ensures that when the chat is saved again (e.g., after revert), 
+        // the correct UI messages are preserved
+        if (chat.uiMessages && chat.uiMessages.length > 0) {
+            this.uiMessageHistory = chat.uiMessages.map(msg => ({
+                id: msg.id,
+                role: msg.role,
+                content: msg.content,
+                timestamp: msg.timestamp ? new Date(msg.timestamp) : undefined,
+                toolExecution: msg.toolExecution,
+                shouldStream: false,
+                isCommandMode: msg.isCommandMode,
+                tool_call_id: msg.tool_call_id,
+                tool_calls: msg.tool_calls,
+                thinkingDuration: msg.thinkingDuration,
+                taskCompletion: msg.taskCompletion,
+                planAccepted: msg.planAccepted,
+                connectionStatus: msg.connectionStatus,
+            }));
+        }
         // Set the current chat ID to continue the conversation
         this.currentChatId = chatId;
         this.conversationStarted = true;
         // Restore CWD if saved (important for commands to run in correct directory)
-        if (chat.cwd) {
+        // Allow skipping when preserving an active remote session
+        if (chat.cwd && !options?.preserveCwd) {
             this.cwd = chat.cwd;
             // Update context manager if available
             if (this.contextManager) {
@@ -3993,8 +4337,12 @@ Create once, run many times across different machines.`;
     }
     /**
      * Handle chat picker selection
+     * @param chatId The ID of the chat to load
+     * @param skipLoadedMessage If true, skip showing the "Loaded chat" message (used during revert to avoid React setState race condition)
+     * @param revertSuccessMessage If provided, this message will be appended to restoredMessages before calling onRestoreMessagesCallback (used during revert to include success message in the same React setState call, avoiding race conditions)
+     * @param preserveRemoteSession If true, keep current SSH/WSL/Docker connection and skip reconnect logic (used during revert in a live remote session)
      */
-    async handleChatPickerSelection(chatId) {
+    async handleChatPickerSelection(chatId, skipLoadedMessage = false, revertSuccessMessage, preserveRemoteSession = false) {
         const chat = localChatStorage.loadChat(chatId);
         if (!chat) {
             if (this.onResponseCallback) {
@@ -4005,7 +4353,8 @@ Create once, run many times across different machines.`;
         // IMPORTANT: Clean up current remote session before loading new chat
         // This ensures that switching from a remote chat to a local chat properly resets the state
         const currentContext = this.contextManager.getCurrentContext();
-        if (currentContext.type !== 'local') {
+        const shouldPreserveRemote = preserveRemoteSession && currentContext.type !== 'local';
+        if (!shouldPreserveRemote && currentContext.type !== 'local') {
             // Disconnect from current remote session
             if (currentContext.handler) {
                 try {
@@ -4022,8 +4371,18 @@ Create once, run many times across different machines.`;
             this.connectionCommandStack = [];
         }
         // Load AI context
-        this.loadChat(chatId);
+        this.loadChat(chatId, shouldPreserveRemote ? { preserveCwd: true } : undefined);
+        // If we're preserving a live remote session (revert case), restore the current remote CWD
+        if (shouldPreserveRemote) {
+            this.cwd = currentContext.metadata.workingDirectory;
+            this.contextManager.updateWorkingDirectory(this.cwd);
+            if (this.onCwdChange) {
+                this.onCwdChange(this.cwd);
+            }
+        }
         // Restore UI messages if available
+        quickLog(`[${new Date().toISOString()}] [handleChatPickerSelection] chat.uiMessages exists: ${!!chat.uiMessages}, length: ${chat.uiMessages?.length ?? 0}\n`);
+        quickLog(`[${new Date().toISOString()}] [handleChatPickerSelection] onRestoreMessagesCallback exists: ${!!this.onRestoreMessagesCallback}\n`);
         if (chat.uiMessages && chat.uiMessages.length > 0 && this.onRestoreMessagesCallback) {
             // Convert StoredUIMessage back to Message format
             const restoredMessages = chat.uiMessages.map(msg => ({
@@ -4041,11 +4400,29 @@ Create once, run many times across different machines.`;
                 planAccepted: msg.planAccepted,
                 connectionStatus: msg.connectionStatus, // For SSH/WSL/Docker connection status boxes
             }));
+            // If a revert success message was provided, append it to the restored messages
+            // This ensures the message is part of the same React setState call, avoiding race conditions
+            quickLog(`[${new Date().toISOString()}] [handleChatPickerSelection] revertSuccessMessage provided: ${!!revertSuccessMessage}\n`);
+            if (revertSuccessMessage) {
+                const revertSystemMessage = {
+                    id: `revert-success-${Date.now()}-${Math.random().toString(36).substring(2, 11)}`,
+                    role: 'assistant',
+                    content: revertSuccessMessage,
+                    timestamp: new Date(),
+                    shouldStream: false,
+                };
+                restoredMessages.push(revertSystemMessage);
+                quickLog(`[${new Date().toISOString()}] [handleChatPickerSelection] Appended revert success message, restoredMessages.length: ${restoredMessages.length}\n`);
+            }
+            quickLog(`[${new Date().toISOString()}] [handleChatPickerSelection] Calling onRestoreMessagesCallback with ${restoredMessages.length} messages\n`);
             this.onRestoreMessagesCallback(restoredMessages);
+            quickLog(`[${new Date().toISOString()}] [handleChatPickerSelection] onRestoreMessagesCallback completed\n`);
         }
         // Attempt to restore remote context if chat was saved while in remote environment
         // Support both remoteContextStack (for nested sessions) and single remoteContext (backward compat)
-        const contextStackToRestore = chat.remoteContextStack ?? (chat.remoteContext ? [chat.remoteContext] : null);
+        const contextStackToRestore = shouldPreserveRemote
+            ? null
+            : (chat.remoteContextStack ?? (chat.remoteContext ? [chat.remoteContext] : null));
         if (contextStackToRestore && contextStackToRestore.length > 0) {
             // Get the base local CWD from the first context's localCwdBeforeRemote
             const baseLocalCwd = contextStackToRestore[0].localCwdBeforeRemote;
@@ -4082,12 +4459,15 @@ Create once, run many times across different machines.`;
                     if (!detection) {
                         throw new Error(`Could not detect handler for: ${connectionCommand}`);
                     }
+                    // Create a new instance of the handler to ensure each level of the connection
+                    // has its own state (client, stream, etc.), critical for nested sessions of the same type
+                    const handler = detection.handler.createNew();
                     let context;
                     // Check if this is a nested connection (i > 0 means we're inside a remote session)
                     if (i > 0) {
                         const currentContext = this.contextManager.getCurrentContext();
-                        if (detection.handler.connectFromRemote) {
-                            context = await detection.handler.connectFromRemote(connectionCommand, previousCwd, currentContext);
+                        if (handler.connectFromRemote) {
+                            context = await handler.connectFromRemote(connectionCommand, previousCwd, currentContext);
                         }
                         else {
                             throw new Error(`Nested connections not supported by ${type} handler`);
@@ -4095,7 +4475,7 @@ Create once, run many times across different machines.`;
                     }
                     else {
                         // First level: connect from local
-                        context = await detection.handler.connect(connectionCommand, previousCwd);
+                        context = await handler.connect(connectionCommand, previousCwd);
                     }
                     // Push to stacks
                     this.connectionCommandStack.push(connectionCommand);
@@ -4167,9 +4547,8 @@ Create once, run many times across different machines.`;
             }
             return;
         }
-        // No remote context - show regular confirmation
         // No remote context - show regular confirmation and restore CWD
-        if (chat.cwd && !chat.remoteContext) {
+        if (!shouldPreserveRemote && chat.cwd && !chat.remoteContext) {
             if (fs.existsSync(chat.cwd)) {
                 this.cwd = chat.cwd;
                 this.contextManager.updateWorkingDirectory(chat.cwd);
@@ -4178,7 +4557,8 @@ Create once, run many times across different machines.`;
                 }
             }
         }
-        if (this.onDirectMessageCallback) {
+        // Skip loaded message if called from revert (to avoid React setState race condition)
+        if (this.onDirectMessageCallback && !skipLoadedMessage) {
             const responseMessage = `✅ Loaded chat: "${chat.title}"\n\nYou have ${chat.messageCount} messages in AI context. Continue your conversation!`;
             this.onDirectMessageCallback(responseMessage);
         }
@@ -4210,12 +4590,19 @@ Create once, run many times across different machines.`;
      * Start a new chat session
      */
     startNewChat() {
+        const currentContext = this.contextManager.getCurrentContext();
         this.conversationHistory = [];
         this.currentChatId = null;
         this.conversationStarted = false;
         this.uiMessageHistory = [];
-        this.cwdStack = [];
-        this.connectionCommandStack = [];
+        // Only clear stacks if we're in local context
+        // If we're in a remote context (SSH/WSL/Docker), we want to PRESERVE the connection state
+        // so the new chat continues in the same environment
+        if (currentContext.type === 'local') {
+            this.cwdStack = [];
+            this.connectionCommandStack = [];
+        }
+        // Else: Preserve cwdStack and connectionCommandStack for remote sessions
         // Reset context limit state
         if (this.contextLimitReached) {
             this.contextLimitReached = false;
@@ -4256,6 +4643,12 @@ Create once, run many times across different machines.`;
     setOnRestoreMessagesCallback(callback) {
         this.onRestoreMessagesCallback = callback;
     }
+    /**
+     * Set callback for setting input value (e.g., for revert)
+     */
+    setOnSetInput(callback) {
+        this.onSetInputCallback = callback;
+    }
     /**
      * Get environment context for backend
      * Returns structured environment information to be sent to backend