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:
@@ -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)]
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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(())
|
||||||
|
|||||||
@@ -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());
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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)]
|
||||||
|
|||||||
@@ -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
@@ -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
@@ -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"));
|
||||||
|
|||||||
Reference in New Issue
Block a user