feat: implement subagent tool gating and timeout mechanisms for enhanced security and performance
This commit is contained in:
@@ -894,6 +894,11 @@ fn archive_message(db: &Option<std::sync::Arc<std::sync::Mutex<rusqlite::Connect
|
||||
/// the agent gets stuck in a loop (e.g. an unachievable todo item).
|
||||
const MAX_TURN_STEPS: usize = 10000;
|
||||
|
||||
/// Hard wall-clock timeout per agent turn (5 minutes). Prevents a single
|
||||
/// user turn from running indefinitely even if the step budget isn't
|
||||
/// exhausted (e.g. slow LLM responses, stuck tool calls).
|
||||
const MAX_TURN_TIMEOUT_MS: u64 = 300_000;
|
||||
|
||||
/// Execute one full agent turn: stream the conversation to the LLM,
|
||||
/// handle tool calls, and loop until the LLM produces a non-tool response
|
||||
/// or runs out of unfinished todo items.
|
||||
@@ -924,7 +929,11 @@ fn run_agent_turn(
|
||||
let mut msgs = messages.to_vec();
|
||||
let mut edits_this_turn = 0u32;
|
||||
let mut prev_shaped = false;
|
||||
let turn_start_ms = std::time::Instant::now();
|
||||
|
||||
// Build system prompt components once and cache them for the entire turn
|
||||
// instead of regenerating on every loop iteration (which walks the full
|
||||
// workspace tree and reads all memory files each time).
|
||||
let tree_info = generate_workspace_tree(&tc.workspace_roots);
|
||||
let memory_section = build_memory_section(&tc.ctx.memory_dir);
|
||||
let system_text = format!(
|
||||
@@ -941,6 +950,8 @@ fn run_agent_turn(
|
||||
}
|
||||
|
||||
let mut turn_step = 0usize;
|
||||
let mut todo_retry_count = 0usize;
|
||||
const MAX_TODO_RETRIES: usize = 5;
|
||||
|
||||
loop {
|
||||
turn_step += 1;
|
||||
@@ -951,6 +962,13 @@ fn run_agent_turn(
|
||||
MAX_TURN_STEPS,
|
||||
);
|
||||
}
|
||||
if turn_start_ms.elapsed().as_millis() as u64 > MAX_TURN_TIMEOUT_MS {
|
||||
anyhow::bail!(
|
||||
"turn exceeded maximum duration ({}s) — aborting. \
|
||||
Use /compact or shorter prompts if the model needs more time.",
|
||||
MAX_TURN_TIMEOUT_MS / 1000,
|
||||
);
|
||||
}
|
||||
let total_chars: usize = msgs.iter()
|
||||
.filter_map(|m| m.content.as_deref())
|
||||
.map(|c| c.len())
|
||||
@@ -1049,10 +1067,18 @@ fn run_agent_turn(
|
||||
}
|
||||
}
|
||||
if has_unfinished {
|
||||
todo_retry_count += 1;
|
||||
if todo_retry_count > MAX_TODO_RETRIES {
|
||||
anyhow::bail!(
|
||||
"exhausted {} todo-retries — giving up on unfinished tasks. \
|
||||
Edit todo.md manually or ask me to focus on specific items.",
|
||||
MAX_TODO_RETRIES,
|
||||
);
|
||||
}
|
||||
if let Ok(mut q) = events_q.lock() {
|
||||
q.push_back(TurnEvent::SystemNote {
|
||||
kind: "task_retry".to_string(),
|
||||
message: format!("Network/API error: {}. Auto-retrying to finish tasks...", api_err),
|
||||
message: format!("Network/API error: {}. Auto-retrying to finish tasks... (retry {}/{})", api_err, todo_retry_count, MAX_TODO_RETRIES),
|
||||
});
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_secs(5));
|
||||
@@ -1160,14 +1186,25 @@ fn run_agent_turn(
|
||||
}
|
||||
|
||||
if has_unfinished {
|
||||
let sys_text = "You stopped, but you still have unfinished tasks in todo.md (marked with '- [ ]'). You MUST continue working and use tools to finish them, or edit todo.md to mark them as done if they are finished.";
|
||||
todo_retry_count += 1;
|
||||
if todo_retry_count > MAX_TODO_RETRIES {
|
||||
if let Ok(mut q) = events_q.lock() {
|
||||
q.push_back(TurnEvent::SystemNote {
|
||||
kind: "task_retry".to_string(),
|
||||
message: format!("Giving up after {} retries — some todo items remain unfinished. Edit todo.md manually or ask again.", MAX_TODO_RETRIES),
|
||||
});
|
||||
}
|
||||
break;
|
||||
}
|
||||
let sys_text = format!("You stopped, but you still have unfinished tasks in todo.md (marked with '- [ ]'). You MUST continue working and use tools to finish them, or edit todo.md to mark them as done if they are finished. (Retry {}/{})", todo_retry_count, MAX_TODO_RETRIES);
|
||||
let sys_text_clone = sys_text.clone();
|
||||
let msg = ChatMessage::system(sys_text);
|
||||
archive_message(&tc.db, &tc.session_id, &msg);
|
||||
msgs.push(msg);
|
||||
if let Ok(mut q) = events_q.lock() {
|
||||
q.push_back(TurnEvent::SystemNote {
|
||||
kind: "task_retry".to_string(),
|
||||
message: sys_text.to_string(),
|
||||
message: sys_text_clone,
|
||||
});
|
||||
}
|
||||
continue;
|
||||
|
||||
Reference in New Issue
Block a user