Skip to content

feat(smith): trust provider input token usage #1329

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
edwin-zvs merged 1 commit into from
Sep 10, 2026
Merged

feat(smith): trust provider input token usage #1329

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 38 additions & 15 deletions crates/adapter-smith/src/agent.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -616,6 +616,7 @@ pub async fn run(
let provider = spec.provider;
let mut provider_context_window = None;
let mut provider_context_checked = false;
let mut input_tokens = context::InputTokenCounter::default();
// Per-model learned token limits — adapts on overflow errors
// and bumps upward on successful probe calls. Shared across
// all construct sessions on this machine via state_dir.
Expand Down Expand Up @@ -864,11 +865,11 @@ pub async fn run(
}
let hardcoded_cap = context::context_window_tokens(provider_name, &model);
let learned = limits.get(provider_name, &model);
let est =
fixed_context_tokens.saturating_add(context::estimate_tokens(&messages) as u64);
let preflight_tokens =
context::preflight_input_tokens(&input_tokens, fixed_context_tokens, &messages);
let is_probe = provider_context_window.is_none()
&& learned.is_some()
&& limits.should_probe(provider_name, &model, est, now_ms);
&& limits.should_probe(provider_name, &model, preflight_tokens, now_ms);
let effective_cap = provider_context_window
.or(learned)
.unwrap_or(hardcoded_cap as u64);
Expand All @@ -877,7 +878,7 @@ pub async fn run(
} else {
context::UTILIZATION
};
let budget = context::message_budget(effective_cap, utilization, fixed_context_tokens);
let budget = (effective_cap as f64 * utilization) as u64;
// Auto-compact pass before the destructive prune. Headless
// sessions don't get a `/compact` UI, so this is the only
// way summaries get generated outside of an interactive
Expand All @@ -886,7 +887,7 @@ pub async fn run(
match crate::compact::maybe_auto_compact(
&mut messages,
effective_cap,
fixed_context_tokens,
preflight_tokens,
provider.as_ref(),
&model,
)
Expand Down Expand Up @@ -915,10 +916,18 @@ pub async fn run(
}
}
}
if context::prune_to_budget(&mut messages, budget) > 0 {
if context::prune_to_input_budget(
&mut messages,
budget,
fixed_context_tokens,
&input_tokens,
) > 0
{
reset_context_serve(&tool_ctx);
}

let mut sent_heuristic_tokens =
fixed_context_tokens.saturating_add(context::estimate_tokens(&messages) as u64);
let mut sink = MessageSink { emit: &emit };
let turn = match crate::provider_watchdog::complete(
provider.as_ref(),
Expand All @@ -945,14 +954,18 @@ pub async fn run(
effective_cap,
now_ms,
);
let retry_budget = context::message_budget(
new_limit,
context::UTILIZATION,
let retry_budget = (new_limit as f64 * context::UTILIZATION) as u64;
if context::prune_to_input_budget(
&mut messages,
retry_budget,
fixed_context_tokens,
);
if context::prune_to_budget(&mut messages, retry_budget) > 0 {
&input_tokens,
) > 0
{
reset_context_serve(&tool_ctx);
}
sent_heuristic_tokens = fixed_context_tokens
.saturating_add(context::estimate_tokens(&messages) as u64);
emit.emit(SessionEvent::Status {
state: SessionState::Running,
detail: Some(format!(
Expand Down Expand Up @@ -993,6 +1006,8 @@ pub async fn run(
}
};

input_tokens.record(turn.usage.input_tokens, sent_heuristic_tokens);

// Record the successful call so probe state advances
// (and the learned limit grows on a probe that pushed
// past the prior cap).
Expand All @@ -1007,16 +1022,16 @@ pub async fn run(

emit.emit(SessionEvent::Cost {
usd: turn.usage.usd,
tokens_in: turn.usage.input_tokens,
tokens_in: turn.usage.input_tokens_or_zero(),
tokens_out: turn.usage.output_tokens,
tokens_cached: turn.usage.cached_tokens,
model: Some(current_model_spec.clone()),
});
// Context gauge (spec 0104): prefer a provider-reported runtime
// allocation; otherwise retain Smith's learned/static fallback.
if turn.usage.input_tokens > 0 {
if let Some(reported_input_tokens) = turn.usage.input_tokens {
emit.emit(SessionEvent::ContextUsage {
used_tokens: turn.usage.input_tokens,
used_tokens: reported_input_tokens,
window_tokens: Some(effective_cap),
});
// Per-component detail behind the gauge (spec 0156) — all
Expand Down Expand Up @@ -1838,8 +1853,13 @@ fn build_profile_model(
profile_api_key(profile, name, &["OPENAI_API_KEY"])?,
)?),
provider::routing::Provider::Anthropic => {
Box::new(provider::anthropic::Anthropic::with_config(
Box::new(provider::anthropic::Anthropic::with_options(
base_url,
provider::anthropic::AnthropicOptions {
cache_control: profile.anthropic_cache_control,
betas: profile.anthropic_betas.clone(),
context_window_tokens: profile.anthropic_context_window_tokens,
},
profile_api_key(profile, name, &["ANTHROPIC_API_KEY"])?,
)?)
}
Expand Down Expand Up @@ -2016,6 +2036,9 @@ mod tests {
api_key: Some("test-key".to_string()),
api_key_env: None,
model: model.map(str::to_string),
anthropic_cache_control: None,
anthropic_betas: Vec::new(),
anthropic_context_window_tokens: None,
}
}

Expand Down
64 changes: 50 additions & 14 deletions crates/adapter-smith/src/compact.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -165,19 +165,18 @@ pub async fn compact(
/// the rolling-prune path.
///
/// `effective_cap` is the per-model input-token cap the caller is using for
/// budget math. `fixed_tokens` is the system/tool prefix paid on every call.
/// The trigger compares their sum with the whole-window threshold so a large
/// fixed prefix cannot make Smith believe the context is emptier than it is.
/// budget math. `preflight_input_tokens` is provider-anchored when the prior
/// response reported usage and otherwise is Smith's full char heuristic. The
/// trigger never re-estimates that authoritative baseline here.
pub async fn maybe_auto_compact(
messages: &mut Vec<Message>,
effective_cap: u64,
fixed_tokens: u64,
preflight_input_tokens: u64,
provider: &dyn LlmProvider,
model: &str,
) -> Result<Option<CompactOutcome>> {
let est = fixed_tokens.saturating_add(context::estimate_tokens(messages) as u64);
let trigger = ((effective_cap as f64) * AUTO_COMPACT_RATIO) as u64;
if est < trigger {
if preflight_input_tokens < trigger {
return Ok(None);
}
compact(messages, DEFAULT_KEEP_PAIRS, provider, model).await
Expand Down Expand Up @@ -619,7 +618,8 @@ mod tests {
let provider = StubProvider::new("unused");
// cap = 100k tokens; conversation is ~10 tokens; way under
// threshold.
let outcome = maybe_auto_compact(&mut messages, 100_000, 0, &provider, "stub")
let preflight = context::estimate_tokens(&messages) as u64;
let outcome = maybe_auto_compact(&mut messages, 100_000, preflight, &provider, "stub")
.await
.unwrap();
assert!(outcome.is_none());
Expand All @@ -640,7 +640,8 @@ mod tests {
messages.push(user("recent"));
messages.push(asst("latest"));
let provider = StubProvider::new("auto-summary");
let outcome = maybe_auto_compact(&mut messages, 1000, 0, &provider, "stub")
let preflight = context::estimate_tokens(&messages) as u64;
let outcome = maybe_auto_compact(&mut messages, 1000, preflight, &provider, "stub")
.await
.unwrap()
.expect("should auto-compact");
Expand Down Expand Up @@ -670,10 +671,16 @@ mod tests {
assert!(est < prune_budget);

let provider = StubProvider::new("pre-prune-summary");
let outcome = maybe_auto_compact(&mut messages, effective_cap as u64, 0, &provider, "stub")
.await
.unwrap()
.expect("should auto-compact before rolling prune would fire");
let outcome = maybe_auto_compact(
&mut messages,
effective_cap as u64,
est as u64,
&provider,
"stub",
)
.await
.unwrap()
.expect("should auto-compact before rolling prune would fire");
assert!(outcome.dropped_turn_pairs > 0);
}

Expand All @@ -692,10 +699,39 @@ mod tests {
assert!(message_tokens < (effective_cap as f64 * AUTO_COMPACT_RATIO) as u64);

let provider = StubProvider::new("fixed-overhead-summary");
let outcome = maybe_auto_compact(&mut messages, effective_cap, 6_000, &provider, "stub")
let outcome = maybe_auto_compact(
&mut messages,
effective_cap,
message_tokens + 6_000,
&provider,
"stub",
)
.await
.unwrap()
.expect("fixed overhead should push the total over the trigger");
assert!(outcome.dropped_turn_pairs > 0);
}

#[tokio::test]
async fn auto_compact_obeys_provider_anchored_pressure() {
let mut messages = vec![
user("one"),
asst("one"),
user("two"),
asst("two"),
user("three"),
asst("three"),
user("four"),
asst("four"),
user("five"),
asst("five"),
];
assert!(context::estimate_tokens(&messages) < 100);
let provider = StubProvider::new("provider-pressure-summary");
let outcome = maybe_auto_compact(&mut messages, 1_000, 700, &provider, "stub")
.await
.unwrap()
.expect("fixed overhead should push the total over the trigger");
.expect("provider count above 65% should trigger compaction");
assert!(outcome.dropped_turn_pairs > 0);
}
}
Loading
Loading