diff --git a/crates/buzz-agent/src/handoff.rs b/crates/buzz-agent/src/handoff.rs index 887b9560da..9c27c6606d 100644 --- a/crates/buzz-agent/src/handoff.rs +++ b/crates/buzz-agent/src/handoff.rs @@ -4,6 +4,18 @@ use crate::config::{ }; use crate::types::HistoryItem; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct HandoffTokenCounts { + before: u64, + after: u64, +} + +impl std::fmt::Display for HandoffTokenCounts { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{} -> {} tokens", self.before, self.after) + } +} + pub(crate) enum HandoffOutcome { Performed, Skipped, @@ -28,6 +40,7 @@ impl RunCtx<'_> { return HandoffOutcome::Skipped; } let prompt = self.build_handoff_prompt(); + let tokens_before = self.projected_handoff_input_tokens(); let summary = tokio::select! { biased; _ = self.cancel.changed() => return HandoffOutcome::Cancelled, @@ -81,8 +94,12 @@ impl RunCtx<'_> { self.history.push(HistoryItem::User(prompt)); } *self.handoff_count += 1; + let token_counts = HandoffTokenCounts { + before: tokens_before, + after: estimate_history_tokens(self.history), + }; tracing::info!( - "handoff #{} (history {prior} -> {} items)", + "handoff #{} (history {prior} -> {} items; {token_counts})", *self.handoff_count, self.history.len() ); @@ -90,6 +107,29 @@ impl RunCtx<'_> { } fn should_handoff(&self) -> bool { + match *self.last_request_input_tokens { + Some(_) => { + self.projected_handoff_input_tokens() + >= token_threshold(self.cfg.max_context_tokens, self.cfg.max_output_tokens) + } + None => { + let bytes: usize = self + .history + .iter() + .map(HistoryItem::context_pressure_bytes) + .sum(); + bytes + > byte_fallback_threshold( + self.cfg.max_context_tokens, + self.cfg.max_output_tokens, + self.cfg.max_history_bytes, + ) + } + } + } + + fn projected_handoff_input_tokens(&self) -> u64 { + let current_tokens = estimate_history_tokens(self.history); match *self.last_request_input_tokens { // Token-first: the provider told us exactly how many input tokens // the PREVIOUS request used. But history has grown since that @@ -107,9 +147,7 @@ impl RunCtx<'_> { .map(HistoryItem::context_pressure_bytes) .sum(); let grown = current_bytes.saturating_sub(measured_bytes); - let projected = measured_tokens.saturating_add(estimate_tokens_from_bytes(grown)); - projected - >= token_threshold(self.cfg.max_context_tokens, self.cfg.max_output_tokens) + measured_tokens.saturating_add(estimate_tokens_from_bytes(grown)) } // No usage yet (first request, or just after a handoff reset). // Fall back to the byte heuristic, capped conservatively so a @@ -122,19 +160,7 @@ impl RunCtx<'_> { // Caveat: this can't shrink a single oversized current prompt, // since a handoff re-adds the current prompt verbatim — that is a // prompt-cap concern (MAX_PROMPT_BYTES), not this gate. - None => { - let bytes: usize = self - .history - .iter() - .map(HistoryItem::context_pressure_bytes) - .sum(); - bytes - > byte_fallback_threshold( - self.cfg.max_context_tokens, - self.cfg.max_output_tokens, - self.cfg.max_history_bytes, - ) - } + None => current_tokens, } } @@ -299,6 +325,15 @@ pub(crate) fn clamp_bytes(s: &str, max_bytes: usize) -> String { /// than risk the next request exceeding the window. const CONSERVATIVE_BYTES_PER_TOKEN: u64 = 1; +fn estimate_history_tokens(history: &[HistoryItem]) -> u64 { + estimate_tokens_from_bytes( + history + .iter() + .map(HistoryItem::context_pressure_bytes) + .sum(), + ) +} + /// Estimate tokens from a byte count at the conservative ratio (rounding up, /// so a partial token still counts). At a 1:1 ratio this is just the byte /// count — a guaranteed upper bound on tokens. diff --git a/crates/buzz-agent/tests/regressions.rs b/crates/buzz-agent/tests/regressions.rs index 3c2237ac99..6d25d050f2 100644 --- a/crates/buzz-agent/tests/regressions.rs +++ b/crates/buzz-agent/tests/regressions.rs @@ -7,7 +7,7 @@ use std::collections::VecDeque; use std::process::Stdio; -use std::sync::Arc; +use std::sync::{Arc, Mutex as StdMutex}; use std::time::{Duration, Instant}; use serde_json::{json, Value}; @@ -88,6 +88,7 @@ struct Harness { child: tokio::process::Child, stdin: tokio::process::ChildStdin, stdout: BufReader, + stderr: Arc>, next_id: i64, } @@ -108,15 +109,36 @@ impl Harness { } cmd.stdin(Stdio::piped()) .stdout(Stdio::piped()) - .stderr(Stdio::inherit()) + .stderr(Stdio::piped()) .kill_on_drop(true); let mut child = cmd.spawn().expect("spawn buzz-agent"); let stdin = child.stdin.take().unwrap(); let stdout = BufReader::new(child.stdout.take().unwrap()); + let stderr = child.stderr.take().unwrap(); + let stderr_buf = Arc::new(StdMutex::new(String::new())); + let stderr_out = Arc::clone(&stderr_buf); + tokio::spawn(async move { + let mut reader = BufReader::new(stderr); + let mut line = String::new(); + loop { + line.clear(); + let n = match reader.read_line(&mut line).await { + Ok(n) => n, + Err(_) => break, + }; + if n == 0 { + break; + } + if let Ok(mut out) = stderr_out.lock() { + out.push_str(&line); + } + } + }); Self { child, stdin, stdout, + stderr: stderr_buf, next_id: 1, } } @@ -169,6 +191,10 @@ impl Harness { let _ = tokio::time::timeout(Duration::from_secs(2), self.child.wait()).await; let _ = self.child.start_kill(); } + + fn stderr_text(&self) -> String { + self.stderr.lock().map(|s| s.clone()).unwrap_or_default() + } } fn openai_text(content: &str) -> Value { @@ -1381,6 +1407,15 @@ async fn token_usage_over_budget_triggers_handoff() { "expected handoff summarize() between the two prompts (3 reqs), saw {captured} — \ token gate did not fire on usage over budget" ); + let stderr = h.stderr_text(); + assert!( + stderr.contains("handoff #1 (history"), + "expected handoff log line in stderr, got: {stderr}" + ); + assert!( + stderr.contains(" -> ") && stderr.contains(" tokens"), + "expected handoff log to include before/after token counts, got: {stderr}" + ); h.shutdown().await; }