feat: long-running session improvements

- Hang fix: 120s timeout on compression LLM call + raw_stream
  channel-close handling (None => break instead of spinning)
- Supervisor effective active count: stop counting finished
  JoinHandles as occupying capacity slots
- Universal tool result size cap: truncate_if_needed() on
  ToolResult, applied in eval_tool_calls() after escalation block;
  configurable via max_tool_result_chars in AppConfig/AgentConfig
- Windowed compression: compression_keep_last config param keeps
  the N most recent messages visible after compression
- Fix pre-existing flaky test: add #[serial] to
  handle_list_available_unrestricted_when_no_whitelist so it does
  not race with TestConfigDirGuard-based tests that temporarily
  populate the agents data dir
This commit is contained in:
2026-07-23 12:25:54 -06:00
parent d51bdd3086
commit 54c5079cb7
9 changed files with 118 additions and 19 deletions
+12
View File
@@ -531,6 +531,14 @@ impl Agent {
self.config.compression_threshold self.config.compression_threshold
} }
pub fn max_tool_result_chars(&self) -> Option<usize> {
self.config.max_tool_result_chars
}
pub fn compression_keep_last(&self) -> Option<usize> {
self.config.compression_keep_last
}
pub fn is_dynamic_instructions(&self) -> bool { pub fn is_dynamic_instructions(&self) -> bool {
self.config.dynamic_instructions self.config.dynamic_instructions
} }
@@ -679,6 +687,10 @@ pub struct AgentConfig {
pub memory: Option<bool>, pub memory: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
pub compression_threshold: Option<usize>, pub compression_threshold: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_tool_result_chars: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub compression_keep_last: Option<usize>,
#[serde(default)] #[serde(default)]
pub description: String, pub description: String,
#[serde(default)] #[serde(default)]
+6
View File
@@ -62,8 +62,10 @@ pub struct AppConfig {
pub save_session: Option<bool>, pub save_session: Option<bool>,
pub compression_threshold: usize, pub compression_threshold: usize,
pub compression_keep_last: usize,
pub summarization_prompt: Option<String>, pub summarization_prompt: Option<String>,
pub summary_context_prompt: Option<String>, pub summary_context_prompt: Option<String>,
pub max_tool_result_chars: Option<usize>,
pub memory: Option<bool>, pub memory: Option<bool>,
pub memory_cap_with_tools: Option<usize>, pub memory_cap_with_tools: Option<usize>,
@@ -143,8 +145,10 @@ impl Default for AppConfig {
save_session: None, save_session: None,
compression_threshold: 4000, compression_threshold: 4000,
compression_keep_last: 0,
summarization_prompt: None, summarization_prompt: None,
summary_context_prompt: None, summary_context_prompt: None,
max_tool_result_chars: None,
memory: None, memory: None,
memory_cap_with_tools: None, memory_cap_with_tools: None,
@@ -225,8 +229,10 @@ impl AppConfig {
save_session: config.save_session, save_session: config.save_session,
compression_threshold: config.compression_threshold, compression_threshold: config.compression_threshold,
compression_keep_last: config.compression_keep_last,
summarization_prompt: config.summarization_prompt, summarization_prompt: config.summarization_prompt,
summary_context_prompt: config.summary_context_prompt, summary_context_prompt: config.summary_context_prompt,
max_tool_result_chars: config.max_tool_result_chars,
memory: config.memory, memory: config.memory,
memory_cap_with_tools: config.memory_cap_with_tools, memory_cap_with_tools: config.memory_cap_with_tools,
+4
View File
@@ -239,8 +239,10 @@ pub struct Config {
pub save_session: Option<bool>, pub save_session: Option<bool>,
pub compression_threshold: usize, pub compression_threshold: usize,
pub compression_keep_last: usize,
pub summarization_prompt: Option<String>, pub summarization_prompt: Option<String>,
pub summary_context_prompt: Option<String>, pub summary_context_prompt: Option<String>,
pub max_tool_result_chars: Option<usize>,
pub memory: Option<bool>, pub memory: Option<bool>,
pub memory_cap_with_tools: Option<usize>, pub memory_cap_with_tools: Option<usize>,
@@ -318,8 +320,10 @@ impl Default for Config {
save_session: None, save_session: None,
compression_threshold: 4000, compression_threshold: 4000,
compression_keep_last: 0,
summarization_prompt: None, summarization_prompt: None,
summary_context_prompt: None, summary_context_prompt: None,
max_tool_result_chars: None,
memory: None, memory: None,
memory_cap_with_tools: None, memory_cap_with_tools: None,
+12 -2
View File
@@ -3941,7 +3941,12 @@ impl RequestContext {
.clone() .clone()
.unwrap_or_else(|| SUMMARIZATION_PROMPT.into()); .unwrap_or_else(|| SUMMARIZATION_PROMPT.into());
let input = Input::from_str(self, &prompt, None)?; 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 let summary_context_prompt = self
.app .app
.config .config
@@ -3958,8 +3963,13 @@ impl RequestContext {
String::new() 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() { 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(); self.discontinuous_last_message();
Ok(()) Ok(())
+9 -3
View File
@@ -570,7 +570,7 @@ impl Session {
self.compressing = compressing; 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 let Some(system_prompt) = self.messages.first().and_then(|v| {
if MessageRole::System == v.role { if MessageRole::System == v.role {
let content = v.content.to_text(); let content = v.content.to_text();
@@ -582,11 +582,17 @@ impl Session {
}) { }) {
prompt = format!("{system_prompt}\n\n{prompt}",); 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.compressed_messages.append(&mut self.messages);
self.messages.push(Message::new( self.messages.push(Message::new(
MessageRole::System, MessageRole::System,
MessageContent::Text(prompt), MessageContent::Text(prompt),
)); ));
self.messages.extend(messages_to_keep);
self.dirty = true; self.dirty = true;
self.update_tokens(); self.update_tokens();
} }
@@ -1032,7 +1038,7 @@ mod tests {
assert_eq!(session.messages.len(), 2); assert_eq!(session.messages.len(), 2);
assert!(session.compressed_messages.is_empty()); 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!(!session.compressed_messages.is_empty());
assert_eq!(session.messages.len(), 1); assert_eq!(session.messages.len(), 1);
@@ -1047,7 +1053,7 @@ mod tests {
MessageContent::Text("hello".to_string()), MessageContent::Text("hello".to_string()),
)); ));
session.compress("Summary".to_string()); session.compress("Summary".to_string(), 0);
assert!(!session.is_empty()); assert!(!session.is_empty());
} }
+25
View File
@@ -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) Ok(output)
} }
@@ -219,6 +233,17 @@ impl ToolResult {
thinking: vec![], 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)] #[derive(Debug, Clone, Default)]
+2
View File
@@ -1446,6 +1446,7 @@ mod tests {
use crate::config::{AppState, WorkingMode}; use crate::config::{AppState, WorkingMode};
use crate::supervisor::escalation::{EscalationQueue, EscalationRequest}; use crate::supervisor::escalation::{EscalationQueue, EscalationRequest};
use serde_json::json; use serde_json::json;
use serial_test::serial;
fn default_app_state() -> Arc<AppState> { fn default_app_state() -> Arc<AppState> {
Arc::new(AppState::test_default()) Arc::new(AppState::test_default())
@@ -1537,6 +1538,7 @@ mod tests {
} }
#[test] #[test]
#[serial]
fn handle_list_available_unrestricted_when_no_whitelist() { fn handle_list_available_unrestricted_when_no_whitelist() {
let ctx = ctx_with_supervisor(4, 3); let ctx = ctx_with_supervisor(4, 3);
let result = handle_list_available(&ctx).unwrap(); let result = handle_list_available(&ctx).unwrap();
+14 -11
View File
@@ -42,18 +42,21 @@ pub async fn raw_stream(
if abort_signal.aborted() { if abort_signal.aborted() {
break; break;
} }
if let Some(evt) = rx.recv().await { match rx.recv().await {
if let Some(spinner) = spinner.take() { None => break,
spinner.stop(); Some(evt) => {
} if let Some(spinner) = spinner.take() {
spinner.stop();
match evt {
SseEvent::Text(text) => {
print!("{text}");
stdout().flush()?;
} }
SseEvent::Done => {
break; match evt {
SseEvent::Text(text) => {
print!("{text}");
stdout().flush()?;
}
SseEvent::Done => {
break;
}
} }
} }
} }
+34 -3
View File
@@ -58,6 +58,13 @@ impl Supervisor {
self.handles.len() 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 { pub fn max_concurrent(&self) -> usize {
self.max_concurrent self.max_concurrent
} }
@@ -75,10 +82,10 @@ impl Supervisor {
} }
pub fn register(&mut self, handle: AgentHandle) -> Result<()> { pub fn register(&mut self, handle: AgentHandle) -> Result<()> {
if self.handles.len() >= self.max_concurrent { if self.effective_active_count() >= self.max_concurrent {
bail!( bail!(
"Cannot spawn agent: at capacity ({}/{})", "Cannot spawn agent: at capacity ({}/{})",
self.handles.len(), self.effective_active_count(),
self.max_concurrent self.max_concurrent
); );
} }
@@ -188,8 +195,32 @@ mod tests {
#[test] #[test]
fn supervisor_register_rejects_at_capacity() { 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, anyhow::Error>(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); 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)); let result = sup.register(make_handle("a2", "coder", 1));
assert!(result.is_err()); assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("at capacity")); assert!(result.unwrap_err().to_string().contains("at capacity"));