fix the issue of losing context when continuing a conversation

Author: Blushyes

cli/src/interactive/app.rs

@@ -14,7 +14,8 @@ use coro_core::ResolvedLlmConfig;
 use iocraft::prelude::*;
 use regex::Regex;
 use std::path::{Path, PathBuf};
-use tokio::sync::broadcast;
+use std::sync::Arc;
+use tokio::sync::{broadcast, Mutex};
 
 /// Represents a file reference found in user input
 #[derive(Debug, Clone)]
@@ -269,14 +270,15 @@ pub fn submit_task_with_file_processing(
     llm_config: ResolvedLlmConfig,
     project_path: PathBuf,
     ui_sender: broadcast::Sender<AppMessage>,
+    agent: Arc<Mutex<Option<coro_core::agent::AgentCore>>>,
 ) {
-    use crate::interactive::components::input_section::spawn_ui_agent_task;
     use crate::interactive::message_handler::get_random_status_word;
 
     // Process file references asynchronously and send combined message
     let ui_sender_clone = ui_sender.clone();
     let llm_config_clone = llm_config.clone();
     let project_path_clone = project_path.clone();
+    let agent_clone = agent.clone();
 
     tokio::spawn(async move {
         let input_clone = input.clone();
@@ -297,12 +299,13 @@ pub fn submit_task_with_file_processing(
                     operation: get_random_status_word(),
                 });
 
-                // Use the existing spawn_ui_agent_task with enhanced input
-                spawn_ui_agent_task(
+                // Use the enhanced spawn_ui_agent_task_with_context with enhanced input
+                crate::interactive::components::input_section::spawn_ui_agent_task_with_context(
                     enhanced_input,
                     llm_config_clone,
                     project_path_clone,
                     ui_sender_clone,
+                    agent_clone.clone(),
                 );
             }
             Err(e) => {
@@ -317,25 +320,28 @@ pub fn submit_task_with_file_processing(
                 });
 
                 // Fall back to original input
-                spawn_ui_agent_task(
+                crate::interactive::components::input_section::spawn_ui_agent_task_with_context(
                     input_clone,
                     llm_config_clone,
                     project_path_clone,
                     ui_sender_clone,
+                    agent_clone,
                 );
             }
         }
     });
 }
 
 /// Context for interactive mode - immutable application configuration
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 struct AppContext {
     llm_config: ResolvedLlmConfig,
     project_path: PathBuf,
     ui_sender: broadcast::Sender<AppMessage>,
     ui_anim: UiAnimationConfig,
     debug_model: bool,
+    // Persistent agent instance for conversation continuity
+    agent: Arc<Mutex<Option<coro_core::agent::AgentCore>>>,
 }
 
 impl AppContext {
@@ -353,6 +359,7 @@ impl AppContext {
             ui_sender,
             ui_anim,
             debug_model,
+            agent: Arc::new(Mutex::new(None)),
         }
     }
 }
@@ -542,6 +549,7 @@ fn CoroApp(mut hooks: Hooks) -> impl Into<AnyElement<'static>> {
         llm_config: app_context.llm_config.clone(),
         project_path: app_context.project_path.clone(),
         ui_sender: app_context.ui_sender.clone(),
+        agent: app_context.agent.clone(),
     };
 
     // Create router configuration with main page using new API

cli/src/interactive/components/input_section.rs

@@ -16,8 +16,9 @@ use coro_core::ResolvedLlmConfig;
 use iocraft::prelude::*;
 use std::cmp::min;
 use std::path::PathBuf;
+use std::sync::Arc;
 use std::time::{Duration, Instant};
-use tokio::sync::broadcast;
+use tokio::sync::{broadcast, Mutex};
 use unicode_width::UnicodeWidthStr;
 
 /// Find the nearest character boundary at or before the given byte position
@@ -117,17 +118,19 @@ impl Default for InputSectionProps {
                 ),
                 project_path: PathBuf::new(),
                 ui_sender: tokio::sync::broadcast::channel(1).0,
