1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
//! Token budgeting for high-volume MCP tool responses.
//!
//! An agent can pass `max_tokens` to bound how many tokens a list-returning tool's
//! result list costs. The result list is already best-first (the tools rank as they
//! build), so budgeting is a keep-prefix-that-fits operation: walk the items in order,
//! sum each item's estimated token cost (serialized bytes / 4, the same factor the
//! telemetry estimator in [`super::savings`] uses), and drop the suffix once the running
//! total would exceed the budget.
//!
//! The budget bounds the RESULT LIST, not the whole response envelope — the surrounding
//! struct fields (`total`, `truncated`, cursors) are not counted. Truncation is always
//! explicit: the caller sets a `budgeted: bool` on the response and keeps `next_cursor`
//! working so the agent can page the dropped tail.
use serde::Serialize;
use super::savings::bytes_to_tokens;
/// Outcome of applying a token budget to an already-ranked result list.
pub(super) struct Budgeted<T> {
/// The kept leading prefix of items whose cumulative token cost fits the budget.
pub items: Vec<T>,
/// True when at least one trailing item was dropped to fit the budget. The number
/// dropped is recoverable as `original_len - items.len()` (asserted in the unit tests);
/// callers only ever branch on the boolean, so it is not stored separately.
pub budgeted: bool,
}
/// Keep the longest leading prefix of `items` whose cumulative estimated token cost
/// stays within `max_tokens`, estimating each item's cost as `serialized_bytes / 4`.
///
/// - `max_tokens = None` keeps everything (`budgeted = false`, `dropped = 0`) and never
/// serializes — the zero-overhead path for callers that didn't ask for a budget.
/// - `max_tokens = Some(0)` drops everything except — to guarantee forward progress for
/// pagination — the first item, which is always kept even when it alone exceeds the
/// budget. Without this an agent that set an impossibly small budget would get an empty
/// page with a cursor pointing at the same item forever.
///
/// Each item is serialized exactly once (its bytes feed the running total directly), so
/// there is no per-item re-serialize in the accumulation loop.
pub(super) fn apply_budget<T: Serialize>(items: Vec<T>, max_tokens: Option<u32>) -> Budgeted<T> {
let Some(budget) = max_tokens.map(u64::from) else {
return Budgeted { items, budgeted: false };
};
let total = items.len();
let mut kept: Vec<T> = Vec::with_capacity(total);
let mut used: u64 = 0;
for item in items {
// Deliberately the cheap `bytes/4` heuristic, not the real tokenizer: this loop runs
// per-item on every budgeted tool call (a hot path) and only ranks/fits items — it is
// not user-facing telemetry, so a tokenizer call per item would be an unjustified cost.
let cost = match serde_json::to_vec(&item) {
Ok(bytes) => bytes_to_tokens(bytes.len() as u64),
// A serialize failure here is unexpected (every response item is plain JSON),
// but rather than drop the whole page treat the item as zero-cost and keep it —
// the worst case is a slight budget overshoot, never a silent data loss.
Err(_) => 0,
};
// Always admit the first item so a tiny budget still makes forward progress; after
// that, stop as soon as adding this item would push the running total over budget.
if kept.is_empty() || used.saturating_add(cost) <= budget {
used = used.saturating_add(cost);
kept.push(item);
} else {
break;
}
}
Budgeted {
budgeted: kept.len() < total,
items: kept,
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde::Serialize;
/// A fixed-size item: serializes to a predictable byte count so token math is exact.
#[derive(Serialize)]
struct Item {
/// 10-char value → the serialized object is a stable width.
v: String,
}
fn item(tag: char) -> Item {
Item {
v: std::iter::repeat_n(tag, 10).collect(),
}
}
/// Serialized form of one `Item` is `{"v":"xxxxxxxxxx"}` = 18 bytes → 18/4 = 4 tokens.
const ITEM_TOKENS: u64 = 4;
#[test]
fn none_budget_keeps_everything_without_serializing() {
let items: Vec<Item> = ('a'..='e').map(item).collect();
let input_len = 5;
let out = apply_budget(items, None);
assert_eq!(out.items.len(), 5);
assert!(!out.budgeted);
assert_eq!(input_len - out.items.len(), 0, "nothing dropped");
}
#[test]
fn keeps_exact_prefix_that_fits_and_reports_dropped_count() {
// 5 items × 4 tokens each = 20 tokens total. Budget of 9 admits exactly 2 (8
// tokens); the 3rd would reach 12 > 9, so it and the rest are dropped.
let items: Vec<Item> = ('a'..='e').map(item).collect();
let input_len = items.len();
let budget = ITEM_TOKENS * 2 + 1; // 9: fits 2 (8), not 3 (12).
let out = apply_budget(items, Some(budget as u32));
assert_eq!(out.items.len(), 2, "exactly the 2-item prefix fits");
assert_eq!(
out.items.iter().map(|i| i.v.clone()).collect::<Vec<_>>(),
vec!["aaaaaaaaaa".to_string(), "bbbbbbbbbb".to_string()],
"kept items must be the leading prefix, in order"
);
assert!(out.budgeted);
assert_eq!(input_len - out.items.len(), 3, "dropped count");
}
#[test]
fn budget_exactly_at_boundary_keeps_all() {
// Budget exactly equal to the total cost keeps every item, drops none.
let items: Vec<Item> = ('a'..='c').map(item).collect();
let out = apply_budget(items, Some((ITEM_TOKENS * 3) as u32));
assert_eq!(out.items.len(), 3);
assert!(!out.budgeted);
}
#[test]
fn zero_budget_still_keeps_first_item_for_forward_progress() {
let items: Vec<Item> = ('a'..='c').map(item).collect();
let input_len = items.len();
let out = apply_budget(items, Some(0));
assert_eq!(out.items.len(), 1, "first item always admitted");
assert_eq!(out.items[0].v, "aaaaaaaaaa");
assert!(out.budgeted);
assert_eq!(input_len - out.items.len(), 2, "dropped count");
}
#[test]
fn empty_input_is_never_budgeted() {
let out = apply_budget(Vec::<Item>::new(), Some(0));
assert!(out.items.is_empty());
assert!(!out.budgeted);
}
}