diff --git a/src/config/agent.rs b/src/config/agent.rs index 8edcd3b..41a50d6 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -531,6 +531,14 @@ impl Agent { self.config.compression_threshold } + pub fn max_tool_result_chars(&self) -> Option { + self.config.max_tool_result_chars + } + + pub fn compression_keep_last(&self) -> Option { + self.config.compression_keep_last + } + pub fn is_dynamic_instructions(&self) -> bool { self.config.dynamic_instructions } @@ -679,6 +687,10 @@ pub struct AgentConfig { pub memory: Option, #[serde(skip_serializing_if = "Option::is_none")] pub compression_threshold: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub max_tool_result_chars: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub compression_keep_last: Option, #[serde(default)] pub description: String, #[serde(default)] diff --git a/src/config/app_config.rs b/src/config/app_config.rs index 28a0a4f..c3f671e 100644 --- a/src/config/app_config.rs +++ b/src/config/app_config.rs @@ -62,8 +62,10 @@ pub struct AppConfig { pub save_session: Option, pub compression_threshold: usize, + pub compression_keep_last: usize, pub summarization_prompt: Option, pub summary_context_prompt: Option, + pub max_tool_result_chars: Option, pub memory: Option, pub memory_cap_with_tools: Option, @@ -143,8 +145,10 @@ impl Default for AppConfig { save_session: None, compression_threshold: 4000, + compression_keep_last: 0, summarization_prompt: None, summary_context_prompt: None, + max_tool_result_chars: None, memory: None, memory_cap_with_tools: None, @@ -225,8 +229,10 @@ impl AppConfig { save_session: config.save_session, compression_threshold: config.compression_threshold, + compression_keep_last: config.compression_keep_last, summarization_prompt: config.summarization_prompt, summary_context_prompt: config.summary_context_prompt, + max_tool_result_chars: config.max_tool_result_chars, memory: config.memory, memory_cap_with_tools: config.memory_cap_with_tools, diff --git a/src/config/mod.rs b/src/config/mod.rs index ef0432c..bfe6d52 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -239,8 +239,10 @@ pub struct Config { pub save_session: Option, pub compression_threshold: usize, + pub compression_keep_last: usize, pub summarization_prompt: Option, pub summary_context_prompt: Option, + pub max_tool_result_chars: Option, pub memory: Option, pub memory_cap_with_tools: Option, @@ -318,8 +320,10 @@ impl Default for Config { save_session: None, compression_threshold: 4000, + compression_keep_last: 0, summarization_prompt: None, summary_context_prompt: None, + max_tool_result_chars: None, memory: None, memory_cap_with_tools: None, diff --git a/src/config/request_context.rs b/src/config/request_context.rs index c8ccbd5..227b868 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -3941,7 +3941,12 @@ impl RequestContext { .clone() .unwrap_or_else(|| SUMMARIZATION_PROMPT.into()); let input = Input::from_str(self, &prompt, None)?; - let summary = input.fetch_chat_text().await?; + let summary = tokio::time::timeout( + std::time::Duration::from_secs(120), + input.fetch_chat_text(), + ) + .await + .map_err(|_| anyhow::anyhow!("Compression LLM call timed out after 120 s"))??; let summary_context_prompt = self .app .config @@ -3958,8 +3963,13 @@ impl RequestContext { String::new() }; + let keep_last = self + .agent + .as_ref() + .and_then(|a| a.compression_keep_last()) + .unwrap_or(self.app.config.compression_keep_last); if let Some(session) = self.session.as_mut() { - session.compress(format!("{todo_prefix}{summary_context_prompt}{summary}")); + session.compress(format!("{todo_prefix}{summary_context_prompt}{summary}"), keep_last); } self.discontinuous_last_message(); Ok(()) diff --git a/src/config/session.rs b/src/config/session.rs index 02379f1..98cc304 100644 --- a/src/config/session.rs +++ b/src/config/session.rs @@ -570,7 +570,7 @@ impl Session { self.compressing = compressing; } - pub fn compress(&mut self, mut prompt: String) { + pub fn compress(&mut self, mut prompt: String, keep_last: usize) { if let Some(system_prompt) = self.messages.first().and_then(|v| { if MessageRole::System == v.role { let content = v.content.to_text(); @@ -582,11 +582,17 @@ impl Session { }) { prompt = format!("{system_prompt}\n\n{prompt}",); } + let messages_to_keep = if keep_last > 0 && keep_last < self.messages.len() { + self.messages.split_off(self.messages.len() - keep_last) + } else { + vec![] + }; self.compressed_messages.append(&mut self.messages); self.messages.push(Message::new( MessageRole::System, MessageContent::Text(prompt), )); + self.messages.extend(messages_to_keep); self.dirty = true; self.update_tokens(); } @@ -1032,7 +1038,7 @@ mod tests { assert_eq!(session.messages.len(), 2); assert!(session.compressed_messages.is_empty()); - session.compress("Summary of conversation".to_string()); + session.compress("Summary of conversation".to_string(), 0); assert!(!session.compressed_messages.is_empty()); assert_eq!(session.messages.len(), 1); @@ -1047,7 +1053,7 @@ mod tests { MessageContent::Text("hello".to_string()), )); - session.compress("Summary".to_string()); + session.compress("Summary".to_string(), 0); assert!(!session.is_empty()); } diff --git a/src/function/mod.rs b/src/function/mod.rs index 079a2f4..778cc0b 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -184,6 +184,20 @@ pub async fn eval_tool_calls( } } + { + let max_chars = ctx + .agent + .as_ref() + .and_then(|a| a.max_tool_result_chars()) + .or_else(|| ctx.app.config.max_tool_result_chars); + if let Some(max_chars) = max_chars.filter(|&n| n > 0) { + output = output + .into_iter() + .map(|r| r.truncate_if_needed(max_chars)) + .collect(); + } + } + Ok(output) } @@ -219,6 +233,17 @@ impl ToolResult { thinking: vec![], } } + + pub fn truncate_if_needed(mut self, max_chars: usize) -> Self { + let s = self.output.to_string(); + if s.len() > max_chars { + let prefix = s.get(..max_chars).unwrap_or(s.as_str()); + self.output = json!(format!( + "[truncated: tool output exceeded {max_chars} chars]\n{prefix}" + )); + } + self + } } #[derive(Debug, Clone, Default)] diff --git a/src/function/supervisor.rs b/src/function/supervisor.rs index d0a1ea2..d392821 100644 --- a/src/function/supervisor.rs +++ b/src/function/supervisor.rs @@ -1446,6 +1446,7 @@ mod tests { use crate::config::{AppState, WorkingMode}; use crate::supervisor::escalation::{EscalationQueue, EscalationRequest}; use serde_json::json; + use serial_test::serial; fn default_app_state() -> Arc { Arc::new(AppState::test_default()) @@ -1537,6 +1538,7 @@ mod tests { } #[test] + #[serial] fn handle_list_available_unrestricted_when_no_whitelist() { let ctx = ctx_with_supervisor(4, 3); let result = handle_list_available(&ctx).unwrap(); diff --git a/src/render/stream.rs b/src/render/stream.rs index fffa161..689990a 100644 --- a/src/render/stream.rs +++ b/src/render/stream.rs @@ -42,18 +42,21 @@ pub async fn raw_stream( if abort_signal.aborted() { break; } - if let Some(evt) = rx.recv().await { - if let Some(spinner) = spinner.take() { - spinner.stop(); - } - - match evt { - SseEvent::Text(text) => { - print!("{text}"); - stdout().flush()?; + match rx.recv().await { + None => break, + Some(evt) => { + if let Some(spinner) = spinner.take() { + spinner.stop(); } - SseEvent::Done => { - break; + + match evt { + SseEvent::Text(text) => { + print!("{text}"); + stdout().flush()?; + } + SseEvent::Done => { + break; + } } } } diff --git a/src/supervisor/mod.rs b/src/supervisor/mod.rs index e1f86ea..c3ee65f 100644 --- a/src/supervisor/mod.rs +++ b/src/supervisor/mod.rs @@ -58,6 +58,13 @@ impl Supervisor { self.handles.len() } + pub fn effective_active_count(&self) -> usize { + self.handles + .values() + .filter(|h| !h.join_handle.is_finished()) + .count() + } + pub fn max_concurrent(&self) -> usize { self.max_concurrent } @@ -75,10 +82,10 @@ impl Supervisor { } pub fn register(&mut self, handle: AgentHandle) -> Result<()> { - if self.handles.len() >= self.max_concurrent { + if self.effective_active_count() >= self.max_concurrent { bail!( "Cannot spawn agent: at capacity ({}/{})", - self.handles.len(), + self.effective_active_count(), self.max_concurrent ); } @@ -188,8 +195,32 @@ mod tests { #[test] fn supervisor_register_rejects_at_capacity() { + // Keep the runtime alive in this scope so the spawned task is never + // polled (current_thread only polls inside block_on), keeping + // join_handle.is_finished() == false and the slot occupied. + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let join_handle = rt.spawn(async { + Ok::(AgentResult { + id: "done".into(), + agent_name: "test".into(), + output: "result".into(), + exit_status: AgentExitStatus::Completed, + }) + }); + let running_handle = AgentHandle { + id: "a1".to_string(), + agent_name: "explore".to_string(), + depth: 1, + inbox: Arc::new(Inbox::new()), + abort_signal: create_abort_signal(), + join_handle, + child_supervisor: None, + }; let mut sup = Supervisor::new(1, 3); - sup.register(make_handle("a1", "explore", 1)).unwrap(); + sup.register(running_handle).unwrap(); let result = sup.register(make_handle("a2", "coder", 1)); assert!(result.is_err()); assert!(result.unwrap_err().to_string().contains("at capacity"));