+                agent: Arc::new(Mutex::new(None)),
             },
         }
     }
 }
 
 /// Context for the input section component
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 pub struct InputSectionContext {
     pub llm_config: ResolvedLlmConfig,
     pub project_path: PathBuf,
     pub ui_sender: broadcast::Sender<AppMessage>,
+    pub agent: Arc<Mutex<Option<coro_core::agent::AgentCore>>>,
 }
 
 /// Enhanced text input component that wraps iocraft's TextInput with submit handling
@@ -912,6 +915,69 @@ pub fn spawn_ui_agent_task(
     });
 }
 
+/// Spawn agent task execution with persistent agent for conversation continuity
+pub fn spawn_ui_agent_task_with_context(
+    input: String,
+    llm_config: ResolvedLlmConfig,
+    project_path: PathBuf,
+    ui_sender: broadcast::Sender<AppMessage>,
+    agent: Arc<Mutex<Option<coro_core::agent::AgentCore>>>,
+) {
+    use crate::interactive::message_handler::get_random_status_word;
+    use crate::interactive::task_executor::execute_agent_task_with_context;
+
+    // Start with a random status word
+    let _ = ui_sender.send(AppMessage::AgentTaskStarted {
+        operation: get_random_status_word(),
+    });
+
+    // Create a cancellation token for the timer
+    let (cancel_sender, mut cancel_receiver) = tokio::sync::oneshot::channel::<()>();
+
+    // Change status word once after 1 second (unless cancelled)
+    let ui_sender_timer = ui_sender.clone();
+    tokio::spawn(async move {
+        tokio::select! {
+            _ = tokio::time::sleep(tokio::time::Duration::from_secs(1)) => {
+                let _ = ui_sender_timer.send(AppMessage::AgentTaskStarted {
+                    operation: get_random_status_word(),
+                });
+            }
+            _ = &mut cancel_receiver => {
+                // Timer cancelled, do nothing
+            }
+        }
+    });
+
+    // Execute agent task with persistent context
+    tokio::spawn(async move {
+        match execute_agent_task_with_context(
+            input,
+            llm_config,
+            project_path,
+            ui_sender.clone(),
+            agent,
+        )
+        .await
+        {
+            Ok(_) => {
+                let _ = cancel_sender.send(()); // Cancel the timer
+                let _ = ui_sender.send(AppMessage::AgentExecutionCompleted);
+            }
+            Err(e) => {
+                let _ = cancel_sender.send(()); // Cancel the timer
+                                                // Check if it's an interruption error
+                if e.to_string().contains("Task interrupted by user") {
+                    // Don't show error message for user interruptions
+                } else {
+                    let _ = ui_sender.send(AppMessage::SystemMessage(format!("Error: {}", e)));
+                }
+                let _ = ui_sender.send(AppMessage::AgentExecutionCompleted);
+            }
+        }
+    });
+}
+
 /// Input Section Component - Fixed bottom area for input and status
 #[component]
 pub fn InputSection(mut hooks: Hooks, props: &InputSectionProps) -> impl Into<AnyElement<'static>> {
@@ -1105,6 +1171,7 @@ pub fn InputSection(mut hooks: Hooks, props: &InputSectionProps) -> impl Into<An
                     let ui_sender = ui_sender.clone();
                     let llm_config = llm_config.clone();
                     let project_path = project_path.clone();
+                    let agent = context.agent.clone();
                     move |input: String| {
                         if input.trim().is_empty() {
                             return;
@@ -1137,6 +1204,7 @@ pub fn InputSection(mut hooks: Hooks, props: &InputSectionProps) -> impl Into<An
                             llm_config.clone(),
                             project_path.clone(),
                             ui_sender.clone(),
+                            agent.clone(),
                         );
                     }
                 },

cli/src/interactive/pages/main_page.rs

@@ -31,6 +31,7 @@ impl Default for MainPageProps {
                 ),
                 project_path: std::path::PathBuf::from("."),
                 ui_sender: tokio::sync::broadcast::channel(1).0,
+                agent: std::sync::Arc::new(tokio::sync::Mutex::new(None)),
             },
         }
     }

cli/src/interactive/task_executor.rs

