pr patch: task tracking for more aggressive compression

crates/goose-cli/src/commands/web.rs

@@ -448,8 +448,8 @@ async fn process_message_streaming(
     // Create a user message
     let user_message = GooseMessage::user().with_text(content.clone());
 
-    // Get existing messages from session and add the new user message
-    let mut messages = {
+    // Messages will be auto-compacted in agent.reply() if needed
+    let messages = {
         let mut session_msgs = session_messages.lock().await;
         session_msgs.push(user_message.clone());
         session_msgs.clone()
@@ -618,14 +618,41 @@ async fn process_message_streaming(
                                     // TODO: Implement proper UI for context handling
                                     let (summarized_messages, _) =
                                         agent.summarize_context(&messages).await?;
-                                    messages = summarized_messages;
+                                    {
+                                        let mut session_msgs = session_messages.lock().await;
+                                        *session_msgs = summarized_messages;
+                                    }
                                 }
                                 _ => {
                                     // Handle other message types as needed
                                 }
                             }
                         }
                     }
+                    Ok(AgentEvent::HistoryReplaced(new_messages)) => {
+                        // Replace the session's message history with the compacted messages
+                        {
+                            let mut session_msgs = session_messages.lock().await;
+                            *session_msgs = new_messages;
+                        }
+
+                        // Persist the updated messages to the JSONL file
+                        let current_messages = {
+                            let session_msgs = session_messages.lock().await;
+                            session_msgs.clone()
+                        };
+
+                        if let Err(e) = session::persist_messages(
+                            &session_file,
+                            &current_messages,
+                            None, // No provider needed for persisting
+                            working_dir.clone(),
+                        )
+                        .await
+                        {
+                            error!("Failed to persist compacted messages: {}", e);
+                        }
+                    }
                     Ok(AgentEvent::McpNotification(_notification)) => {
                         // Handle MCP notifications if needed
                         // For now, we'll just log them

crates/goose-cli/src/session/mod.rs

@@ -841,6 +841,7 @@ impl Session {
     }
 
     async fn process_agent_response(&mut self, interactive: bool) -> Result<()> {
+        // Messages will be auto-compacted in agent.reply() if needed
         let cancel_token = CancellationToken::new();
         let cancel_token_clone = cancel_token.clone();
 
@@ -1135,6 +1136,25 @@ impl Session {
                                 _ => (),
                             }
                         }
+                        Some(Ok(AgentEvent::HistoryReplaced(new_messages))) => {
+                            // Replace the session's message history with the compacted messages
+                            self.messages = new_messages;
+
+                            // Persist the updated messages to the session file
+                            if let Some(session_file) = &self.session_file {
+                                let provider = self.agent.provider().await.ok();
+                                let working_dir = std::env::current_dir().ok();
+                                if let Err(e) = session::persist_messages_with_schedule_id(
+                                    session_file,
+                                    &self.messages,
+                                    provider,
+                                    self.scheduled_job_id.clone(),
+                                    working_dir,
+                                ).await {
+                                    eprintln!("Failed to persist compacted messages: {}", e);
+                                }
+                            }
+                        }
                         Some(Ok(AgentEvent::ModelChange { model, mode })) => {
                             // Log model change if in debug mode
                             if self.debug {

crates/goose-ffi/src/lib.rs

@@ -269,7 +269,9 @@ pub unsafe extern "C" fn goose_agent_send_message(
                 Ok(AgentEvent::ModelChange { .. }) => {
                     // Model change events are informational, just continue
                 }
-
+                Ok(AgentEvent::HistoryReplaced(_)) => {
+                    // Handle history replacement events if needed
+                }
                 Err(e) => {
                     full_response.push_str(&format!("\nError in message stream: {}", e));
                 }

crates/goose-server/src/routes/reply.rs

@@ -159,8 +159,15 @@ async fn reply_handler(
             retry_config: None,
         };
 
+        // Messages will be auto-compacted in agent.reply() if needed
+        let messages_to_process = messages.clone();
+
         let mut stream = match agent
-            .reply(&messages, Some(session_config), Some(task_cancel.clone()))
+            .reply(
+                &messages_to_process,
+                Some(session_config),
+                Some(task_cancel.clone()),
+            )
             .await
         {
             Ok(stream) => stream,
@@ -215,6 +222,12 @@ async fn reply_handler(
                                             break;
                                         }
                                     }
+                                    Ok(Some(Ok(AgentEvent::HistoryReplaced(new_messages)))) => {
+                                        // Replace the message history with the compacted messages
+                                        all_messages = new_messages;
+                                        // Note: We don't send this as a stream event since it's an internal operation
+                                        // The client will see the compaction notification message that was sent before this event
+                                    }
                                     Ok(Some(Ok(AgentEvent::ModelChange { model, mode }))) => {
                                         if let Err(e) = stream_event(MessageEvent::ModelChange { model, mode }, &tx).await {
                                             tracing::error!("Error sending model change through channel: {}", e);

crates/goose/src/agents/agent.rs

@@ -29,11 +29,13 @@ use crate::agents::subagent_execution_tool::subagent_execute_task_tool::{
     self, SUBAGENT_EXECUTE_TASK_TOOL_NAME,
 };
 use crate::agents::subagent_execution_tool::tasks_manager::TasksManager;
+use crate::agents::task_tracker_tool::{TaskTracker, TASK_TRACKER_TOOL_NAME};
 use crate::agents::tool_router_index_manager::ToolRouterIndexManager;
 use crate::agents::tool_vectordb::generate_table_id;
 use crate::agents::types::SessionConfig;
 use crate::agents::types::{FrontendTool, ToolResultReceiver};
 use crate::config::{Config, ExtensionConfigManager, PermissionManager};
+use crate::context_mgmt::auto_compact;
 use crate::message::{push_message, Message};
 use crate::permission::permission_judge::check_tool_permissions;
 use crate::permission::PermissionConfirmation;
@@ -77,13 +79,15 @@ pub struct Agent {
     pub(super) router_tool_selector: Mutex<Option<Arc<Box<dyn RouterToolSelector>>>>,
     pub(super) scheduler_service: Mutex<Option<Arc<dyn SchedulerTrait>>>,
     pub(super) retry_manager: RetryManager,
+    pub(super) task_tracker: TaskTracker,
 }
 
 #[derive(Clone, Debug)]
 pub enum AgentEvent {
     Message(Message),
     McpNotification((String, ServerNotification)),
     ModelChange { model: String, mode: String },
+    HistoryReplaced(Vec<Message>),
 }
 
 impl Default for Agent {
@@ -152,6 +156,7 @@ impl Agent {
             router_tool_selector: Mutex::new(None),
             scheduler_service: Mutex::new(None),
             retry_manager,
+            task_tracker: TaskTracker::new(),
         }
     }
 
@@ -367,6 +372,8 @@ impl Agent {
             ToolCallResult::from(Err(ToolError::ExecutionError(
                 "Frontend tool execution required".to_string(),
             )))
+        } else if tool_call.name == TASK_TRACKER_TOOL_NAME {
+            ToolCallResult::from(self.task_tracker.execute(tool_call.arguments.clone()).await)
         } else if tool_call.name == ROUTER_VECTOR_SEARCH_TOOL_NAME
             || tool_call.name == ROUTER_LLM_SEARCH_TOOL_NAME
         {
@@ -622,6 +629,9 @@ impl Agent {
                 prefixed_tools.push(final_output_tool.tool());
             }
             prefixed_tools.push(subagent_execute_task_tool::create_subagent_execute_task_tool());
+
+            // Add task tracker tool
+            prefixed_tools.push(TaskTracker::tool());
         }
 
         prefixed_tools
@@ -712,11 +722,63 @@ impl Agent {
         messages: &[Message],
         session: Option<SessionConfig>,
         cancel_token: Option<CancellationToken>,
+    ) -> Result<BoxStream<'_, Result<AgentEvent>>> {
+        let mut messages = messages.to_vec();
+        let _initial_messages = messages.clone();
+        let _reply_span = tracing::Span::current();
+        self.reset_retry_attempts().await;
+        let _config = Config::global();
+
+        // Handle auto-compaction before processing
+        let compact_result =
+            auto_compact::check_and_compact_messages(self, &messages, None).await?;
+        if compact_result.compacted {
+            messages = compact_result.messages;
+
+            // Create compaction notification message
+            let compaction_msg = if let (Some(before), Some(after)) =
+                (compact_result.tokens_before, compact_result.tokens_after)
+            {
+                format!(
+                    "Auto-compacted context: {} → {} tokens ({:.0}% reduction)",
+                    before,
+                    after,
+                    (1.0 - (after as f64 / before as f64)) * 100.0
+                )
+            } else {
+                "Auto-compacted context to reduce token usage".to_string()
+            };
+
+            // Create a new AgentEvent type to signal message history replacement
+            return Ok(Box::pin(async_stream::try_stream! {
+                yield AgentEvent::Message(Message::assistant().with_text(compaction_msg));
+
+                // Yield a special event to indicate the session should replace its message history
+                // with the compacted messages
+                yield AgentEvent::HistoryReplaced(messages.clone());
+
+                // Continue with normal reply processing using compacted messages
+                let mut reply_stream = self.reply_internal(&messages, session, cancel_token).await?;
+                while let Some(event) = reply_stream.next().await {
+                    yield event?;
+                }
+            }));
+        }
+
+        // No compaction needed, proceed with normal processing
+        self.reply_internal(&messages, session, cancel_token).await
+    }
+
+    /// Internal reply method that handles the actual agent processing
+    async fn reply_internal(
+        &self,
+        messages: &[Message],
+        session: Option<SessionConfig>,
+        cancel_token: Option<CancellationToken>,
     ) -> Result<BoxStream<'_, Result<AgentEvent>>> {
         let mut messages = messages.to_vec();
         let initial_messages = messages.clone();
         let reply_span = tracing::Span::current();
-        self.reset_retry_attempts().await;
         let config = Config::global();
 
         let (mut tools, mut toolshim_tools, mut system_prompt) =

crates/goose/src/agents/mod.rs

@@ -17,6 +17,7 @@ pub mod subagent;
 pub mod subagent_execution_tool;
 pub mod subagent_handler;
 mod subagent_task_config;
+pub mod task_tracker_tool;
 mod tool_execution;
 mod tool_router_index_manager;
 pub(crate) mod tool_vectordb;