feat: present options to user when context length is exceeded (#2207)

Co-authored-by: Yingjie He <yingjiehe@squareup.com>

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

@@ -6,6 +6,7 @@ mod prompt;
 mod thinking;
 
 pub use builder::{build_session, SessionBuilderConfig};
+use console::Color;
 use goose::permission::permission_confirmation::PrincipalType;
 use goose::permission::Permission;
 use goose::permission::PermissionConfirmation;
@@ -592,7 +593,7 @@ impl Session {
             .reply(
                 &self.messages,
                 Some(SessionConfig {
-                    id: session_id,
+                    id: session_id.clone(),
                     working_dir: std::env::current_dir()
                         .expect("failed to get current session working directory"),
                 }),
@@ -622,6 +623,54 @@ impl Session {
                                     principal_type: PrincipalType::Tool,
                                     permission,
                                 },).await;
+                            } else if let Some(MessageContent::ContextLengthExceeded(_)) = message.content.first() {
+                                output::hide_thinking();
+
+                                let prompt = "The model's context length is maxed out. You will need to reduce the # msgs. Do you want to?".to_string();
+                                let selected = cliclack::select(prompt)
+                                    .item("clear", "Clear Session", "Removes all messages from Goose's memory")
+                                    .item("truncate", "Truncate Messages", "Removes old messages till context is within limits")
+                                    .item("summarize", "Summarize Session", "Summarize the session to reduce context length")
+                                    .interact()?;
+
+                                match selected {
+                                    "clear" => {
+                                        self.messages.clear();
+                                        let msg = format!("Session cleared.\n{}", "-".repeat(50));
+                                        output::render_text(&msg, Some(Color::Yellow), true);
+                                        break;  // exit the loop to hand back control to the user
+                                    }
+                                    "truncate" => {
+                                        // Truncate messages to fit within context length
+                                        let (truncated_messages, _) = self.agent.truncate_context(&self.messages).await?;
+                                        let msg = format!("Context maxed out\n{}\nGoose tried its best to truncate messages for you.", "-".repeat(50));
+                                        output::render_text("", Some(Color::Yellow), true);
+                                        output::render_text(&msg, Some(Color::Yellow), true);
+                                        self.messages = truncated_messages;
+                                    }
+                                    "summarize" => {
+                                        // Summarize messages to fit within context length
+                                        let (summarized_messages, _) = self.agent.summarize_context(&self.messages).await?;
+                                        let msg = format!("Context maxed out\n{}\nGoose summarized messages for you.", "-".repeat(50));
+                                        output::render_text(&msg, Some(Color::Yellow), true);
+                                        self.messages = summarized_messages;
+                                    }
+                                    _ => {
+                                        unreachable!()
+                                    }
+                                }
+                                // Restart the stream after handling ContextLengthExceeded
+                                stream = self
+                                    .agent
+                                    .reply(
+                                        &self.messages,
+                                        Some(SessionConfig {
+                                            id: session_id.clone(),
+                                            working_dir: std::env::current_dir()
+                                                .expect("failed to get current session working directory"),
+                                        }),
+                                    )
+                                    .await?;
                             }
                             // otherwise we have a model/tool to render
                             else {

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

@@ -1,5 +1,5 @@
 use bat::WrappingMode;
-use console::style;
+use console::{style, Color};
 use goose::config::Config;
 use goose::message::{Message, MessageContent, ToolRequest, ToolResponse};
 use mcp_core::prompt::PromptArgument;
@@ -143,6 +143,19 @@ pub fn render_message(message: &Message, debug: bool) {
     println!();
 }
 
+pub fn render_text(text: &str, color: Option<Color>, dim: bool) {
+    let mut styled_text = style(text);
+    if dim {
+        styled_text = styled_text.dim();
+    }
+    if let Some(color) = color {
+        styled_text = styled_text.fg(color);
+    } else {
+        styled_text = styled_text.green();
+    }
+    println!("\n{}\n", styled_text);
+}
+
 pub fn render_enter_plan_mode() {
     println!(
         "\n{} {}\n",