From 24b454b145eb8595c0a58b4d4b4abb4e8d8e970e Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Sun, 1 Mar 2026 17:39:13 -0500 Subject: [PATCH] Fix retro error handling and add Session wall-clock timeout - Add wall_clock_timeout field to SessionConfig that spawns a tokio timer to cancel the session after a duration, reusing existing Aborted path - Set 120s wall-clock timeout on retro agent to prevent unbounded runs - Replace silent `let _ =` with logged warnings for checkpoint load and retro save failures - Reuse LLM client from initial from_env() call instead of creating a second one for the retro agent - Extract retro generation into generate_retro() helper and call it from both run_command and run_from_branch (resume path) - Tolerate mutex poisoning in retro agent with unwrap_or_else Co-Authored-By: Claude Opus 4.6 --- crates/arc-agent/src/config.rs | 7 + crates/arc-agent/src/session.rs | 105 ++++++++++-- crates/arc-workflows/src/cli/run.rs | 206 +++++++++++++++++------- crates/arc-workflows/src/retro_agent.rs | 6 +- 4 files changed, 256 insertions(+), 68 deletions(-) diff --git a/crates/arc-agent/src/config.rs b/crates/arc-agent/src/config.rs index 8364279ac..cedd89201 100644 --- a/crates/arc-agent/src/config.rs +++ b/crates/arc-agent/src/config.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; use std::sync::Arc; +use std::time::Duration; use arc_mcp::config::McpServerConfig; @@ -32,6 +33,9 @@ pub struct SessionConfig { pub skill_dirs: Option>, /// MCP server configurations to connect to on session startup. pub mcp_servers: Vec, + /// Wall-clock timeout for the entire `process_input` call. + /// When set, the session's cancel token is triggered after this duration. + pub wall_clock_timeout: Option, } impl std::fmt::Debug for SessionConfig { @@ -65,6 +69,7 @@ impl std::fmt::Debug for SessionConfig { .field("compaction_preserve_turns", &self.compaction_preserve_turns) .field("skill_dirs", &self.skill_dirs) .field("mcp_servers", &self.mcp_servers.len()) + .field("wall_clock_timeout", &self.wall_clock_timeout) .finish() } } @@ -91,6 +96,7 @@ impl Default for SessionConfig { compaction_preserve_turns: 6, skill_dirs: None, mcp_servers: Vec::new(), + wall_clock_timeout: None, } } } @@ -114,6 +120,7 @@ mod tests { assert_eq!(config.max_subagent_depth, 1); assert!(config.user_instructions.is_none()); assert!(config.mcp_servers.is_empty()); + assert!(config.wall_clock_timeout.is_none()); } #[test] diff --git a/crates/arc-agent/src/session.rs b/crates/arc-agent/src/session.rs index 5006647bd..183c06958 100644 --- a/crates/arc-agent/src/session.rs +++ b/crates/arc-agent/src/session.rs @@ -287,21 +287,44 @@ impl Session { return Err(AgentError::SessionClosed); } + // Spawn wall-clock timeout task if configured + let timer_handle = self.config.wall_clock_timeout.map(|duration| { + let token = self.cancel_token.clone(); + tokio::spawn(async move { + tokio::time::sleep(duration).await; + token.cancel(); + }) + }); + // Process the initial input, then drain any followups - self.run_single_input(input).await?; - loop { - let followup = self - .followup_queue - .lock() - .expect("followup queue lock poisoned") - .pop_front(); - let Some(followup) = followup else { break }; - self.run_single_input(&followup).await?; + let mut result = self.run_single_input(input).await; + + if result.is_ok() { + loop { + let followup = self + .followup_queue + .lock() + .expect("followup queue lock poisoned") + .pop_front(); + let Some(followup) = followup else { break }; + result = self.run_single_input(&followup).await; + if result.is_err() { + break; + } + } } - self.state = SessionState::Idle; + // Abort the timer so it doesn't fire after we're done + if let Some(handle) = timer_handle { + handle.abort(); + } - Ok(()) + // Only transition to Idle if the session wasn't closed by an error + if self.state != SessionState::Closed { + self.state = SessionState::Idle; + } + + result } async fn run_single_input(&mut self, input: &str) -> Result<(), AgentError> { @@ -2028,4 +2051,64 @@ mod tests { "ToolCallCompleted should be emitted for MCP tool" ); } + + #[tokio::test] + async fn wall_clock_timeout_aborts_session() { + // Register a tool that loops until the cancel token fires + let slow_tool = RegisteredTool { + definition: ToolDefinition { + name: "slow_tool".into(), + description: "Waits until cancelled".into(), + parameters: serde_json::json!({"type": "object"}), + }, + executor: Arc::new(|_args, ctx| { + Box::pin(async move { + ctx.cancel.cancelled().await; + Ok("cancelled".to_string()) + }) + }), + }; + let mut registry = ToolRegistry::new(); + registry.register(slow_tool); + + // LLM will call the slow tool, then (if it ever gets there) respond with text + let responses = vec![ + tool_call_response("slow_tool", "call_1", serde_json::json!({})), + text_response("Should not reach this"), + ]; + + let config = SessionConfig { + wall_clock_timeout: Some(std::time::Duration::from_millis(10)), + enable_loop_detection: false, + ..Default::default() + }; + + let mut session = make_session_with_tools_and_config(responses, registry, config).await; + let result = session.process_input("Do something slow").await; + + assert!( + matches!(result, Err(AgentError::Aborted)), + "expected Aborted, got {result:?}" + ); + assert_eq!(session.state(), SessionState::Closed); + } + + #[tokio::test] + async fn wall_clock_timeout_does_not_fire_when_session_completes_in_time() { + let responses = vec![text_response("Fast response")]; + + let config = SessionConfig { + wall_clock_timeout: Some(std::time::Duration::from_secs(10)), + ..Default::default() + }; + + let mut session = make_session_with_config(responses, config).await; + let result = session.process_input("Hello").await; + + assert!(result.is_ok()); + assert_eq!(session.state(), SessionState::Idle); + let turns = session.history().turns(); + assert_eq!(turns.len(), 2); + assert!(matches!(&turns[1], Turn::Assistant { content, .. } if content == "Fast response")); + } } diff --git a/crates/arc-workflows/src/cli/run.rs b/crates/arc-workflows/src/cli/run.rs index 849a08932..cead128e9 100644 --- a/crates/arc-workflows/src/cli/run.rs +++ b/crates/arc-workflows/src/cli/run.rs @@ -443,8 +443,8 @@ pub async fn run_command(args: RunArgs, styles: &'static Styles) -> anyhow::Resu } // 6. Resolve backend, model, and provider - let dry_run_mode = if args.dry_run { - true + let (dry_run_mode, llm_client) = if args.dry_run { + (true, None) } else { match arc_llm::client::Client::from_env().await { Ok(c) if c.provider_names().is_empty() => { @@ -453,15 +453,15 @@ pub async fn run_command(args: RunArgs, styles: &'static Styles) -> anyhow::Resu yellow = styles.yellow, reset = styles.reset, ); - true + (true, None) } - Ok(_) => false, + Ok(c) => (false, Some(c)), Err(e) => { eprintln!( "{yellow}Warning:{reset} Failed to initialize LLM client: {e}. Running in dry-run mode.", yellow = styles.yellow, reset = styles.reset, ); - true + (true, None) } } }; @@ -607,56 +607,22 @@ pub async fn run_command(args: RunArgs, styles: &'static Styles) -> anyhow::Resu ), Err(e) => (true, Some(e.to_string())), }; - if let Ok(cp) = Checkpoint::load(&logs_dir.join("checkpoint.json")) { - let stage_durations = crate::retro::extract_stage_durations(&logs_dir); - let mut retro = crate::retro::derive_retro( - &config.run_id, - &graph.name, - graph.goal(), - &cp, - failed, - failure_reason.as_deref(), - run_duration_ms, - &stage_durations, - ); - let _ = retro.save(&logs_dir); - - // Run retro agent session (execution_env still alive via _cleanup_guard) - let narrative_result = if dry_run_mode { - Ok(crate::retro_agent::dry_run_narrative()) - } else if let Ok(client) = arc_llm::client::Client::from_env().await { - crate::retro_agent::run_retro_agent( - &execution_env, - &logs_dir, - &client, - provider_enum, - &model, - ) - .await - } else { - Err(anyhow::anyhow!("No LLM client available")) - }; - - match narrative_result { - Ok(narrative) => { - retro.apply_narrative(narrative); - let _ = retro.save(&logs_dir); - eprintln!( - "{dim}Retro saved to {}/retro.json{reset}", - logs_dir.display(), - dim = styles.dim, - reset = styles.reset, - ); - } - Err(e) => { - eprintln!( - "{dim}Retro agent skipped: {e}{reset}", - dim = styles.dim, - reset = styles.reset, - ); - } - } - } + generate_retro( + &config.run_id, + &graph.name, + graph.goal(), + &logs_dir, + failed, + failure_reason.as_deref(), + run_duration_ms, + dry_run_mode, + llm_client.as_ref(), + &execution_env, + provider_enum, + &model, + styles, + ) + .await; } let outcome = engine_result?; @@ -955,6 +921,40 @@ async fn run_from_branch( let _ = std::env::set_current_dir(&original_cwd); let _ = crate::git::remove_worktree(&original_cwd, &worktree_path); + // Auto-derive retro + { + let (failed, failure_reason) = match &engine_result { + Ok(o) => ( + o.status == StageStatus::Fail, + o.failure_reason.clone(), + ), + Err(e) => (true, Some(e.to_string())), + }; + + let llm_client = if dry_run_mode { + None + } else { + arc_llm::client::Client::from_env().await.ok() + }; + + generate_retro( + &config.run_id, + &graph.name, + graph.goal(), + &logs_dir, + failed, + failure_reason.as_deref(), + run_duration_ms, + dry_run_mode, + llm_client.as_ref(), + &execution_env, + provider_enum, + &model, + styles, + ) + .await; + } + let outcome = engine_result?; eprintln!( @@ -988,6 +988,102 @@ async fn run_from_branch( } } +/// Generate a retro report for a completed pipeline run. +/// +/// Derives a basic retro from the checkpoint, then optionally runs the retro agent +/// for a richer narrative. Errors are logged as warnings rather than propagated. +#[allow(clippy::too_many_arguments)] +async fn generate_retro( + run_id: &str, + pipeline_name: &str, + goal: &str, + logs_dir: &std::path::Path, + failed: bool, + failure_reason: Option<&str>, + run_duration_ms: u64, + dry_run_mode: bool, + llm_client: Option<&arc_llm::client::Client>, + execution_env: &Arc, + provider_enum: Provider, + model: &str, + styles: &'static Styles, +) { + let cp = match Checkpoint::load(&logs_dir.join("checkpoint.json")) { + Ok(cp) => cp, + Err(e) => { + eprintln!( + "{yellow}Warning:{reset} Could not load checkpoint, skipping retro: {e}", + yellow = styles.yellow, + reset = styles.reset, + ); + return; + } + }; + + let stage_durations = crate::retro::extract_stage_durations(logs_dir); + let mut retro = crate::retro::derive_retro( + run_id, + pipeline_name, + goal, + &cp, + failed, + failure_reason, + run_duration_ms, + &stage_durations, + ); + + match retro.save(logs_dir) { + Ok(()) => {} + Err(e) => { + eprintln!( + "{yellow}Warning:{reset} Failed to save initial retro: {e}", + yellow = styles.yellow, + reset = styles.reset, + ); + } + } + + // Run retro agent session + let narrative_result = if dry_run_mode { + Ok(crate::retro_agent::dry_run_narrative()) + } else if let Some(client) = llm_client { + crate::retro_agent::run_retro_agent(execution_env, logs_dir, client, provider_enum, model) + .await + } else { + Err(anyhow::anyhow!("No LLM client available")) + }; + + match narrative_result { + Ok(narrative) => { + retro.apply_narrative(narrative); + match retro.save(logs_dir) { + Ok(()) => { + eprintln!( + "{dim}Retro saved to {}/retro.json{reset}", + logs_dir.display(), + dim = styles.dim, + reset = styles.reset, + ); + } + Err(e) => { + eprintln!( + "{yellow}Warning:{reset} Failed to save retro with narrative: {e}", + yellow = styles.yellow, + reset = styles.reset, + ); + } + } + } + Err(e) => { + eprintln!( + "{dim}Retro agent skipped: {e}{reset}", + dim = styles.dim, + reset = styles.reset, + ); + } + } +} + #[cfg(test)] mod tests { #[test] diff --git a/crates/arc-workflows/src/retro_agent.rs b/crates/arc-workflows/src/retro_agent.rs index 1b76a980a..1a0bba1cb 100644 --- a/crates/arc-workflows/src/retro_agent.rs +++ b/crates/arc-workflows/src/retro_agent.rs @@ -1,5 +1,6 @@ use std::path::Path; use std::sync::{Arc, Mutex}; +use std::time::Duration; use arc_agent::{ AnthropicProfile, ExecutionEnvironment, GeminiProfile, OpenAiProfile, ProviderProfile, @@ -141,7 +142,7 @@ pub async fn run_retro_agent( Box::pin(async move { let narrative: RetroNarrative = serde_json::from_value(args) .map_err(|e| format!("Invalid retro submission: {e}"))?; - *captured.lock().unwrap() = Some(narrative); + *captured.lock().unwrap_or_else(|e| e.into_inner()) = Some(narrative); Ok("Retrospective submitted successfully.".to_string()) }) }), @@ -152,6 +153,7 @@ pub async fn run_retro_agent( let config = SessionConfig { max_tool_rounds_per_input: 10, + wall_clock_timeout: Some(Duration::from_secs(120)), // Disable features not needed for retro analysis enable_context_compaction: false, skill_dirs: Some(vec![]), @@ -183,7 +185,7 @@ pub async fn run_retro_agent( // Extract the captured narrative let narrative = captured .lock() - .unwrap() + .unwrap_or_else(|e| e.into_inner()) .take() .ok_or_else(|| anyhow::anyhow!("Retro agent did not call submit_retro"))?;