Log handoff token counts (#1954)

Signed-off-by: npub12gtutshhh76rx0jx697f32f9tffd4hhp3hx58fp4x6u4uemkm7sqf8f757 <5217c5c2f7bfb4333e46d17c98a9255a52dadee18dcd43a43536b95e6776dfa0@sprout-oss.stage.blox.sqprod.co>
Co-authored-by: npub12gtutshhh76rx0jx697f32f9tffd4hhp3hx58fp4x6u4uemkm7sqf8f757 <5217c5c2f7bfb4333e46d17c98a9255a52dadee18dcd43a43536b95e6776dfa0@sprout-oss.stage.blox.sqprod.co>

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.

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<tokio::process::ChildStdout>,
+    stderr: Arc<StdMutex<String>>,
     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;
 }