Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 52 additions & 17 deletions crates/buzz-agent/src/handoff.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -81,15 +94,42 @@ 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()
);
HandoffOutcome::Performed
}

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
Expand All @@ -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
Expand All @@ -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,
}
}

Expand Down Expand Up @@ -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.
Expand Down
39 changes: 37 additions & 2 deletions crates/buzz-agent/tests/regressions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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,
}

Expand All @@ -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,
}
}
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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;
}

Expand Down