feat: add ability to stop Claude execution mid-way using loading icon as cancel button

2025-06-23 00:45:58 +05:30
parent abe0891b0b
commit 3dc741fd6b
5 changed files with 257 additions and 97 deletions
--- a/src-tauri/src/commands/claude.rs
+++ b/src-tauri/src/commands/claude.rs
@@ -6,10 +6,25 @@ use std::time::SystemTime;
 use std::io::{BufRead, BufReader};
 use std::process::Stdio;
 use tauri::{AppHandle, Emitter, Manager};
-use tokio::process::Command;
+use tokio::process::{Command, Child};
+use tokio::sync::Mutex;
+use std::sync::Arc;
 use crate::process::ProcessHandle;
 use crate::checkpoint::{CheckpointResult, CheckpointDiff, SessionTimeline, Checkpoint};

+/// Global state to track current Claude process
+pub struct ClaudeProcessState {
+    pub current_process: Arc<Mutex<Option<Child>>>,
+}
+
+impl Default for ClaudeProcessState {
+    fn default() -> Self {
+        Self {
+            current_process: Arc::new(Mutex::new(None)),
+        }
+    }
+}
+
 /// Represents a project in the ~/.claude/projects directory
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Project {
@@ -925,6 +940,41 @@ pub async fn resume_claude_code(
    spawn_claude_process(app, cmd).await
 }

+/// Cancel the currently running Claude Code execution
+#[tauri::command]
+pub async fn cancel_claude_execution(app: AppHandle) -> Result<(), String> {
+    log::info!("Cancelling Claude Code execution");
+    
+    let claude_state = app.state::<ClaudeProcessState>();
+    let mut current_process = claude_state.current_process.lock().await;
+    
+    if let Some(mut child) = current_process.take() {
+        // Try to get the PID before killing
+        let pid = child.id();
+        log::info!("Attempting to kill Claude process with PID: {:?}", pid);
+        
+        // Kill the process
+        match child.kill().await {
+            Ok(_) => {
+                log::info!("Successfully killed Claude process");
+                // Emit cancellation event
+                let _ = app.emit("claude-cancelled", true);
+                // Also emit complete with false to indicate failure
+                tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+                let _ = app.emit("claude-complete", false);
+                Ok(())
+            }
+            Err(e) => {
+                log::error!("Failed to kill Claude process: {}", e);
+                Err(format!("Failed to kill Claude process: {}", e))
+            }
+        }
+    } else {
+        log::warn!("No active Claude process to cancel");
+        Ok(())
+    }
+}
+
 /// Helper function to check if sandboxing should be used based on settings
 fn should_use_sandbox(app: &AppHandle) -> Result<bool, String> {
    // First check if sandboxing is even available on this platform
@@ -1097,10 +1147,21 @@ async fn spawn_claude_process(app: AppHandle, mut cmd: Command) -> Result<(), St
    let stdout = child.stdout.take().ok_or("Failed to get stdout")?;
    let stderr = child.stderr.take().ok_or("Failed to get stderr")?;
    
+    // Get the child PID for logging
+    let pid = child.id();
+    log::info!("Spawned Claude process with PID: {:?}", pid);
+    
    // Create readers
    let stdout_reader = BufReader::new(stdout);
    let stderr_reader = BufReader::new(stderr);
    
+    // Store the child process in the global state
+    let claude_state = app.state::<ClaudeProcessState>();
+    {
+        let mut current_process = claude_state.current_process.lock().await;
+        *current_process = Some(child);
+    }
+    
    // Spawn tasks to read stdout and stderr
    let app_handle = app.clone();
    let stdout_task = tokio::spawn(async move {
@@ -1123,24 +1184,33 @@ async fn spawn_claude_process(app: AppHandle, mut cmd: Command) -> Result<(), St
    });
    
    // Wait for the process to complete
+    let app_handle_wait = app.clone();
+    let claude_state_wait = claude_state.current_process.clone();
    tokio::spawn(async move {
        let _ = stdout_task.await;
        let _ = stderr_task.await;
        
-        match child.wait().await {
-            Ok(status) => {
-                log::info!("Claude process exited with status: {}", status);
-                // Add a small delay to ensure all messages are processed
-                tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
-                let _ = app.emit("claude-complete", status.success());
-            }
-            Err(e) => {
-                log::error!("Failed to wait for Claude process: {}", e);
-                // Add a small delay to ensure all messages are processed
-                tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
-                let _ = app.emit("claude-complete", false);
+        // Get the child from the state to wait on it
+        let mut current_process = claude_state_wait.lock().await;
+        if let Some(mut child) = current_process.take() {
+            match child.wait().await {
+                Ok(status) => {
+                    log::info!("Claude process exited with status: {}", status);
+                    // Add a small delay to ensure all messages are processed
+                    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+                    let _ = app_handle_wait.emit("claude-complete", status.success());
+                }
+                Err(e) => {
+                    log::error!("Failed to wait for Claude process: {}", e);
+                    // Add a small delay to ensure all messages are processed
+                    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+                    let _ = app_handle_wait.emit("claude-complete", false);
+                }
            }
        }
+        
+        // Clear the process from state
+        *current_process = None;
    });
    
    Ok(())
--- a/src-tauri/src/main.rs
+++ b/src-tauri/src/main.rs
@@ -17,7 +17,7 @@ use commands::claude::{
    get_session_timeline, update_checkpoint_settings, get_checkpoint_diff,
    track_checkpoint_message, track_session_messages, check_auto_checkpoint, cleanup_old_checkpoints,
    get_checkpoint_settings, clear_checkpoint_manager, get_checkpoint_state_stats,
-    get_recently_modified_files,
+    get_recently_modified_files, cancel_claude_execution, ClaudeProcessState,
 };
 use commands::agents::{
    init_database, list_agents, create_agent, update_agent, delete_agent, 
@@ -93,6 +93,9 @@ fn main() {
            // Initialize process registry
            app.manage(ProcessRegistryState::default());
            
+            // Initialize Claude process state
+            app.manage(ClaudeProcessState::default());
+            
            Ok(())
        })
        .invoke_handler(tauri::generate_handler![
@@ -111,6 +114,7 @@ fn main() {
            execute_claude_code,
            continue_claude_code,
            resume_claude_code,
+            cancel_claude_execution,
            list_directory_contents,
            search_files,
            create_checkpoint,
--- a/src/components/ClaudeCodeSession.tsx
+++ b/src/components/ClaudeCodeSession.tsx
@@ -9,7 +9,8 @@ import {
  ChevronDown,
  GitBranch,
  Settings,
-  Globe
+  Globe,
+  Square
 } from "lucide-react";
 import { Button } from "@/components/ui/button";
 import { Input } from "@/components/ui/input";
@@ -84,6 +85,7 @@ export const ClaudeCodeSession: React.FC<ClaudeCodeSessionProps> = ({
  const [showForkDialog, setShowForkDialog] = useState(false);
  const [forkCheckpointId, setForkCheckpointId] = useState<string | null>(null);
  const [forkSessionName, setForkSessionName] = useState("");
+  const [isCancelling, setIsCancelling] = useState(false);
  
  // New state for preview feature
  const [showPreview, setShowPreview] = useState(false);
@@ -278,6 +280,7 @@ export const ClaudeCodeSession: React.FC<ClaudeCodeSessionProps> = ({
      const completeUnlisten = await listen<boolean>("claude-complete", async (event) => {
        console.log('[ClaudeCodeSession] Received claude-complete:', event.payload);
        setIsLoading(false);
+        setIsCancelling(false);
        hasActiveSessionRef.current = false;
        if (!event.payload) {
          setError("Claude execution failed");
@@ -437,6 +440,40 @@ export const ClaudeCodeSession: React.FC<ClaudeCodeSessionProps> = ({
    setTimelineVersion((v) => v + 1);
  };

+  const handleCancelExecution = async () => {
+    if (!isLoading || isCancelling) return;
+    
+    try {
+      setIsCancelling(true);
+      
+      // Cancel the Claude execution
+      await api.cancelClaudeExecution();
+      
+      // Clean up listeners
+      unlistenRefs.current.forEach(unlisten => unlisten());
+      unlistenRefs.current = [];
+      
+      // Add a system message indicating cancellation
+      const cancelMessage: ClaudeStreamMessage = {
+        type: "system",
+        subtype: "cancelled",
+        result: "Execution cancelled by user",
+        timestamp: new Date().toISOString()
+      };
+      setMessages(prev => [...prev, cancelMessage]);
+      
+      // Reset states
+      setIsLoading(false);
+      hasActiveSessionRef.current = false;
+      setError(null);
+    } catch (err) {
+      console.error("Failed to cancel execution:", err);
+      setError("Failed to cancel execution");
+    } finally {
+      setIsCancelling(false);
+    }
+  };
+
  const handleFork = (checkpointId: string) => {
    setForkCheckpointId(checkpointId);
    setForkSessionName(`Fork-${new Date().toISOString().slice(0, 10)}`);
@@ -817,6 +854,42 @@ export const ClaudeCodeSession: React.FC<ClaudeCodeSessionProps> = ({
              {messagesList}
            </div>
          )}
+
+          {isLoading && enhancedMessages.length === 0 && (
+            <div className="flex items-center justify-center h-full">
+              <div className="flex items-center gap-3">
+                <Loader2 className="h-6 w-6 animate-spin" />
+                <span className="text-sm text-muted-foreground">
+                  {session ? "Loading session history..." : "Initializing Claude Code..."}
+                </span>
+              </div>
+            </div>
+          )}
+
+          <AnimatePresence>
+            {enhancedMessages.map((message, index) => (
+              <motion.div
+                key={index}
+                initial={{ opacity: 0, y: 10 }}
+                animate={{ opacity: 1, y: 0 }}
+                transition={{ duration: 0.2 }}
+              >
+                <ErrorBoundary>
+                  <StreamMessage message={message} streamMessages={enhancedMessages} />
+                </ErrorBoundary>
+              </motion.div>
+            ))}
+          </AnimatePresence>
+          
+          {/* Show loading indicator when processing, even if there are messages */}
+          {isLoading && enhancedMessages.length > 0 && (
+            <div className="flex items-center gap-2 p-4">
+              <Loader2 className="h-4 w-4 animate-spin" />
+              <span className="text-sm text-muted-foreground">
+                {isCancelling ? "Cancelling..." : "Processing..."}
+              </span>
+            </div>
+          )}
        </div>

        {/* Floating Prompt Input - Always visible */}
@@ -824,6 +897,7 @@ export const ClaudeCodeSession: React.FC<ClaudeCodeSessionProps> = ({
          <FloatingPromptInput
            ref={floatingPromptRef}
            onSend={handleSendPrompt}
+            onCancel={handleCancelExecution}
            isLoading={isLoading}
            disabled={!projectPath}
            projectPath={projectPath}
@@ -844,14 +918,6 @@ export const ClaudeCodeSession: React.FC<ClaudeCodeSessionProps> = ({
        )}
      </div>

-      {/* Preview Prompt Dialog */}
-      <PreviewPromptDialog
-        isOpen={showPreviewPrompt}
-        url={detectedUrl}
-        onConfirm={handleOpenPreview}
-        onCancel={() => setShowPreviewPrompt(false)}
-      />
-
      {/* Fork Dialog */}
      <Dialog open={showForkDialog} onOpenChange={setShowForkDialog}>
        <DialogContent>
--- a/src/components/FloatingPromptInput.tsx
+++ b/src/components/FloatingPromptInput.tsx
@@ -6,7 +6,8 @@ import {
  Minimize2,
  ChevronUp,
  Sparkles,
-  Zap
+  Zap,
+  Square
 } from "lucide-react";
 import { cn } from "@/lib/utils";
 import { Button } from "@/components/ui/button";
@@ -42,6 +43,10 @@ interface FloatingPromptInputProps {
   * Optional className for styling
   */
  className?: string;
+  /**
+   * Callback when cancel is clicked (only during loading)
+   */
+  onCancel?: () => void;
 }

 export interface FloatingPromptInputRef {
@@ -81,14 +86,18 @@ const MODELS: Model[] = [
 *   isLoading={false}
 * />
 */
-export const FloatingPromptInput = React.forwardRef<FloatingPromptInputRef, FloatingPromptInputProps>(({
-  onSend,
-  isLoading = false,
-  disabled = false,
-  defaultModel = "sonnet",
-  projectPath,
-  className,
-}, ref) => {
+const FloatingPromptInputInner = (
+  {
+    onSend,
+    isLoading = false,
+    disabled = false,
+    defaultModel = "sonnet",
+    projectPath,
+    className,
+    onCancel,
+  }: FloatingPromptInputProps,
+  ref: React.Ref<FloatingPromptInputRef>,
+) => {
  const [prompt, setPrompt] = useState("");
  const [selectedModel, setSelectedModel] = useState<"sonnet" | "opus">(defaultModel);
  const [isExpanded, setIsExpanded] = useState(false);
@@ -444,21 +453,18 @@ export const FloatingPromptInput = React.forwardRef<FloatingPromptInputRef, Floa
                  </Button>
                </div>

-                {isLoading ? (
-                  <div className="flex items-center justify-center min-w-[60px] h-10">
-                    <div className="rotating-symbol text-primary text-2xl"></div>
-                  </div>
-                ) : (
-                  <Button
-                    onClick={handleSend}
-                    disabled={!prompt.trim() || disabled}
-                    size="sm"
-                    className="min-w-[80px]"
-                  >
-                    <Send className="mr-2 h-4 w-4" />
-                    Send
-                  </Button>
-                )}
+                <Button
+                  onClick={handleSend}
+                  disabled={!prompt.trim() || isLoading || disabled}
+                  size="default"
+                  className="min-w-[60px]"
+                >
+                  {isLoading ? (
+                    <div className="rotating-symbol text-primary-foreground" />
+                  ) : (
+                    <Send className="h-4 w-4" />
+                  )}
+                </Button>
              </div>
            </motion.div>
          </motion.div>
@@ -574,21 +580,23 @@ export const FloatingPromptInput = React.forwardRef<FloatingPromptInputRef, Floa
                </AnimatePresence>
              </div>

-              {/* Send Button */}
-              {isLoading ? (
-                <div className="flex items-center justify-center min-w-[60px] h-10">
-                  <div className="rotating-symbol text-primary text-2xl"></div>
-                </div>
-              ) : (
-                <Button
-                  onClick={handleSend}
-                  disabled={!prompt.trim() || disabled}
-                  size="default"
-                  className="min-w-[60px]"
-                >
+              {/* Send/Stop Button */}
+              <Button
+                onClick={isLoading ? onCancel : handleSend}
+                disabled={isLoading ? false : (!prompt.trim() || disabled)}
+                variant={isLoading ? "destructive" : "default"}
+                size="default"
+                className="min-w-[60px]"
+              >
+                {isLoading ? (
+                  <>
+                    <Square className="h-4 w-4 mr-1" />
+                    Stop
+                  </>
+                ) : (
                  <Send className="h-4 w-4" />
-                </Button>
-              )}
+                )}
+              </Button>
            </div>

            <div className="mt-2 text-xs text-muted-foreground">
@@ -599,6 +607,11 @@ export const FloatingPromptInput = React.forwardRef<FloatingPromptInputRef, Floa
      </div>
    </>
  );
-});
+};
+
+export const FloatingPromptInput = React.forwardRef<
+  FloatingPromptInputRef,
+  FloatingPromptInputProps
+>(FloatingPromptInputInner);

 FloatingPromptInput.displayName = 'FloatingPromptInput';
--- a/src/lib/api.ts
+++ b/src/lib/api.ts
@@ -918,6 +918,13 @@ export const api = {
    return invoke("resume_claude_code", { projectPath, sessionId, prompt, model });
  },

+  /**
+   * Cancels the currently running Claude Code execution
+   */
+  async cancelClaudeExecution(): Promise<void> {
+    return invoke("cancel_claude_execution");
+  },
+
  /**
   * Lists files and directories in a given path
   */