From a0020b22a5f438c279095f52931b741c9ae2f6fe Mon Sep 17 00:00:00 2001 From: Zaki Date: Sat, 28 Mar 2026 11:20:01 -0700 Subject: [PATCH] fix(routines): address review feedback on retry loop - Remove outer retry for LlmFailed errors since RetryProvider already handles transient LLM failures with its own bounded budget, preventing multiplicative retry counts - Preserve Option semantics for token accumulation: None means "unknown/not tracked" rather than converting to Some(0) via unwrap_or(0), so downstream API/UI correctly distinguishes null from zero - Add partial_tokens field to EmptyResponse and TruncatedResponse variants so token usage from those failed attempts is captured in the retry accumulator - Persist accumulated token total on final failure path so usage from earlier retry attempts is not silently discarded Co-Authored-By: Claude Opus 4.6 (1M context) --- src/agent/routine_engine.rs | 86 ++++++++++++++++++++++++++++--------- src/error.rs | 26 ++++++++--- 2 files changed, 86 insertions(+), 26 deletions(-) diff --git a/src/agent/routine_engine.rs b/src/agent/routine_engine.rs index 7a9b84db..6ec80f8e 100644 --- a/src/agent/routine_engine.rs +++ b/src/agent/routine_engine.rs @@ -1095,9 +1095,13 @@ async fn execute_routine(ctx: EngineContext, routine: Routine, run: RoutineRun) let is_lightweight = matches!(routine.action, RoutineAction::Lightweight { .. }); - let result = { + // The retry block returns both the execution result and any accumulated + // token count so that usage is preserved even on final failure. + let (result, accumulated_tokens) = { let mut attempt = 0u32; - let mut accumulated_tokens: i32 = 0; + // Track accumulated tokens as Option to preserve None semantics: + // None = no attempt reported tokens; Some(n) = at least one attempt did. + let mut accumulated_tokens: Option = None; let uses_tools = matches!( routine.action, RoutineAction::Lightweight { @@ -1106,6 +1110,33 @@ async fn execute_routine(ctx: EngineContext, routine: Routine, run: RoutineRun) } ) && ctx.config.lightweight_tools_enabled; + /// Extract partial_tokens from any RoutineError variant that carries them. + fn extract_partial_tokens(e: &RoutineError) -> Option { + match e { + RoutineError::LlmFailed { + partial_tokens: Some(t), + .. + } + | RoutineError::EmptyResponse { + partial_tokens: Some(t), + } + | RoutineError::TruncatedResponse { + partial_tokens: Some(t), + } => Some(*t), + _ => None, + } + } + + /// Merge an optional partial token count into the accumulator, + /// only materializing Some when at least one source had Some. + fn accumulate(acc: Option, partial: Option) -> Option { + match (acc, partial) { + (Some(a), Some(p)) => Some(a.saturating_add(p)), + (Some(a), None) => Some(a), + (None, p) => p, + } + } + loop { let execution_result = match &routine.action { RoutineAction::Lightweight { @@ -1142,24 +1173,22 @@ async fn execute_routine(ctx: EngineContext, routine: Routine, run: RoutineRun) match execution_result { Ok((status, summary, tokens)) => { - // Accumulate tokens from this attempt with any previous attempts - let total = Some(accumulated_tokens.saturating_add(tokens.unwrap_or(0))); - break Ok((status, summary, total)); + // Merge tokens: only produce Some when at least one source had Some. + let total = accumulate(accumulated_tokens, tokens); + break (Ok((status, summary, total)), accumulated_tokens); } Err(ref e) if is_lightweight && !uses_tools && e.is_retryable() + // Skip outer retry for LlmFailed — RetryProvider already + // retries transient LLM errors with its own budget. Retrying + // here would create a multiplicative retry count. + && !matches!(e, RoutineError::LlmFailed { .. }) && attempt < MAX_RETRIES => { - // Accumulate any partial tokens from the failed attempt - if let RoutineError::LlmFailed { - partial_tokens: Some(t), - .. - } = e - { - accumulated_tokens = accumulated_tokens.saturating_add(*t); - } + // Accumulate partial tokens from the failed attempt. + accumulated_tokens = accumulate(accumulated_tokens, extract_partial_tokens(e)); attempt += 1; @@ -1169,7 +1198,11 @@ async fn execute_routine(ctx: EngineContext, routine: Routine, run: RoutineRun) tracing::event!(target: "transient_routine_errors", tracing::Level::WARN, routine = %routine.name, attempt = attempt, max_retries = MAX_RETRIES, delay_ms = delay.as_millis() as u64, "Transient routine error, retrying: {}", e); tokio::time::sleep(delay).await; } - Err(e) => break Err(e), + Err(e) => { + // Accumulate tokens from the final failed attempt. + accumulated_tokens = accumulate(accumulated_tokens, extract_partial_tokens(&e)); + break (Err(e), accumulated_tokens); + } } } }; @@ -1177,12 +1210,13 @@ async fn execute_routine(ctx: EngineContext, routine: Routine, run: RoutineRun) // Decrement running count ctx.running_count.fetch_sub(1, Ordering::Relaxed); - // Process result + // Process result — on failure, preserve accumulated token total from + // earlier retry attempts so usage reporting stays accurate. let (status, summary, tokens) = match result { Ok(execution) => execution, Err(e) => { tracing::error!(routine = %routine.name, "Execution failed: {}", e); - (RunStatus::Failed, Some(e.to_string()), None) + (RunStatus::Failed, Some(e.to_string()), accumulated_tokens) } }; @@ -1591,12 +1625,18 @@ fn handle_text_response( ) -> Result<(RunStatus, Option, Option), RoutineError> { let content = content.trim(); - // Empty content guard + // Empty content guard — carry consumed tokens so the retry loop can + // accumulate them even when the response shape is invalid. if content.is_empty() { + let consumed = Some((total_input_tokens + total_output_tokens) as i32); return if finish_reason == FinishReason::Length { - Err(RoutineError::TruncatedResponse) + Err(RoutineError::TruncatedResponse { + partial_tokens: consumed, + }) } else { - Err(RoutineError::EmptyResponse) + Err(RoutineError::EmptyResponse { + partial_tokens: consumed, + }) }; } @@ -2591,8 +2631,12 @@ mod tests { partial_tokens: Some(42), retryable: true, }, - RoutineError::EmptyResponse, - RoutineError::TruncatedResponse, + RoutineError::EmptyResponse { + partial_tokens: None, + }, + RoutineError::TruncatedResponse { + partial_tokens: Some(100), + }, ]; for err in &transient_errors { assert!(err.is_retryable(), "{} should be retryable", err); diff --git a/src/error.rs b/src/error.rs index 5da9068f..79b10d4f 100644 --- a/src/error.rs +++ b/src/error.rs @@ -411,10 +411,16 @@ pub enum RoutineError { JobDispatchFailed { reason: String }, #[error("LLM returned empty content")] - EmptyResponse, + EmptyResponse { + /// Tokens consumed by the call that produced the empty response. + partial_tokens: Option, + }, #[error("LLM response truncated (finish_reason=length) with no content")] - TruncatedResponse, + TruncatedResponse { + /// Tokens consumed by the call that produced the truncated response. + partial_tokens: Option, + }, } impl RoutineError { @@ -428,7 +434,7 @@ impl RoutineError { pub fn is_retryable(&self) -> bool { match self { RoutineError::LlmFailed { retryable, .. } => *retryable, - RoutineError::EmptyResponse | RoutineError::TruncatedResponse => true, + RoutineError::EmptyResponse { .. } | RoutineError::TruncatedResponse { .. } => true, _ => false, } } @@ -561,8 +567,18 @@ mod tests { } .is_retryable() ); - assert!(RoutineError::EmptyResponse.is_retryable()); - assert!(RoutineError::TruncatedResponse.is_retryable()); + assert!( + RoutineError::EmptyResponse { + partial_tokens: None + } + .is_retryable() + ); + assert!( + RoutineError::TruncatedResponse { + partial_tokens: None + } + .is_retryable() + ); // Hard failures should NOT be retryable assert!(