feat: refactor agent step limits and enhance workflow orchestration with new findings tool

This commit is contained in:
asepharyana
2026-07-13 14:39:39 +07:00
parent 0d6f558b2b
commit 3b660e09a8
9 changed files with 229 additions and 100 deletions
-26
View File
@@ -957,16 +957,6 @@ fn archive_message(db: Option<&std::sync::Arc<std::sync::Mutex<rusqlite::Connect
}
}
/// Maximum number of LLM call + tool-execution iterations per single
/// agent turn before bailing. Prevents runaway token consumption when
/// 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;
/// Maximum number of auto inline reviews spawned per single agent turn.
/// After N edits, the inline review is skipped to keep the turn fast;
/// background subagents still fire at the end of the turn.
@@ -1006,7 +996,6 @@ fn run_agent_turn(
let mut edited_paths: Vec<String> = Vec::new();
let mut inline_reviews_count: usize = 0;
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
@@ -1141,24 +1130,9 @@ fn run_agent_turn(
return Ok(());
}
let mut turn_step = 0usize;
let mut todo_retry_count = 0usize;
loop {
turn_step += 1;
if turn_step > MAX_TURN_STEPS {
anyhow::bail!(
"turn exceeded maximum steps ({MAX_TURN_STEPS}) — possible runaway loop. \
aborting to prevent excessive token usage",
);
}
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(str::len)
+4 -11
View File
@@ -37,12 +37,6 @@ const SKIP_REVIEW_FILES: &[&str] = &[
".gitignore", ".env", ".env.example",
];
/// Maximum LLM steps for a quick-review subagent. Keeps reviews fast.
const QUICK_REVIEW_MAX_STEPS: usize = 2;
/// Maximum LLM steps for background subagents (test gen, arch, security).
const BG_SUBAGENT_MAX_STEPS: usize = 8;
/// ─── Helpers ───
///
/// Check whether a file path is worth auto-reviewing (not config/lock/data).
@@ -114,8 +108,7 @@ pub fn spawn_quick_review(
"quick-reviewer".to_string(),
"reviewer".to_string(),
)
.with_system_prompt(prompt)
.with_max_steps(QUICK_REVIEW_MAX_STEPS);
.with_system_prompt(prompt);
let mut ctx = build_subagent_context(&def);
ctx.session_dir = session_dir.to_path_buf();
@@ -189,7 +182,7 @@ pub fn spawn_background_test_gen(
"coder".to_string(), // needs write access
)
.with_system_prompt(prompt)
.with_max_steps(BG_SUBAGENT_MAX_STEPS);
;
let mut ctx = build_subagent_context(&def);
ctx.session_dir = sd;
@@ -269,7 +262,7 @@ pub fn spawn_background_arch_review(
"reviewer".to_string(),
)
.with_system_prompt(prompt)
.with_max_steps(BG_SUBAGENT_MAX_STEPS);
;
let mut ctx = build_subagent_context(&def);
ctx.session_dir = sd;
@@ -355,7 +348,7 @@ pub fn spawn_background_security_review(
"reviewer".to_string(),
)
.with_system_prompt(prompt)
.with_max_steps(BG_SUBAGENT_MAX_STEPS);
;
let mut ctx = build_subagent_context(&def);
ctx.session_dir = sd;
+1 -1
View File
@@ -45,7 +45,7 @@ pub fn build_subagent_context(def: &AgentDefinition) -> SubagentContext {
Vec::new()
}
});
let max_steps = def.max_steps.unwrap_or(25);
let max_steps = def.max_steps.unwrap_or(usize::MAX);
SubagentContext {
system_prompt: String::new(),
allowed_tools,
+5 -5
View File
@@ -43,7 +43,6 @@ pub fn strategy_division() -> AgentDefinition {
roles::STRATEGY.to_string(),
)
.with_system_prompt(crate::resources::DIVISION_PLANNER_PROMPT.to_string())
.with_max_steps(15)
.with_allowed_tools(vec![
"read".to_string(),
"grep".to_string(),
@@ -57,6 +56,7 @@ pub fn strategy_division() -> AgentDefinition {
"lsp_hover".to_string(),
"lsp_definition".to_string(),
"lsp_references".to_string(),
"read_findings".to_string(),
])
}
@@ -70,7 +70,6 @@ pub fn engineering_division() -> AgentDefinition {
roles::ENGINEERING.to_string(),
)
.with_system_prompt(crate::resources::DIVISION_IMPLEMENTER_PROMPT.to_string())
.with_max_steps(50)
.with_allowed_tools(vec![
"read".to_string(),
"write".to_string(),
@@ -90,6 +89,7 @@ pub fn engineering_division() -> AgentDefinition {
"lsp_disconnect".to_string(),
"todowrite".to_string(),
"todofinish".to_string(),
"read_findings".to_string(),
])
}
@@ -103,7 +103,6 @@ pub fn quality_division() -> AgentDefinition {
roles::QUALITY.to_string(),
)
.with_system_prompt(crate::resources::DIVISION_TESTER_PROMPT.to_string())
.with_max_steps(30)
.with_allowed_tools(vec![
"read".to_string(),
"write".to_string(),
@@ -119,6 +118,7 @@ pub fn quality_division() -> AgentDefinition {
"lsp_hover".to_string(),
"lsp_definition".to_string(),
"lsp_references".to_string(),
"read_findings".to_string(),
])
}
@@ -132,7 +132,6 @@ pub fn security_division() -> AgentDefinition {
roles::SECURITY.to_string(),
)
.with_system_prompt(crate::resources::SECURITY_REVIEWER_PROMPT.to_string())
.with_max_steps(15)
.with_allowed_tools(vec![
"read".to_string(),
"grep".to_string(),
@@ -146,6 +145,7 @@ pub fn security_division() -> AgentDefinition {
"lsp_hover".to_string(),
"lsp_definition".to_string(),
"lsp_references".to_string(),
"read_findings".to_string(),
])
}
@@ -159,7 +159,6 @@ pub fn documentation_division() -> AgentDefinition {
roles::DOCUMENTATION.to_string(),
)
.with_system_prompt(crate::resources::DIVISION_DOCUMENTER_PROMPT.to_string())
.with_max_steps(15)
.with_allowed_tools(vec![
"read".to_string(),
"write".to_string(),
@@ -168,6 +167,7 @@ pub fn documentation_division() -> AgentDefinition {
"glob".to_string(),
"recall".to_string(),
"remember".to_string(),
"read_findings".to_string(),
])
}
+1
View File
@@ -30,6 +30,7 @@ impl AgentDefinition {
}
/// Builder method: limit this agent to at most `steps` LLM calls.
#[allow(dead_code)]
pub fn with_max_steps(mut self, steps: usize) -> Self {
self.max_steps = Some(steps);
self
+176 -55
View File
@@ -1,6 +1,6 @@
//! Company-style workflow orchestrator: runs the complete division pipeline
//! (Strategy → Engineering → Quality Security Documentation) with
//! findings flowing between stages, then returns a consolidated executive
//! (Strategy → Engineering → [Quality || Security || Documentation] in parallel)
//! with findings flowing between stages, then returns a consolidated executive
//! summary to the CEO (main agent).
//!
//! Flow:
@@ -9,32 +9,114 @@
//! │ delegates to run_company_pipeline(request)
//! ▼
//! ┌──────────────────────────────────────────────────┐
//! │ Strategy Division — plan + mermaid diagrams │
//! │ Engineering Division — implement per plan │
//! Quality Division — review + write tests
//! │ Security Division — vulnerability audit │
//! │ Documentation Div — update docs │
//! └─────────────────────────────────────────────────┘
//! returns consolidated summary
//!
//! CEO Main Agent delivers to user
//! │ Strategy Division — plan + mermaid diagrams │ (runs sequentially first)
//! └─────────────────────────┬────────────────────────┘
//!
//! ┌──────────────────────────────────────────────────┐
//! │ Engineering Division — implement per plan │ (runs sequentially second)
//! └─────────────────────────────────────────────────┘
//!
//! ┌────────────┼────────────┐
//! ▼ ▼ ▼
//! ┌───────────┐┌───────────┐┌───────────┐
//! │ Quality ││ Security ││ Docs │ (run concurrently in parallel)
//! └───────────┘└───────────┘└───────────┘
//! │ │ │
//! └────────────┼────────────┘
//! ▼
//! CEO Main Agent delivers consolidated summary to user
//! ```
use std::collections::HashMap;
use std::fmt::Write;
use std::sync::{Arc, Mutex, atomic::AtomicBool};
use crate::app::workflow::engine::{execute_primitive, LiveStateFn, AgentStatus};
use crate::app::workflow::script::{ScriptPrimitive, ScriptOptions, WorkflowScript};
use crate::app::workflow::engine::{execute_primitive, LiveStateFn, AgentStatus};
use crate::app::subagent::division;
/// Construct the 4 specialized agents for a division.
fn make_division_specialists(
div: &division::Division,
user_request: &str,
) -> Vec<ScriptPrimitive> {
let div_prompt = div.agent_def.system_prompt.as_deref().unwrap_or("");
let specializations = match div.name {
"Strategy" => vec![
("Architectural Analysis", "Focus on component tree, file layout, and module structure."),
("Data Flow Planning", "Focus on sequence of calls, interface definitions, and APIs."),
("Task Breakdown", "Focus on step-by-step TODO lists and implementation order."),
("Risk Evaluation", "Focus on edge cases, compatibility, and system constraints."),
],
"Engineering" => vec![
("Core Logic", "Focus on core algorithms, mathematical processing, and backend logic."),
("Interface & Endpoints", "Focus on implementing routes, IPC handlers, and struct mappings."),
("Error Handling & Logs", "Focus on implementing robust error handling, try/catch, tracing, and Result wrappers."),
("Utility & Helpers", "Focus on filesystem helpers, input sanitization, and parsing utilities."),
],
"Quality" => vec![
("Code Reviewer", "Focus on checking coding style, naming standards, and coding conventions."),
("Unit Testing", "Focus on writing and running unit tests for individual functions and modules."),
("Integration Testing", "Focus on writing and running integration tests for system interactions and state flows."),
("Performance Analyst", "Focus on efficiency check, bottleneck analysis, and time complexity."),
],
"Security" => vec![
("Dependency Auditor", "Focus on auditing cargo lock and checking dependencies for vulnerabilities."),
("Input Sanitizer", "Focus on auditing input validation, injection prevention, path traversal, and shell safety."),
("Access Control", "Focus on auditing authorization, filesystem access permissions, and API scopes."),
("Secrets Auditor", "Focus on auditing secrets leakage, credentials safety, and log auditing."),
],
"Documentation" => vec![
("README & Setup", "Focus on updating README, installation guides, usage examples, and high-level setup."),
("API Reference", "Focus on updating API reference, parameter details, and traits/functions documentation."),
("Changelog & Architecture", "Focus on updating CHANGELOG and describing system architecture/diagrams."),
("Inline Comments", "Focus on adding explanatory inline comments and documentation comments inside source files."),
],
_ => vec![
("Specialist 1", "Focus on general tasks and responsibilities of this division."),
("Specialist 2", "Focus on code review and validation."),
("Specialist 3", "Focus on error handling and reporting."),
("Specialist 4", "Focus on documentation and testing."),
],
};
specializations
.into_iter()
.map(|(label, focus)| {
// Prepend [Division Name: Specialist Label] so the first 40 chars
// of the prompt become the agent_name in spawn_single_agent.
let prompt = format!(
"[{}: {}]\n\n{}\n\n{}\n\nUser request: {}\n\nFindings from previous divisions: {{findings}}",
div.name,
label,
focus,
div_prompt,
user_request,
);
ScriptPrimitive::Agent(prompt)
})
.collect()
}
/// Construct a named Phase wrapper containing a Parallel block of division specialists.
fn make_division_phase(
div: &division::Division,
user_request: &str,
) -> ScriptPrimitive {
let specialists = make_division_specialists(div, user_request);
ScriptPrimitive::Phase {
name: div.name.to_string(),
script: Box::new(ScriptPrimitive::Parallel(specialists)),
}
}
/// Run the full company-style pipeline for a given user request.
///
/// This orchestrates all five divisions in sequence:
/// 1. **Strategy** — create plan with diagrams
/// 2. **Engineering** — implement code
/// 3. **Quality** — review + write tests
/// 4. **Security** — audit
/// 5. **Documentation** — update docs
/// This orchestrates all five divisions, running Strategy and Engineering
/// sequentially, followed by Quality, Security, and Documentation in parallel:
/// 1. **Strategy** — create plan with diagrams (4 parallel subagents)
/// 2. **Engineering** — implement code per the plan (4 parallel subagents)
/// 3. **Quality** || **Security** || **Documentation** (in parallel, up to 10 concurrent subagents total)
///
/// Each division receives findings from all previous divisions, enabling
/// context to flow through the pipeline.
@@ -48,35 +130,39 @@ pub fn run_company_pipeline(
abort_flag: &Option<Arc<AtomicBool>>,
) -> anyhow::Result<String> {
let divisions = division::all_divisions();
let mut pipeline_scripts: Vec<ScriptPrimitive> = Vec::with_capacity(divisions.len());
for div in &divisions {
let div_prompt = div.agent_def.system_prompt.as_deref().unwrap_or("");
// Prepend [Division Name] so the first 40 chars of the prompt
// become the agent_name in spawn_single_agent, making the TUI
// panel show division names instead of UUID fragments.
let prompt = format!(
"[{}]\n\n{}\n\nUser request: {}\n\nFindings from previous divisions: {{findings}}",
div.name,
div_prompt,
user_request,
);
pipeline_scripts.push(ScriptPrimitive::Agent(prompt));
}
let strategy_phase = make_division_phase(&divisions[0], user_request);
let engineering_phase = make_division_phase(&divisions[1], user_request);
let quality_phase = make_division_phase(&divisions[2], user_request);
let security_phase = make_division_phase(&divisions[3], user_request);
let documentation_phase = make_division_phase(&divisions[4], user_request);
let parallel_divisions = ScriptPrimitive::Parallel(vec![
quality_phase,
security_phase,
documentation_phase,
]);
let pipeline_primitive = ScriptPrimitive::Pipeline(vec![
strategy_phase,
engineering_phase,
parallel_divisions,
]);
let wf = WorkflowScript {
name: "company-pipeline".to_string(),
description: "Company Pipeline (full): Strategy → Engineering → Quality Security Documentation".to_string(),
script: ScriptPrimitive::Pipeline(pipeline_scripts),
description: "Company Pipeline: Strategy → Engineering → (Quality || Security || Documentation)".to_string(),
script: pipeline_primitive,
options: ScriptOptions {
max_concurrency: 1, // sequential by design
max_concurrency: 10, // Max concurrent agents in execution
continue_on_error: true, // one division failing shouldn't block the rest
timeout_ms: None,
},
};
// Build a live callback for TUI updates if turn_events is available.
// Uses agent_name (division name) for the display label in the panel.
// Uses agent_name (division + specialist name) for the display label in the panel.
let live: Option<LiveStateFn> = turn_events.map(|events| {
let events = events.clone();
let f: LiveStateFn = Arc::new(move |_agent_id: String, agent_name: String, status: AgentStatus| {
@@ -101,7 +187,7 @@ pub fn run_company_pipeline(
let results = execute_primitive(
&wf.script,
&args,
1,
wf.options.max_concurrency,
true,
abort_flag,
live_ref,
@@ -134,24 +220,22 @@ pub fn run_company_pipeline_quick(
// Only use first 3 divisions for quick pipeline: Strategy, Engineering, Quality
let quick_divisions = &divisions[..3];
let mut pipeline_scripts: Vec<ScriptPrimitive> = Vec::with_capacity(quick_divisions.len());
for div in quick_divisions {
let div_prompt = div.agent_def.system_prompt.as_deref().unwrap_or("");
let prompt = format!(
"[{}]\n\n{}\n\nUser request: {}\n\nFindings from previous divisions: {{findings}}",
div.name,
div_prompt,
user_request,
);
pipeline_scripts.push(ScriptPrimitive::Agent(prompt));
}
let strategy_phase = make_division_phase(&quick_divisions[0], user_request);
let engineering_phase = make_division_phase(&quick_divisions[1], user_request);
let quality_phase = make_division_phase(&quick_divisions[2], user_request);
let pipeline_primitive = ScriptPrimitive::Pipeline(vec![
strategy_phase,
engineering_phase,
quality_phase,
]);
let wf = WorkflowScript {
name: "company-pipeline-quick".to_string(),
description: "Company Pipeline (quick): Strategy → Engineering → Quality".to_string(),
script: ScriptPrimitive::Pipeline(pipeline_scripts),
script: pipeline_primitive,
options: ScriptOptions {
max_concurrency: 1,
max_concurrency: 10,
continue_on_error: true,
timeout_ms: None,
},
@@ -176,7 +260,7 @@ pub fn run_company_pipeline_quick(
let findings: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
let results = execute_primitive(
&wf.script, &args, 1, true,
&wf.script, &args, wf.options.max_concurrency, true,
abort_flag, live.as_ref(), session_dir, workspaces, &findings, None,
)?;
@@ -202,10 +286,20 @@ fn build_executive_summary(
writeln!(summary, "Pipeline for: {request}").unwrap();
for (i, div) in divisions.iter().enumerate() {
let verdict = results.get(i).map_or_else(|| "".to_string(), |r| {
r.lines().next().unwrap_or(r)
.chars().take(100).collect::<String>()
});
let mut division_verdicts = Vec::new();
for offset in 0..4 {
if let Some(r) = results.get(4 * i + offset) {
let first_line = r.lines().next().unwrap_or(r);
let trimmed = first_line.chars().take(40).collect::<String>();
division_verdicts.push(trimmed);
}
}
let verdict = if division_verdicts.is_empty() {
"".to_string()
} else {
division_verdicts.join(" | ")
};
writeln!(summary, " {}: {}", div.name, verdict).unwrap();
}
@@ -265,3 +359,30 @@ pub fn is_complex_request(request: &str) -> bool {
];
complexity_keywords.iter().any(|k| lower.contains(k))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_is_complex_request_too_short() {
assert!(!is_complex_request("abc"));
}
#[test]
fn test_is_complex_request_simple_keywords() {
assert!(!is_complex_request("just a simple update to the readme"));
assert!(!is_complex_request("minor typo fix in main.rs"));
}
#[test]
fn test_is_complex_request_multi_sentence() {
assert!(is_complex_request("This is sentence one. This is sentence two. This is sentence three."));
}
#[test]
fn test_is_complex_request_complex_keywords() {
assert!(is_complex_request("implement user authentication endpoint"));
assert!(is_complex_request("refactor the whole engine module"));
}
}
+1 -2
View File
@@ -134,8 +134,7 @@ fn spawn_single_agent(
);
}
let def = AgentDefinition::new(agent_name.to_string(), "coder".to_string())
.with_max_steps(50);
let def = AgentDefinition::new(agent_name.to_string(), "coder".to_string());
let mut ctx = build_subagent_context(&def);
ctx.session_dir = session_dir.to_path_buf();
ctx.workspaces = workspaces.to_vec();
+1
View File
@@ -162,6 +162,7 @@ pub fn all_tools() -> Vec<Box<dyn Tool>> {
Box::new(super::tool::plan::PlanReady),
Box::new(super::tool::workflow::WorkflowRun),
Box::new(super::tool::workflow::NoteFinding),
Box::new(super::tool::workflow::ReadFindings),
Box::new(super::tool::workflow::CompanyPipeline),
Box::new(super::tool::spawn::SpawnAgents),
Box::new(super::tool::spawn::SpawnPipeline),
+40
View File
@@ -209,3 +209,43 @@ impl Tool for CompanyPipeline {
}
}
}
/// Tool that retrieves all findings shared by sibling agents in the current workflow run.
pub struct ReadFindings;
impl Tool for ReadFindings {
fn name(&self) -> &'static str {
"read_findings"
}
fn description(&self) -> &'static str {
"Retrieve all findings shared by sibling agents in the current workflow run. Use this to get real-time context updates from other divisions/subagents working in parallel."
}
fn parameters(&self) -> Value {
json!({
"type": "object",
"properties": {}
})
}
fn run(&self, ctx: &ToolCtx, _args: &Value) -> Result<String> {
if let Some(ref findings) = ctx.workflow_findings {
let f = findings.lock().map_err(|e| anyhow!("poisoned lock: {e}"))?;
if f.is_empty() {
Ok("No findings recorded yet in this workflow run.".to_string())
} else {
let formatted = f
.iter()
.enumerate()
.map(|(i, f)| format!("{}. {}", i + 1, f))
.collect::<Vec<_>>()
.join("\n");
Ok(format!("Findings in this workflow run:\n{}", formatted))
}
} else {
Ok("No findings database available (called outside a workflow run).".to_string())
}
}
}