mirror of
https://github.com/block/buzz.git
synced 2026-10-27 12:40:30 +00:00
fix(agent): gate handoff on provider token usage, not byte estimate (#821)
Signed-off-by: npub1qyvc0c5kl4gqv2fd97fsk46tu378sqgy35vc83rvgfwne90sel7s0ed67d <011987e296fd5006292d2f930b574be47c7801048d1983c46c425d3c95f0cffd@sprout-oss.stage.blox.sqprod.co> Signed-off-by: npub1mprnacetjua2xx3p5eddmhxyk6wv929ymm5py8kd2xfxurxahspqqlgyta <d8473ee32b973aa31a21a65adddcc4b69cc2a8a4dee8121ecd51926e0cddbc02@sprout-oss.stage.blox.sqprod.co> Co-authored-by: npub1qyvc0c5kl4gqv2fd97fsk46tu378sqgy35vc83rvgfwne90sel7s0ed67d <011987e296fd5006292d2f930b574be47c7801048d1983c46c425d3c95f0cffd@sprout-oss.stage.blox.sqprod.co> Co-authored-by: npub1mprnacetjua2xx3p5eddmhxyk6wv929ymm5py8kd2xfxurxahspqqlgyta <d8473ee32b973aa31a21a65adddcc4b69cc2a8a4dee8121ecd51926e0cddbc02@sprout-oss.stage.blox.sqprod.co>
This commit is contained in:
tlongwell-block
@@ -145,6 +145,8 @@ Everything is environment variables. No flags, no config files. (We are a subpro
|
||||
| `SPROUT_AGENT_SYSTEM_PROMPT_FILE` | — | File path. Mutually exclusive with the above. |
|
||||
| `SPROUT_AGENT_MAX_ROUNDS` | `0` | Tool-loop iteration cap. 0 = unlimited. |
|
||||
| `SPROUT_AGENT_MAX_OUTPUT_TOKENS` | `32768` | Per LLM call. Headroom for large tool-call inputs (e.g. file writes via heredoc); Sonnet 4 / Opus 4 cap at 64K. |
|
||||
| `SPROUT_AGENT_MAX_CONTEXT_TOKENS` | `200000` | Provider context window used by the handoff gate. |
|
||||
| `SPROUT_AGENT_MAX_HANDOFFS` | `10` | Max context handoffs per session before falling back to truncation. |
|
||||
| `SPROUT_AGENT_LLM_TIMEOUT_SECS` | `120` | |
|
||||
| `SPROUT_AGENT_TOOL_TIMEOUT_SECS` | `660` | Per-tool call timeout in seconds |
|
||||
| `SPROUT_AGENT_MAX_PARALLEL_TOOLS` | `8` | Max concurrent tool calls per turn (1 = sequential) |
|
||||
|
||||
@@ -34,6 +34,18 @@ pub struct RunCtx<'a> {
|
||||
/// `cfg.stop_max_rejections` we stop calling `_Stop` for that
|
||||
/// session — a runaway hook can't burn rejections on every prompt.
|
||||
pub stop_rejections: &'a mut u32,
|
||||
/// Cache-summed input tokens reported by the provider on this session's
|
||||
/// most recent request (persists across `session/prompt` calls), or `None`
|
||||
/// before the first response and immediately after a handoff resets the
|
||||
/// context. The handoff gate reads this to compare against the token
|
||||
/// budget; falls back to the byte heuristic when `None`.
|
||||
pub last_request_input_tokens: &'a mut Option<u64>,
|
||||
/// History byte size at the moment `last_request_input_tokens` was
|
||||
/// measured. Paired with it so the gate can add a conservative token
|
||||
/// estimate of history that has grown since (tool results, next prompt),
|
||||
/// which the exact-but-stale token count would otherwise miss. Cleared and
|
||||
/// preserved in lockstep with `last_request_input_tokens`.
|
||||
pub last_request_history_bytes: &'a mut Option<usize>,
|
||||
}
|
||||
|
||||
impl RunCtx<'_> {
|
||||
@@ -63,7 +75,15 @@ impl RunCtx<'_> {
|
||||
}
|
||||
match self.maybe_handoff().await {
|
||||
HandoffOutcome::Cancelled => return Ok(StopReason::Cancelled),
|
||||
HandoffOutcome::Performed => {}
|
||||
// Context was just reset — the prior request's token count no
|
||||
// longer describes the (now much smaller) history. Clear both
|
||||
// the token count and its byte baseline so a stale over-
|
||||
// threshold reading can't immediately re-fire the handoff
|
||||
// before the next response reports fresh usage.
|
||||
HandoffOutcome::Performed => {
|
||||
*self.last_request_input_tokens = None;
|
||||
*self.last_request_history_bytes = None;
|
||||
}
|
||||
HandoffOutcome::Skipped => {
|
||||
truncate_history(self.history, self.cfg.max_history_bytes)
|
||||
}
|
||||
@@ -77,6 +97,21 @@ impl RunCtx<'_> {
|
||||
r = self.llm.complete(self.cfg, self.system_prompt, self.history, &tools) => r?,
|
||||
};
|
||||
|
||||
// Record provider-reported input usage so the next loop iteration's
|
||||
// handoff gate can compare it against the token budget. We capture
|
||||
// it together with the history byte size AT THIS MOMENT — which is
|
||||
// exactly the history that was just sent to `complete()` (the
|
||||
// assistant response is appended below, after this point). Pairing
|
||||
// them lets the gate add a conservative estimate for any history
|
||||
// appended before the next request. Preserve both when a response
|
||||
// omits usage (`None`) rather than clobbering — a one-off missing
|
||||
// field shouldn't blind the gate or zero the growth baseline.
|
||||
if let Some(tokens) = response.input_tokens {
|
||||
*self.last_request_input_tokens = Some(tokens);
|
||||
*self.last_request_history_bytes =
|
||||
Some(self.history.iter().map(HistoryItem::estimated_bytes).sum());
|
||||
}
|
||||
|
||||
if !response.text.is_empty() {
|
||||
wire::send(
|
||||
self.wire,
|
||||
|
||||
@@ -6,9 +6,6 @@ pub const MAX_PROMPT_BYTES: usize = 1024 * 1024;
|
||||
pub const MAX_TOOL_RESULT_BYTES: usize = 8 * 1024 * 1024;
|
||||
pub const MAX_TOOL_CALLS_PER_TURN: usize = 64;
|
||||
|
||||
/// Leaves headroom for the summary call.
|
||||
pub const HANDOFF_THRESHOLD: f64 = 0.75;
|
||||
|
||||
pub const HANDOFF_MAX_OUTPUT_TOKENS: u32 = 8192;
|
||||
|
||||
pub const HANDOFF_TAIL_ITEMS: usize = 5;
|
||||
@@ -59,6 +56,13 @@ pub struct Config {
|
||||
pub max_sessions: usize,
|
||||
pub max_line_bytes: usize,
|
||||
pub max_history_bytes: usize,
|
||||
/// Provider context window in tokens used to gate handoff. The handoff
|
||||
/// fires when the previous request's (cache-summed) input tokens cross the
|
||||
/// handoff threshold for this budget, before the next request can exceed
|
||||
/// the window and 400. Default 200_000 — matching Claude 4.x windows;
|
||||
/// operators lower/raise it for other models. Set via
|
||||
/// `SPROUT_AGENT_MAX_CONTEXT_TOKENS`.
|
||||
pub max_context_tokens: u64,
|
||||
pub max_handoffs: usize,
|
||||
pub max_parallel_tools: usize,
|
||||
pub hook_timeout: Duration,
|
||||
@@ -149,7 +153,8 @@ impl Config {
|
||||
max_sessions: parse_env("SPROUT_AGENT_MAX_SESSIONS", usize::MAX)?,
|
||||
max_line_bytes: parse_env("SPROUT_AGENT_MAX_LINE_BYTES", 4 * 1024 * 1024)?,
|
||||
max_history_bytes: parse_env("SPROUT_AGENT_MAX_HISTORY_BYTES", 16 * 1024 * 1024)?,
|
||||
max_handoffs: parse_env("SPROUT_AGENT_MAX_HANDOFFS", 5)?,
|
||||
max_context_tokens: parse_env("SPROUT_AGENT_MAX_CONTEXT_TOKENS", 200_000u64)?,
|
||||
max_handoffs: parse_env("SPROUT_AGENT_MAX_HANDOFFS", 10)?,
|
||||
max_parallel_tools: parse_env("SPROUT_AGENT_MAX_PARALLEL_TOOLS", 8usize)?,
|
||||
hook_timeout: Duration::from_millis(parse_env(
|
||||
"SPROUT_AGENT_HOOK_TIMEOUT_MS",
|
||||
@@ -171,6 +176,12 @@ impl Config {
|
||||
if self.max_output_tokens < 1 {
|
||||
return Err("config: SPROUT_AGENT_MAX_OUTPUT_TOKENS must be >= 1".into());
|
||||
}
|
||||
if self.max_context_tokens <= u64::from(self.max_output_tokens) {
|
||||
return Err(format!(
|
||||
"config: SPROUT_AGENT_MAX_CONTEXT_TOKENS ({}) must be > SPROUT_AGENT_MAX_OUTPUT_TOKENS ({}) — the context window must leave room for the response",
|
||||
self.max_context_tokens, self.max_output_tokens
|
||||
));
|
||||
}
|
||||
if self.max_history_bytes < MIN_HISTORY_BYTES {
|
||||
return Err(format!(
|
||||
"config: SPROUT_AGENT_MAX_HISTORY_BYTES must be >= {MIN_HISTORY_BYTES}"
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use crate::agent::{push_hook_outputs_as_tool_results, RunCtx};
|
||||
use crate::config::{
|
||||
HANDOFF_MAX_OUTPUT_TOKENS, HANDOFF_MAX_TOOL_NAMES, HANDOFF_ORIGINAL_TASK_MAX_BYTES,
|
||||
HANDOFF_PROMPT_MAX_BYTES, HANDOFF_TAIL_ITEMS, HANDOFF_THRESHOLD,
|
||||
HANDOFF_PROMPT_MAX_BYTES, HANDOFF_TAIL_ITEMS,
|
||||
};
|
||||
use crate::types::HistoryItem;
|
||||
|
||||
@@ -90,9 +90,45 @@ impl RunCtx<'_> {
|
||||
}
|
||||
|
||||
fn should_handoff(&self) -> bool {
|
||||
let usage: usize = self.history.iter().map(HistoryItem::estimated_bytes).sum();
|
||||
let threshold = (self.cfg.max_history_bytes as f64 * HANDOFF_THRESHOLD) as usize;
|
||||
usage > threshold
|
||||
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
|
||||
// measurement — new assistant text, tool results, and the next
|
||||
// user prompt are appended before the next `complete()`. The exact
|
||||
// count alone would miss "previous request was under threshold, but
|
||||
// newly appended content pushes the next one over" (the stale-usage
|
||||
// cousin of the original stale-bytes bug). So we add a conservative
|
||||
// token estimate of the bytes added since the measurement.
|
||||
Some(measured_tokens) => {
|
||||
let measured_bytes = self.last_request_history_bytes.unwrap_or(0);
|
||||
let current_bytes: usize =
|
||||
self.history.iter().map(HistoryItem::estimated_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)
|
||||
}
|
||||
// No usage yet (first request, or just after a handoff reset).
|
||||
// Fall back to the byte heuristic, capped conservatively so a
|
||||
// single pre-usage request can't blow the window. We map the token
|
||||
// threshold to bytes using a deliberately LOW bytes/token ratio:
|
||||
// a low ratio implies more tokens per byte, so the byte cap is
|
||||
// small and the handoff fires early rather than late. Never raise
|
||||
// the cap above the configured byte budget.
|
||||
//
|
||||
// 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::estimated_bytes).sum();
|
||||
bytes
|
||||
> byte_fallback_threshold(
|
||||
self.cfg.max_context_tokens,
|
||||
self.cfg.max_output_tokens,
|
||||
self.cfg.max_history_bytes,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn build_handoff_prompt(&self) -> String {
|
||||
@@ -208,3 +244,94 @@ pub(crate) fn clamp_bytes(s: &str, max_bytes: usize) -> String {
|
||||
}
|
||||
format!("{}…", &s[..cut])
|
||||
}
|
||||
|
||||
/// Conservative bytes-per-token ratio used when estimating tokens from raw
|
||||
/// history bytes. We use 1: a token is always at least one byte, so treating
|
||||
/// every byte as a whole token is an unconditional UPPER bound on the true
|
||||
/// token count — it can never undercount, regardless of content density (even
|
||||
/// the densest real content sits at ~1.4 bytes/token). That over-estimate is
|
||||
/// exactly what a fail-early preflight gate wants: it hands off sooner rather
|
||||
/// than risk the next request exceeding the window.
|
||||
const CONSERVATIVE_BYTES_PER_TOKEN: u64 = 1;
|
||||
|
||||
/// 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.
|
||||
fn estimate_tokens_from_bytes(bytes: usize) -> u64 {
|
||||
(bytes as u64).div_ceil(CONSERVATIVE_BYTES_PER_TOKEN)
|
||||
}
|
||||
|
||||
/// Input-token count at which to hand off. Caps at the configured fraction of
|
||||
/// the window and also leaves room for `max_output_tokens`, so input + output
|
||||
/// can't together exceed the window. Free function so the policy math is unit
|
||||
/// testable without constructing a [`RunCtx`].
|
||||
fn token_threshold(max_context_tokens: u64, max_output_tokens: u32) -> u64 {
|
||||
// Integer math: handoff threshold is 90%, i.e. window * 9 / 10.
|
||||
let fractional = max_context_tokens / 10 * 9;
|
||||
let output_reserved = max_context_tokens.saturating_sub(u64::from(max_output_tokens));
|
||||
fractional.min(output_reserved)
|
||||
}
|
||||
|
||||
/// Conservative byte cap used only before any usage is known. Maps the token
|
||||
/// threshold to bytes at the conservative bytes/token ratio (so the cap is
|
||||
/// small and the handoff fires early), clamped to the configured byte budget
|
||||
/// so it can only ever be more conservative than the old byte-only behavior.
|
||||
fn byte_fallback_threshold(
|
||||
max_context_tokens: u64,
|
||||
max_output_tokens: u32,
|
||||
max_history_bytes: usize,
|
||||
) -> usize {
|
||||
let derived = token_threshold(max_context_tokens, max_output_tokens)
|
||||
.saturating_mul(CONSERVATIVE_BYTES_PER_TOKEN);
|
||||
let byte_cap = max_history_bytes / 10 * 9;
|
||||
usize::try_from(derived).unwrap_or(usize::MAX).min(byte_cap)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{byte_fallback_threshold, estimate_tokens_from_bytes, token_threshold};
|
||||
|
||||
#[test]
|
||||
fn token_threshold_uses_fraction_when_output_is_small() {
|
||||
// 200k window, 1k output. fractional = 0.9*200000 = 180000;
|
||||
// output_reserved = 200000-1000 = 199000; min = 180000.
|
||||
assert_eq!(token_threshold(200_000, 1_000), 180_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_threshold_reserves_output_headroom() {
|
||||
// Large output relative to window: the output-reserve term dominates,
|
||||
// keeping input+output within the window.
|
||||
// 100k window, 40k output: fractional=90k, reserved=60k -> 60k.
|
||||
assert_eq!(token_threshold(100_000, 40_000), 60_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_threshold_saturates_when_output_exceeds_window() {
|
||||
// Degenerate (config validation forbids this, but math must not panic):
|
||||
// reserved saturates to 0, so threshold is 0 -> always hand off.
|
||||
assert_eq!(token_threshold(1000, 5000), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn byte_fallback_is_conservative_and_capped() {
|
||||
// Derived = token_threshold * 1 (1 byte/token upper bound). For
|
||||
// 200k/1k: 180000 bytes, well under a 16 MiB byte budget, so derived
|
||||
// wins (early handoff).
|
||||
let t = byte_fallback_threshold(200_000, 1_000, 16 * 1024 * 1024);
|
||||
assert_eq!(t, 180_000);
|
||||
// With a tiny byte budget the cap wins -> never exceeds it (window*90%).
|
||||
let capped = byte_fallback_threshold(200_000, 1_000, 8192);
|
||||
assert_eq!(capped, 8192 / 10 * 9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn estimate_tokens_is_upper_bound_on_tokens() {
|
||||
// 1 byte/token: a token is always >= 1 byte, so byte count is an
|
||||
// unconditional upper bound on the true token count.
|
||||
assert_eq!(estimate_tokens_from_bytes(0), 0);
|
||||
assert_eq!(estimate_tokens_from_bytes(1), 1);
|
||||
assert_eq!(estimate_tokens_from_bytes(4), 4);
|
||||
assert_eq!(estimate_tokens_from_bytes(5), 5);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -42,6 +42,14 @@ struct Session {
|
||||
original_task: Option<String>,
|
||||
handoff_count: usize,
|
||||
stop_rejections: u32,
|
||||
/// Cache-summed input tokens the provider reported for this session's most
|
||||
/// recent request, or `None` before the first response (or after a handoff
|
||||
/// resets the context). Drives the token-based handoff gate; see
|
||||
/// [`RunCtx::should_handoff`].
|
||||
last_request_input_tokens: Option<u64>,
|
||||
/// History byte size when `last_request_input_tokens` was measured, paired
|
||||
/// with it so the gate can account for history appended since.
|
||||
last_request_history_bytes: Option<usize>,
|
||||
effective_system_prompt: Arc<str>,
|
||||
}
|
||||
|
||||
@@ -296,6 +304,8 @@ async fn session_new(app: &Arc<App>, id: Value, params: Value, wire_tx: &WireSen
|
||||
original_task: None,
|
||||
handoff_count: 0,
|
||||
stop_rejections: 0,
|
||||
last_request_input_tokens: None,
|
||||
last_request_history_bytes: None,
|
||||
effective_system_prompt,
|
||||
},
|
||||
);
|
||||
@@ -335,6 +345,8 @@ async fn run_prompt(app: Arc<App>, id: Value, params: Value, wire_tx: WireSender
|
||||
mut original_task,
|
||||
mut handoff_count,
|
||||
mut stop_rejections,
|
||||
mut last_request_input_tokens,
|
||||
mut last_request_history_bytes,
|
||||
mut cancel_rx,
|
||||
effective_system_prompt,
|
||||
) = match acquire_session(&app, &p.session_id).await {
|
||||
@@ -361,6 +373,8 @@ async fn run_prompt(app: Arc<App>, id: Value, params: Value, wire_tx: WireSender
|
||||
original_task: &mut original_task,
|
||||
handoff_count: &mut handoff_count,
|
||||
stop_rejections: &mut stop_rejections,
|
||||
last_request_input_tokens: &mut last_request_input_tokens,
|
||||
last_request_history_bytes: &mut last_request_history_bytes,
|
||||
};
|
||||
let result = ctx.run(p.prompt).await;
|
||||
if let Some(s) = app.sessions.lock().await.get_mut(&sid) {
|
||||
@@ -369,6 +383,8 @@ async fn run_prompt(app: Arc<App>, id: Value, params: Value, wire_tx: WireSender
|
||||
s.original_task = original_task;
|
||||
s.handoff_count = handoff_count;
|
||||
s.stop_rejections = stop_rejections;
|
||||
s.last_request_input_tokens = last_request_input_tokens;
|
||||
s.last_request_history_bytes = last_request_history_bytes;
|
||||
}
|
||||
match result {
|
||||
Ok(stop) => {
|
||||
@@ -393,6 +409,8 @@ async fn acquire_session(
|
||||
Option<String>,
|
||||
usize,
|
||||
u32,
|
||||
Option<u64>,
|
||||
Option<usize>,
|
||||
watch::Receiver<bool>,
|
||||
Arc<str>,
|
||||
),
|
||||
@@ -413,6 +431,8 @@ async fn acquire_session(
|
||||
s.original_task.take(),
|
||||
s.handoff_count,
|
||||
s.stop_rejections,
|
||||
s.last_request_input_tokens,
|
||||
s.last_request_history_bytes,
|
||||
rx,
|
||||
Arc::clone(&s.effective_system_prompt),
|
||||
))
|
||||
|
||||
@@ -568,10 +568,12 @@ fn parse_responses(v: Value) -> Result<LlmResponse, AgentError> {
|
||||
Some("completed") => ProviderStop::EndTurn,
|
||||
_ => ProviderStop::Other,
|
||||
};
|
||||
let input_tokens = sum_usage(&v, &["input_tokens"]);
|
||||
Ok(LlmResponse {
|
||||
text,
|
||||
tool_calls,
|
||||
stop,
|
||||
input_tokens,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -585,6 +587,54 @@ fn map_stop(s: Option<&str>) -> ProviderStop {
|
||||
}
|
||||
}
|
||||
|
||||
/// Sum a set of `usage` token fields, returning `None` only when the `usage`
|
||||
/// object is absent or carries none of the requested fields. A field that is
|
||||
/// present is added; a field that is missing contributes 0. This keeps the
|
||||
/// result an inclusive total (so cached tokens are never silently dropped)
|
||||
/// while still distinguishing "no usage reported" from "usage was zero".
|
||||
fn sum_usage(v: &Value, fields: &[&str]) -> Option<u64> {
|
||||
let usage = v.get("usage")?;
|
||||
let mut total: u64 = 0;
|
||||
let mut saw_any = false;
|
||||
for f in fields {
|
||||
if let Some(n) = usage.get(*f).and_then(Value::as_u64) {
|
||||
total = total.saturating_add(n);
|
||||
saw_any = true;
|
||||
}
|
||||
}
|
||||
saw_any.then_some(total)
|
||||
}
|
||||
|
||||
/// Input-token total for Anthropic / Databricks (Anthropic-style) responses.
|
||||
/// `input_tokens` alone EXCLUDES cached tokens, so we sum it with the two
|
||||
/// cache fields to get the inclusive total the context budget must gate on.
|
||||
fn anthropic_input_tokens(v: &Value) -> Option<u64> {
|
||||
sum_usage(
|
||||
v,
|
||||
&[
|
||||
"input_tokens",
|
||||
"cache_read_input_tokens",
|
||||
"cache_creation_input_tokens",
|
||||
],
|
||||
)
|
||||
}
|
||||
|
||||
/// Input-token total for OpenAI Chat Completions and Databricks responses.
|
||||
/// OpenAI's `prompt_tokens` is already inclusive. Databricks uses the same
|
||||
/// `prompt_tokens` wire field but ALSO reports Anthropic-style cache fields
|
||||
/// alongside it, so we sum them; the cache fields are simply absent (and
|
||||
/// contribute 0) for vanilla OpenAI.
|
||||
fn openai_chat_input_tokens(v: &Value) -> Option<u64> {
|
||||
sum_usage(
|
||||
v,
|
||||
&[
|
||||
"prompt_tokens",
|
||||
"cache_read_input_tokens",
|
||||
"cache_creation_input_tokens",
|
||||
],
|
||||
)
|
||||
}
|
||||
|
||||
fn str_field(v: &Value, key: &str) -> String {
|
||||
v.get(key).and_then(Value::as_str).unwrap_or("").to_owned()
|
||||
}
|
||||
@@ -610,10 +660,12 @@ fn parse_anthropic(v: Value) -> Result<LlmResponse, AgentError> {
|
||||
}
|
||||
}
|
||||
}
|
||||
let input_tokens = anthropic_input_tokens(&v);
|
||||
Ok(LlmResponse {
|
||||
text,
|
||||
tool_calls,
|
||||
stop,
|
||||
input_tokens,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -644,10 +696,12 @@ fn parse_openai(v: Value) -> Result<LlmResponse, AgentError> {
|
||||
)?);
|
||||
}
|
||||
}
|
||||
let input_tokens = openai_chat_input_tokens(&v);
|
||||
Ok(LlmResponse {
|
||||
text,
|
||||
tool_calls,
|
||||
stop,
|
||||
input_tokens,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -870,6 +924,7 @@ mod tests {
|
||||
max_sessions: 1,
|
||||
max_line_bytes: 1024 * 1024,
|
||||
max_history_bytes: 16 * 1024 * 1024,
|
||||
max_context_tokens: 200_000,
|
||||
max_handoffs: 1,
|
||||
max_parallel_tools: 1,
|
||||
hook_timeout: Duration::from_secs(1),
|
||||
@@ -1297,4 +1352,112 @@ mod tests {
|
||||
accepts.load(Ordering::SeqCst)
|
||||
);
|
||||
}
|
||||
|
||||
// ---- usage / input-token extraction -------------------------------------
|
||||
|
||||
#[test]
|
||||
fn parse_anthropic_sums_input_and_cache_tokens() {
|
||||
// input_tokens alone excludes cached tokens; the inclusive total must
|
||||
// sum all three so a cache-heavy turn can't undercount the budget.
|
||||
let v = serde_json::json!({
|
||||
"stop_reason": "end_turn",
|
||||
"content": [{"type": "text", "text": "hi"}],
|
||||
"usage": {
|
||||
"input_tokens": 100,
|
||||
"cache_read_input_tokens": 900,
|
||||
"cache_creation_input_tokens": 50,
|
||||
"output_tokens": 7
|
||||
}
|
||||
});
|
||||
let r = parse_anthropic(v).unwrap();
|
||||
assert_eq!(r.input_tokens, Some(1050));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_anthropic_input_tokens_only() {
|
||||
let v = serde_json::json!({
|
||||
"stop_reason": "end_turn",
|
||||
"content": [{"type": "text", "text": "hi"}],
|
||||
"usage": {"input_tokens": 42, "output_tokens": 3}
|
||||
});
|
||||
assert_eq!(parse_anthropic(v).unwrap().input_tokens, Some(42));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_anthropic_missing_usage_is_none() {
|
||||
let v = serde_json::json!({
|
||||
"stop_reason": "end_turn",
|
||||
"content": [{"type": "text", "text": "hi"}]
|
||||
});
|
||||
assert_eq!(parse_anthropic(v).unwrap().input_tokens, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_openai_uses_prompt_tokens() {
|
||||
let v = serde_json::json!({
|
||||
"choices": [{"finish_reason": "stop", "message": {"content": "hi"}}],
|
||||
"usage": {"prompt_tokens": 123, "completion_tokens": 4, "total_tokens": 127}
|
||||
});
|
||||
assert_eq!(parse_openai(v).unwrap().input_tokens, Some(123));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_openai_databricks_sums_cache_fields() {
|
||||
// Databricks uses the OpenAI chat wire format (prompt_tokens) but also
|
||||
// reports Anthropic-style cache fields; the inclusive total sums them.
|
||||
let v = serde_json::json!({
|
||||
"choices": [{"finish_reason": "stop", "message": {"content": "hi"}}],
|
||||
"usage": {
|
||||
"prompt_tokens": 200,
|
||||
"completion_tokens": 4,
|
||||
"total_tokens": 204,
|
||||
"cache_read_input_tokens": 800,
|
||||
"cache_creation_input_tokens": 0
|
||||
}
|
||||
});
|
||||
assert_eq!(parse_openai(v).unwrap().input_tokens, Some(1000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_openai_missing_usage_is_none() {
|
||||
let v = serde_json::json!({
|
||||
"choices": [{"finish_reason": "stop", "message": {"content": "hi"}}]
|
||||
});
|
||||
assert_eq!(parse_openai(v).unwrap().input_tokens, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_responses_uses_input_tokens() {
|
||||
let v = serde_json::json!({
|
||||
"status": "completed",
|
||||
"output": [{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "hi"}]
|
||||
}],
|
||||
"usage": {"input_tokens": 321, "output_tokens": 9, "total_tokens": 330}
|
||||
});
|
||||
assert_eq!(parse_responses(v).unwrap().input_tokens, Some(321));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_responses_missing_usage_is_none() {
|
||||
let v = serde_json::json!({
|
||||
"status": "completed",
|
||||
"output": [{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "hi"}]
|
||||
}]
|
||||
});
|
||||
assert_eq!(parse_responses(v).unwrap().input_tokens, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sum_usage_empty_object_is_none() {
|
||||
// A `usage` object present but carrying none of the requested fields
|
||||
// is "no usable reading" -> None, not Some(0).
|
||||
let v = serde_json::json!({"usage": {"output_tokens": 5}});
|
||||
assert_eq!(sum_usage(&v, &["input_tokens", "prompt_tokens"]), None);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -95,6 +95,13 @@ pub struct LlmResponse {
|
||||
pub text: String,
|
||||
pub tool_calls: Vec<ToolCall>,
|
||||
pub stop: ProviderStop,
|
||||
/// Total input tokens the provider reported for this request, or `None`
|
||||
/// if the response carried no usage. For Anthropic/Databricks this is the
|
||||
/// inclusive sum `input_tokens + cache_read_input_tokens +
|
||||
/// cache_creation_input_tokens` (plain `input_tokens` excludes cached
|
||||
/// tokens, so reading it alone would undercount). Used to gate handoff on
|
||||
/// the real token budget rather than a byte estimate.
|
||||
pub input_tokens: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
|
||||
@@ -10,6 +10,8 @@
|
||||
//! FAKE_MCP_DESC_SIZE=N — every tool description is N bytes (overrides HUGE_DESC)
|
||||
//! FAKE_MCP_TOOL_DELAY=N — `tools/call` sleeps N seconds before replying
|
||||
//! (use a large value, e.g. 999, to simulate hang)
|
||||
//! FAKE_MCP_RESULT_SIZE=N — `tools/call` returns an N-byte text result
|
||||
//! (default: the literal "ok"); grows history
|
||||
//! FAKE_MCP_PID_FILE=path — write the child PID to `path` on startup
|
||||
//! (for tests that want to verify the child died)
|
||||
//! FAKE_MCP_SPAWN_GRANDCHILD=1
|
||||
@@ -123,6 +125,9 @@ fn main() {
|
||||
"fake tool".to_owned()
|
||||
};
|
||||
let tool_delay_secs = env_u64("FAKE_MCP_TOOL_DELAY", 0);
|
||||
// Tool-call result text size in bytes (default: the literal "ok"). Lets a
|
||||
// test grow session history by a controlled amount via a tool result.
|
||||
let result_size = env_u64("FAKE_MCP_RESULT_SIZE", 0) as usize;
|
||||
let stop_hook = env_flag("FAKE_MCP_STOP_HOOK");
|
||||
let stop_text = std::env::var("FAKE_MCP_STOP_TEXT").unwrap_or_else(|_| "keep going".to_owned());
|
||||
let stop_delay_secs = env_u64("FAKE_MCP_STOP_DELAY", 0);
|
||||
@@ -266,10 +271,15 @@ fn main() {
|
||||
if tool_delay_secs > 0 {
|
||||
std::thread::sleep(std::time::Duration::from_secs(tool_delay_secs));
|
||||
}
|
||||
let result_text = if result_size > 0 {
|
||||
"x".repeat(result_size)
|
||||
} else {
|
||||
"ok".to_owned()
|
||||
};
|
||||
write_response(
|
||||
id,
|
||||
json!({
|
||||
"content": [{ "type": "text", "text": "ok" }],
|
||||
"content": [{ "type": "text", "text": result_text }],
|
||||
"isError": false,
|
||||
}),
|
||||
);
|
||||
|
||||
@@ -186,6 +186,19 @@ fn openai_text(content: &str) -> Value {
|
||||
})
|
||||
}
|
||||
|
||||
/// Like [`openai_text`] but attaches a `usage` block so tests can drive the
|
||||
/// token-based handoff gate. `prompt_tokens` is the input-token count the
|
||||
/// agent will read and compare against the configured context budget.
|
||||
fn openai_text_with_usage(content: &str, prompt_tokens: u64) -> Value {
|
||||
let mut v = openai_text(content);
|
||||
v["usage"] = json!({
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": 1,
|
||||
"total_tokens": prompt_tokens + 1,
|
||||
});
|
||||
v
|
||||
}
|
||||
|
||||
fn openai_tool_call(id: &str, name: &str, args: Value) -> Value {
|
||||
json!({
|
||||
"id": "cc-2", "object": "chat.completion", "model": "fake-model",
|
||||
@@ -1015,7 +1028,7 @@ async fn hook_post_compact_injects_after_handoff() {
|
||||
// Sequence of canned LLM responses consumed in order:
|
||||
// 1-3. Three `session/prompt` rounds returning short text. Each
|
||||
// prompt body is ~300 KB, so by the 4th prompt we'll be over
|
||||
// the 75% (= 768 KB) threshold of a 1 MB budget.
|
||||
// the 90% (= ~922 KB) threshold of a 1 MB budget.
|
||||
// 4. Handoff `summarize()` call returns the summary text.
|
||||
// 5. Next regular `complete()` call after the handoff returns
|
||||
// a final "done" message; we inspect this request's body.
|
||||
@@ -1027,7 +1040,7 @@ async fn hook_post_compact_injects_after_handoff() {
|
||||
openai_text("done"),
|
||||
])
|
||||
.await;
|
||||
// 1 MB budget = MIN allowed. Threshold = 768 KB. Each ~300 KB prompt
|
||||
// 1 MB budget = MIN allowed. Threshold = ~922 KB. Each ~300 KB prompt
|
||||
// fills the budget on the 4th turn, triggering handoff.
|
||||
let mut h = Harness::spawn_with_env(
|
||||
&llm.url,
|
||||
@@ -1112,6 +1125,151 @@ async fn hook_post_compact_injects_after_handoff() {
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
/// Regression for the original bug: context fills, the provider 400s on the
|
||||
/// next request, and the handoff never fires because the old gate measured
|
||||
/// BYTES (16 MiB threshold) while the limit is in TOKENS. The fix gates on
|
||||
/// provider-reported input tokens. Here the prompts are tiny (bytes nowhere
|
||||
/// near any byte threshold), but the fake LLM reports `usage.prompt_tokens`
|
||||
/// over the configured token budget — so the handoff MUST fire on the token
|
||||
/// signal alone, before the next normal `complete()`.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn token_usage_over_budget_triggers_handoff() {
|
||||
// Context window 1000 tokens, output 100 -> threshold = min(900, 900) = 900.
|
||||
// First response reports 950 input tokens (> 900). The agent stores that;
|
||||
// the next prompt's pre-flight gate sees 950 >= 900 and hands off, which
|
||||
// inserts an extra summarize() call we didn't issue.
|
||||
// req 1: prompt #0 -> text + usage(950)
|
||||
// req 2: summarize() (the handoff) -> summary text
|
||||
// req 3: prompt #1's actual complete() -> done
|
||||
let llm = spawn_capturing_llm(vec![
|
||||
openai_text_with_usage("ack-0", 950),
|
||||
openai_text("handoff summary text"),
|
||||
openai_text_with_usage("done", 10),
|
||||
])
|
||||
.await;
|
||||
let mut h = Harness::spawn_with_env(
|
||||
&llm.url,
|
||||
&[
|
||||
("SPROUT_AGENT_MAX_CONTEXT_TOKENS", "1000"),
|
||||
("SPROUT_AGENT_MAX_OUTPUT_TOKENS", "100"),
|
||||
("SPROUT_AGENT_MAX_HANDOFFS", "3"),
|
||||
// Huge byte budget so the byte path can NOT be what fires — only
|
||||
// the token gate can explain a handoff on these tiny prompts.
|
||||
(
|
||||
"SPROUT_AGENT_MAX_HISTORY_BYTES",
|
||||
&(16 * 1024 * 1024).to_string(),
|
||||
),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
let sid = init_session(&mut h, json!([])).await;
|
||||
|
||||
// Prompt #0: small body; response carries usage(950) -> over threshold.
|
||||
let p0 = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({"sessionId": sid, "prompt": [{"type":"text","text":"hello 0"}]}),
|
||||
)
|
||||
.await;
|
||||
let _ = h.recv_until(|v| v["id"] == json!(p0)).await;
|
||||
assert_eq!(
|
||||
llm.captured.lock().await.len(),
|
||||
1,
|
||||
"first prompt should produce exactly one LLM request (no handoff yet)"
|
||||
);
|
||||
|
||||
// Prompt #1: also small. The pre-flight gate sees the stored 800 tokens
|
||||
// and hands off BEFORE issuing this prompt's complete() -> an extra
|
||||
// summarize request appears (3 total, not 2).
|
||||
let p1 = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({"sessionId": sid, "prompt": [{"type":"text","text":"hello 1"}]}),
|
||||
)
|
||||
.await;
|
||||
let _ = h.recv_until(|v| v["id"] == json!(p1)).await;
|
||||
let captured = llm.captured.lock().await.len();
|
||||
assert_eq!(
|
||||
captured, 3,
|
||||
"expected handoff summarize() between the two prompts (3 reqs), saw {captured} — \
|
||||
token gate did not fire on usage over budget"
|
||||
);
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
/// Regression for the stale-usage gap (caught in review): the exact token
|
||||
/// count describes the PREVIOUS request, but history grows afterward (tool
|
||||
/// results, next prompt). If the gate trusted only the stale `Some(tokens)`
|
||||
/// and skipped the byte signal, a previously-under-threshold session could
|
||||
/// still 400 once a large tool result lands. The fix adds a conservative
|
||||
/// token estimate of the bytes grown since the measurement. Here usage is
|
||||
/// reported UNDER threshold, then a large tool result grows history enough
|
||||
/// that the projection crosses — so the handoff must fire.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn stale_usage_plus_history_growth_triggers_handoff() {
|
||||
// window 10_000, output 1_000 -> threshold = min(9_000, 9_000) = 9_000.
|
||||
// req1 reports usage 8_500 (UNDER 9_000). Its response is a tool_call;
|
||||
// the fake MCP returns a ~6 KB result, appended to history. At the
|
||||
// conservative 1 byte/token estimate that's ~6_000 projected tokens, so
|
||||
// projected ~14_500 >= 9_000 -> the next loop iteration hands off before
|
||||
// the follow-up complete().
|
||||
// req1: tool_call + usage(8500)
|
||||
// (tool result ~6KB appended)
|
||||
// req2: summarize() (handoff)
|
||||
// req3: final text
|
||||
let llm = spawn_capturing_llm(vec![
|
||||
{
|
||||
let mut v = openai_tool_call("tc1", "fake__tool_0", json!({}));
|
||||
v["usage"] =
|
||||
json!({"prompt_tokens": 8500, "completion_tokens": 1, "total_tokens": 8501});
|
||||
v
|
||||
},
|
||||
openai_text("handoff summary text"),
|
||||
openai_text("done"),
|
||||
])
|
||||
.await;
|
||||
let mut h = Harness::spawn_with_env(
|
||||
&llm.url,
|
||||
&[
|
||||
("SPROUT_AGENT_MAX_CONTEXT_TOKENS", "10000"),
|
||||
("SPROUT_AGENT_MAX_OUTPUT_TOKENS", "1000"),
|
||||
("SPROUT_AGENT_MAX_HANDOFFS", "3"),
|
||||
// Huge byte budget so the None-path byte fallback can't be what
|
||||
// fires — only the token-mode growth estimate can explain it.
|
||||
(
|
||||
"SPROUT_AGENT_MAX_HISTORY_BYTES",
|
||||
&(16 * 1024 * 1024).to_string(),
|
||||
),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
let sid = init_session_with_fake_mcp(
|
||||
&mut h,
|
||||
&[
|
||||
("FAKE_MCP_TOOL_COUNT", "1"),
|
||||
("FAKE_MCP_RESULT_SIZE", "6000"),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
|
||||
let p = h
|
||||
.send(
|
||||
"session/prompt",
|
||||
json!({"sessionId": sid, "prompt": [{"type":"text","text":"go"}]}),
|
||||
)
|
||||
.await;
|
||||
let _ = h.recv_until(|v| v["id"] == json!(p)).await;
|
||||
// req1 (tool_call) + summarize (handoff) + req2 (done) = 3. Without the
|
||||
// growth estimate we'd see only 2 (stale 8500 < 9000, no handoff).
|
||||
let captured = llm.captured.lock().await.len();
|
||||
assert_eq!(
|
||||
captured, 3,
|
||||
"expected handoff after history grew past threshold (3 reqs), saw {captured} — \
|
||||
stale under-threshold usage skipped the growth estimate"
|
||||
);
|
||||
h.shutdown().await;
|
||||
}
|
||||
|
||||
/// `_Stop` hook that takes longer than `SPROUT_AGENT_HOOK_TIMEOUT_MS`
|
||||
/// must be treated as no-objection (fail-open). Agent stops normally.
|
||||
///
|
||||
|
||||
Reference in New Issue
Block a user