Skip to content

Commit 28c9603

Browse files
committed
fix(agent): use ground-truth prompt_tokens in the compaction trigger
The auto-compaction trigger sized the context with a bytes/4 heuristic (`estimate_context_tokens`). That under-counts the code/JSON/non-English payloads this agent generates, so a context that has really blown past the model's window can read as "under threshold" and the next request overflows — an unrecoverable failure once the provider rejects it. The provider already returns the exact input size in `usage.prompt_tokens`, but it was only emitted as telemetry and discarded. Track the most recent value per turn and feed it into the trigger via `effective_context_tokens(estimate, last)` = `max(estimate, last_prompt_tokens)`. At the end-of-turn compaction check, `last_prompt_tokens` covers nearly the whole current context (only the just- produced final assistant message is newer), so the max corrects the heuristic's under-count and compaction fires when it should. Self-correcting by construction: after any compaction the next request re-counts a smaller context (overflow recovery `continue`s the loop; the threshold path ends the turn), and `last_prompt_tokens` resets to `None` per inbound — so a stale large value can't cause a spurious re-trigger. Out of scope (follow-up): an OpenAI context-window override for `provider.context_window_tokens()` so the window-aware threshold tightens there too. Tests: `effective_context_tokens_prefers_ground_truth` (fallback, ground-truth wins, estimate wins, zero never lowers).
1 parent e99e8b3 commit 28c9603

1 file changed

Lines changed: 39 additions & 1 deletion

File tree

src/agent/mod.rs

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,14 @@ fn estimate_context_tokens(context: &[crate::utils::ChatMessage]) -> usize {
151151
context.iter().map(estimate_message_tokens).sum()
152152
}
153153

154+
/// Best available context-size estimate for the compaction trigger: the larger of the bytes/4
155+
/// heuristic and the last LLM call's exact `usage.prompt_tokens`. The heuristic under-counts the
156+
/// code/JSON/non-English payloads this agent produces, while `prompt_tokens` is the provider's
157+
/// ground-truth input count — so the max guards against silently overflowing the context window.
158+
fn effective_context_tokens(estimate: usize, last_prompt_tokens: Option<u32>) -> usize {
159+
estimate.max(last_prompt_tokens.unwrap_or(0) as usize)
160+
}
161+
154162
/// Trim context from the front (oldest messages) to stay within a token budget.
155163
/// Preserves the system message at index 0 and never splits tool_call/tool pairs.
156164
/// A marker message is inserted only when messages were actually removed.
@@ -2212,6 +2220,11 @@ impl AgentLogic {
22122220
let mut consecutive_doom_detections: usize = 0;
22132221
// Set when the doom loop persists past the nudge budget; branches the terminal message.
22142222
let mut doom_loop_stuck = false;
2223+
// Ground-truth input size from the most recent LLM call's `usage.prompt_tokens` (exact,
2224+
// server-counted). The bytes/4 heuristic under-counts code/JSON/non-English — exactly what
2225+
// this agent generates — so the compaction trigger uses `max(estimate, last_prompt_tokens)`
2226+
// to avoid silently overflowing the context window. Updated after each provider response.
2227+
let mut last_prompt_tokens: Option<u32> = None;
22152228

22162229
while iterations < max_iterations {
22172230
if cancel_token.is_cancelled() {
@@ -2620,6 +2633,10 @@ impl AgentLogic {
26202633

26212634
// Log USAGE telemetry
26222635
if let Some(usage) = &response.usage {
2636+
// Remember the exact server-counted input size for the compaction trigger.
2637+
if usage.prompt_tokens > 0 {
2638+
last_prompt_tokens = Some(usage.prompt_tokens);
2639+
}
26232640
let usage_evt = TelemetryEvent::AgentUsage {
26242641
chat_id: inbound.chat_id.clone(),
26252642
model: "llm_provider".to_string(),
@@ -2984,7 +3001,15 @@ impl AgentLogic {
29843001
// Auto-compaction check
29853002
let current_context = mem.get_context_since_reflection().await?;
29863003
let user_turns = current_context.iter().filter(|m| m.role == "user").count();
2987-
let approx_tokens: usize = estimate_context_tokens(&current_context);
3004+
// Prefer the ground truth: `last_prompt_tokens` is the exact input size the provider
3005+
// counted for the most recent request, which (at this end-of-turn point) covers
3006+
// nearly the entire current context — only the just-produced final assistant message
3007+
// is newer. Taking the max with the bytes/4 estimate corrects the heuristic's
3008+
// under-count on code/JSON-heavy contexts so a real overflow triggers compaction.
3009+
let approx_tokens: usize = effective_context_tokens(
3010+
estimate_context_tokens(&current_context),
3011+
last_prompt_tokens,
3012+
);
29883013

29893014
// PR-3: pull the model's context window from the provider; if known,
29903015
// tighten the absolute token threshold to whichever is smaller of:
@@ -4056,6 +4081,19 @@ mod tests {
40564081
assert_eq!(super::estimate_context_tokens(std::slice::from_ref(&msg)), 100);
40574082
}
40584083

4084+
#[test]
4085+
fn effective_context_tokens_prefers_ground_truth() {
4086+
// No usage yet -> fall back to the estimate.
4087+
assert_eq!(super::effective_context_tokens(1000, None), 1000);
4088+
// Provider's exact count exceeds the bytes/4 under-estimate -> use the ground truth so a
4089+
// real overflow still triggers compaction.
4090+
assert_eq!(super::effective_context_tokens(1000, Some(9000)), 9000);
4091+
// Estimate larger (e.g. messages added since the last call) -> keep the estimate.
4092+
assert_eq!(super::effective_context_tokens(9000, Some(1000)), 9000);
4093+
// A zero ground-truth never lowers the estimate.
4094+
assert_eq!(super::effective_context_tokens(1000, Some(0)), 1000);
4095+
}
4096+
40594097
#[tokio::test]
40604098
async fn run_reasoning_loop_persists_terminal_message_on_cancel() {
40614099
let (result, context) =

0 commit comments

Comments
 (0)