From 928ce4d5c6f4ac21fa23fb296b3808d63ce7efe1 Mon Sep 17 00:00:00 2001 From: rainbowpuffpuff <83903147+rainbowpuffpuff@users.noreply.github.com> Date: Tue, 30 Jun 2026 17:23:31 +0300 Subject: [PATCH] fix(tokenjuice): attribute compaction savings to the per-turn model (#4122) (#4229) Co-authored-by: M3gA-Mind --- .../agent/harness/session/turn/core.rs | 11 ++ .../harness/subagent_runner/ops/loop_.rs | 54 ++++---- src/openhuman/tokenjuice/savings.rs | 118 +++++++++++++++--- 3 files changed, 142 insertions(+), 41 deletions(-) diff --git a/src/openhuman/agent/harness/session/turn/core.rs b/src/openhuman/agent/harness/session/turn/core.rs index 565705918..962ea5ccd 100644 --- a/src/openhuman/agent/harness/session/turn/core.rs +++ b/src/openhuman/agent/harness/session/turn/core.rs @@ -815,6 +815,15 @@ impl Agent { user_message, ); let (outcome_result, subagent_usage_entries) = + crate::openhuman::tokenjuice::savings::with_turn_model( + effective_model.clone(), + // Box the per-turn context chain onto the heap so the added + // `with_turn_model` scope does not deepen the worker stack — + // the same stack-accumulation guard the sub-agent path uses + // around `run_turn_engine`. Without this the cron agent-job + // lib test overflows its stack under llvm-cov instrumentation + // (issue #4122 review). + Box::pin( super::super::super::turn_subagent_usage::with_turn_collector( super::super::super::turn_attachments_context::with_current_turn_image_placeholders( turn_image_placeholders, @@ -843,6 +852,8 @@ impl Agent { )), ), ), + ) + ), ) .await; let outcome = outcome_result?; diff --git a/src/openhuman/agent/harness/subagent_runner/ops/loop_.rs b/src/openhuman/agent/harness/subagent_runner/ops/loop_.rs index 1b2e108a4..d48a8be3b 100644 --- a/src/openhuman/agent/harness/subagent_runner/ops/loop_.rs +++ b/src/openhuman/agent/harness/subagent_runner/ops/loop_.rs @@ -198,29 +198,37 @@ pub(super) async fn run_inner_loop( // `nested_subagent_dispatch_runs_on_a_constrained_worker_stack`; // the deep end-to-end catcher is the `chat-harness-subagent` // Playwright spec. - let outcome = super::super::super::model_vision_context::with_current_model_vision( - model_vision, - Box::pin(super::super::super::engine::run_turn_engine( - provider, - history, - &mut tool_source, - &progress, - &mut observer, - &checkpoint, - &parser, - "subagent", - model, - temperature, - true, // silent — sub-agents never echo to stdout - &crate::openhuman::config::MultimodalConfig::default(), - &crate::openhuman::config::MultimodalFileConfig::default(), - max_iterations, - max_output_tokens, - None, // sub-agents don't stream a draft - &["ask_user_clarification"], - run_queue, // steering channel for `steer_subagent` (None = non-steerable) - autocompact.as_ref(), - )), + let outcome = crate::openhuman::tokenjuice::savings::with_turn_model( + model.to_string(), + // Box the context chain so the added `with_turn_model` scope keeps the + // nested sub-agent future off the constrained worker stack (mirrors the + // existing `Box::pin` guard below; issue #4122 review). + Box::pin( + super::super::super::model_vision_context::with_current_model_vision( + model_vision, + Box::pin(super::super::super::engine::run_turn_engine( + provider, + history, + &mut tool_source, + &progress, + &mut observer, + &checkpoint, + &parser, + "subagent", + model, + temperature, + true, // silent — sub-agents never echo to stdout + &crate::openhuman::config::MultimodalConfig::default(), + &crate::openhuman::config::MultimodalFileConfig::default(), + max_iterations, + max_output_tokens, + None, // sub-agents don't stream a draft + &["ask_user_clarification"], + run_queue, // steering channel for `steer_subagent` (None = non-steerable) + autocompact.as_ref(), + )), + ), + ), ) .await?; diff --git a/src/openhuman/tokenjuice/savings.rs b/src/openhuman/tokenjuice/savings.rs index add57c9e7..4106011cc 100644 --- a/src/openhuman/tokenjuice/savings.rs +++ b/src/openhuman/tokenjuice/savings.rs @@ -52,6 +52,24 @@ pub struct SavingsAggregate { pub by_compressor: HashMap, } +impl SavingsAggregate { + /// Fold one compaction's savings into the aggregate, attributed to `model`. + /// Caller guarantees `original > compacted`. Pure (no global state) so it is + /// unit-testable without touching the process-global aggregate. + fn record_saving(&mut self, model: &str, compressor: &str, original: u64, compacted: u64) { + let cost = cost_saved_usd(model, original.saturating_sub(compacted)); + self.total.add(original, compacted, cost); + self.by_model + .entry(model.to_string()) + .or_default() + .add(original, compacted, cost); + self.by_compressor + .entry(compressor.to_string()) + .or_default() + .add(original, compacted, cost); + } +} + struct State { aggregate: SavingsAggregate, /// Model used to price the saved input tokens (the configured default). @@ -75,6 +93,38 @@ fn state() -> &'static Mutex { STATE.get_or_init(|| Mutex::new(State::default())) } +tokio::task_local! { + /// The model actually running the current turn/sub-agent, scoped by the + /// agent loop around `run_turn_engine` (mirrors + /// [`crate::openhuman::agent::harness::model_vision_context`]). When set, + /// compaction savings are priced against *this* model instead of the + /// process-global configured default (issue #4122). Unset ⇒ fall back to + /// the configured default, so non-harness callers and tests are unaffected + /// — strictly additive. + pub static TURN_MODEL: String; +} + +/// Run `future` with `model` installed as the per-turn attribution model used +/// to price compaction savings. Intended call site is around each +/// `run_turn_engine` invocation, alongside the other per-turn `*_context` +/// scopes (issue #4122). +pub async fn with_turn_model(model: String, future: F) -> R +where + F: std::future::Future, +{ + TURN_MODEL.scope(model, future).await +} + +/// The model to attribute savings to: the per-turn [`TURN_MODEL`] when scoped +/// and non-empty, otherwise the process-global configured `default`. +fn resolve_attribution_model(default: &str) -> String { + TURN_MODEL + .try_with(|m| m.clone()) + .ok() + .filter(|m| !m.trim().is_empty()) + .unwrap_or_else(|| default.to_string()) +} + /// Install the attribution model and snapshot location, loading any prior /// snapshot. Called once at startup from [`crate::openhuman::tokenjuice::install_config`]. pub fn configure(attribution_model: String, workspace_dir: &std::path::Path) { @@ -104,25 +154,17 @@ pub fn record( if original_tokens <= compacted_tokens { return; } - let saved = original_tokens - compacted_tokens; - let mut st = state().lock().unwrap_or_else(|p| p.into_inner()); - let model = st.attribution_model.clone(); - let cost = cost_saved_usd(&model, saved); - - st.aggregate - .total - .add(original_tokens, compacted_tokens, cost); - st.aggregate - .by_model - .entry(model) - .or_default() - .add(original_tokens, compacted_tokens, cost); - st.aggregate - .by_compressor - .entry(compressor.as_str().to_string()) - .or_default() - .add(original_tokens, compacted_tokens, cost); + // Attribute the saving to the per-turn model the agent loop scoped via + // `with_turn_model` (issue #4122); fall back to the configured default when + // unscoped (non-harness callers, tests). + let model = resolve_attribution_model(&st.attribution_model); + st.aggregate.record_saving( + &model, + compressor.as_str(), + original_tokens, + compacted_tokens, + ); let _ = content_kind; // reserved for a future by-kind breakdown persist(&st); @@ -211,4 +253,44 @@ mod tests { record(ContentKind::Json, CompressorKind::SmartCrusher, 50, 100); assert_eq!(stats().total.events, before, "no-op when not smaller"); } + + #[test] + fn record_saving_attributes_to_given_model() { + // Pure aggregation on a LOCAL aggregate — no process-global state, so it + // cannot race the other tests in this module. + let mut agg = SavingsAggregate::default(); + agg.record_saving("turn-model-x", "smartcrusher", 2000, 1000); + assert_eq!(agg.total.tokens_saved, 1000); + assert!( + agg.by_model.contains_key("turn-model-x"), + "saving must be attributed to the supplied model" + ); + assert!(agg.by_model["turn-model-x"].cost_saved_usd > 0.0); + } + + #[tokio::test] + async fn attribution_model_falls_back_to_default_when_unscoped() { + assert_eq!(resolve_attribution_model("default-model"), "default-model"); + } + + #[tokio::test] + async fn attribution_model_prefers_scoped_turn_model() { + let got = with_turn_model("turn-model".to_string(), async { + resolve_attribution_model("default-model") + }) + .await; + assert_eq!( + got, "turn-model", + "scoped per-turn model wins (issue #4122)" + ); + } + + #[tokio::test] + async fn blank_turn_model_falls_back_to_default() { + let got = with_turn_model(" ".to_string(), async { + resolve_attribution_model("default-model") + }) + .await; + assert_eq!(got, "default-model", "blank scoped model is ignored"); + } }