1use crate::value::VmDictExt;
2use rust_decimal::Decimal;
3use std::cell::RefCell;
4use std::collections::BTreeMap;
5use std::str::FromStr;
6
7use crate::value::{categorized_error, ErrorCategory, VmError, VmValue};
8use crate::vm::{Vm, VmBuiltinArity, VmBuiltinMetadata};
9
10thread_local! {
11 static LLM_BUDGET: RefCell<Option<f64>> = const { RefCell::new(None) };
12 static LLM_ACCUMULATED_COST: RefCell<f64> = const { RefCell::new(0.0) };
13 static LLM_TOKEN_BUDGET: RefCell<Option<u64>> = const { RefCell::new(None) };
14 static LLM_ACCUMULATED_TOKENS: RefCell<u64> = const { RefCell::new(0) };
15}
16
17pub(crate) fn reset_cost_state() {
19 LLM_BUDGET.with(|b| *b.borrow_mut() = None);
20 LLM_ACCUMULATED_COST.with(|a| *a.borrow_mut() = 0.0);
21 LLM_TOKEN_BUDGET.with(|b| *b.borrow_mut() = None);
22 LLM_ACCUMULATED_TOKENS.with(|a| *a.borrow_mut() = 0);
23}
24
25pub fn peek_total_cost() -> f64 {
26 LLM_ACCUMULATED_COST.with(|acc| *acc.borrow())
27}
28
29#[must_use = "dropping the guard immediately restores the prior LLM cost budget"]
35pub struct LlmBudgetGuard {
36 previous_budget: Option<f64>,
37 previous_accumulated: f64,
38}
39
40impl Drop for LlmBudgetGuard {
41 fn drop(&mut self) {
42 LLM_BUDGET.with(|b| *b.borrow_mut() = self.previous_budget);
43 LLM_ACCUMULATED_COST.with(|a| *a.borrow_mut() = self.previous_accumulated);
44 }
45}
46
47pub fn install_llm_cost_budget(max_cost_usd: f64) -> LlmBudgetGuard {
53 let previous_budget = LLM_BUDGET.with(|b| b.borrow().to_owned());
54 let previous_accumulated = LLM_ACCUMULATED_COST.with(|a| *a.borrow());
55 LLM_BUDGET.with(|b| *b.borrow_mut() = Some(max_cost_usd.max(0.0)));
56 LLM_ACCUMULATED_COST.with(|a| *a.borrow_mut() = 0.0);
57 LlmBudgetGuard {
58 previous_budget,
59 previous_accumulated,
60 }
61}
62
63#[must_use = "dropping the guard immediately restores the prior LLM token budget"]
67pub struct LlmTokenBudgetGuard {
68 previous_budget: Option<u64>,
69 previous_accumulated: u64,
70}
71
72impl Drop for LlmTokenBudgetGuard {
73 fn drop(&mut self) {
74 LLM_TOKEN_BUDGET.with(|b| *b.borrow_mut() = self.previous_budget);
75 LLM_ACCUMULATED_TOKENS.with(|a| *a.borrow_mut() = self.previous_accumulated);
76 }
77}
78
79pub fn install_llm_token_budget(max_tokens: u64) -> LlmTokenBudgetGuard {
85 let previous_budget = LLM_TOKEN_BUDGET.with(|b| *b.borrow());
86 let previous_accumulated = LLM_ACCUMULATED_TOKENS.with(|a| *a.borrow());
87 LLM_TOKEN_BUDGET.with(|b| *b.borrow_mut() = Some(max_tokens));
88 LLM_ACCUMULATED_TOKENS.with(|a| *a.borrow_mut() = 0);
89 LlmTokenBudgetGuard {
90 previous_budget,
91 previous_accumulated,
92 }
93}
94
95pub fn peek_total_tokens() -> u64 {
96 LLM_ACCUMULATED_TOKENS.with(|acc| *acc.borrow())
97}
98
99pub fn set_llm_cost_budget(max_cost_usd: Option<f64>) {
111 LLM_BUDGET.with(|b| *b.borrow_mut() = max_cost_usd.map(|max| max.max(0.0)));
112}
113
114pub fn set_llm_token_budget(max_tokens: Option<u64>) {
118 LLM_TOKEN_BUDGET.with(|b| *b.borrow_mut() = max_tokens);
119}
120
121pub fn peek_llm_cost_budget() -> Option<f64> {
125 LLM_BUDGET.with(|b| *b.borrow())
126}
127
128pub fn peek_llm_token_budget() -> Option<u64> {
131 LLM_TOKEN_BUDGET.with(|b| *b.borrow())
132}
133
134#[derive(Clone, Debug, Default, PartialEq, serde::Serialize)]
139pub(crate) struct LlmBudgetEnvelope {
140 pub max_cost_usd: Option<f64>,
141 pub total_budget_usd: Option<f64>,
142 pub max_input_tokens: Option<i64>,
143 pub max_output_tokens: Option<i64>,
144}
145
146impl LlmBudgetEnvelope {
147 pub(crate) fn is_empty(&self) -> bool {
148 self.max_cost_usd.is_none()
149 && self.total_budget_usd.is_none()
150 && self.max_input_tokens.is_none()
151 && self.max_output_tokens.is_none()
152 }
153}
154
155#[derive(Clone, Debug)]
156pub(crate) struct LlmBudgetProjection {
157 pub provider: String,
158 pub model: String,
159 pub projected_input_tokens: i64,
160 pub projected_output_tokens: i64,
161 pub projected_cost_usd: f64,
162 pub session_cost_usd: f64,
163}
164
165#[derive(Clone, Copy, Debug, PartialEq, Eq)]
166pub(crate) enum BudgetLimitKind {
167 PerCallCost,
168 TotalCost,
169 InputTokens,
170 OutputTokens,
171}
172
173impl BudgetLimitKind {
174 fn as_str(self) -> &'static str {
175 match self {
176 BudgetLimitKind::PerCallCost => "max_cost_usd",
177 BudgetLimitKind::TotalCost => "total_budget_usd",
178 BudgetLimitKind::InputTokens => "max_input_tokens",
179 BudgetLimitKind::OutputTokens => "max_output_tokens",
180 }
181 }
182}
183
184fn numeric_value(value: &VmValue, key: &str) -> Result<f64, VmError> {
185 let value = match value {
186 VmValue::Float(f) => *f,
187 VmValue::Int(n) => *n as f64,
188 _ => {
189 return Err(VmError::Thrown(VmValue::String(arcstr::ArcStr::from(
190 format!("budget.{key}: expected a non-negative number"),
191 ))));
192 }
193 };
194 if !value.is_finite() || value < 0.0 {
195 return Err(VmError::Thrown(VmValue::String(arcstr::ArcStr::from(
196 format!("budget.{key}: expected a non-negative finite number"),
197 ))));
198 }
199 Ok(value)
200}
201
202fn integer_value(value: &VmValue, key: &str) -> Result<i64, VmError> {
203 let value = match value {
204 VmValue::Int(n) => *n,
205 VmValue::Float(f) if f.is_finite() && f.fract() == 0.0 => *f as i64,
206 _ => {
207 return Err(VmError::Thrown(VmValue::String(arcstr::ArcStr::from(
208 format!("budget.{key}: expected a non-negative integer"),
209 ))));
210 }
211 };
212 if value < 0 {
213 return Err(VmError::Thrown(VmValue::String(arcstr::ArcStr::from(
214 format!("budget.{key}: expected a non-negative integer"),
215 ))));
216 }
217 Ok(value)
218}
219
220fn parse_budget_fields(
221 fields: &crate::value::DictMap,
222 envelope: &mut LlmBudgetEnvelope,
223) -> Result<(), VmError> {
224 if let Some(value) = fields.get("max_cost_usd") {
225 envelope.max_cost_usd = Some(numeric_value(value, "max_cost_usd")?);
226 }
227 if let Some(value) = fields.get("total_budget_usd") {
228 envelope.total_budget_usd = Some(numeric_value(value, "total_budget_usd")?);
229 }
230 if let Some(value) = fields.get("max_input_tokens") {
231 envelope.max_input_tokens = Some(integer_value(value, "max_input_tokens")?);
232 }
233 if let Some(value) = fields.get("max_output_tokens") {
234 envelope.max_output_tokens = Some(integer_value(value, "max_output_tokens")?);
235 }
236 Ok(())
237}
238
239pub(crate) fn parse_budget_envelope(
240 options: Option<&crate::value::DictMap>,
241) -> Result<Option<LlmBudgetEnvelope>, VmError> {
242 let Some(options) = options else {
243 return Ok(None);
244 };
245 let mut envelope = LlmBudgetEnvelope::default();
246 if let Some(value) = options.get("budget") {
247 match value {
248 VmValue::Nil => {}
249 VmValue::Dict(fields) => parse_budget_fields(fields, &mut envelope)?,
250 _ => {
251 return Err(VmError::Thrown(VmValue::String(arcstr::ArcStr::from(
252 "budget: expected a dict {max_cost_usd?, total_budget_usd?, max_input_tokens?, max_output_tokens?}",
253 ))));
254 }
255 }
256 }
257 parse_budget_fields(options, &mut envelope)?;
258 Ok((!envelope.is_empty()).then_some(envelope))
259}
260
261fn estimate_json_tokens(value: &serde_json::Value, model: &str) -> i64 {
262 match value {
263 serde_json::Value::Null | serde_json::Value::Bool(_) | serde_json::Value::Number(_) => 1,
264 serde_json::Value::String(s) => estimate_text_tokens_for_model(s, model),
265 serde_json::Value::Array(items) => items
266 .iter()
267 .map(|item| estimate_json_tokens(item, model))
268 .sum(),
269 serde_json::Value::Object(map) => map
270 .iter()
271 .map(|(key, value)| {
272 estimate_text_tokens_for_model(key, model) + estimate_json_tokens(value, model)
273 })
274 .sum(),
275 }
276}
277
278fn estimate_text_tokens_for_model(text: &str, model: &str) -> i64 {
279 super::token_count::estimate_text_tokens(text, Some(model)).tokens
280}
281
282pub(crate) fn project_llm_call_tokens(opts: &super::api::LlmCallOptions) -> (i64, i64) {
283 let system_tokens = opts
284 .system
285 .as_deref()
286 .map(|system| estimate_text_tokens_for_model(system, &opts.model))
287 .unwrap_or(0);
288 let message_tokens: i64 = opts
289 .messages
290 .iter()
291 .map(|message| estimate_json_tokens(message, &opts.model))
292 .sum();
293 let tool_tokens: i64 = opts
294 .native_tools
295 .as_ref()
296 .map(|tools| {
297 tools
298 .iter()
299 .map(|tool| {
300 estimate_text_tokens_for_model(
301 &serde_json::to_string(tool).unwrap_or_default(),
302 &opts.model,
303 )
304 })
305 .sum()
306 })
307 .unwrap_or(0);
308 let projected_input_tokens = system_tokens
309 .saturating_add(message_tokens)
310 .saturating_add(tool_tokens);
311 let projected_output_tokens = opts.max_tokens.max(0);
312 (projected_input_tokens, projected_output_tokens)
313}
314
315pub(crate) fn project_llm_call_context_tokens(opts: &super::api::LlmCallOptions) -> u64 {
316 let (input, output) = project_llm_call_tokens(opts);
317 input.max(0) as u64 + output.max(0) as u64
318}
319
320pub(crate) fn project_llm_call_cost(
321 opts: &super::api::LlmCallOptions,
322 session_cost_usd: f64,
323) -> LlmBudgetProjection {
324 let (projected_input_tokens, projected_output_tokens) = project_llm_call_tokens(opts);
325 let projected_cost_usd = calculate_cost_for_provider(
326 &opts.provider,
327 &opts.model,
328 projected_input_tokens,
329 projected_output_tokens,
330 );
331 LlmBudgetProjection {
332 provider: opts.provider.clone(),
333 model: opts.model.clone(),
334 projected_input_tokens,
335 projected_output_tokens,
336 projected_cost_usd,
337 session_cost_usd,
338 }
339}
340
341pub(crate) fn budget_exceeded_error(
342 projection: &LlmBudgetProjection,
343 limit_kind: BudgetLimitKind,
344 limit_value: f64,
345) -> VmError {
346 let mut dict = BTreeMap::new();
347 dict.put_str("category", "budget_exceeded");
348 dict.put_str("kind", "terminal");
349 dict.put_str("reason", "budget_exceeded");
350 dict.put_str("limit", limit_kind.as_str());
351 dict.insert("limit_value".to_string(), VmValue::Float(limit_value));
352 dict.insert(
353 "projected_cost_usd".to_string(),
354 VmValue::Float(projection.projected_cost_usd),
355 );
356 dict.insert(
357 "session_cost_usd".to_string(),
358 VmValue::Float(projection.session_cost_usd),
359 );
360 dict.insert(
361 "projected_input_tokens".to_string(),
362 VmValue::Int(projection.projected_input_tokens),
363 );
364 dict.insert(
365 "projected_output_tokens".to_string(),
366 VmValue::Int(projection.projected_output_tokens),
367 );
368 dict.put_str("provider", projection.provider.clone());
369 dict.put_str("model", projection.model.clone());
370 dict.put_str(
371 "message",
372 format!(
373 "LLM budget exceeded before provider call: {} would exceed {}",
374 match limit_kind {
375 BudgetLimitKind::PerCallCost =>
376 format!("projected cost ${:.6}", projection.projected_cost_usd),
377 BudgetLimitKind::TotalCost => format!(
378 "projected session cost ${:.6}",
379 projection.session_cost_usd + projection.projected_cost_usd
380 ),
381 BudgetLimitKind::InputTokens => format!(
382 "projected input tokens {}",
383 projection.projected_input_tokens
384 ),
385 BudgetLimitKind::OutputTokens => format!(
386 "projected output tokens {}",
387 projection.projected_output_tokens
388 ),
389 },
390 limit_kind.as_str(),
391 ),
392 );
393 VmError::Thrown(VmValue::dict(dict))
394}
395
396pub(crate) fn budget_exceeded_limit(
397 envelope: &LlmBudgetEnvelope,
398 projection: &LlmBudgetProjection,
399) -> Option<(BudgetLimitKind, f64)> {
400 if let Some(max) = envelope.max_input_tokens {
401 if projection.projected_input_tokens > max {
402 return Some((BudgetLimitKind::InputTokens, max as f64));
403 }
404 }
405 if let Some(max) = envelope.max_output_tokens {
406 if projection.projected_output_tokens > max {
407 return Some((BudgetLimitKind::OutputTokens, max as f64));
408 }
409 }
410 if let Some(max) = envelope.max_cost_usd {
411 if projection.projected_cost_usd > max {
412 return Some((BudgetLimitKind::PerCallCost, max));
413 }
414 }
415 if let Some(max) = envelope.total_budget_usd {
416 if projection.session_cost_usd + projection.projected_cost_usd > max {
417 return Some((BudgetLimitKind::TotalCost, max));
418 }
419 }
420 None
421}
422
423pub(crate) fn check_budget_envelope(
424 envelope: &LlmBudgetEnvelope,
425 projection: &LlmBudgetProjection,
426) -> Result<(), VmError> {
427 if let Some((kind, limit)) = budget_exceeded_limit(envelope, projection) {
428 return Err(budget_exceeded_error(projection, kind, limit));
429 }
430 Ok(())
431}
432
433pub(crate) fn check_llm_preflight_budget(
434 opts: &super::api::LlmCallOptions,
435) -> Result<LlmBudgetProjection, VmError> {
436 let session_cost_usd = peek_total_cost();
437 let projection = project_llm_call_cost(opts, session_cost_usd);
438 if let Some(envelope) = opts.budget.as_ref() {
439 check_budget_envelope(envelope, &projection)?;
440 }
441 LLM_BUDGET.with(|budget| {
442 if let Some(max) = *budget.borrow() {
443 if session_cost_usd + projection.projected_cost_usd > max {
444 return Err(budget_exceeded_error(
445 &projection,
446 BudgetLimitKind::TotalCost,
447 max,
448 ));
449 }
450 }
451 Ok(())
452 })?;
453 Ok(projection)
454}
455
456#[derive(Clone, Copy, Debug, PartialEq)]
460pub(crate) struct PricingDetail {
461 pub input_per_1k: f64,
462 pub output_per_1k: f64,
463 pub cache_read_per_1k: Option<f64>,
464 pub cache_write_per_1k: Option<f64>,
465 pub source: PricingSource,
466}
467
468#[derive(Clone, Copy, Debug, PartialEq, Eq)]
469pub(crate) enum PricingSource {
470 CatalogModel,
472 CatalogServingTier,
475 ProviderEconomics,
477}
478
479impl PricingSource {
480 pub(crate) fn as_str(self) -> &'static str {
481 match self {
482 PricingSource::CatalogModel => "catalog_model",
483 PricingSource::CatalogServingTier => "catalog_serving_tier",
484 PricingSource::ProviderEconomics => "provider_economics",
485 }
486 }
487}
488
489pub(crate) fn pricing_detail_for(provider: &str, model: &str) -> Option<PricingDetail> {
494 if let Some(pricing) = crate::llm_config::model_pricing_per_mtok(model) {
495 return Some(PricingDetail {
496 input_per_1k: pricing.input_per_mtok / 1000.0,
497 output_per_1k: pricing.output_per_mtok / 1000.0,
498 cache_read_per_1k: pricing.cache_read_per_mtok.map(|rate| rate / 1000.0),
499 cache_write_per_1k: pricing.cache_write_per_mtok.map(|rate| rate / 1000.0),
500 source: PricingSource::CatalogModel,
501 });
502 }
503 let (input, output, _) = crate::llm_config::provider_economics(provider);
504 match (input, output) {
505 (Some(input_per_1k), Some(output_per_1k)) => Some(PricingDetail {
506 input_per_1k,
507 output_per_1k,
508 cache_read_per_1k: None,
509 cache_write_per_1k: None,
510 source: PricingSource::ProviderEconomics,
511 }),
512 _ => None,
513 }
514}
515
516pub(crate) fn pricing_per_1k_for(provider: &str, model: &str) -> Option<(f64, f64)> {
517 pricing_detail_for(provider, model).map(|p| (p.input_per_1k, p.output_per_1k))
518}
519
520pub(crate) fn pricing_detail_for_tier(
526 provider: &str,
527 model: &str,
528 served_fast: bool,
529) -> Option<PricingDetail> {
530 if served_fast {
531 if let Some(pricing) = crate::llm_config::model_serving_tier_pricing_per_mtok(
532 model,
533 crate::llm::serving_tiers::FAST_TIER_ID,
534 ) {
535 return Some(PricingDetail {
536 input_per_1k: pricing.input_per_mtok / 1000.0,
537 output_per_1k: pricing.output_per_mtok / 1000.0,
538 cache_read_per_1k: pricing.cache_read_per_mtok.map(|rate| rate / 1000.0),
539 cache_write_per_1k: pricing.cache_write_per_mtok.map(|rate| rate / 1000.0),
540 source: PricingSource::CatalogServingTier,
541 });
542 }
543 }
544 pricing_detail_for(provider, model)
545}
546
547pub(crate) fn latency_p50_ms_for(provider: &str) -> Option<u64> {
548 let (_, _, latency) = crate::llm_config::provider_economics(provider);
549 latency
550}
551
552fn authored_rate_decimal(rate: f64) -> Decimal {
565 Decimal::from_str(&format!("{rate}"))
566 .ok()
567 .or_else(|| Decimal::from_f64_retain(rate))
568 .unwrap_or(Decimal::ZERO)
569}
570
571pub fn calculate_cost_decimal(model: &str, input_tokens: i64, output_tokens: i64) -> Decimal {
579 let Some(pricing) = crate::llm_config::model_pricing_per_mtok(model) else {
580 return Decimal::ZERO;
581 };
582 let gross = Decimal::from(input_tokens) * authored_rate_decimal(pricing.input_per_mtok)
583 + Decimal::from(output_tokens) * authored_rate_decimal(pricing.output_per_mtok);
584 gross / Decimal::from(1_000_000i64)
585}
586
587pub fn calculate_cost_for_provider(
591 provider: &str,
592 model: &str,
593 input_tokens: i64,
594 output_tokens: i64,
595) -> f64 {
596 let Some(detail) = pricing_detail_for(provider, model) else {
597 return 0.0;
598 };
599 (input_tokens as f64 * detail.input_per_1k + output_tokens as f64 * detail.output_per_1k)
600 / 1000.0
601}
602
603pub fn pricing_aware_call_cost(
611 provider: &str,
612 model: &str,
613 input_tokens: i64,
614 output_tokens: i64,
615) -> Option<f64> {
616 let detail = pricing_detail_for(provider, model)?;
617 Some(
618 (input_tokens as f64 * detail.input_per_1k + output_tokens as f64 * detail.output_per_1k)
619 / 1000.0,
620 )
621}
622
623pub(crate) fn cache_hit_ratio(
624 input_tokens: i64,
625 cache_read_tokens: i64,
626 cache_write_tokens: i64,
627) -> f64 {
628 let input_tokens = input_tokens.max(0);
629 let cache_read_tokens = cache_read_tokens.max(0);
630 let cache_write_tokens = cache_write_tokens.max(0);
631 let reported_cache_tokens = cache_read_tokens.saturating_add(cache_write_tokens);
632 let total_prompt_tokens = if reported_cache_tokens <= input_tokens {
633 input_tokens
634 } else {
635 input_tokens.saturating_add(reported_cache_tokens)
636 };
637 if total_prompt_tokens == 0 {
638 0.0
639 } else {
640 cache_read_tokens as f64 / total_prompt_tokens as f64
641 }
642}
643
644pub(crate) fn cache_savings_usd_for_provider(
645 provider: &str,
646 model: &str,
647 cache_read_tokens: i64,
648 cache_write_tokens: i64,
649) -> f64 {
650 let Some(detail) = pricing_detail_for(provider, model) else {
651 return 0.0;
652 };
653 let input_rate = detail.input_per_1k;
654 let cache_read_rate = detail.cache_read_per_1k.unwrap_or(input_rate);
655 let cache_write_rate = detail.cache_write_per_1k.unwrap_or(input_rate);
656 let cache_read_savings =
657 cache_read_tokens.max(0) as f64 * (input_rate - cache_read_rate) / 1000.0;
658 let cache_write_savings =
659 cache_write_tokens.max(0) as f64 * (input_rate - cache_write_rate) / 1000.0;
660 cache_read_savings + cache_write_savings
661}
662
663pub(crate) fn accumulate_cost_for_provider(
664 provider: &str,
665 model: &str,
666 input_tokens: i64,
667 output_tokens: i64,
668 served_fast: bool,
669) -> Result<(), VmError> {
670 let cost = pricing_detail_for_tier(provider, model, served_fast)
671 .map(|detail| {
672 (input_tokens as f64 * detail.input_per_1k
673 + output_tokens as f64 * detail.output_per_1k)
674 / 1000.0
675 })
676 .unwrap_or(0.0);
677 crate::step_runtime::record_step_llm_usage(model, input_tokens, output_tokens, cost)?;
681 let total_tokens = input_tokens.max(0) as u64 + output_tokens.max(0) as u64;
682 if total_tokens > 0 {
683 LLM_ACCUMULATED_TOKENS.with(|acc| {
684 let mut slot = acc.borrow_mut();
685 *slot = slot.saturating_add(total_tokens);
686 });
687 LLM_TOKEN_BUDGET.with(|budget| {
688 if let Some(max) = *budget.borrow() {
689 let total = LLM_ACCUMULATED_TOKENS.with(|acc| *acc.borrow());
690 if total > max {
691 return Err(categorized_error(
692 format!("LLM token budget exceeded: spent {total} of {max} tokens"),
693 ErrorCategory::BudgetExceeded,
694 ));
695 }
696 }
697 Ok(())
698 })?;
699 }
700 if cost == 0.0 {
701 return Ok(());
702 }
703 LLM_ACCUMULATED_COST.with(|acc| {
704 *acc.borrow_mut() += cost;
705 });
706 LLM_BUDGET.with(|budget| {
707 if let Some(max) = *budget.borrow() {
708 let total = LLM_ACCUMULATED_COST.with(|acc| *acc.borrow());
709 if total > max {
710 return Err(categorized_error(
711 format!("LLM budget exceeded: spent ${total:.4} of ${max:.4} budget"),
712 ErrorCategory::BudgetExceeded,
713 ));
714 }
715 }
716 Ok(())
717 })
718}
719
720pub(crate) fn record_llm_usage_for_provider(
721 provider: &str,
722 model: &str,
723 input_tokens: i64,
724 output_tokens: i64,
725 served_fast: bool,
726) -> Result<(), VmError> {
727 accumulate_cost_for_provider(provider, model, input_tokens, output_tokens, served_fast)
728}
729
730pub(crate) fn register_cost_builtins(vm: &mut Vm) {
731 vm.register_builtin("llm_cost", |args, _out| {
732 let model = args.first().map(|a| a.display()).unwrap_or_default();
733 let input_tokens = args.get(1).and_then(|a| a.as_int()).unwrap_or(0);
734 let output_tokens = args.get(2).and_then(|a| a.as_int()).unwrap_or(0);
735 let cost = calculate_cost_decimal(&model, input_tokens, output_tokens);
740 Ok(VmValue::decimal(cost))
741 });
742
743 vm.register_builtin_with_metadata(
744 VmBuiltinMetadata::sync_static("llm_pricing")
745 .signature_static("llm_pricing(model_or_dict, model?)")
746 .arity(VmBuiltinArity::Range { min: 1, max: 2 })
747 .category_static("llm.economics")
748 .doc_static(
749 "Return catalog pricing for a model: \
750 {input_per_mtok, output_per_mtok, cache_read_per_mtok, cache_write_per_mtok, \
751 provider, model, source} or nil if the model has no priced entry.",
752 ),
753 llm_pricing_builtin,
754 );
755
756 vm.register_builtin_with_metadata(
757 VmBuiltinMetadata::sync_static("llm_format_usd")
758 .signature_static("llm_format_usd(amount, options?)")
759 .arity(VmBuiltinArity::Range { min: 1, max: 2 })
760 .category_static("llm.economics")
761 .doc_static(
762 "Format a USD amount as a string. Default precision auto-scales: 6 decimals \
763 under $1, 4 decimals under $100, 2 decimals otherwise; pass {precision: N} to override.",
764 ),
765 llm_format_usd_builtin,
766 );
767
768 vm.register_builtin_with_metadata(
769 VmBuiltinMetadata::sync_static("llm_compare_costs")
770 .signature_static("llm_compare_costs(candidates, opts)")
771 .arity(VmBuiltinArity::Exact(2))
772 .category_static("llm.economics")
773 .doc_static(
774 "Project a per-call cost across a list of {provider?, model} candidates given \
775 {input_tokens, output_tokens, cache_read_tokens?, cache_write_tokens?, calls?}. \
776 Returns a list sorted ascending by projected cost (unknown pricing trails).",
777 ),
778 llm_compare_costs_builtin,
779 );
780
781 vm.register_builtin("llm_session_cost", |_args, _out| {
782 let (total_input, total_output, _duration, call_count) = super::trace::peek_trace_summary();
783 let total_cost = LLM_ACCUMULATED_COST.with(|acc| *acc.borrow());
784 let mut result = BTreeMap::new();
785 result.insert("total_cost".to_string(), VmValue::Float(total_cost));
786 result.insert("input_tokens".to_string(), VmValue::Int(total_input));
787 result.insert("output_tokens".to_string(), VmValue::Int(total_output));
788 result.insert("call_count".to_string(), VmValue::Int(call_count));
789 Ok(VmValue::dict(result))
790 });
791
792 vm.register_builtin("llm_budget", |args, _out| {
793 let max_cost = match args.first() {
794 Some(VmValue::Float(f)) => *f,
795 Some(VmValue::Int(n)) => *n as f64,
796 _ => {
797 return Err(VmError::Thrown(VmValue::String(arcstr::ArcStr::from(
798 "llm_budget: requires a numeric argument",
799 ))));
800 }
801 };
802 set_llm_cost_budget(Some(max_cost));
803 Ok(VmValue::Nil)
804 });
805
806 vm.register_builtin("llm_budget_remaining", |_args, _out| {
807 let remaining = LLM_BUDGET.with(|budget| {
808 budget.borrow().map(|max| {
809 let spent = LLM_ACCUMULATED_COST.with(|acc| *acc.borrow());
810 max - spent
811 })
812 });
813 match remaining {
814 Some(r) => Ok(VmValue::Float(r)),
815 None => Ok(VmValue::Nil),
816 }
817 });
818
819 vm.register_builtin_with_metadata(
820 VmBuiltinMetadata::sync_static("tiktoken_count_tokens")
821 .signature_static("tiktoken_count_tokens(text, model)")
822 .arity(VmBuiltinArity::Exact(2))
823 .category_static("llm.budget")
824 .doc_static("Count text tokens with the tiktoken encoder selected for a model."),
825 |args, _out| {
826 let text = args.first().map(|arg| arg.display()).unwrap_or_default();
827 let model = args.get(1).map(|arg| arg.display()).unwrap_or_default();
828 if model.trim().is_empty() {
829 return Err(VmError::Runtime(
830 "tiktoken_count_tokens: model is required".to_string(),
831 ));
832 }
833 let estimate = super::token_count::tiktoken_count_text(&text, &model)
834 .map_err(|error| VmError::Runtime(format!("tiktoken_count_tokens: {error}")))?;
835 Ok(VmValue::Int(estimate.tokens))
836 },
837 );
838
839 vm.register_builtin_with_metadata(
840 VmBuiltinMetadata::sync_static("tiktoken_tokenizer_info")
841 .signature_static("tiktoken_tokenizer_info(model)")
842 .arity(VmBuiltinArity::Exact(1))
843 .category_static("llm.budget")
844 .doc_static("Return the tiktoken encoder metadata used for a model token count."),
845 |args, _out| {
846 let model = args.first().map(|arg| arg.display()).unwrap_or_default();
847 Ok(tokenizer_info_to_vm_value(
848 &model,
849 super::token_count::tokenizer_info_for_model(&model),
850 ))
851 },
852 );
853}
854
855fn pricing_detail_to_vm_value(provider: &str, model: &str, detail: &PricingDetail) -> VmValue {
856 let mut dict = BTreeMap::new();
857 dict.put_str("provider", provider);
858 dict.put_str("model", model);
859 dict.insert(
860 "input_per_mtok".to_string(),
861 VmValue::Float(detail.input_per_1k * 1000.0),
862 );
863 dict.insert(
864 "output_per_mtok".to_string(),
865 VmValue::Float(detail.output_per_1k * 1000.0),
866 );
867 dict.insert(
868 "cache_read_per_mtok".to_string(),
869 detail
870 .cache_read_per_1k
871 .map(|rate| VmValue::Float(rate * 1000.0))
872 .unwrap_or(VmValue::Nil),
873 );
874 dict.insert(
875 "cache_write_per_mtok".to_string(),
876 detail
877 .cache_write_per_1k
878 .map(|rate| VmValue::Float(rate * 1000.0))
879 .unwrap_or(VmValue::Nil),
880 );
881 dict.put_str("source", detail.source.as_str());
882 VmValue::dict(dict)
883}
884
885fn resolve_pricing_args(args: &[VmValue]) -> (String, String) {
886 if let Some(VmValue::Dict(dict)) = args.first() {
887 let provider = dict
888 .get("provider")
889 .map(|value| value.display())
890 .unwrap_or_default();
891 let model = dict
892 .get("model")
893 .map(|value| value.display())
894 .unwrap_or_default();
895 if !provider.is_empty() && !model.is_empty() {
896 return (provider, model);
897 }
898 if !model.is_empty() {
899 let resolved = crate::llm_config::resolve_model_info(&model);
900 return (resolved.provider, resolved.id);
901 }
902 }
903 let first = args.first().map(|a| a.display()).unwrap_or_default();
904 let second = args.get(1).map(|a| a.display()).unwrap_or_default();
905 match (first.is_empty(), second.is_empty()) {
906 (false, false) => (first, second),
907 (false, true) => {
908 let resolved = crate::llm_config::resolve_model_info(&first);
909 (resolved.provider, resolved.id)
910 }
911 _ => (String::new(), String::new()),
912 }
913}
914
915fn llm_pricing_builtin(args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
916 let (provider, model) = resolve_pricing_args(args);
917 if model.trim().is_empty() {
918 return Err(VmError::Runtime(
919 "llm_pricing: model is required".to_string(),
920 ));
921 }
922 Ok(pricing_detail_for(&provider, &model)
923 .map(|detail| pricing_detail_to_vm_value(&provider, &model, &detail))
924 .unwrap_or(VmValue::Nil))
925}
926
927fn llm_format_usd_builtin(args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
928 let amount = match args.first() {
929 Some(VmValue::Float(value)) => *value,
930 Some(VmValue::Int(value)) => *value as f64,
931 Some(VmValue::Decimal(value)) => {
936 use rust_decimal::prelude::ToPrimitive;
937 value.to_f64().unwrap_or(0.0)
938 }
939 Some(VmValue::Nil) | None => 0.0,
940 Some(other) => {
941 return Err(VmError::Runtime(format!(
942 "llm_format_usd: amount must be a number (got {})",
943 other.type_name(),
944 )))
945 }
946 };
947 let options = args.get(1).and_then(|v| v.as_dict());
948 let explicit_precision = options
949 .and_then(|opts| opts.get("precision"))
950 .and_then(|value| match value {
951 VmValue::Int(n) if *n >= 0 => Some(*n as usize),
952 VmValue::Float(f) if f.is_finite() && *f >= 0.0 => Some(*f as usize),
953 _ => None,
954 });
955 let sign_always = options
956 .and_then(|opts| opts.get("sign"))
957 .and_then(|value| match value {
958 VmValue::Bool(b) => Some(*b),
959 _ => None,
960 })
961 .unwrap_or(false);
962 let formatted = format_usd_amount(amount, explicit_precision, sign_always);
963 Ok(VmValue::String(arcstr::ArcStr::from(formatted)))
964}
965
966fn format_usd_amount(amount: f64, precision: Option<usize>, sign_always: bool) -> String {
967 if !amount.is_finite() {
968 return "$NaN".to_string();
969 }
970 let precision = precision.unwrap_or_else(|| {
971 let abs = amount.abs();
972 if abs == 0.0 || abs >= 100.0 {
973 2
974 } else if abs >= 1.0 {
975 4
976 } else {
977 6
978 }
979 });
980 let sign = if amount < 0.0 {
981 "-"
982 } else if sign_always {
983 "+"
984 } else {
985 ""
986 };
987 let rounded = format!("{:.*}", precision, amount.abs());
990 let (whole_str, frac_part) = match rounded.find('.') {
991 Some(idx) => (&rounded[..idx], &rounded[idx + 1..]),
992 None => (rounded.as_str(), ""),
993 };
994 let mut grouped = String::new();
995 for (idx, ch) in whole_str.chars().enumerate() {
996 if idx > 0 && (whole_str.len() - idx) % 3 == 0 {
997 grouped.push(',');
998 }
999 grouped.push(ch);
1000 }
1001 if precision == 0 || frac_part.is_empty() {
1002 format!("{sign}${grouped}")
1003 } else {
1004 format!("{sign}${grouped}.{frac_part}")
1005 }
1006}
1007
1008fn llm_compare_costs_builtin(args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
1009 let candidates = match args.first() {
1010 Some(VmValue::List(items)) => items.clone(),
1011 _ => {
1012 return Err(VmError::Runtime(
1013 "llm_compare_costs: candidates must be a list".to_string(),
1014 ))
1015 }
1016 };
1017 let opts = match args.get(1) {
1018 Some(VmValue::Dict(dict)) => dict.clone(),
1019 _ => {
1020 return Err(VmError::Runtime(
1021 "llm_compare_costs: options dict is required".to_string(),
1022 ))
1023 }
1024 };
1025 let input_tokens = opts
1026 .get("input_tokens")
1027 .and_then(|v| v.as_int())
1028 .unwrap_or(0)
1029 .max(0);
1030 let output_tokens = opts
1031 .get("output_tokens")
1032 .and_then(|v| v.as_int())
1033 .unwrap_or(0)
1034 .max(0);
1035 let cache_read_tokens = opts
1036 .get("cache_read_tokens")
1037 .and_then(|v| v.as_int())
1038 .unwrap_or(0)
1039 .max(0);
1040 let cache_write_tokens = opts
1041 .get("cache_write_tokens")
1042 .and_then(|v| v.as_int())
1043 .unwrap_or(0)
1044 .max(0);
1045 let calls = opts
1046 .get("calls")
1047 .and_then(|v| v.as_int())
1048 .unwrap_or(1)
1049 .max(1);
1050
1051 let mut rows: Vec<(Option<f64>, VmValue)> = Vec::with_capacity(candidates.len());
1052 for candidate in candidates.iter() {
1053 let (provider, model) = match candidate {
1054 VmValue::Dict(dict) => {
1055 let provider = dict
1056 .get("provider")
1057 .map(|v| v.display())
1058 .unwrap_or_default();
1059 let model = dict.get("model").map(|v| v.display()).unwrap_or_default();
1060 if model.is_empty() {
1061 return Err(VmError::Runtime(
1062 "llm_compare_costs: each candidate dict must include `model`".to_string(),
1063 ));
1064 }
1065 if provider.is_empty() {
1066 let resolved = crate::llm_config::resolve_model_info(&model);
1067 (resolved.provider, resolved.id)
1068 } else {
1069 (provider, model)
1070 }
1071 }
1072 VmValue::String(s) => {
1073 let resolved = crate::llm_config::resolve_model_info(s);
1074 (resolved.provider, resolved.id)
1075 }
1076 _ => {
1077 return Err(VmError::Runtime(format!(
1078 "llm_compare_costs: candidates must be strings or dicts (got {})",
1079 candidate.type_name(),
1080 )))
1081 }
1082 };
1083 let detail = pricing_detail_for(&provider, &model);
1084 let projection = detail.map(|d| {
1085 project_call_cost(
1086 &d,
1087 input_tokens,
1088 output_tokens,
1089 cache_read_tokens,
1090 cache_write_tokens,
1091 ) * calls as f64
1092 });
1093 let mut row = BTreeMap::new();
1094 row.put_str("provider", provider.clone());
1095 row.put_str("model", model.clone());
1096 row.insert(
1097 "pricing".to_string(),
1098 detail
1099 .as_ref()
1100 .map(|d| pricing_detail_to_vm_value(&provider, &model, d))
1101 .unwrap_or(VmValue::Nil),
1102 );
1103 row.insert(
1104 "cost_usd".to_string(),
1105 projection.map(VmValue::Float).unwrap_or(VmValue::Nil),
1106 );
1107 row.insert("calls".to_string(), VmValue::Int(calls));
1108 row.insert("pricing_known".to_string(), VmValue::Bool(detail.is_some()));
1109 rows.push((projection, VmValue::dict(row)));
1110 }
1111
1112 rows.sort_by(|left, right| match (left.0, right.0) {
1113 (Some(a), Some(b)) => a.partial_cmp(&b).unwrap_or(std::cmp::Ordering::Equal),
1114 (Some(_), None) => std::cmp::Ordering::Less,
1115 (None, Some(_)) => std::cmp::Ordering::Greater,
1116 (None, None) => std::cmp::Ordering::Equal,
1117 });
1118 Ok(VmValue::List(std::sync::Arc::new(
1119 rows.into_iter().map(|(_, value)| value).collect(),
1120 )))
1121}
1122
1123pub(crate) fn project_call_cost(
1124 detail: &PricingDetail,
1125 input_tokens: i64,
1126 output_tokens: i64,
1127 cache_read_tokens: i64,
1128 cache_write_tokens: i64,
1129) -> f64 {
1130 let cache_read_rate = detail.cache_read_per_1k.unwrap_or(detail.input_per_1k);
1131 let cache_write_rate = detail.cache_write_per_1k.unwrap_or(detail.input_per_1k);
1132 let billable_input = (input_tokens - cache_read_tokens - cache_write_tokens).max(0);
1133 (billable_input as f64 * detail.input_per_1k
1134 + output_tokens as f64 * detail.output_per_1k
1135 + cache_read_tokens as f64 * cache_read_rate
1136 + cache_write_tokens as f64 * cache_write_rate)
1137 / 1000.0
1138}
1139
1140fn tokenizer_info_to_vm_value(model: &str, info: super::token_count::TokenizerInfo) -> VmValue {
1141 let mut result = BTreeMap::new();
1142 result.put_str("model", model);
1143 result.put_str("model_family", info.model_family);
1144 result.put_str("source", info.source.as_str());
1145 result.insert("exact".to_string(), VmValue::Bool(info.exact));
1146 result.insert(
1147 "known_model_family".to_string(),
1148 VmValue::Bool(info.known_model_family),
1149 );
1150 result.insert(
1151 "encoder".to_string(),
1152 info.encoder
1153 .map(|encoder| VmValue::String(arcstr::ArcStr::from(encoder)))
1154 .unwrap_or(VmValue::Nil),
1155 );
1156 VmValue::dict(result)
1157}
1158
1159#[cfg(test)]
1160mod tests {
1161 use super::*;
1162
1163 #[test]
1164 fn calculate_cost_uses_catalog_model_pricing() {
1165 let _guard = crate::llm::env_guard();
1166 let mut overlay = crate::llm_config::ProvidersConfig::default();
1167 overlay.models.insert(
1168 "gpt-4o-mini".to_string(),
1169 crate::llm_config::ModelDef {
1170 name: "Test GPT-4o Mini".to_string(),
1171 provider: "openai".to_string(),
1172 context_window: 128_000,
1173 logical_model: None,
1174 equivalence_group: None,
1175 served_variant: None,
1176 wire_model: None,
1177 api_dialect: None,
1178 rate_limits: None,
1179 performance: None,
1180 architecture: None,
1181 local_memory: None,
1182 runtime_context_window: None,
1183 stream_timeout: None,
1184 capabilities: Vec::new(),
1185 pricing: Some(crate::llm_config::ModelPricing {
1186 input_per_mtok: 10.0,
1187 output_per_mtok: 20.0,
1188 cache_read_per_mtok: None,
1189 cache_write_per_mtok: None,
1190 }),
1191 deprecated: false,
1192 deprecation_note: None,
1193 superseded_by: None,
1194 serving_tiers: Vec::new(),
1195 quality_tags: Vec::new(),
1196 availability: crate::llm_config::ModelAvailability::default(),
1197 tier: None,
1198 open_weight: None,
1199 strengths: Vec::new(),
1200 benchmarks: std::collections::BTreeMap::new(),
1201 family: None,
1202 lineage: None,
1203 complementary_with: Vec::new(),
1204 avoid_as_reviewer_for: Vec::new(),
1205 },
1206 );
1207 crate::llm_config::set_user_overrides(Some(overlay));
1208
1209 assert_eq!(
1211 calculate_cost_decimal("gpt-4o-mini", 1000, 1000),
1212 Decimal::from_str("0.03").unwrap()
1213 );
1214
1215 crate::llm_config::clear_user_overrides();
1216 }
1217
1218 #[test]
1219 fn calculate_cost_is_zero_for_unknown_model() {
1220 let _guard = crate::llm::env_guard();
1221 crate::llm_config::clear_user_overrides();
1222 assert_eq!(
1223 calculate_cost_decimal("definitely-unpriced-model", 1_000, 1_000),
1224 Decimal::ZERO
1225 );
1226 }
1227
1228 #[test]
1229 fn authored_rate_decimal_recovers_the_written_literal_not_float_noise() {
1230 for (raw, written) in [
1235 (0.15_f64, "0.15"),
1236 (0.8, "0.8"),
1237 (0.08, "0.08"),
1238 (4.0, "4"),
1239 (0.0, "0"),
1240 (3.75, "3.75"),
1241 ] {
1242 let recovered = authored_rate_decimal(raw);
1243 assert_eq!(
1244 recovered,
1245 Decimal::from_str(written).unwrap(),
1246 "rate {raw} should recover as {written}"
1247 );
1248 }
1249 assert_ne!(
1252 authored_rate_decimal(0.1),
1253 Decimal::from_f64_retain(0.1).unwrap()
1254 );
1255 }
1256
1257 #[test]
1258 fn calculate_cost_decimal_is_exact_for_inexact_catalog_rates() {
1259 let _guard = crate::llm::env_guard();
1260 let mut overlay = crate::llm_config::ProvidersConfig::default();
1261 overlay.models.insert(
1262 "gpt-4o-mini".to_string(),
1263 crate::llm_config::ModelDef {
1264 name: "Test GPT-4o Mini".to_string(),
1265 provider: "openai".to_string(),
1266 context_window: 128_000,
1267 logical_model: None,
1268 equivalence_group: None,
1269 served_variant: None,
1270 wire_model: None,
1271 api_dialect: None,
1272 rate_limits: None,
1273 performance: None,
1274 architecture: None,
1275 local_memory: None,
1276 runtime_context_window: None,
1277 stream_timeout: None,
1278 capabilities: Vec::new(),
1279 pricing: Some(crate::llm_config::ModelPricing {
1281 input_per_mtok: 0.15,
1282 output_per_mtok: 0.60,
1283 cache_read_per_mtok: None,
1284 cache_write_per_mtok: None,
1285 }),
1286 deprecated: false,
1287 deprecation_note: None,
1288 superseded_by: None,
1289 serving_tiers: Vec::new(),
1290 quality_tags: Vec::new(),
1291 availability: crate::llm_config::ModelAvailability::default(),
1292 tier: None,
1293 open_weight: None,
1294 strengths: Vec::new(),
1295 benchmarks: std::collections::BTreeMap::new(),
1296 family: None,
1297 lineage: None,
1298 complementary_with: Vec::new(),
1299 avoid_as_reviewer_for: Vec::new(),
1300 },
1301 );
1302 crate::llm_config::set_user_overrides(Some(overlay));
1303
1304 assert_eq!(
1306 calculate_cost_decimal("gpt-4o-mini", 1000, 500),
1307 Decimal::from_str("0.00045").unwrap()
1308 );
1309
1310 crate::llm_config::clear_user_overrides();
1311 }
1312
1313 #[test]
1314 fn calculate_cost_for_provider_falls_back_to_provider_economics() {
1315 let _guard = crate::llm::env_guard();
1316 crate::llm_config::clear_user_overrides();
1317 let cost =
1318 calculate_cost_for_provider("openai", "some-bespoke-openai-deployment", 1_000, 1_000);
1319 let (input_per_1k, output_per_1k, _) = crate::llm_config::provider_economics("openai");
1320 let expected =
1321 (1_000.0 * input_per_1k.unwrap() + 1_000.0 * output_per_1k.unwrap()) / 1_000.0;
1322 assert!(
1323 (cost - expected).abs() < 1e-9,
1324 "cost={cost}, expected={expected}"
1325 );
1326 }
1327
1328 #[test]
1329 fn pricing_detail_reports_source() {
1330 let _guard = crate::llm::env_guard();
1331 crate::llm_config::clear_user_overrides();
1332 let exact = pricing_detail_for("anthropic", "claude-sonnet-4-20250514").unwrap();
1333 assert_eq!(exact.source, PricingSource::CatalogModel);
1334 assert!(exact.cache_read_per_1k.is_some());
1335
1336 let provider_only = pricing_detail_for("openai", "some-bespoke-openai-deployment").unwrap();
1337 assert_eq!(provider_only.source, PricingSource::ProviderEconomics);
1338 assert!(provider_only.cache_read_per_1k.is_none());
1339
1340 assert!(pricing_detail_for("local", "no-such-local-model").is_some()); assert!(pricing_detail_for("nonexistent_provider", "ghost-model").is_none());
1342 }
1343
1344 #[test]
1345 fn pricing_aware_call_cost_distinguishes_unpriced_from_zero() {
1346 let _guard = crate::llm::env_guard();
1347 crate::llm_config::clear_user_overrides();
1348
1349 let priced = pricing_aware_call_cost("anthropic", "claude-sonnet-4-20250514", 1_000, 1_000);
1351 let expected =
1352 calculate_cost_for_provider("anthropic", "claude-sonnet-4-20250514", 1_000, 1_000);
1353 assert!(priced.is_some());
1354 assert!((priced.unwrap() - expected).abs() < 1e-9);
1355
1356 assert_eq!(
1361 pricing_aware_call_cost("nonexistent_provider", "ghost-model", 1_000, 1_000),
1362 None
1363 );
1364 assert_eq!(
1365 calculate_cost_for_provider("nonexistent_provider", "ghost-model", 1_000, 1_000),
1366 0.0
1367 );
1368 }
1369
1370 #[test]
1371 fn format_usd_amount_auto_precision_and_grouping() {
1372 assert_eq!(format_usd_amount(0.000_045, None, false), "$0.000045");
1373 assert_eq!(format_usd_amount(1.234_5, None, false), "$1.2345");
1374 assert_eq!(format_usd_amount(1234.5, None, false), "$1,234.50");
1375 assert_eq!(format_usd_amount(-1234.5, None, false), "-$1,234.50");
1376 assert_eq!(format_usd_amount(1234.5, None, true), "+$1,234.50");
1377 assert_eq!(format_usd_amount(0.123_456_789, Some(2), false), "$0.12");
1378 assert_eq!(format_usd_amount(1.0, Some(0), false), "$1");
1379 }
1380
1381 #[test]
1382 fn format_usd_handles_fractional_carry_into_whole() {
1383 let amount = 0.000_27_f64 * 300_000.0;
1387 assert!((amount - 81.0).abs() < 1e-6);
1388 assert_eq!(format_usd_amount(amount, None, false), "$81.0000");
1389 }
1390
1391 #[test]
1392 fn fast_tier_bills_premium_pricing_when_served_fast() {
1393 let _guard = crate::llm::env_guard();
1394 crate::llm_config::clear_user_overrides();
1395
1396 let standard = pricing_detail_for_tier("anthropic", "claude-opus-4-8", false).unwrap();
1398 let fast = pricing_detail_for_tier("anthropic", "claude-opus-4-8", true).unwrap();
1399 assert_eq!(standard.source, PricingSource::CatalogModel);
1400 assert_eq!(fast.source, PricingSource::CatalogServingTier);
1401 assert!((fast.input_per_1k - 2.0 * standard.input_per_1k).abs() < 1e-9);
1402 assert!((fast.output_per_1k - 2.0 * standard.output_per_1k).abs() < 1e-9);
1403
1404 let no_fast =
1406 pricing_detail_for_tier("anthropic", "claude-sonnet-4-20250514", true).unwrap();
1407 assert_eq!(no_fast.source, PricingSource::CatalogModel);
1408 }
1409
1410 #[test]
1411 fn project_call_cost_excludes_cached_input_from_full_rate() {
1412 let detail = pricing_detail_for("anthropic", "claude-sonnet-4-20250514").unwrap();
1413 let with_cache = project_call_cost(&detail, 10_000, 500, 8_000, 0);
1414 let no_cache = project_call_cost(&detail, 10_000, 500, 0, 0);
1415 assert!(with_cache < no_cache);
1416 }
1417
1418 #[test]
1419 fn cache_savings_uses_catalog_cache_pricing() {
1420 let _guard = crate::llm::env_guard();
1421 crate::llm_config::clear_user_overrides();
1422
1423 let savings =
1424 cache_savings_usd_for_provider("anthropic", "claude-sonnet-4-20250514", 1000, 0);
1425 assert!((savings - 0.0027).abs() < 0.0000001);
1426
1427 let write_delta =
1428 cache_savings_usd_for_provider("anthropic", "claude-sonnet-4-20250514", 0, 1000);
1429 assert!((write_delta + 0.00075).abs() < 0.0000001);
1430
1431 crate::llm_config::clear_user_overrides();
1432 }
1433
1434 #[test]
1435 fn cache_hit_ratio_handles_subset_and_separate_anthropic_counts() {
1436 assert!((cache_hit_ratio(1000, 250, 0) - 0.25).abs() < f64::EPSILON);
1437 assert!((cache_hit_ratio(100, 900, 0) - 0.9).abs() < f64::EPSILON);
1438 assert_eq!(cache_hit_ratio(0, 0, 0), 0.0);
1439 }
1440
1441 #[test]
1442 fn token_budget_guard_restores_prior_state_on_drop() {
1443 let _guard_outer = crate::llm::env_guard();
1444 reset_cost_state();
1445
1446 let outer = install_llm_token_budget(100);
1447 assert_eq!(peek_total_tokens(), 0);
1448 LLM_ACCUMULATED_TOKENS.with(|a| *a.borrow_mut() = 50);
1450
1451 {
1453 let _inner = install_llm_token_budget(10);
1454 assert_eq!(peek_total_tokens(), 0);
1455 LLM_ACCUMULATED_TOKENS.with(|a| *a.borrow_mut() = 5);
1456 }
1457
1458 assert_eq!(peek_total_tokens(), 50);
1460 drop(outer);
1461 assert_eq!(peek_total_tokens(), 0);
1462
1463 reset_cost_state();
1464 }
1465
1466 #[test]
1467 fn set_budget_rearms_in_place_without_resetting_accumulation() {
1468 let _guard_outer = crate::llm::env_guard();
1469 reset_cost_state();
1470
1471 let _budget = install_llm_cost_budget(1.0);
1473 LLM_ACCUMULATED_COST.with(|a| *a.borrow_mut() = 0.60);
1474
1475 set_llm_cost_budget(Some(0.50));
1478 assert!((peek_total_cost() - 0.60).abs() < f64::EPSILON);
1479 LLM_BUDGET.with(|b| assert_eq!(*b.borrow(), Some(0.50)));
1480
1481 set_llm_cost_budget(Some(2.0));
1483 assert!((peek_total_cost() - 0.60).abs() < f64::EPSILON);
1484 LLM_BUDGET.with(|b| assert_eq!(*b.borrow(), Some(2.0)));
1485
1486 set_llm_cost_budget(None);
1488 LLM_BUDGET.with(|b| assert_eq!(*b.borrow(), None));
1489
1490 set_llm_cost_budget(Some(-5.0));
1492 LLM_BUDGET.with(|b| assert_eq!(*b.borrow(), Some(0.0)));
1493
1494 reset_cost_state();
1495 }
1496
1497 #[test]
1498 fn set_token_budget_rearms_in_place_without_resetting_accumulation() {
1499 let _guard_outer = crate::llm::env_guard();
1500 reset_cost_state();
1501
1502 let _budget = install_llm_token_budget(100);
1503 LLM_ACCUMULATED_TOKENS.with(|a| *a.borrow_mut() = 60);
1504
1505 set_llm_token_budget(Some(50));
1506 assert_eq!(peek_total_tokens(), 60);
1507 LLM_TOKEN_BUDGET.with(|b| assert_eq!(*b.borrow(), Some(50)));
1508
1509 set_llm_token_budget(None);
1510 assert_eq!(peek_total_tokens(), 60);
1511 LLM_TOKEN_BUDGET.with(|b| assert_eq!(*b.borrow(), None));
1512
1513 reset_cost_state();
1514 }
1515
1516 #[test]
1517 fn token_budget_raises_categorized_error_when_exhausted() {
1518 let _guard_outer = crate::llm::env_guard();
1519 reset_cost_state();
1520 let _budget = install_llm_token_budget(10);
1521
1522 let first =
1524 accumulate_cost_for_provider("anthropic", "claude-sonnet-4-20250514", 5, 0, false);
1525 assert!(first.is_ok());
1526
1527 let second =
1529 accumulate_cost_for_provider("anthropic", "claude-sonnet-4-20250514", 8, 0, false);
1530 match second {
1531 Err(VmError::CategorizedError { category, message }) => {
1532 assert_eq!(category, ErrorCategory::BudgetExceeded);
1533 assert!(message.contains("token budget"), "got: {message}");
1534 }
1535 other => panic!("expected BudgetExceeded, got {other:?}"),
1536 }
1537
1538 reset_cost_state();
1539 }
1540}