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
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
//! Pre-dispatch spend reservations for per-token `max_tokens` budgets.
//!
//! A spend cap cannot be enforced from recorded usage alone. Usage is only known
//! once a response completes, so checking `used_tokens < max_tokens` at admission
//! lets a single large answer push the persisted total arbitrarily far past the
//! cap — the router observed `39273/30000` after one tool loop (issue #195).
//!
//! The router therefore reserves budget *before* dispatching. Every protocol
//! declares, or is given, a maximum output size; that figure plus an estimate of
//! the prompt is reserved against the cap at admission and released when the real
//! usage arrives. Because the reservation happens inside the same locked
//! read-modify-write that increments the request counters, concurrent requests
//! cannot collectively overshoot either.
//!
//! # Enforcement contract
//!
//! * A request is admitted only when `used + reserved + estimate <= max_tokens`.
//! * A request whose declared output budget cannot fit is rejected up front with
//! `token_budget_exceeded` rather than being truncated mid-answer.
//! * Actual usage is always recorded in full, even when it exceeds the estimate,
//! so the persisted total never understates real spend.
//! * The bound is therefore exact against *declared* budgets. Providers that
//! report more tokens than the caller's declared maximum (hidden reasoning
//! tokens, for example) can still land above the cap by that provider-side
//! excess; the overshoot is bounded by one request's unreported surplus rather
//! than being unbounded, and both figures are visible to administrators.
use serde_json::Value;
/// Characters per token used by the router's local estimator.
///
/// Matches [`crate::output_limit`] and
/// [`crate::anthropic_bridge::count_tokens_estimate`] so one request is not
/// measured two different ways.
const CHARS_PER_TOKEN: u64 = 4;
/// Output budget assumed when a request declares none.
///
/// Protocol adapters already default to this figure when shaping upstream
/// requests, so reserving it keeps the reservation aligned with what the
/// upstream is actually allowed to return.
pub const DEFAULT_MAX_OUTPUT_TOKENS: u64 = 4096;
/// The spend a request could report, used to reserve budget before dispatch.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct RequestBudget {
/// Estimated prompt tokens.
pub input: u64,
/// Declared (or defaulted) maximum output tokens.
pub output: u64,
}
impl RequestBudget {
/// Total tokens to reserve for this request.
#[must_use]
pub const fn total(self) -> u64 {
self.input.saturating_add(self.output)
}
}
/// Estimate the spend a request body could report.
///
/// Reads whichever output-cap field the request's protocol uses:
/// `max_output_tokens` (Responses), `max_completion_tokens`/`max_tokens` (Chat
/// Completions and Anthropic Messages), and
/// `generationConfig.maxOutputTokens` (Gemini). When none is present the
/// protocol default applies, matching what the adapters forward upstream.
#[must_use]
pub fn estimate(body: &Value) -> RequestBudget {
RequestBudget {
input: estimate_input_tokens(body),
output: declared_output_tokens(body).unwrap_or(DEFAULT_MAX_OUTPUT_TOKENS),
}
}
/// The caller's declared output cap, in whichever field its protocol uses.
#[must_use]
pub fn declared_output_tokens(body: &Value) -> Option<u64> {
for field in ["max_output_tokens", "max_completion_tokens", "max_tokens"] {
if let Some(value) = body.get(field).and_then(Value::as_u64) {
return Some(value);
}
}
body.get("generationConfig")
.and_then(|config| config.get("maxOutputTokens"))
.and_then(Value::as_u64)
}
/// Estimate prompt tokens from the request's text payload.
///
/// The router has no upstream tokenizer, so it measures the serialized textual
/// content with the same ~4-characters-per-token heuristic used elsewhere. This
/// only has to be good enough to keep a reservation in the right order of
/// magnitude; the real figure replaces it at settlement.
#[must_use]
pub fn estimate_input_tokens(body: &Value) -> u64 {
let mut characters = 0u64;
for field in ["messages", "input", "contents", "system", "prompt", "tools"] {
if let Some(value) = body.get(field) {
characters = characters.saturating_add(text_length(value));
}
}
characters.div_ceil(CHARS_PER_TOKEN)
}
/// Total length of every string in a JSON value.
fn text_length(value: &Value) -> u64 {
match value {
Value::String(text) => text.len() as u64,
Value::Array(items) => items
.iter()
.map(text_length)
.fold(0u64, u64::saturating_add),
Value::Object(fields) => fields
.values()
.map(text_length)
.fold(0u64, u64::saturating_add),
_ => 0,
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn reads_the_responses_output_cap() {
let body = json!({"max_output_tokens": 1500});
assert_eq!(declared_output_tokens(&body), Some(1500));
assert_eq!(estimate(&body).output, 1500);
}
#[test]
fn reads_the_chat_completions_output_caps() {
assert_eq!(
declared_output_tokens(&json!({"max_completion_tokens": 700})),
Some(700)
);
assert_eq!(
declared_output_tokens(&json!({"max_tokens": 800})),
Some(800)
);
}
#[test]
fn reads_the_gemini_output_cap() {
let body = json!({"generationConfig": {"maxOutputTokens": 256}});
assert_eq!(declared_output_tokens(&body), Some(256));
}
#[test]
fn falls_back_to_the_protocol_default_when_no_cap_is_declared() {
let body = json!({"messages": []});
assert_eq!(declared_output_tokens(&body), None);
assert_eq!(estimate(&body).output, DEFAULT_MAX_OUTPUT_TOKENS);
}
#[test]
fn estimates_prompt_tokens_from_message_text() {
// 40 characters of content over the 4-characters-per-token heuristic.
let body = json!({"messages": [{"role": "user", "content": "a".repeat(40)}]});
// "user" (4) + 40 characters = 44 characters => 11 tokens.
assert_eq!(estimate_input_tokens(&body), 11);
}
#[test]
fn counts_gemini_contents_and_anthropic_system_prompts() {
let gemini = json!({"contents": [{"parts": [{"text": "b".repeat(8)}]}]});
assert!(estimate_input_tokens(&gemini) > 0);
let anthropic = json!({"system": "c".repeat(16), "messages": []});
assert!(estimate_input_tokens(&anthropic) >= 4);
}
#[test]
fn total_combines_input_and_output() {
let budget = RequestBudget {
input: 10,
output: 20,
};
assert_eq!(budget.total(), 30);
}
#[test]
fn saturates_rather_than_overflowing_on_absurd_declared_budgets() {
let budget = RequestBudget {
input: u64::MAX,
output: u64::MAX,
};
assert_eq!(budget.total(), u64::MAX);
}
#[test]
fn ignores_non_textual_payload_fields() {
assert_eq!(estimate_input_tokens(&json!({"temperature": 0.7})), 0);
}
}