@@ -75,6 +75,113 @@ impl coro_core::output::AgentOutput for TokenTrackingOutputHandler {
     }
 }
 
+/// Execute agent task with persistent agent to maintain conversation context
+pub async fn execute_agent_task_with_context(
+    task: String,
+    llm_config: ResolvedLlmConfig,
+    project_path: PathBuf,
+    ui_sender: broadcast::Sender<AppMessage>,
+    agent: std::sync::Arc<tokio::sync::Mutex<Option<coro_core::agent::AgentCore>>>,
+) -> Result<()> {
+    // Create a receiver to listen for interruption signals
+    let mut interrupt_receiver = ui_sender.subscribe();
+    use crate::tools::StatusReportToolFactory;
+
+    // Create channel for InteractiveMessage and forward to AppMessage
+    let (interactive_sender, mut interactive_receiver) = mpsc::unbounded_channel();
+    let ui_sender_clone = ui_sender.clone();
+
+    // Forward InteractiveMessage to AppMessage
+    tokio::spawn(async move {
+        while let Some(interactive_msg) = interactive_receiver.recv().await {
+            let _ = ui_sender_clone.send(AppMessage::InteractiveUpdate(interactive_msg));
+        }
+    });
+
+    // Lock the agent for the duration of this task
+    let mut agent_guard = agent.lock().await;
+
+    // If no agent exists, create one
+    if agent_guard.is_none() {
+        // Create agent configuration with CLI tools and status_report tool for interactive mode
+        let mut agent_config = coro_core::AgentConfig {
+            tools: crate::tools::get_default_cli_tools(),
+            ..Default::default()
+        };
+        if !agent_config.tools.contains(&"status_report".to_string()) {
+            agent_config.tools.push("status_report".to_string());
+        }
+
+        // Create TokenTrackingOutputHandler with UI integration
+        let interactive_config = InteractiveOutputConfig {
+            realtime_updates: true,
+            show_tool_details: true,
+        };
+        let token_tracking_output = Box::new(TokenTrackingOutputHandler::new(
+            interactive_config,
+            interactive_sender,
+            ui_sender.clone(),
+        ));
+
+        // Create CLI tool registry with status_report tool for interactive mode
+        let mut tool_registry = crate::tools::create_cli_tool_registry();
+        tool_registry.register_factory(Box::new(StatusReportToolFactory::with_ui_sender(
+            ui_sender.clone(),
+        )));
+
+        // Create new agent
+        let new_agent = coro_core::agent::AgentCore::new_with_output_and_registry(
+            agent_config,
+            llm_config,
+            token_tracking_output,
+            tool_registry,
+        )
+        .await?;
+
+        *agent_guard = Some(new_agent);
+    }
+
+    // Get mutable reference to the agent
+    let agent_ref = agent_guard.as_mut().unwrap();
+
+    // Execute task with conversation continuation
+    let task_future = agent_ref.continue_conversation(&task, &project_path);
+
+    // Listen for interruption signals
+    let interrupt_future = async {
+        loop {
+            match interrupt_receiver.recv().await {
+                Ok(AppMessage::AgentExecutionInterrupted { .. }) => {
+                    tracing::warn!("Task interrupted by user");
+                    return Err(anyhow::anyhow!("Task interrupted by user"));
+                }
+                Ok(_) => continue, // Ignore other messages
+                Err(_) => break,   // Channel closed
+            }
+        }
+        Ok(())
+    };
+
+    // Add timeout to prevent hanging
+    let timeout_future = tokio::time::sleep(tokio::time::Duration::from_secs(300)); // 5 minutes timeout
+
+    // Race between task execution, interruption, and timeout
+    tokio::select! {
+        result = task_future => {
+            result?;
+        }
+        interrupt_result = interrupt_future => {
+            interrupt_result?;
+        }
+        _ = timeout_future => {
+            tracing::error!("Task execution timed out after 5 minutes");
+            return Err(anyhow::anyhow!("Task execution timed out"));
+        }
+    }
+
+    Ok(())
+}
+
 /// Execute agent task asynchronously and send updates to UI
 pub async fn execute_agent_task(
     task: String,