diff --git a/crates/attractor/src/cli/backend.rs b/crates/attractor/src/cli/backend.rs index bc033df2c..f5aaaba85 100644 --- a/crates/attractor/src/cli/backend.rs +++ b/crates/attractor/src/cli/backend.rs @@ -280,20 +280,8 @@ impl CodergenBackend for AgentBackend { })?; // Aggregate token usage from all assistant turns. - let (mut turn_count, mut tool_call_count, mut input_tokens, mut output_tokens) = - (0usize, 0usize, 0i64, 0i64); - let (mut cache_read_tokens, mut cache_write_tokens, mut reasoning_tokens): ( - Option, - Option, - Option, - ) = (None, None, None); - - fn add_optional(acc: &mut Option, val: Option) { - if let Some(v) = val { - *acc = Some(acc.unwrap_or(0) + v); - } - } - + let (mut turn_count, mut tool_call_count) = (0usize, 0usize); + let mut total_usage = llm::types::Usage::default(); for turn in session.history().turns() { if let Turn::Assistant { tool_calls, usage, .. @@ -301,26 +289,22 @@ impl CodergenBackend for AgentBackend { { turn_count += 1; tool_call_count += tool_calls.len(); - input_tokens += usage.input_tokens; - output_tokens += usage.output_tokens; - add_optional(&mut cache_read_tokens, usage.cache_read_tokens); - add_optional(&mut cache_write_tokens, usage.cache_write_tokens); - add_optional(&mut reasoning_tokens, usage.reasoning_tokens); + total_usage = total_usage + usage.clone(); } } let stage_usage = StageUsage { model: self.model.clone(), - input_tokens, - output_tokens, - cache_read_tokens, - cache_write_tokens, - reasoning_tokens, + input_tokens: total_usage.input_tokens, + output_tokens: total_usage.output_tokens, + cache_read_tokens: total_usage.cache_read_tokens, + cache_write_tokens: total_usage.cache_write_tokens, + reasoning_tokens: total_usage.reasoning_tokens, }; // Print session summary to stderr. if self.verbose >= 1 { - let total_tokens = input_tokens + output_tokens; + let total_tokens = total_usage.input_tokens + total_usage.output_tokens; let token_str = if total_tokens >= 1000 { format!("{}k tokens", total_tokens / 1000) } else { diff --git a/crates/attractor/src/cli/mod.rs b/crates/attractor/src/cli/mod.rs index 181e3a5c6..e7778c167 100644 --- a/crates/attractor/src/cli/mod.rs +++ b/crates/attractor/src/cli/mod.rs @@ -178,11 +178,12 @@ pub fn format_event_summary(event: &PipelineEvent, styles: &Styles) -> String { PipelineEvent::PipelineFailed { error, duration_ms } => { format!("[PIPELINE_FAILED] error=\"{error}\" duration={duration_ms}ms") } - PipelineEvent::StageStarted { name, index, handler_type } => { + PipelineEvent::StageStarted { name, index, handler_type, attempt, max_attempts } => { let mut s = format!("[STAGE_STARTED] name={name} index={index}"); if let Some(ht) = handler_type { s.push_str(&format!(" handler_type={ht}")); } + s.push_str(&format!(" attempt={attempt}/{max_attempts}")); s } PipelineEvent::StageCompleted { @@ -196,6 +197,8 @@ pub fn format_event_summary(event: &PipelineEvent, styles: &Styles) -> String { failure_reason, notes, files_touched, + attempt, + max_attempts, } => { let mut s = format!("[STAGE_COMPLETED] name={name} index={index} duration={duration_ms}ms status={status}"); if let Some(label) = preferred_label { @@ -222,6 +225,7 @@ pub fn format_event_summary(event: &PipelineEvent, styles: &Styles) -> String { if !files_touched.is_empty() { s.push_str(&format!(" files_touched={}", files_touched.len())); } + s.push_str(&format!(" attempt={attempt}/{max_attempts}")); s } PipelineEvent::StageFailed { @@ -243,10 +247,11 @@ pub fn format_event_summary(event: &PipelineEvent, styles: &Styles) -> String { name, index, attempt, + max_attempts, delay_ms, } => { format!( - "[STAGE_RETRYING] name={name} index={index} attempt={attempt} delay={delay_ms}ms" + "[STAGE_RETRYING] name={name} index={index} attempt={attempt}/{max_attempts} delay={delay_ms}ms" ) } PipelineEvent::ParallelStarted { branch_count, join_policy, error_policy } => { @@ -290,6 +295,19 @@ pub fn format_event_summary(event: &PipelineEvent, styles: &Styles) -> String { PipelineEvent::CheckpointSaved { node_id } => { format!("[CHECKPOINT_SAVED] node={node_id}") } + PipelineEvent::EdgeSelected { from_node, to_node, label, condition } => { + let mut s = format!("[EDGE_SELECTED] from={from_node} to={to_node}"); + if let Some(l) = label { + s.push_str(&format!(" label=\"{l}\"")); + } + if let Some(c) = condition { + s.push_str(&format!(" condition=\"{c}\"")); + } + s + } + PipelineEvent::LoopRestart { from_node, to_node } => { + format!("[LOOP_RESTART] from={from_node} to={to_node}") + } PipelineEvent::Prompt { stage, text } => { let truncated = if text.len() > 80 { &text[..80] } else { text }; format!("[PROMPT] stage={stage} text=\"{truncated}\"") @@ -384,13 +402,14 @@ pub fn format_event_detail(event: &PipelineEvent, styles: &Styles) -> String { PipelineEvent::PipelineFailed { error, duration_ms } => { format!("{d}── PIPELINE_FAILED ──────────────────────────{r}\n {d}error:{r} {error}\n {d}duration_ms:{r} {duration_ms}\n") } - PipelineEvent::StageStarted { name, index, handler_type } => { + PipelineEvent::StageStarted { name, index, handler_type, attempt, max_attempts } => { let mut s = format!( "{d}── STAGE_STARTED ────────────────────────────{r}\n {d}name:{r} {name}\n {d}index:{r} {index}\n" ); if let Some(ht) = handler_type { s.push_str(&format!(" {d}handler_type:{r} {ht}\n")); } + s.push_str(&format!(" {d}attempt:{r} {attempt}/{max_attempts}\n")); s } PipelineEvent::StageCompleted { @@ -404,6 +423,8 @@ pub fn format_event_detail(event: &PipelineEvent, styles: &Styles) -> String { failure_reason, notes, files_touched, + attempt, + max_attempts, } => { let mut s = format!("{d}── STAGE_COMPLETED ──────────────────────────{r}\n {d}name:{r} {name}\n {d}index:{r} {index}\n {d}duration_ms:{r} {duration_ms}\n {d}status:{r} {status}\n"); if let Some(label) = preferred_label { @@ -442,6 +463,7 @@ pub fn format_event_detail(event: &PipelineEvent, styles: &Styles) -> String { if let Some(n) = notes { s.push_str(&format!(" {d}notes:{r} {n}\n")); } + s.push_str(&format!(" {d}attempt:{r} {attempt}/{max_attempts}\n")); s } PipelineEvent::StageFailed { @@ -461,9 +483,10 @@ pub fn format_event_detail(event: &PipelineEvent, styles: &Styles) -> String { name, index, attempt, + max_attempts, delay_ms, } => { - format!("{d}── STAGE_RETRYING ───────────────────────────{r}\n {d}name:{r} {name}\n {d}index:{r} {index}\n {d}attempt:{r} {attempt}\n {d}delay_ms:{r} {delay_ms}\n") + format!("{d}── STAGE_RETRYING ───────────────────────────{r}\n {d}name:{r} {name}\n {d}index:{r} {index}\n {d}attempt:{r} {attempt}/{max_attempts}\n {d}delay_ms:{r} {delay_ms}\n") } PipelineEvent::ParallelStarted { branch_count, join_policy, error_policy } => { format!("{d}── PARALLEL_STARTED ─────────────────────────{r}\n {d}branch_count:{r} {branch_count}\n {d}join_policy:{r} {join_policy}\n {d}error_policy:{r} {error_policy}\n") @@ -508,6 +531,19 @@ pub fn format_event_detail(event: &PipelineEvent, styles: &Styles) -> String { "{d}── CHECKPOINT_SAVED ─────────────────────────{r}\n {d}node_id:{r} {node_id}\n" ) } + PipelineEvent::EdgeSelected { from_node, to_node, label, condition } => { + let mut s = format!("{d}── EDGE_SELECTED ────────────────────────────{r}\n {d}from:{r} {from_node}\n {d}to:{r} {to_node}\n"); + if let Some(l) = label { + s.push_str(&format!(" {d}label:{r} {l}\n")); + } + if let Some(c) = condition { + s.push_str(&format!(" {d}condition:{r} {c}\n")); + } + s + } + PipelineEvent::LoopRestart { from_node, to_node } => { + format!("{d}── LOOP_RESTART ─────────────────────────────{r}\n {d}from:{r} {from_node}\n {d}to:{r} {to_node}\n") + } PipelineEvent::Prompt { stage, text } => { format!("{d}── PROMPT ───────────────────────────────────{r}\n {d}stage:{r} {stage}\n {d}text:{r}\n{text}\n") } diff --git a/crates/attractor/src/engine.rs b/crates/attractor/src/engine.rs index 8476ba73e..8a3966546 100644 --- a/crates/attractor/src/engine.rs +++ b/crates/attractor/src/engine.rs @@ -572,6 +572,7 @@ impl PipelineEngine { name: node.label().to_string(), index: stage_index, attempt: usize::try_from(attempt).unwrap_or(usize::MAX), + max_attempts: usize::try_from(policy.max_attempts).unwrap_or(usize::MAX), delay_ms: millis_u64(delay), }); tokio::time::sleep(delay).await; @@ -595,6 +596,7 @@ impl PipelineEngine { name: node.label().to_string(), index: stage_index, attempt: usize::try_from(attempt).unwrap_or(usize::MAX), + max_attempts: usize::try_from(policy.max_attempts).unwrap_or(usize::MAX), delay_ms: millis_u64(delay), }); tokio::time::sleep(delay).await; @@ -811,6 +813,8 @@ impl PipelineEngine { name: node.label().to_string(), index: stage_index, handler_type: node.handler_type().map(String::from), + attempt: 1, + max_attempts: usize::try_from(retry_policy.max_attempts).unwrap_or(usize::MAX), }); if node.handler_type() != Some("wait.human") { self.inform( @@ -866,6 +870,8 @@ impl PipelineEngine { failure_reason: outcome.failure_reason.clone(), notes: outcome.notes.clone(), files_touched: outcome.files_touched.clone(), + attempt: usize::try_from(attempts_used).unwrap_or(usize::MAX), + max_attempts: usize::try_from(retry_policy.max_attempts).unwrap_or(usize::MAX), }); self.inform( &format!("Stage completed: {}", node.label()), @@ -891,6 +897,14 @@ impl PipelineEngine { // Step 5: Select next edge (done before checkpoint so we can store next_node_id) let next_edge = select_edge(&node.id, &outcome, &context, graph); + if let Some(edge) = next_edge { + self.services.emitter.emit(&PipelineEvent::EdgeSelected { + from_node: node.id.clone(), + to_node: edge.to.clone(), + label: edge.label().map(String::from), + condition: edge.condition().map(String::from), + }); + } let next_node_id_for_checkpoint = next_edge.map(|e| e.to.clone()); // Step 6: Save checkpoint with all state @@ -939,6 +953,10 @@ impl PipelineEngine { incoming_edge = Some(edge); // Gap #6: Handle loop_restart by recursively running from the target if edge.loop_restart() { + self.services.emitter.emit(&PipelineEvent::LoopRestart { + from_node: node.id.clone(), + to_node: edge.to.clone(), + }); return Box::pin(self.run_internal( graph, config, diff --git a/crates/attractor/src/event.rs b/crates/attractor/src/event.rs index 7143f107a..e2b38e6c9 100644 --- a/crates/attractor/src/event.rs +++ b/crates/attractor/src/event.rs @@ -22,6 +22,8 @@ pub enum PipelineEvent { name: String, index: usize, handler_type: Option, + attempt: usize, + max_attempts: usize, }, StageCompleted { name: String, @@ -34,6 +36,8 @@ pub enum PipelineEvent { failure_reason: Option, notes: Option, files_touched: Vec, + attempt: usize, + max_attempts: usize, }, StageFailed { name: String, @@ -46,6 +50,7 @@ pub enum PipelineEvent { name: String, index: usize, attempt: usize, + max_attempts: usize, delay_ms: u64, }, ParallelStarted { @@ -86,6 +91,16 @@ pub enum PipelineEvent { CheckpointSaved { node_id: String, }, + EdgeSelected { + from_node: String, + to_node: String, + label: Option, + condition: Option, + }, + LoopRestart { + from_node: String, + to_node: String, + }, Prompt { stage: String, text: String, @@ -218,17 +233,23 @@ mod tests { name: "plan".to_string(), index: 0, handler_type: Some("codergen".to_string()), + attempt: 1, + max_attempts: 3, }; let json = serde_json::to_string(&event).unwrap(); assert!(json.contains("StageStarted")); assert!(json.contains("plan")); assert!(json.contains("\"handler_type\":\"codergen\"")); + assert!(json.contains("\"attempt\":1")); + assert!(json.contains("\"max_attempts\":3")); // None handler_type serializes as null let event_none = PipelineEvent::StageStarted { name: "plan".to_string(), index: 0, handler_type: None, + attempt: 1, + max_attempts: 1, }; let json_none = serde_json::to_string(&event_none).unwrap(); assert!(json_none.contains("\"handler_type\":null")); @@ -324,11 +345,15 @@ mod tests { failure_reason: Some("lint errors remain".to_string()), notes: Some("fixed 3 of 5 issues".to_string()), files_touched: vec!["src/main.rs".to_string()], + attempt: 2, + max_attempts: 3, }; let json = serde_json::to_string(&event).unwrap(); assert!(json.contains("\"failure_reason\":\"lint errors remain\"")); assert!(json.contains("\"notes\":\"fixed 3 of 5 issues\"")); assert!(json.contains("src/main.rs")); + assert!(json.contains("\"attempt\":2")); + assert!(json.contains("\"max_attempts\":3")); let event_none = PipelineEvent::StageCompleted { name: "plan".to_string(), @@ -341,6 +366,8 @@ mod tests { failure_reason: None, notes: None, files_touched: vec![], + attempt: 1, + max_attempts: 1, }; let json_none = serde_json::to_string(&event_none).unwrap(); assert!(json_none.contains("\"failure_reason\":null")); @@ -438,4 +465,68 @@ mod tests { let deserialized: PipelineEvent = serde_json::from_str(&json).unwrap(); assert!(matches!(deserialized, PipelineEvent::CompactionCompleted { stage, .. } if stage == "code")); } + + #[test] + fn edge_selected_event_serialization() { + let event = PipelineEvent::EdgeSelected { + from_node: "plan".to_string(), + to_node: "code".to_string(), + label: Some("success".to_string()), + condition: Some("outcome == 'success'".to_string()), + }; + let json = serde_json::to_string(&event).unwrap(); + assert!(json.contains("EdgeSelected")); + assert!(json.contains("\"from_node\":\"plan\"")); + assert!(json.contains("\"to_node\":\"code\"")); + assert!(json.contains("\"label\":\"success\"")); + assert!(json.contains("\"condition\":\"outcome == 'success'\"")); + + let deserialized: PipelineEvent = serde_json::from_str(&json).unwrap(); + assert!(matches!(deserialized, PipelineEvent::EdgeSelected { from_node, to_node, .. } if from_node == "plan" && to_node == "code")); + + // None label/condition + let event_none = PipelineEvent::EdgeSelected { + from_node: "a".to_string(), + to_node: "b".to_string(), + label: None, + condition: None, + }; + let json_none = serde_json::to_string(&event_none).unwrap(); + assert!(json_none.contains("\"label\":null")); + assert!(json_none.contains("\"condition\":null")); + } + + #[test] + fn loop_restart_event_serialization() { + let event = PipelineEvent::LoopRestart { + from_node: "review".to_string(), + to_node: "code".to_string(), + }; + let json = serde_json::to_string(&event).unwrap(); + assert!(json.contains("LoopRestart")); + assert!(json.contains("\"from_node\":\"review\"")); + assert!(json.contains("\"to_node\":\"code\"")); + + let deserialized: PipelineEvent = serde_json::from_str(&json).unwrap(); + assert!(matches!(deserialized, PipelineEvent::LoopRestart { from_node, to_node } if from_node == "review" && to_node == "code")); + } + + #[test] + fn stage_retrying_event_serialization() { + let event = PipelineEvent::StageRetrying { + name: "lint".to_string(), + index: 2, + attempt: 3, + max_attempts: 5, + delay_ms: 400, + }; + let json = serde_json::to_string(&event).unwrap(); + assert!(json.contains("StageRetrying")); + assert!(json.contains("\"attempt\":3")); + assert!(json.contains("\"max_attempts\":5")); + assert!(json.contains("\"delay_ms\":400")); + + let deserialized: PipelineEvent = serde_json::from_str(&json).unwrap(); + assert!(matches!(deserialized, PipelineEvent::StageRetrying { max_attempts: 5, .. })); + } }