diff --git a/lib/components/fabro-workflow/src/interview_runtime.rs b/lib/components/fabro-workflow/src/interview_runtime.rs index dc2855864..69f3e0899 100644 --- a/lib/components/fabro-workflow/src/interview_runtime.rs +++ b/lib/components/fabro-workflow/src/interview_runtime.rs @@ -152,7 +152,10 @@ pub(crate) struct WorkflowAgentQuestionRuntime { interviewer: Arc, emitter: Arc, stage_scope: StageScope, - stage_id: String, + /// Graph node id, reported as the `stage` on interview events. Distinct + /// from `stage_scope.stage_id()`, which is the visit-qualified `StageId` + /// used to key block state. + node_id: String, blocker: Arc, } @@ -162,14 +165,14 @@ impl WorkflowAgentQuestionRuntime { interviewer: Arc, emitter: Arc, stage_scope: StageScope, - stage_id: impl Into, + node_id: impl Into, blocker: Arc, ) -> Self { Self { interviewer, emitter, stage_scope, - stage_id: stage_id.into(), + node_id: node_id.into(), blocker, } } @@ -183,7 +186,7 @@ struct PreparedQuestion { struct PendingAgentQuestionBatch { emitter: Arc, stage_scope: StageScope, - stage_id: String, + node_id: String, questions: Vec<(String, String)>, started_at: Instant, guard: Option, @@ -193,7 +196,7 @@ impl PendingAgentQuestionBatch { fn new( emitter: Arc, stage_scope: StageScope, - stage_id: String, + node_id: String, prepared: &[PreparedQuestion], guard: RunInterviewGuard, started_at: Instant, @@ -201,7 +204,7 @@ impl PendingAgentQuestionBatch { Self { emitter, stage_scope, - stage_id, + node_id, questions: prepared .iter() .map(|prepared_question| { @@ -237,7 +240,7 @@ impl Drop for PendingAgentQuestionBatch { }), question_id: question_id.clone(), question: question.clone(), - stage: self.stage_id.clone(), + stage: self.node_id.clone(), reason: "interrupted".to_string(), duration_ms, }, @@ -274,7 +277,7 @@ impl AgentQuestionRuntime for WorkflowAgentQuestionRuntime { &Event::InterviewStarted { question_id: question.id.clone(), question: question.text.clone(), - stage: self.stage_id.clone(), + stage: self.node_id.clone(), question_type: question.question_type.to_string(), options: question.options.clone(), allow_freeform: question.allow_freeform, @@ -290,7 +293,7 @@ impl AgentQuestionRuntime for WorkflowAgentQuestionRuntime { let cleanup = PendingAgentQuestionBatch::new( Arc::clone(&self.emitter), self.stage_scope.clone(), - self.stage_id.clone(), + self.node_id.clone(), &prepared, self.blocker .block(Arc::clone(&self.emitter), self.stage_scope.stage_id()), @@ -361,7 +364,7 @@ impl WorkflowAgentQuestionRuntime { question.id = internal_question_id(&self.stage_scope, tool_call_id, index); question.options.clone_from(&agent_question.options); question.allow_freeform = agent_question.allow_freeform; - question.stage.clone_from(&self.stage_id); + question.stage.clone_from(&self.node_id); question.metadata.insert( "agent.tool_call_id".to_string(), serde_json::json!(tool_call_id), @@ -401,7 +404,7 @@ impl WorkflowAgentQuestionRuntime { }), question_id: prepared.question.id.clone(), question: prepared.question.text.clone(), - stage: self.stage_id.clone(), + stage: self.node_id.clone(), duration_ms, }, &self.stage_scope, @@ -444,7 +447,7 @@ impl WorkflowAgentQuestionRuntime { actor, question_id: prepared.question.id.clone(), question: prepared.question.text.clone(), - stage: self.stage_id.clone(), + stage: self.node_id.clone(), reason: reason.to_string(), duration_ms, }, diff --git a/lib/components/fabro-workflow/src/pipeline/execute.rs b/lib/components/fabro-workflow/src/pipeline/execute.rs index 60d07b070..3b0fb81f9 100644 --- a/lib/components/fabro-workflow/src/pipeline/execute.rs +++ b/lib/components/fabro-workflow/src/pipeline/execute.rs @@ -5,6 +5,7 @@ use fabro_core::executor::ExecutorBuilder; use fabro_core::handler::NodeHandler; use fabro_core::state::ExecutionState; use tokio::sync::watch; +use tokio::task::JoinHandle; use tokio::time::{Instant as TokioInstant, sleep_until}; use tokio_util::sync::CancellationToken; @@ -30,6 +31,50 @@ fn seed_context_from_checkpoint(checkpoint: Option<&Checkpoint>) -> Context { context } +/// Background watchdog that cancels a run which stops emitting events. +struct StallWatchdog { + /// Cancelled by the monitor once the run stalls. Handed to the executor. + stall_token: CancellationToken, + /// Cancelled by us to stop the monitor once the run finishes. + shutdown: CancellationToken, + task: JoinHandle<()>, +} + +impl StallWatchdog { + fn spawn( + stall_timeout: Duration, + emitter: Arc, + interview_blocks: watch::Receiver, + ) -> Self { + let stall_token = CancellationToken::new(); + let shutdown = CancellationToken::new(); + emitter.touch(); + let task = tokio::spawn(monitor_for_stall( + stall_timeout, + stall_token.clone(), + shutdown.clone(), + emitter, + interview_blocks, + )); + Self { + stall_token, + shutdown, + task, + } + } + + fn stall_token(&self) -> CancellationToken { + self.stall_token.clone() + } + + async fn stop(self) { + self.shutdown.cancel(); + if let Err(error) = self.task.await { + tracing::error!(error = ?error, "stall watchdog task failed"); + } + } +} + /// Cancels `stall_token` once the run goes `stall_timeout` without emitting an /// event. Waiting on human input suspends the timer, and the first unblock /// starts a fresh full deadline. @@ -228,31 +273,19 @@ pub async fn execute(init: Initialized) -> Executed { None }; - let stall_timeout_opt = graph.stall_timeout(); - let stall_token = stall_timeout_opt.map(|_| CancellationToken::new()); - let stall_watchdog = - if let (Some(stall_timeout), Some(ref token)) = (stall_timeout_opt, &stall_token) { - let shutdown = CancellationToken::new(); - let emitter = Arc::clone(&engine.run.emitter); - let interview_blocks = engine.run.interview_blocker.subscribe(); - emitter.touch(); - let task = tokio::spawn(monitor_for_stall( - stall_timeout, - token.clone(), - shutdown.clone(), - emitter, - interview_blocks, - )); - Some((shutdown, task)) - } else { - None - }; + let stall_watchdog = graph.stall_timeout().map(|stall_timeout| { + StallWatchdog::spawn( + stall_timeout, + Arc::clone(&engine.run.emitter), + engine.run.interview_blocker.subscribe(), + ) + }); let mut builder = ExecutorBuilder::new(handler as Arc>) .lifecycle(Box::new(lifecycle)); builder = builder.cancel_token(run_options.cancel_token.clone()); - if let Some(token) = stall_token.clone() { + if let Some(token) = stall_watchdog.as_ref().map(StallWatchdog::stall_token) { builder = builder.stall_token(token); } if let Some(limit) = max_node_visits { @@ -262,11 +295,8 @@ pub async fn execute(init: Initialized) -> Executed { let executor = builder.build(); let result = executor.run(&wf_graph, state).await; - if let Some((shutdown, task)) = stall_watchdog { - shutdown.cancel(); - if let Err(error) = task.await { - tracing::error!(error = ?error, "stall watchdog task failed"); - } + if let Some(watchdog) = stall_watchdog { + watchdog.stop().await; } let (outcome, final_context) = match result { diff --git a/lib/components/fabro-workflow/src/pipeline/execute/tests.rs b/lib/components/fabro-workflow/src/pipeline/execute/tests.rs index 8b77767ce..af31798e3 100644 --- a/lib/components/fabro-workflow/src/pipeline/execute/tests.rs +++ b/lib/components/fabro-workflow/src/pipeline/execute/tests.rs @@ -1460,12 +1460,14 @@ async fn stall_watchdog_starts_a_fresh_deadline_after_human_input() { #[tokio::test] async fn stall_watchdog_suspends_while_run_waits_for_human_input() { let dir = tempfile::tempdir().unwrap(); - let graph = interview_wait_graph(Duration::from_millis(50), None); + // The blocked wait outruns the stall timeout, so this only passes if the + // watchdog stays suspended and then restarts on a fresh deadline. + let graph = interview_wait_graph(Duration::from_millis(300), None); let mut registry = make_registry(); registry.register( "interview_wait", Box::new(InterviewWaitHandler { - wait_ms: 150, + wait_ms: 500, active_ms: 10, }), ); @@ -1486,12 +1488,14 @@ async fn stall_watchdog_suspends_while_run_waits_for_human_input() { #[tokio::test] async fn node_timeout_excludes_human_input_wait() { let dir = tempfile::tempdir().unwrap(); - let graph = interview_wait_graph(Duration::ZERO, Some(Duration::from_millis(50))); + // The blocked wait outruns the node timeout, but the active work is well + // inside it, so this only fails if the interview wait is being charged. + let graph = interview_wait_graph(Duration::ZERO, Some(Duration::from_millis(300))); let mut registry = make_registry(); registry.register( "interview_wait", Box::new(InterviewWaitHandler { - wait_ms: 150, + wait_ms: 500, active_ms: 20, }), );