use serde_json::Value;
use crate::SharedMessage;
pub const SAFETY_MARGIN_PERCENT: u64 = 15;
pub const EXPECTED_RESPONSE_TOKENS: u64 = 16_384;
pub const ZERO_BUDGET_THINKING_RESERVE_TOKENS: u64 =
crate::core::models::DEFAULT_LEGACY_ADAPTIVE_FALLBACK as u64;
pub const MIN_COMPACTION_MESSAGES: usize = 4;
pub const IMAGE_TOKEN_ESTIMATE: u64 = 1_600;
const PER_MESSAGE_FRAMING_TOKENS: u64 = 4;
const TOOL_RESULT_BYTES_PER_TOKEN: u64 = 3;
pub fn conservative_token_estimate(text: &str) -> u64 {
let mut fifths: u64 = 0;
for c in text.chars() {
if c.is_ascii() {
fifths += 2;
} else {
fifths += 5 * c.len_utf8() as u64;
}
}
fifths.div_ceil(5)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TokenEstimator {
Conservative,
}
impl TokenEstimator {
pub fn estimate(&self, text: &str) -> u64 {
match self {
TokenEstimator::Conservative => conservative_token_estimate(text),
}
}
}
pub fn estimator_for_model(_model: &str) -> TokenEstimator {
TokenEstimator::Conservative
}
#[derive(Debug, Clone, Copy)]
pub struct ContextBudgetInputs<'a> {
pub model: &'a str,
pub provider_window: u64,
pub system_prompt: Option<&'a str>,
pub tools_schema: &'a [Value],
pub messages: &'a [SharedMessage],
pub skill_contents: &'a [&'a str],
pub memory_contents: &'a [&'a str],
pub thinking_budget_tokens: u64,
pub next_tool_result_bytes: u64,
pub max_output_tokens: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ContextBreakdown {
pub system_tokens: u64,
pub tool_schema_tokens: u64,
pub history_tokens: u64,
pub framing_tokens: u64,
pub skill_tokens: u64,
pub memory_tokens: u64,
}
impl ContextBreakdown {
pub fn total(&self) -> u64 {
self.system_tokens
+ self.tool_schema_tokens
+ self.history_tokens
+ self.framing_tokens
+ self.skill_tokens
+ self.memory_tokens
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ThinkingAccounting {
InsideOutputBudget,
SeparateFromOutput,
}
pub fn thinking_accounting_for_model(model: &str) -> ThinkingAccounting {
use crate::runtime::openai::WireProtocol;
match crate::runtime::openai::resolve_route(model).map(|route| route.wire) {
Some(
WireProtocol::AnthropicMessages
| WireProtocol::OpenAiChatCompletions
| WireProtocol::OpenAiResponses
| WireProtocol::CodexResponses,
) => ThinkingAccounting::InsideOutputBudget,
Some(WireProtocol::GoogleGeminiCodeAssist) | None => ThinkingAccounting::SeparateFromOutput,
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ContextReserves {
pub thinking_tokens: u64,
pub tool_result_tokens: u64,
pub output_tokens: u64,
pub safety_margin_tokens: u64,
pub thinking_accounting: ThinkingAccounting,
}
impl ContextReserves {
pub fn total(&self) -> u64 {
self.thinking_tokens
+ self.tool_result_tokens
+ self.output_tokens
+ self.safety_margin_tokens
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ContextAssessment {
pub provider_window: u64,
pub breakdown: ContextBreakdown,
pub reserves: ContextReserves,
pub history_messages: usize,
pub anchored_tokens: Option<u64>,
}
impl ContextAssessment {
pub fn used_tokens(&self) -> u64 {
self.anchored_tokens
.unwrap_or_else(|| self.estimated_tokens())
}
pub fn estimated_tokens(&self) -> u64 {
self.breakdown.total()
}
pub fn fixed_tokens(&self) -> u64 {
let b = &self.breakdown;
b.total() - b.history_tokens - b.framing_tokens
}
pub fn anchored(mut self, anchor: Option<&UsageAnchor>, messages: &[SharedMessage]) -> Self {
self.anchored_tokens =
anchor.and_then(|a| a.footprint(self.fixed_tokens(), messages));
self
}
pub fn budget_tokens(&self) -> u64 {
self.provider_window.saturating_sub(self.reserves.total())
}
pub fn remaining_tokens(&self) -> u64 {
self.budget_tokens().saturating_sub(self.used_tokens())
}
pub fn should_compact(&self) -> bool {
self.history_messages >= MIN_COMPACTION_MESSAGES
&& self.used_tokens() >= self.budget_tokens()
}
}
pub fn assess(inputs: &ContextBudgetInputs<'_>) -> ContextAssessment {
let estimator = estimator_for_model(inputs.model);
let system_tokens = inputs
.system_prompt
.map(|s| estimator.estimate(s))
.unwrap_or(0);
let tool_schema_tokens: u64 = inputs
.tools_schema
.iter()
.map(|schema| estimate_serialized(&estimator, schema))
.sum();
let history_tokens: u64 = inputs
.messages
.iter()
.map(|msg| estimate_message(&estimator, msg))
.sum();
let framing_tokens = PER_MESSAGE_FRAMING_TOKENS * inputs.messages.len() as u64;
let skill_tokens: u64 = inputs
.skill_contents
.iter()
.map(|s| estimator.estimate(s))
.sum();
let memory_tokens: u64 = inputs
.memory_contents
.iter()
.map(|s| estimator.estimate(s))
.sum();
ContextAssessment {
provider_window: inputs.provider_window,
breakdown: ContextBreakdown {
system_tokens,
tool_schema_tokens,
history_tokens,
framing_tokens,
skill_tokens,
memory_tokens,
},
reserves: {
let accounting = thinking_accounting_for_model(inputs.model);
let thinking = match inputs.thinking_budget_tokens {
0 => ZERO_BUDGET_THINKING_RESERVE_TOKENS,
budget => budget,
};
let cap = inputs.max_output_tokens;
let (thinking_tokens, output_tokens) = match accounting {
ThinkingAccounting::InsideOutputBudget => (
0,
cap.min(thinking.saturating_add(EXPECTED_RESPONSE_TOKENS)),
),
ThinkingAccounting::SeparateFromOutput => {
(thinking, cap.min(EXPECTED_RESPONSE_TOKENS))
}
};
ContextReserves {
thinking_tokens,
tool_result_tokens: inputs
.next_tool_result_bytes
.div_ceil(TOOL_RESULT_BYTES_PER_TOKEN),
output_tokens,
safety_margin_tokens: inputs.provider_window * SAFETY_MARGIN_PERCENT / 100,
thinking_accounting: accounting,
}
},
history_messages: inputs.messages.len(),
anchored_tokens: None,
}
}
#[derive(Clone, Debug)]
pub struct UsageAnchor {
pub provider_tokens: u64,
pub fixed_tokens: u64,
pub messages: Vec<SharedMessage>,
}
impl UsageAnchor {
pub fn footprint(&self, fixed_tokens: u64, messages: &[SharedMessage]) -> Option<u64> {
let shared = self
.messages
.iter()
.zip(messages)
.take_while(|(a, b)| super::continuation::same_on_wire(a, b))
.count();
(shared + 1 >= self.messages.len()).then(|| {
self.provider_tokens
.saturating_add(estimate_history(&messages[shared..]))
.saturating_add(fixed_tokens.saturating_sub(self.fixed_tokens))
})
}
}
fn estimate_serialized<T: serde::Serialize + ?Sized>(estimator: &TokenEstimator, value: &T) -> u64 {
match serde_json::to_string(value) {
Ok(s) => estimator.estimate(&s),
Err(_) => u64::MAX / 1024,
}
}
pub(crate) fn estimate_history(messages: &[SharedMessage]) -> u64 {
messages.iter().fold(0u64, |sum, message| {
sum.saturating_add(estimate_message(&TokenEstimator::Conservative, message))
.saturating_add(PER_MESSAGE_FRAMING_TOKENS)
})
}
fn estimate_message(estimator: &TokenEstimator, msg: &Value) -> u64 {
let wire = super::continuation::WireMessage(msg);
let Some(blocks) = msg["content"].as_array() else {
return estimate_serialized(estimator, &wire);
};
if !blocks.iter().any(|b| {
is_base64_image(b)
|| is_reasoning_block(b)
|| b["content"]
.as_array()
.is_some_and(|a| a.iter().any(is_base64_image))
}) {
return estimate_serialized(estimator, &wire);
}
let mut total = 0u64;
for block in blocks {
if is_base64_image(block) {
total += IMAGE_TOKEN_ESTIMATE;
continue;
}
if is_reasoning_block(block) {
total += estimate_reasoning_block(estimator, block);
continue;
}
if block["type"] == "tool_result" {
if let Some(inner) = block["content"].as_array() {
total += 8; for b in inner {
total += if is_base64_image(b) {
IMAGE_TOKEN_ESTIMATE
} else {
estimate_serialized(estimator, b)
};
}
continue;
}
}
total += estimate_serialized(estimator, block);
}
total
}
const REASONING_SIGNATURE_CHARS_PER_TOKEN: u64 = 4;
const REASONING_BLOCK_FRAMING_TOKENS: u64 = 8;
fn is_reasoning_block(b: &Value) -> bool {
b["type"] == "thinking" || b["type"] == "redacted_thinking"
}
fn estimate_reasoning_block(estimator: &TokenEstimator, block: &Value) -> u64 {
let text = block["thinking"].as_str().map_or(0, |t| estimator.estimate(t));
let opaque = block["signature"]
.as_str()
.or_else(|| block["data"].as_str())
.map_or(0, |s| s.len() as u64 / REASONING_SIGNATURE_CHARS_PER_TOKEN);
text.max(opaque) + REASONING_BLOCK_FRAMING_TOKENS
}
fn is_base64_image(b: &Value) -> bool {
b["type"] == "image" && b["source"]["type"] == "base64"
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use std::sync::Arc;
fn history_tokens_of(messages: &[SharedMessage]) -> u64 {
let inputs = ContextBudgetInputs {
model: "claude-sonnet-4-6",
provider_window: 200_000,
system_prompt: None,
tools_schema: &[],
messages,
skill_contents: &[],
memory_contents: &[],
thinking_budget_tokens: 0,
next_tool_result_bytes: 0,
max_output_tokens: 0,
};
assess(&inputs).breakdown.history_tokens
}
fn image_block(len: usize) -> Value {
json!({
"type": "image",
"source": {"type": "base64", "media_type": "image/png", "data": "A".repeat(len)}
})
}
fn shape_fixture(json: &str) -> (u64, Vec<SharedMessage>, u64) {
let fixture: Value = serde_json::from_str(json).unwrap();
let n = |v: &Value| v.as_u64().unwrap() as usize;
let filler = |len: usize| "x".repeat(len);
let base64 = |len: usize| {
const B64: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
(0..len).map(|i| B64[(i * 7 + i / 64) % 64] as char).collect::<String>()
};
let messages = fixture["messages"]
.as_array()
.unwrap()
.iter()
.map(|m| {
let role = m[0].as_str().unwrap();
let blocks = m[1].as_array().unwrap();
let content = match blocks[0][0].as_str().unwrap() {
"s" => json!(filler(n(&blocks[0][1]))),
_ => Value::Array(
blocks
.iter()
.map(|b| match b[0].as_str().unwrap() {
"b" => json!({"type": "text", "text": filler(n(&b[1]).saturating_sub(25))}),
"k" => json!({
"type": "thinking",
"thinking": filler(n(&b[1])),
"signature": base64(n(&b[2])),
}),
"i" => image_block(1_000_000),
other => panic!("block kind {other}"),
})
.collect(),
),
};
Arc::new(json!({"role": role, "content": content})) as SharedMessage
})
.collect();
let fixed = conservative_token_estimate(&filler(
n(&fixture["system_bytes"]) + n(&fixture["tool_schema_bytes"]),
));
(fixed, messages, fixture["billed"].as_u64().unwrap())
}
#[test]
fn xhigh_reasoning_history_estimate_tracks_its_billed_input() {
let (fixed, messages, billed) =
shape_fixture(include_str!("../../tests/fixtures/context_estimate/xhigh_worker_197k.json"));
assert_eq!(billed, 197_253);
assert!(messages.iter().any(|m| m["content"]
.as_array()
.is_some_and(|b| b.iter().any(|b| b["signature"].as_str().is_some_and(|s| s.len() > 1_000)))));
let estimate = fixed + estimate_history(&messages);
let ratio = estimate as f64 / billed as f64;
assert!((0.9..=1.1).contains(&ratio), "estimate {estimate} vs billed {billed}: {ratio:.3}");
let serialized: u64 = messages
.iter()
.map(|m| {
estimate_serialized(&TokenEstimator::Conservative, &**m) + PER_MESSAGE_FRAMING_TOKENS
})
.sum();
let old = (fixed + serialized) as f64 / billed as f64;
assert!(old > 1.25, "the old charge read {old:.3}×");
}
#[test]
fn reasoning_blocks_charge_the_larger_of_text_and_signature() {
let block = |text: usize, sig: usize| {
json!({"type": "thinking", "thinking": "t".repeat(text), "signature": "S".repeat(sig)})
};
let one = |b: Value| {
estimate_history(&[Arc::new(json!({"role": "assistant", "content": [b]}))])
- PER_MESSAGE_FRAMING_TOKENS
- REASONING_BLOCK_FRAMING_TOKENS
};
assert_eq!(one(block(10_000, 4_000)), 4_000);
assert_eq!(one(block(0, 40_000)), 10_000);
assert_eq!(one(block(2_500, 40_000)), 10_000);
assert_eq!(
one(json!({"type": "redacted_thinking", "data": "D".repeat(8_000)})),
2_000
);
let mixed = estimate_history(&[Arc::new(json!({"role": "assistant", "content": [
block(0, 400),
{"type": "text", "text": "a".repeat(2_500)},
]}))]);
assert!((1_100..1_130).contains(&mixed), "{mixed}");
}
#[test]
fn host_only_keys_are_not_charged() {
let plain = json!({"role": "user", "content": "hello there"});
let mut stamped = plain.clone();
stamped["_synaps_at"] = json!("2026-10-30T11:00:00.000Z");
stamped["_synaps_turn"] = json!({"id": "x".repeat(400)});
assert_eq!(
estimate_history(&[Arc::new(stamped.clone())]),
estimate_history(&[Arc::new(plain)])
);
let mut reasoning = json!({"role": "assistant", "content": [
{"type": "thinking", "thinking": "plan", "signature": "S".repeat(40)},
]});
let bare = estimate_history(&[Arc::new(reasoning.clone())]);
reasoning["_synaps_turn"] = json!({"id": "x".repeat(400)});
assert_eq!(estimate_history(&[Arc::new(reasoning)]), bare);
}
fn msg(role: &str, text: &str) -> SharedMessage {
Arc::new(json!({"role": role, "content": text}))
}
fn assessment_of(system: &str, messages: &[SharedMessage]) -> ContextAssessment {
assess(&ContextBudgetInputs {
model: "claude-opus-5-5",
provider_window: 1_000_000,
system_prompt: Some(system),
tools_schema: &[],
messages,
skill_contents: &[],
memory_contents: &[],
thinking_budget_tokens: 0,
next_tool_result_bytes: 0,
max_output_tokens: 128_000,
})
}
#[test]
fn footprint_is_anchored_to_the_last_reported_usage() {
let system = "s".repeat(2_500);
let earlier = vec![msg("user", "task"), msg("assistant", &"a".repeat(250_000))];
let anchor = UsageAnchor {
provider_tokens: 60_000,
fixed_tokens: assessment_of(&system, &earlier).fixed_tokens(),
messages: earlier.clone(),
};
let appended = [msg("user", &"r".repeat(25_000))];
let mut next = earlier.clone();
next.extend(appended.iter().cloned());
let assessment = assessment_of(&system, &next).anchored(Some(&anchor), &next);
assert_eq!(
assessment.used_tokens(),
60_000 + estimate_history(&appended),
"provider count + appended estimate"
);
assert!(assessment.estimated_tokens() > 100_000);
assert_eq!(assessment.anchored_tokens, Some(assessment.used_tokens()));
let mut stamped = next.clone();
let mut first = (*stamped[0]).clone();
first["_synaps_at"] = json!("2026-10-30T11:00:00.000Z");
stamped[0] = Arc::new(first);
assert_eq!(
assessment_of(&system, &stamped).anchored(Some(&anchor), &stamped).used_tokens(),
assessment.used_tokens()
);
let grown = "s".repeat(5_000);
assert_eq!(
assessment_of(&grown, &next).anchored(Some(&anchor), &next).used_tokens(),
assessment.used_tokens() + 1_000
);
assert_eq!(
assessment_of("", &next).anchored(Some(&anchor), &next).used_tokens(),
assessment.used_tokens()
);
}
#[test]
fn a_changed_request_tail_keeps_the_anchor_and_errs_high() {
let base = vec![msg("user", "task"), msg("assistant", &"a".repeat(25_000))];
let mut injected = base.clone();
injected.push(msg("user", &format!("next{}", "i".repeat(5_000))));
let anchor = UsageAnchor {
provider_tokens: 30_000,
fixed_tokens: 0,
messages: injected,
};
let mut next = base.clone();
next.push(msg("user", "next"));
next.push(msg("assistant", &"b".repeat(2_500)));
assert_eq!(
anchor.footprint(0, &next),
Some(30_000 + estimate_history(&next[2..]))
);
assert_eq!(anchor.footprint(0, &base), Some(30_000));
}
#[test]
fn footprint_falls_back_to_the_estimate_without_an_anchor() {
let system = "s".repeat(2_500);
let history = vec![
msg("user", "task"),
msg("assistant", &"a".repeat(25_000)),
msg("user", "more"),
msg("assistant", &"b".repeat(25_000)),
];
let plain = assessment_of(&system, &history);
assert_eq!(plain.anchored_tokens, None);
assert_eq!(plain.used_tokens(), plain.estimated_tokens());
assert_eq!(plain.anchored(None, &history).used_tokens(), plain.estimated_tokens());
let anchor = UsageAnchor {
provider_tokens: 5_000,
fixed_tokens: plain.fixed_tokens(),
messages: history.clone(),
};
let successor = vec![history[0].clone(), msg("user", "[host note]"), history[3].clone()];
let rolled = assessment_of(&system, &successor);
assert_eq!(
rolled.anchored(Some(&anchor), &successor).used_tokens(),
rolled.estimated_tokens()
);
}
#[test]
fn image_block_in_tool_result_charged_flat() {
let messages: Vec<SharedMessage> = vec![Arc::new(json!({
"role": "user",
"content": [{
"type": "tool_result",
"tool_use_id": "toolu_1",
"content": [
{"type": "text", "text": "Image: /tmp/x.png (1024x1536, image/png, 2292 KB)"},
image_block(3_000_000),
]
}]
}))];
let tokens = history_tokens_of(&messages);
assert!(tokens >= IMAGE_TOKEN_ESTIMATE, "{tokens}");
assert!(tokens <= IMAGE_TOKEN_ESTIMATE + 200, "{tokens}");
assert!(tokens < 10_000, "{tokens}");
}
#[test]
fn user_message_image_charged_flat() {
let messages: Vec<SharedMessage> = vec![Arc::new(json!({
"role": "user",
"content": [
{"type": "text", "text": "what is this?"},
image_block(3_000_000),
]
}))];
let tokens = history_tokens_of(&messages);
assert!(tokens >= IMAGE_TOKEN_ESTIMATE, "{tokens}");
assert!(tokens <= IMAGE_TOKEN_ESTIMATE + 200, "{tokens}");
}
#[test]
fn string_content_unchanged() {
let msg = json!({"role": "user", "content": "plain string content, no blocks here"});
let estimator = estimator_for_model("claude-sonnet-4-6");
assert_eq!(
estimate_message(&estimator, &msg),
estimate_serialized(&estimator, &msg)
);
}
#[test]
fn array_content_without_images_unchanged() {
let estimator = estimator_for_model("claude-sonnet-4-6");
for msg in [
json!({"role": "assistant", "content": [
{"type": "text", "text": "let me look"},
{"type": "tool_use", "id": "toolu_1", "name": "read", "input": {"path": "x"}}
]}),
json!({"role": "user", "content": [
{"type": "tool_result", "tool_use_id": "toolu_1", "content": "1\tfoo\n2\tbar"}
]}),
json!({"role": "user", "content": [
{"type": "tool_result", "tool_use_id": "toolu_1", "content": [
{"type": "text", "text": "nested text only"}
]}
]}),
] {
assert_eq!(
estimate_message(&estimator, &msg),
estimate_serialized(&estimator, &msg),
"{msg}"
);
}
}
#[test]
fn non_image_blocks_still_charged_by_payload() {
let msg = json!({"role": "user", "content": [
{"type": "text", "text": "A".repeat(100_000)}
]});
let estimator = estimator_for_model("claude-sonnet-4-6");
assert!(estimate_message(&estimator, &msg) > 10_000);
}
#[test]
fn ascii_is_charged_at_the_dense_bpe_rate() {
assert_eq!(conservative_token_estimate(&"a".repeat(25)), 10);
assert_eq!(conservative_token_estimate(""), 0);
}
#[test]
fn non_ascii_is_charged_per_utf8_byte() {
assert_eq!(conservative_token_estimate("你好"), 6);
assert_eq!(conservative_token_estimate("😀"), 4);
assert_eq!(conservative_token_estimate("да"), 4);
}
#[test]
fn tiny_window_never_underflows_and_fails_toward_compaction() {
let messages: Vec<SharedMessage> = (0..MIN_COMPACTION_MESSAGES)
.map(|i| Arc::new(json!({"role": "user", "content": format!("m{i}")})) as SharedMessage)
.collect();
let assessment = assess(&ContextBudgetInputs {
model: "claude-sonnet-4-6",
provider_window: 10,
system_prompt: None,
tools_schema: &[],
messages: &messages,
skill_contents: &[],
memory_contents: &[],
thinking_budget_tokens: 1_000,
next_tool_result_bytes: 0,
max_output_tokens: 64_000,
});
assert_eq!(assessment.budget_tokens(), 0);
assert_eq!(assessment.remaining_tokens(), 0);
assert!(assessment.should_compact());
}
#[test]
fn anthropic_wire_counts_thinking_exactly_once_inside_the_output_reserve() {
let messages: Vec<SharedMessage> =
vec![Arc::new(json!({"role": "user", "content": "hi"})) as SharedMessage];
let window = 200_000;
let margin = window * SAFETY_MARGIN_PERCENT / 100;
let inputs = ContextBudgetInputs {
model: "claude-sonnet-4-6",
provider_window: window,
system_prompt: None,
tools_schema: &[],
messages: &messages,
skill_contents: &[],
memory_contents: &[],
thinking_budget_tokens: 8_000,
next_tool_result_bytes: 0,
max_output_tokens: 64_000,
};
let assessment = assess(&inputs);
assert_eq!(
assessment.reserves.thinking_accounting,
ThinkingAccounting::InsideOutputBudget
);
assert_eq!(
assessment.reserves.thinking_tokens, 0,
"thinking must not be reserved a second time on the Anthropic wire"
);
let envelope = 8_000 + EXPECTED_RESPONSE_TOKENS;
assert_eq!(assessment.reserves.output_tokens, envelope);
assert_eq!(
assessment.budget_tokens(),
window - (envelope + margin),
"budget must subtract the output/thinking envelope exactly once"
);
let big_thinking = ContextBudgetInputs {
thinking_budget_tokens: 100_000,
..inputs
};
let assessment = assess(&big_thinking);
assert_eq!(assessment.reserves.output_tokens, 64_000);
assert_eq!(assessment.reserves.thinking_tokens, 0);
let adaptive = assess(&ContextBudgetInputs {
thinking_budget_tokens: 0,
..inputs
});
assert_eq!(
adaptive.reserves.output_tokens,
ZERO_BUDGET_THINKING_RESERVE_TOKENS + EXPECTED_RESPONSE_TOKENS
);
}
#[test]
fn the_output_reserve_is_a_realistic_need_not_the_max_tokens_ceiling() {
let hard_line = |model: &str, window: u64, thinking: u64| {
let (assessment, budget) = crate::runtime::stream::pre_request_context_budget(
model, window, None, &[], &[], thinking, 30_000, None,
);
(assessment.reserves.output_tokens, window - budget.required_next_round_tokens)
};
for thinking in [0, 16_384] {
for model in ["claude-opus-5-5", "claude-sonnet-4-6"] {
assert_eq!(hard_line(model, 200_000, thinking), (32_768, 126_720), "{model}");
assert_eq!(hard_line(model, 1_000_000, thinking), (32_768, 806_720), "{model}");
}
}
assert_eq!(hard_line("claude-opus-5-5", 200_000, 4_096), (20_480, 139_008));
assert_eq!(hard_line("claude-opus-5-5", 200_000, 32_768), (49_152, 110_336));
assert_eq!(hard_line("claude-sonnet-4-6", 200_000, 100_000).0, 64_000);
}
#[test]
fn openai_wires_account_reasoning_inside_the_output_reserve() {
for model in ["openai-codex/gpt-5.2-codex", "xai-auth/grok-4"] {
let accounting = thinking_accounting_for_model(model);
if model.starts_with("xai-auth") && accounting == ThinkingAccounting::SeparateFromOutput
{
continue;
}
assert_eq!(
accounting,
ThinkingAccounting::InsideOutputBudget,
"{model}"
);
}
}
#[test]
fn gemini_and_unknown_models_reserve_thinking_separately() {
let messages: Vec<SharedMessage> =
vec![Arc::new(json!({"role": "user", "content": "hi"})) as SharedMessage];
let window = 200_000;
let margin = window * SAFETY_MARGIN_PERCENT / 100;
for model in ["google-gemini/gemini-2.5-pro", "totally/unroutable-model"] {
let assessment = assess(&ContextBudgetInputs {
model,
provider_window: window,
system_prompt: None,
tools_schema: &[],
messages: &messages,
skill_contents: &[],
memory_contents: &[],
thinking_budget_tokens: 8_000,
next_tool_result_bytes: 0,
max_output_tokens: 64_000,
});
assert_eq!(
assessment.reserves.thinking_accounting,
ThinkingAccounting::SeparateFromOutput,
"{model}"
);
assert_eq!(assessment.reserves.thinking_tokens, 8_000, "{model}");
assert_eq!(
assessment.reserves.output_tokens, EXPECTED_RESPONSE_TOKENS,
"{model}"
);
assert_eq!(
assessment.budget_tokens(),
window - (8_000 + EXPECTED_RESPONSE_TOKENS + margin),
"{model}: separate wires reserve both envelopes"
);
}
}
#[tokio::test]
async fn runtime_assessment_matches_engine_calculation() {
let mut runtime = crate::Runtime::new_headless();
runtime.set_system_prompt("shared system prompt for budget parity".into());
runtime.set_context_window(Some(200_000));
let messages: Vec<SharedMessage> = vec![
Arc::new(json!({"role": "user", "content": "hello 你好 😀"})),
Arc::new(json!({"role": "assistant", "content": [
{"type": "text", "text": "fn main() { println!(\"hi\"); }"}
]})),
];
let via_runtime = runtime.assess_context(&messages).await;
let system = runtime.effective_system_prompt().await;
let schema = runtime.tools_shared().read().await.tools_schema();
let direct = assess(&ContextBudgetInputs {
model: runtime.model(),
provider_window: runtime.context_window(),
system_prompt: system.as_deref(),
tools_schema: &schema,
messages: &messages,
skill_contents: &[],
memory_contents: &[],
thinking_budget_tokens: runtime.thinking_budget() as u64,
next_tool_result_bytes: runtime.max_tool_output() as u64,
max_output_tokens: crate::runtime::helpers::HelperMethods::max_tokens_for_model(
runtime.model(),
),
});
assert_eq!(
via_runtime, direct,
"Runtime surface must not fork the budget math"
);
assert!(
via_runtime.reserves.output_tokens > 0,
"output reserve must be model-derived"
);
assert!(
via_runtime.breakdown.system_tokens > 0,
"actual system segment must be accounted"
);
}
fn synthetic_memory_contribution(
rendered: &str,
) -> crate::runtime::memory_context::MemoryContextContribution {
use crate::runtime::memory_context as mc;
use agent_core::BoundedText;
mc::MemoryContextContribution {
schema: mc::ContributionSchemaVersion::parse("contribution/1").expect("valid schema"),
provider_id: mc::ContextProviderId::parse("axel-memory").expect("valid provider"),
project_id: crate::runtime::memory_project_id(),
records: vec![mc::MemoryContributionRecord {
memory_id: mc::MemoryId::parse("mem-0001").expect("valid memory id"),
source: mc::MemorySource::ChatHistory,
timestamp: std::time::SystemTime::UNIX_EPOCH,
rank_reason: vec![mc::RankReason::ExactTopic],
sensitivity: agent_core::core::disclosure::DisclosureClass::ModelVisible,
retention: mc::RetentionClass::Standard,
content: BoundedText::new("the project uses session-scoped authorization", 2048),
truncated: false,
supersedes: None,
}],
rendered: BoundedText::new(rendered, 16 * 1024),
accounting: mc::ContributionAccounting::default(),
}
}
#[tokio::test]
async fn memory_contribution_charges_only_the_memory_lane_and_never_reserves() {
let mut runtime = crate::Runtime::new_headless();
runtime.set_system_prompt("system prompt for the memory budget test".into());
runtime.set_context_window(Some(200_000));
let messages: Vec<SharedMessage> = vec![
Arc::new(json!({"role": "user", "content": "what auth model do we use?"})),
Arc::new(json!({"role": "assistant", "content": "let me check"})),
];
let without = runtime.assess_context(&messages).await;
assert_eq!(without.breakdown.memory_tokens, 0);
let rendered = "[Axel memory — lower-authority project data; verify before relying]\n\
1. mem-0001 — Decision — the project uses session-scoped \
authorization rather than persisted grants.";
runtime
.hold_memory_contribution(synthetic_memory_contribution(rendered))
.expect("in-bounds contribution for the runtime's own project is accepted");
let with = runtime.assess_context(&messages).await;
let expected_memory = conservative_token_estimate(rendered);
assert!(expected_memory > 0, "test contribution must be non-empty");
assert_eq!(
with.breakdown.memory_tokens, expected_memory,
"memory lane must carry exactly the estimated rendered tokens"
);
assert_eq!(
with.reserves, without.reserves,
"reserves must be COMPLETELY unaffected by memory content"
);
assert_eq!(
with.budget_tokens(),
without.budget_tokens(),
"the admissible budget (window minus reserves) must not move"
);
assert_eq!(
with.breakdown.system_tokens,
without.breakdown.system_tokens
);
assert_eq!(
with.breakdown.tool_schema_tokens,
without.breakdown.tool_schema_tokens
);
assert_eq!(
with.breakdown.history_tokens,
without.breakdown.history_tokens
);
assert_eq!(
with.breakdown.framing_tokens,
without.breakdown.framing_tokens
);
assert_eq!(with.breakdown.skill_tokens, without.breakdown.skill_tokens);
assert_eq!(with.history_messages, without.history_messages);
assert_eq!(with.used_tokens(), without.used_tokens() + expected_memory);
assert_eq!(
with.remaining_tokens(),
without.remaining_tokens() - expected_memory
);
assert_eq!(
with.should_compact(),
with.history_messages >= MIN_COMPACTION_MESSAGES
&& with.used_tokens() >= with.budget_tokens()
);
runtime.clear_memory_contribution();
let cleared = runtime.assess_context(&messages).await;
assert_eq!(
cleared, without,
"no held contribution ⇒ assessment identical to the empty lane"
);
}
#[tokio::test]
async fn foreign_project_contribution_is_refused_and_never_budgeted() {
use crate::runtime::memory_context as mc;
let mut runtime = crate::Runtime::new_headless();
runtime.set_context_window(Some(200_000));
let mut contribution = synthetic_memory_contribution("foreign project memory");
contribution.project_id =
mc::ProjectId::parse("project-not-this-host").expect("valid project id");
assert_eq!(
runtime.hold_memory_contribution(contribution),
Err(mc::MemoryContextError::ContributionProjectMismatch)
);
let messages: Vec<SharedMessage> =
vec![Arc::new(json!({"role": "user", "content": "hi"})) as SharedMessage];
let assessment = runtime.assess_context(&messages).await;
assert_eq!(
assessment.breakdown.memory_tokens, 0,
"a refused contribution must never be budgeted"
);
}
}