use super::*;
use crate::{runtime::config::AgentLoopConfig, ticket_token_budget::NativeWorkTokenBudget};
use bamboo_llm::{provider::ProviderModelInfo, LLMChunk, LLMError, LLMProvider, LLMStream};
use bamboo_tickets::{ExecutionTokenUsage, WorkTokenLimits};
use futures::StreamExt;
use std::sync::atomic::{AtomicUsize, Ordering};
mod terminal;
include!("token_budget_outer_guard_fixture.rs");
enum Script {
ReplacementInvalidAfterDone,
ReplacementOverCapWhileOpen,
ReplacementInputOverCapWhileOpen,
Known,
KnownText,
DoneError,
DoneOverCap,
DonePending,
RepeatedDone,
NoDone,
Partial,
BootstrapError,
StreamError,
Pending,
OverCap,
PartialOverCap,
InconsistentOverCap,
LegacyCachedInputOverCap,
LegacyCompositeInputOverCap,
LegacyOutputOverCap,
LegacyThinkingOverCap,
LegacyOverflowOverCap,
LegacyWithinCap,
MixedKnown,
CanonicalCacheOnlyOverCap,
CanonicalReasoningOnlyOverCap,
CanonicalTotalOnlyOverCap,
LegacySplitInputOverCap,
CanonicalSplitCacheOverCap,
SplitTotalMismatch,
}
struct Provider {
script: Script,
calls: AtomicUsize,
}
impl Provider {
fn new(script: Script) -> Arc<Self> {
Arc::new(Self {
script,
calls: AtomicUsize::new(0),
})
}
}
fn usage(input: Option<u64>, output: Option<u64>) -> LLMChunk {
LLMChunk::ProviderUsage {
input_tokens: input,
output_tokens: output,
total_tokens: None,
reasoning_tokens: Some(5),
cache_creation_input_tokens: None,
cache_read_input_tokens: Some(10),
cache_write_input_tokens: None,
}
}
#[async_trait::async_trait]
impl LLMProvider for Provider {
async fn vision_support_override(&self, model: &str) -> Option<bool> {
(model == "auxiliary-model").then_some(true)
}
async fn list_model_info(&self) -> bamboo_llm::provider::Result<Vec<ProviderModelInfo>> {
Ok(vec![
ProviderModelInfo {
id: "fixture-model".into(),
max_context_tokens: Some(100),
max_output_tokens: Some(40),
},
ProviderModelInfo {
id: "auxiliary-model".into(),
max_context_tokens: Some(8192),
max_output_tokens: Some(2048),
},
])
}
async fn chat_stream(
&self,
_: &[bamboo_domain::Message],
_: &[bamboo_domain::ToolSchema],
cap: Option<u32>,
model: &str,
) -> bamboo_llm::provider::Result<LLMStream> {
self.calls.fetch_add(1, Ordering::SeqCst);
if model == "auxiliary-model" {
assert!(cap.is_some_and(|cap| cap > 0 && cap <= 2048));
} else {
assert_eq!(cap, Some(40));
}
let mut complete = usage(Some(20), Some(u64::from(cap.unwrap_or(40).min(10))));
if let LLMChunk::ProviderUsage {
reasoning_tokens,
cache_read_input_tokens,
..
} = &mut complete
{
*reasoning_tokens = None;
*cache_read_input_tokens = None;
}
let chunks = match self.script {
Script::ReplacementInvalidAfterDone => {
let first = bamboo_llm::types::ProviderUsageReplacement {
input_tokens: Some(20),
output_tokens: Some(10),
total_tokens: Some(30),
cache_read_input_tokens: Some(5),
..Default::default()
};
let invalid = bamboo_llm::types::ProviderUsageReplacement {
cache_read_input_tokens: None,
invalid_fields: bamboo_llm::types::ProviderUsageInvalidFields {
cache_read_input_tokens: true,
..Default::default()
},
..first
};
vec![
Ok(LLMChunk::ProviderUsageReplace(first)),
Ok(first.legacy_chunk().unwrap()),
Ok(LLMChunk::Done),
Ok(LLMChunk::ProviderUsageReplace(invalid)),
Ok(invalid.legacy_chunk().unwrap()),
]
}
Script::ReplacementOverCapWhileOpen => {
let first = bamboo_llm::types::ProviderUsageReplacement {
input_tokens: Some(20),
output_tokens: Some(10),
total_tokens: Some(30),
..Default::default()
};
let over = bamboo_llm::types::ProviderUsageReplacement {
output_tokens: None,
output_floor: Some(41),
total_tokens: None,
..first
};
let prefix = futures::stream::iter(vec![
Ok(LLMChunk::ProviderUsageReplace(first)),
Ok(LLMChunk::ProviderUsageReplace(over)),
]);
return Ok(Box::pin(prefix.chain(futures::stream::pending())));
}
Script::ReplacementInputOverCapWhileOpen => {
let first = bamboo_llm::types::ProviderUsageReplacement {
input_tokens: Some(20),
output_tokens: Some(10),
total_tokens: Some(30),
..Default::default()
};
let over = bamboo_llm::types::ProviderUsageReplacement {
input_tokens: None,
input_floor: Some(101),
total_tokens: None,
..first
};
let prefix = futures::stream::iter(vec![
Ok(LLMChunk::ProviderUsageReplace(first)),
Ok(LLMChunk::ProviderUsageReplace(over)),
]);
return Ok(Box::pin(prefix.chain(futures::stream::pending())));
}
Script::KnownText => vec![
Ok(LLMChunk::Token("Fixture result".into())),
Ok(complete.clone()),
Ok(LLMChunk::Done),
Ok(LLMChunk::Token("FORBIDDEN_LATE_TEXT".into())),
Ok(LLMChunk::ReasoningToken("FORBIDDEN_LATE_REASONING".into())),
Ok(LLMChunk::ReasoningSignature(
"FORBIDDEN_LATE_SIGNATURE".into(),
)),
Ok(LLMChunk::ResponseId("FORBIDDEN_LATE_ID".into())),
Ok(LLMChunk::ResponsesEvent {
event_type: "late".into(),
data: Box::new(serde_json::json!({"late":true})),
}),
Ok(LLMChunk::ToolCalls(vec![])),
Ok(complete),
Ok(LLMChunk::Done),
],
Script::DoneError => vec![
Ok(complete),
Ok(LLMChunk::Done),
Err(LLMError::Stream("late failure".into())),
],
Script::DoneOverCap => vec![
Ok(complete),
Ok(LLMChunk::Done),
Ok(usage(Some(101), Some(10))),
],
Script::NoDone => vec![
Ok(LLMChunk::Token("Incomplete response".into())),
Ok(complete),
],
Script::DonePending | Script::RepeatedDone => {
let prefix = futures::stream::iter(vec![
Ok(LLMChunk::Token("Fixture result".into())),
Ok(complete),
Ok(LLMChunk::Done),
]);
return if matches!(self.script, Script::DonePending) {
Ok(Box::pin(prefix.chain(futures::stream::pending())))
} else {
Ok(Box::pin(prefix.chain(futures::stream::repeat_with(|| {
Ok(LLMChunk::Done)
}))))
};
}
Script::BootstrapError => {
return Err(LLMError::Api("controlled bootstrap error".into()))
}
Script::Pending => return std::future::pending().await,
Script::StreamError => vec![
Ok(usage(Some(20), Some(10))),
Err(LLMError::Stream("controlled stream error".into())),
],
Script::Partial => vec![Ok(LLMChunk::Done), Ok(usage(None, Some(10)))],
Script::OverCap => vec![Ok(usage(Some(101), Some(10)))],
Script::PartialOverCap => vec![Ok(usage(Some(101), None))],
Script::LegacySplitInputOverCap => vec![
Ok(LLMChunk::CacheUsage {
input_tokens: 80,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 10,
}),
Ok(LLMChunk::CacheUsage {
input_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 40,
}),
],
Script::CanonicalSplitCacheOverCap => {
let mut first = usage(None, None);
let mut second = usage(None, None);
if let LLMChunk::ProviderUsage {
cache_creation_input_tokens,
..
} = &mut first
{
*cache_creation_input_tokens = Some(80);
}
if let LLMChunk::ProviderUsage {
cache_read_input_tokens,
..
} = &mut second
{
*cache_read_input_tokens = Some(40);
}
vec![Ok(first), Ok(second)]
}
Script::SplitTotalMismatch => {
let mut total = usage(None, None);
if let LLMChunk::ProviderUsage { total_tokens, .. } = &mut total {
*total_tokens = Some(100);
}
vec![Ok(usage(Some(20), Some(10))), Ok(LLMChunk::Done), Ok(total)]
}
Script::CanonicalCacheOnlyOverCap
| Script::CanonicalReasoningOnlyOverCap
| Script::CanonicalTotalOnlyOverCap => {
let mut chunk = usage(None, None);
if let LLMChunk::ProviderUsage {
cache_read_input_tokens,
reasoning_tokens,
total_tokens,
..
} = &mut chunk
{
match self.script {
Script::CanonicalCacheOnlyOverCap => *cache_read_input_tokens = Some(101),
Script::CanonicalReasoningOnlyOverCap => *reasoning_tokens = Some(41),
Script::CanonicalTotalOnlyOverCap => *total_tokens = Some(141),
_ => unreachable!(),
}
}
vec![Ok(chunk)]
}
Script::LegacyCachedInputOverCap => vec![Ok(LLMChunk::CacheUsage {
input_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 101,
})],
Script::LegacyCompositeInputOverCap => vec![Ok(LLMChunk::CacheUsage {
input_tokens: 40,
cache_creation_input_tokens: 40,
cache_read_input_tokens: 40,
})],
Script::LegacyOverflowOverCap => vec![Ok(LLMChunk::CacheUsage {
input_tokens: u64::MAX,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 1,
})],
Script::LegacyOutputOverCap => vec![Ok(LLMChunk::UsageSummary {
output_tokens: 41,
thinking_tokens: 0,
})],
Script::LegacyThinkingOverCap => vec![Ok(LLMChunk::UsageSummary {
output_tokens: 0,
thinking_tokens: 41,
})],
Script::LegacyWithinCap => vec![
Ok(LLMChunk::CacheUsage {
input_tokens: 60,
cache_creation_input_tokens: 10,
cache_read_input_tokens: 30,
}),
Ok(LLMChunk::UsageSummary {
output_tokens: 40,
thinking_tokens: 5,
}),
Ok(LLMChunk::Done),
],
Script::MixedKnown => {
let mut canonical = usage(Some(20), Some(10));
if let LLMChunk::ProviderUsage {
cache_creation_input_tokens,
cache_read_input_tokens,
..
} = &mut canonical
{
*cache_creation_input_tokens = Some(10);
*cache_read_input_tokens = Some(5);
}
vec![
Ok(LLMChunk::CacheUsage {
input_tokens: 5,
cache_creation_input_tokens: 10,
cache_read_input_tokens: 5,
}),
Ok(LLMChunk::UsageSummary {
output_tokens: 10,
thinking_tokens: 5,
}),
Ok(canonical),
Ok(LLMChunk::Done),
]
}
Script::InconsistentOverCap => {
let mut chunk = usage(Some(101), Some(10));
if let LLMChunk::ProviderUsage { total_tokens, .. } = &mut chunk {
*total_tokens = Some(1);
}
vec![Ok(chunk)]
}
Script::Known => vec![
Ok(LLMChunk::Done),
Ok(usage(Some(20), Some(10))),
Ok(LLMChunk::Done),
],
};
Ok(Box::pin(futures::stream::iter(chunks)))
}
}
async fn admitted() -> (Fixture, crate::session_activation::SessionRunRegistration) {
let mut f = Fixture::with_project(true).await;
let registration = f
.runtime
.router
.register_run(&f.child.id, "budget-host-run")
.await
.unwrap();
admit_ticket_run(&f.service, &f.runtime, &mut f.child, "budget-host-run")
.await
.unwrap();
(f, registration)
}
fn configure(f: &Fixture, cap: u64) {
let user = Authority::from_verified_host(
f.spec.binding.clone(),
Principal::User {
user_id: "budget-owner".into(),
},
);
f.service
.configure_work_token_budget(
&user,
0,
WorkTokenLimits {
work_tokens: std::collections::BTreeMap::from([(f.spec.work_id.clone(), cap)]),
utc_day_tokens: cap,
total_tokens: cap,
},
)
.unwrap();
}
fn ledger(f: &Fixture) -> bamboo_tickets::WorkTokenBudget {
f.service
.work_token_budget(&Authority::from_verified_host(
f.spec.binding.clone(),
Principal::Runtime,
))
.unwrap()
.unwrap()
}
async fn provider(f: &Fixture, inner: Arc<Provider>) -> Arc<dyn LLMProvider> {
NativeWorkTokenBudget::capture(Some(f.service.clone()), f.storage.clone(), &f.child)
.await
.unwrap()
.unwrap()
.wrap(
inner,
&AgentLoopConfig {
app_data_dir: Some(f.dir.path().to_path_buf()),
..Default::default()
},
)
}
#[tokio::test]
async fn token_budget_decorator_forwards_known_fast_refusal_without_reservation() {
let routing = crate::app_context::model_routing_test_support::RoutingContext::new(None);
let mut config = routing.config.read().await.clone();
config.provider_instances.get_mut("fast-b").unwrap().enabled = false;
let (_, rejected) = crate::model_config_helper::resolve_fast_model_for_runner(
Some(&config),
"chat-a",
&routing.registry,
);
let (fixture, _registration) = admitted().await;
configure(&fixture, 140);
let budget = NativeWorkTokenBudget::capture(
Some(fixture.service.clone()),
fixture.storage.clone(),
&fixture.child,
)
.await
.unwrap()
.unwrap();
let decorated = budget.wrap(rejected.unwrap(), &AgentLoopConfig::default());
assert!(matches!(
decorated.known_pre_dispatch_error(),
Some(LLMError::Auth(_))
));
assert_eq!(ledger(&fixture).total_charged, 0);
assert!(routing.requests("fast-b").is_empty());
assert!(routing.requests("chat-a").is_empty());
}
#[tokio::test]
async fn token_budget_reserves_before_provider_and_settles_only_after_trailing_usage_eof() {
let (f, _registration) = admitted().await;
configure(&f, 140);
let inner = Provider::new(Script::Known);
let decorated = provider(&f, inner.clone()).await;
assert_eq!(
decorated.resource_allocation_identity(),
Some(Arc::as_ptr(&inner) as *const () as usize)
);
let mut stream = decorated
.chat_stream(&[], &[], None, "fixture-model")
.await
.unwrap();
assert_eq!(inner.calls.load(Ordering::SeqCst), 1);
assert_eq!(ledger(&f).total_charged, 140);
assert!(matches!(
stream.next().await.unwrap().unwrap(),
LLMChunk::ProviderUsage { .. }
));
assert_eq!(ledger(&f).total_charged, 140);
assert!(matches!(
stream.next().await.unwrap().unwrap(),
LLMChunk::Done
));
let budget = ledger(&f);
assert_eq!(budget.total_charged, 30);
let final_seq = f.service.published().unwrap().1.seq;
assert!(stream.next().await.is_none());
assert!(stream.next().await.is_none());
assert_eq!(f.service.published().unwrap().1.seq, final_seq);
let record = budget.reservations.values().next().unwrap();
assert_eq!(
record.request.attempt.as_ref().unwrap().assignment_id,
f.spec.assignment_id
);
assert_eq!(record.request.project_id.as_str(), "budget-fixture-project");
assert_eq!(
record.usage,
Some(ExecutionTokenUsage::Known {
input_tokens: 20,
output_tokens: 10
})
);
assert!(decorated
.chat_stream(&[], &[], None, "fixture-model")
.await
.is_err());
assert_eq!(
inner.calls.load(Ordering::SeqCst),
1,
"exhaustion blocks actual provider entry"
);
}
#[tokio::test]
async fn token_budget_missing_error_cancel_and_dropped_stream_keep_the_full_charge() {
for script in [
Script::Partial,
Script::BootstrapError,
Script::StreamError,
Script::Pending,
Script::Known,
] {
let (f, _registration) = admitted().await;
configure(&f, 140);
let inner = Provider::new(script);
let decorated = provider(&f, inner.clone()).await;
let result = tokio::time::timeout(
std::time::Duration::from_millis(100),
decorated.chat_stream(&[], &[], None, "fixture-model"),
)
.await;
if let Ok(Ok(mut stream)) = result {
if matches!(inner.script, Script::Known) {
drop(stream);
} else {
while let Some(chunk) = stream.next().await {
if chunk.is_err() {
break;
}
}
}
}
assert_eq!(inner.calls.load(Ordering::SeqCst), 1);
assert_eq!(ledger(&f).total_charged, 140);
assert!(decorated
.chat_stream(&[], &[], None, "fixture-model")
.await
.is_err());
assert_eq!(inner.calls.load(Ordering::SeqCst), 1);
}
}
#[tokio::test]
async fn token_budget_over_cap_report_halts_subsequent_provider_entry() {
for script in [
Script::OverCap,
Script::PartialOverCap,
Script::InconsistentOverCap,
Script::LegacyCachedInputOverCap,
Script::LegacyCompositeInputOverCap,
Script::LegacyOutputOverCap,
Script::LegacyThinkingOverCap,
Script::LegacyOverflowOverCap,
Script::CanonicalCacheOnlyOverCap,
Script::CanonicalReasoningOnlyOverCap,
Script::CanonicalTotalOnlyOverCap,
Script::LegacySplitInputOverCap,
Script::CanonicalSplitCacheOverCap,
] {
let (f, _registration) = admitted().await;
configure(&f, 500);
let inner = Provider::new(script);
let decorated = provider(&f, inner.clone()).await;
let mut stream = decorated
.chat_stream(&[], &[], None, "fixture-model")
.await
.unwrap();
let mut saw_error = false;
while let Some(chunk) = stream.next().await {
if chunk.is_err() {
saw_error = true;
break;
}
}
assert!(
saw_error,
"overrun must fail while consuming its real stream"
);
let budget = ledger(&f);
assert_eq!(budget.total_charged, 140);
assert!(budget
.reservations
.values()
.next()
.unwrap()
.over_cap_usage
.is_some());
assert!(decorated
.chat_stream(&[], &[], None, "fixture-model")
.await
.is_err());
assert_eq!(inner.calls.load(Ordering::SeqCst), 1);
}
}
#[tokio::test]
async fn token_budget_incomplete_or_inconsistent_usage_keeps_caps_without_double_counting() {
for (script, charged, known) in [
(Script::LegacyWithinCap, 140, false),
(Script::MixedKnown, 30, true),
(Script::SplitTotalMismatch, 140, false),
] {
let (f, _registration) = admitted().await;
configure(&f, 500);
let decorated = provider(&f, Provider::new(script)).await;
let mut stream = decorated
.chat_stream(&[], &[], None, "fixture-model")
.await
.unwrap();
while let Some(chunk) = stream.next().await {
chunk.unwrap();
}
let budget = ledger(&f);
assert_eq!(budget.total_charged, charged);
let receipt = budget.reservations.values().next().unwrap();
assert!(receipt.over_cap_usage.is_none());
assert_eq!(
matches!(receipt.usage, Some(ExecutionTokenUsage::Known { .. })),
known
);
}
}
#[tokio::test]
async fn token_budget_dynamic_auxiliary_and_vision_cannot_escape_shared_exhaustion() {
use crate::runtime::config::AuxiliaryModelConfig;
let (f, _registration) = admitted().await;
configure(&f, 140);
let inner = Provider::new(Script::Known);
let budget =
NativeWorkTokenBudget::capture(Some(f.service.clone()), f.storage.clone(), &f.child)
.await
.unwrap()
.unwrap();
let dynamic: Arc<dyn LLMProvider> = inner.clone();
let vision: Arc<dyn LLMProvider> = inner.clone();
let mut config = AgentLoopConfig {
app_data_dir: Some(f.dir.path().to_path_buf()),
fast_model_provider: Some(inner.clone()),
summarization_model_provider: Some(inner.clone()),
auxiliary_model_resolver: Some(Arc::new(move || AuxiliaryModelConfig {
background_model_provider: Some(dynamic.clone()),
..Default::default()
})),
vision_model_resolver: Some(Arc::new(move || {
let vision = vision.clone();
Box::pin(async move { Ok(bamboo_llm::ResolvedModel::new(vision, "fixture-model")) })
})),
..Default::default()
};
let main = budget.decorate(&mut config, inner.clone());
drop(
main.chat_stream(&[], &[], None, "fixture-model")
.await
.unwrap(),
);
let late_aux = config.auxiliary_model_resolver.as_ref().unwrap()();
let late_vision = config.vision_model_resolver.as_ref().unwrap()()
.await
.unwrap();
for provider in [
config.fast_model_provider.unwrap(),
config.summarization_model_provider.unwrap(),
late_aux.background_model_provider.unwrap(),
late_vision.provider,
] {
assert!(provider
.chat_stream(&[], &[], None, "fixture-model")
.await
.is_err());
}
assert_eq!(inner.calls.load(Ordering::SeqCst), 1);
assert_eq!(ledger(&f).total_charged, 140);
}
#[tokio::test]
async fn token_budget_background_only_override_cannot_escape_shared_exhaustion() {
let (f, _registration) = admitted().await;
configure(&f, 140);
let inner = Provider::new(Script::Known);
let budget =
NativeWorkTokenBudget::capture(Some(f.service.clone()), f.storage.clone(), &f.child)
.await
.unwrap()
.unwrap();
let dynamic: Arc<dyn LLMProvider> = inner.clone();
let mut config = AgentLoopConfig {
app_data_dir: Some(f.dir.path().to_path_buf()),
auxiliary_model_resolver: None,
background_model_resolver: Some(Arc::new(move |_| {
Some(bamboo_llm::ResolvedModel::new(
dynamic.clone(),
"fixture-model",
))
})),
..Default::default()
};
let main = budget.decorate(&mut config, inner.clone());
drop(
main.chat_stream(&[], &[], None, "fixture-model")
.await
.unwrap(),
);
let background = config.resolve_auxiliary_models();
assert_eq!(
background.background_model_name.as_deref(),
Some("fixture-model")
);
assert!(background
.background_model_provider
.unwrap()
.chat_stream(&[], &[], None, "fixture-model")
.await
.is_err());
assert_eq!(inner.calls.load(Ordering::SeqCst), 1);
assert_eq!(ledger(&f).total_charged, 140);
}
#[tokio::test]
async fn token_budget_unknown_model_and_revoked_root_block_before_provider() {
let (f, _registration) = admitted().await;
configure(&f, 500);
let inner = Provider::new(Script::Known);
let decorated = provider(&f, inner.clone()).await;
assert!(decorated
.chat_stream(&[], &[], None, "unknown-budget-model")
.await
.is_err());
assert!(decorated
.chat_stream(&[], &[], Some(41), "fixture-model")
.await
.is_err());
let foreign_storage = SessionStoreV2::new(f.dir.path().join("foreign-host"))
.await
.unwrap();
let bootstrap = foreign_storage
.get_or_create_default_supervisor("fixture-model")
.await
.unwrap();
let (mut foreign_binding, _) = verified_scope_binding(&foreign_storage, &bootstrap.session_id)
.await
.unwrap();
foreign_binding.binding_revision = 7;
assert_eq!(
foreign_binding.supervisor_session_id,
f.spec.binding.supervisor_session_id
);
assert_ne!(foreign_binding.scope_id, f.spec.binding.scope_id);
let foreign = Arc::new(
TicketService::open(f.dir.path().join("foreign-tickets"), foreign_binding).unwrap(),
);
match NativeWorkTokenBudget::capture(Some(foreign), f.storage.clone(), &f.child).await {
Err(Error::ScopeDenied(reason)) => assert_eq!(reason, "Supervisor scope binding changed"),
_ => panic!("a foreign historical scope cannot authorize Native capture"),
}
f.storage.delete_session(&f.supervisor.id).await.unwrap();
assert!(decorated
.chat_stream(&[], &[], None, "fixture-model")
.await
.is_err());
assert_eq!(inner.calls.load(Ordering::SeqCst), 0);
assert_eq!(ledger(&f).total_charged, 0);
}
#[tokio::test]
async fn token_budget_native_descendant_uses_original_work_and_rejects_lost_ancestry() {
let (f, _registration) = admitted().await;
configure(&f, 500);
let mut descendant = Session::new_child_of(
"native-grandchild",
&f.child,
"fixture-model",
"nested work",
);
f.storage.save_session(&descendant).await.unwrap();
let budget =
NativeWorkTokenBudget::capture(Some(f.service.clone()), f.storage.clone(), &descendant)
.await
.unwrap()
.unwrap();
let inner = Provider::new(Script::Known);
let decorated = budget.wrap(
inner.clone(),
&AgentLoopConfig {
app_data_dir: Some(f.dir.path().to_path_buf()),
..Default::default()
},
);
let mut stream = decorated
.chat_stream(&[], &[], None, "fixture-model")
.await
.unwrap();
while let Some(chunk) = stream.next().await {
chunk.unwrap();
}
assert_eq!(ledger(&f).total_charged, 30);
descendant.parent_session_id = Some(f.supervisor.id.clone());
assert!(f.storage.save_session(&descendant).await.is_err());
f.storage.delete_session(&descendant.id).await.unwrap();
assert!(decorated
.chat_stream(&[], &[], None, "fixture-model")
.await
.is_err());
assert_eq!(inner.calls.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn token_budget_unconfigured_ordinary_child_preserves_legacy_and_later_config_fails_closed() {
let f = Fixture::with_project(true).await;
let mut root = f
.storage
.load_root_authority(&f.supervisor.id)
.await
.unwrap()
.unwrap();
root.set_root_orchestration_only(false).unwrap();
f.storage.save_session(&root).await.unwrap();
let child = Session::new_child_of(
"ordinary-no-budget",
&root,
"fixture-model",
"ordinary child",
);
f.storage.save_session(&child).await.unwrap();
let capability =
NativeWorkTokenBudget::capture(Some(f.service.clone()), f.storage.clone(), &child)
.await
.unwrap()
.unwrap();
let inner = Provider::new(Script::Known);
let decorated = capability.wrap(
inner.clone(),
&AgentLoopConfig {
app_data_dir: Some(f.dir.path().to_path_buf()),
..Default::default()
},
);
let mut stream = decorated
.chat_stream(&[], &[], Some(40), "fixture-model")
.await
.unwrap();
while let Some(chunk) = stream.next().await {
chunk.unwrap();
}
assert!(f.service.published().unwrap().1.execution_budget.is_none());
assert_eq!(inner.calls.load(Ordering::SeqCst), 1);
configure(&f, 500);
assert!(decorated
.chat_stream(&[], &[], Some(40), "fixture-model")
.await
.is_err());
assert_eq!(
inner.calls.load(Ordering::SeqCst),
1,
"late configuration cannot escape the retained capability"
);
assert_eq!(ledger(&f).total_charged, 0);
}
#[tokio::test]
async fn work_budget_and_reasoning_snapshot_share_real_auxiliary_semaphore() {
use bamboo_domain::{ProviderModelRef, ReasoningEffort};
use bamboo_llm::{provider::queued_snapshot_for, ResolvedModel};
use std::time::Duration;
struct Snapshotting {
inner: Arc<dyn LLMProvider>,
}
#[async_trait::async_trait]
impl LLMProvider for Snapshotting {
fn queued_provider_snapshot(
&self,
) -> bamboo_llm::provider::Result<Option<Arc<dyn LLMProvider>>> {
Ok(Some(self.inner.clone()))
}
async fn list_model_info(&self) -> bamboo_llm::provider::Result<Vec<ProviderModelInfo>> {
self.inner.list_model_info().await
}
async fn chat_stream(
&self,
messages: &[bamboo_domain::Message],
tools: &[bamboo_domain::ToolSchema],
cap: Option<u32>,
model: &str,
) -> bamboo_llm::provider::Result<LLMStream> {
self.inner.chat_stream(messages, tools, cap, model).await
}
}
let (fixture, _registration) = admitted().await;
configure(&fixture, 140);
let budget = NativeWorkTokenBudget::capture(
Some(fixture.service.clone()),
fixture.storage.clone(),
&fixture.child,
)
.await
.unwrap()
.unwrap();
let calls = Provider::new(Script::Known);
let inner: Arc<dyn LLMProvider> = Arc::new(Snapshotting {
inner: calls.clone(),
});
let decorated = budget.wrap(
inner,
&AgentLoopConfig {
app_data_dir: Some(fixture.dir.path().to_path_buf()),
..Default::default()
},
);
let reference = ProviderModelRef::new("fixture", "fixture-model")
.with_reasoning_effort(ReasoningEffort::Low);
let original = ResolvedModel::from_ref(decorated, &reference).provider;
let snapshot = queued_snapshot_for(original.clone()).unwrap();
assert_eq!(
snapshot.resource_allocation_identity(),
original.resource_allocation_identity()
);
let permit =
crate::runtime::runner::auxiliary_budget::acquire(&original, "fixture-model", 1).await;
let mut waiting = Box::pin(crate::runtime::runner::auxiliary_budget::acquire(
&snapshot,
"fixture-model",
1,
));
assert!(
tokio::time::timeout(Duration::from_millis(100), waiting.as_mut())
.await
.is_err()
);
assert_eq!(calls.calls.load(Ordering::SeqCst), 0);
drop(permit);
let _permit = tokio::time::timeout(Duration::from_secs(2), waiting)
.await
.unwrap();
let mut stream = snapshot
.chat_stream(&[], &[], Some(40), "fixture-model")
.await
.unwrap();
while let Some(chunk) = stream.next().await {
chunk.unwrap();
}
assert_eq!(calls.calls.load(Ordering::SeqCst), 1);
assert_eq!(
ledger(&fixture).total_charged,
30,
"the simulated usage settles at EOF"
);
}
#[tokio::test]
async fn token_budget_resume_auxiliary_pair_cannot_escape_shared_exhaustion() {
use std::sync::atomic::AtomicBool;
let (f, _registration) = admitted().await;
configure(&f, 140);
let main_inner = Provider::new(Script::Known);
let summary_inner = Provider::new(Script::Known);
let background_inner = Provider::new(Script::Known);
let budget =
NativeWorkTokenBudget::capture(Some(f.service.clone()), f.storage.clone(), &f.child)
.await
.unwrap()
.unwrap();
let want_summary = Arc::new(AtomicBool::new(true));
let callback_want_summary = Arc::clone(&want_summary);
let summary: Arc<dyn LLMProvider> = summary_inner.clone();
let background: Arc<dyn LLMProvider> = background_inner.clone();
let mut config = AgentLoopConfig {
app_data_dir: Some(f.dir.path().to_path_buf()),
resume_auxiliary_resolver: Some(Arc::new(move || {
let summary = callback_want_summary
.load(Ordering::SeqCst)
.then(|| bamboo_llm::ResolvedModel::new(summary.clone(), "fixture-model"));
(
Ok(summary),
Some(bamboo_llm::ResolvedModel::new(
background.clone(),
"fixture-model",
)),
)
})),
..Default::default()
};
let main = budget.decorate(&mut config, main_inner.clone());
drop(
main.chat_stream(&[], &[], None, "fixture-model")
.await
.unwrap(),
);
let before_attempts = ledger(&f);
let pair = config.resolve_auxiliary_models();
assert_eq!(
pair.summarization_model_name.as_deref(),
Some("fixture-model")
);
assert_eq!(pair.background_model_name.as_deref(), Some("fixture-model"));
assert!(pair.summarization_model_error.is_none());
for provider in [
pair.summarization_model_provider.unwrap(),
pair.background_model_provider.unwrap(),
] {
assert!(provider
.chat_stream(&[], &[], None, "fixture-model")
.await
.is_err());
}
want_summary.store(false, Ordering::SeqCst);
let absent = config.resolve_auxiliary_models();
assert!(absent.summarization_model_name.is_none());
assert!(absent.summarization_model_provider.is_none());
assert!(absent.summarization_model_error.is_none());
assert!(absent
.background_model_provider
.unwrap()
.chat_stream(&[], &[], None, "fixture-model")
.await
.is_err());
assert_eq!(main_inner.calls.load(Ordering::SeqCst), 1);
assert_eq!(summary_inner.calls.load(Ordering::SeqCst), 0);
assert_eq!(background_inner.calls.load(Ordering::SeqCst), 0);
let after_attempts = ledger(&f);
assert_eq!(after_attempts.total_charged, 140);
assert_eq!(
after_attempts.reservations.len(),
before_attempts.reservations.len()
);
}
#[tokio::test]
async fn token_budget_resume_auxiliary_rejection_has_no_reservation() {
let (f, _registration) = admitted().await;
configure(&f, 140);
let main_inner = Provider::new(Script::Known);
let background_inner = Provider::new(Script::Known);
let budget =
NativeWorkTokenBudget::capture(Some(f.service.clone()), f.storage.clone(), &f.child)
.await
.unwrap()
.unwrap();
let background: Arc<dyn LLMProvider> = background_inner.clone();
let mut config = AgentLoopConfig {
app_data_dir: Some(f.dir.path().to_path_buf()),
resume_auxiliary_resolver: Some(Arc::new(move || {
(
Err(LLMError::Api("synthetic resume route rejected".into())),
Some(bamboo_llm::ResolvedModel::new(
background.clone(),
"fixture-model",
)),
)
})),
..Default::default()
};
let _main = budget.decorate(&mut config, main_inner.clone());
let before_selection = ledger(&f);
let selected = config.resolve_auxiliary_models();
assert!(selected.summarization_model_name.is_none());
assert!(selected.summarization_model_provider.is_none());
assert!(selected
.summarization_model_error
.as_deref()
.is_some_and(|error| error.contains("synthetic resume route rejected")));
assert!(selected.background_model_provider.is_some());
assert_eq!(main_inner.calls.load(Ordering::SeqCst), 0);
assert_eq!(background_inner.calls.load(Ordering::SeqCst), 0);
let after_selection = ledger(&f);
assert_eq!(after_selection.total_charged, 0);
assert_eq!(
after_selection.reservations.len(),
before_selection.reservations.len()
);
}
struct FixedGoldPolicyWrapper {
inner: Arc<dyn LLMProvider>,
calls: Arc<AtomicUsize>,
}
#[async_trait::async_trait]
impl LLMProvider for FixedGoldPolicyWrapper {
fn fixed_gold_policy_delegate(&self) -> Option<Arc<dyn LLMProvider>> {
Some(self.inner.clone())
}
fn resource_allocation_identity(&self) -> Option<usize> {
self.inner.resource_allocation_identity()
}
async fn chat_stream(
&self,
messages: &[bamboo_domain::Message],
tools: &[bamboo_domain::ToolSchema],
max: Option<u32>,
model: &str,
) -> bamboo_llm::provider::Result<LLMStream> {
self.calls.fetch_add(1, Ordering::SeqCst);
self.inner.chat_stream(messages, tools, max, model).await
}
}
#[tokio::test]
async fn child_resume_gold_delegate_keeps_work_budget_and_selected_provider() {
let (fixture, _registration) = admitted().await;
configure(&fixture, 140);
let budget = NativeWorkTokenBudget::capture(
Some(fixture.service.clone()),
fixture.storage.clone(),
&fixture.child,
)
.await
.unwrap()
.unwrap();
let selected = Provider::new(Script::Known);
let other = Provider::new(Script::Known);
let main: Arc<dyn LLMProvider> = selected.clone();
let other_handle: Arc<dyn LLMProvider> = other.clone();
let mut host_config = bamboo_llm::Config::default();
host_config.provider = "relay".into();
host_config.default_provider_instance = Some("relay".into());
host_config.http_proxy = "http://127.0.0.1:9".into();
host_config.https_proxy.clear();
host_config.proxy_auth = None;
host_config.provider_instances.insert(
"relay".into(),
serde_json::from_value(serde_json::json!({
"provider_type":"openai", "enabled":true, "api_key":"synthetic-only",
"base_url":"http://127.0.0.1:9", "model":"fixture-model",
"runtime_models":["fixture-model"]
}))
.unwrap(),
);
let registry = Arc::new(
bamboo_llm::ProviderRegistry::from_config(&host_config, fixture.dir.path().to_path_buf())
.await
.unwrap(),
);
registry.insert("relay".into(), main.clone());
let (expected, origin) = registry.get_with_origin("relay").unwrap();
let origin = origin.expect("missing binding still has managed origin");
let live = Arc::new(tokio::sync::RwLock::new(host_config));
let resolver: crate::runtime::config::GoldTargetResolver = Arc::new(move |model, delegate| {
bamboo_llm::gold_admission::guarded_gold_target_with_origin(
live.clone(),
origin.clone(),
bamboo_domain::ProviderModelRef::new("relay", model),
expected.clone(),
delegate,
)
});
let mut denied_config = AgentLoopConfig {
app_data_dir: Some(fixture.dir.path().to_path_buf()),
gold_target_resolver: Some(resolver),
..Default::default()
};
let denied_decorated = budget.decorate(&mut denied_config, main.clone());
crate::runtime::runtime::bind_child_gold_delegate(
&mut denied_config,
Some(&main),
&main,
&denied_decorated,
)
.unwrap();
let denied =
denied_config.gold_target_resolver.as_ref().unwrap()("fixture-model", other_handle.clone());
assert!(denied
.chat_stream(&[], &[], None, "fixture-model")
.await
.is_err());
assert!(denied.local_dispatch_refused());
assert_eq!(selected.calls.load(Ordering::SeqCst), 0);
assert_eq!(other.calls.load(Ordering::SeqCst), 0);
assert_eq!(ledger(&fixture).total_charged, 0);
assert!(ledger(&fixture).reservations.is_empty());
let replacement = other_handle.clone();
let inherited: crate::runtime::config::GoldTargetResolver =
Arc::new(move |_, _| replacement.clone());
let managed: crate::runtime::config::GoldTargetResolver = Arc::new(|_, delegate| delegate);
let mut denied_policy = AgentLoopConfig {
app_data_dir: Some(fixture.dir.path().to_path_buf()),
gold_target_resolver: Some(crate::runtime::agent::compose_child_gold_policy(
Some(inherited),
managed,
)),
..Default::default()
};
let policy_decorated = budget.decorate(&mut denied_policy, main.clone());
crate::runtime::runtime::bind_child_gold_delegate(
&mut denied_policy,
Some(&main),
&main,
&policy_decorated,
)
.unwrap();
let refused =
denied_policy.gold_target_resolver.as_ref().unwrap()("fixture-model", other_handle.clone());
assert!(refused.known_pre_dispatch_error().is_some());
assert!(refused
.chat_stream(&[], &[], None, "fixture-model")
.await
.is_err());
assert_eq!(selected.calls.load(Ordering::SeqCst), 0);
assert_eq!(other.calls.load(Ordering::SeqCst), 0);
assert_eq!(ledger(&fixture).total_charged, 0);
assert!(ledger(&fixture).reservations.is_empty());
let permitted = Arc::new(std::sync::atomic::AtomicBool::new(true));
let policy_state = permitted.clone();
let replacement = other_handle.clone();
let changing: crate::runtime::config::GoldTargetResolver = Arc::new(move |_, delegate| {
if policy_state.load(Ordering::SeqCst) {
delegate
} else {
replacement.clone()
}
});
let mut changing_config = AgentLoopConfig {
app_data_dir: Some(fixture.dir.path().to_path_buf()),
gold_target_resolver: Some(crate::runtime::agent::compose_child_gold_policy(
Some(changing),
Arc::new(|_, delegate| delegate),
)),
..Default::default()
};
let changing_delegate = budget.decorate(&mut changing_config, main.clone());
crate::runtime::runtime::bind_child_gold_delegate(
&mut changing_config,
Some(&main),
&main,
&changing_delegate,
)
.unwrap();
let changing_gold =
changing_config.gold_target_resolver.as_ref().unwrap()("fixture-model", main.clone());
assert!(changing_gold.known_pre_dispatch_error().is_none());
permitted.store(false, Ordering::SeqCst);
assert!(changing_gold
.chat_stream(&[], &[], None, "fixture-model")
.await
.is_err());
assert!(changing_gold.local_dispatch_refused());
assert_eq!(selected.calls.load(Ordering::SeqCst), 0);
assert_eq!(other.calls.load(Ordering::SeqCst), 0);
assert_eq!(ledger(&fixture).total_charged, 0);
assert!(ledger(&fixture).reservations.is_empty());
let wrapper_calls = Arc::new(AtomicUsize::new(0));
let policy_calls = wrapper_calls.clone();
let policy: crate::runtime::config::GoldTargetResolver = Arc::new(move |_, delegate| {
Arc::new(FixedGoldPolicyWrapper {
inner: delegate,
calls: policy_calls.clone(),
})
});
let mut config = AgentLoopConfig {
app_data_dir: Some(fixture.dir.path().to_path_buf()),
gold_target_resolver: Some(crate::runtime::agent::compose_child_gold_policy(
Some(policy),
Arc::new(|_model, delegate| delegate),
)),
..Default::default()
};
let decorated = budget.decorate(&mut config, main.clone());
crate::runtime::runtime::bind_child_gold_delegate(&mut config, Some(&main), &main, &decorated)
.unwrap();
let gold = config.gold_target_resolver.as_ref().unwrap()("fixture-model", other_handle);
let mut stream = gold
.chat_stream(&[], &[], None, "fixture-model")
.await
.unwrap();
assert_eq!(selected.calls.load(Ordering::SeqCst), 1);
assert_eq!(other.calls.load(Ordering::SeqCst), 0);
assert_eq!(wrapper_calls.load(Ordering::SeqCst), 1);
assert_eq!(ledger(&fixture).total_charged, 140);
while let Some(chunk) = stream.next().await {
chunk.unwrap();
}
assert_eq!(ledger(&fixture).total_charged, 30);
}