Skip to main content

harn_vm/llm/
cost.rs

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
17/// Reset thread-local cost state. Call between test runs to avoid leaking.
18pub(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/// RAII guard installed by [`install_llm_cost_budget`]. Restores the
30/// prior ceiling (and accumulated total) on drop so nested dispatches
31/// (a handler that re-enters the dispatcher) cannot leak a tighter
32/// budget into the outer scope, or a wider one back into a finished
33/// inner scope.
34#[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
47/// Pin the per-call LLM cost ceiling at `max_cost_usd` for the lifetime
48/// of the returned guard. Sourced from `@budget(llm_cost_usd = …)` on
49/// `.harn` handlers in `harn-serve`; mid-call exhaustion raises a
50/// `BudgetExceeded`-categorised error which adapter codecs render as
51/// HTTP 429.
52pub 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/// RAII guard for [`install_llm_token_budget`]. Pairs the dispatch-level
64/// token cap with the cost-cap guard so `@budget(llm_tokens, llm_cost_usd)`
65/// both restore on drop.
66#[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
79/// Pin the per-dispatch LLM token ceiling (input + output combined) at
80/// `max_tokens` for the lifetime of the returned guard. Sourced from
81/// `@budget(llm_tokens: …)` on `.harn` handlers in `harn-serve`. Like
82/// the cost-cap variant, mid-stream exhaustion raises a
83/// `BudgetExceeded`-categorised error that adapters render as HTTP 429.
84pub 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
99/// Re-arm the live per-thread LLM **cost** ceiling in place, preserving the
100/// accumulated total. `None` clears the cap.
101///
102/// Unlike [`install_llm_cost_budget`] this returns no guard and does not reset
103/// the running total: it mutates the same `LLM_BUDGET` thread-local a dispatch
104/// already consults at preflight ([`check_llm_preflight_budget`]) and after
105/// each call ([`accumulate_cost_for_provider`]). A supervisor on the dispatch
106/// thread can therefore tighten or loosen the ceiling mid-run and have the next
107/// LLM call observe it — the basis for ACP `session/set_budget` re-arm
108/// (burin-labs/burin-code#1561). Callers that want fresh per-scope accounting
109/// (HTTP `@budget`, per-turn guards) keep using `install_*` instead.
110pub 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
114/// Re-arm the live per-thread LLM **token** ceiling in place, preserving the
115/// accumulated total. `None` clears the cap. The token counterpart to
116/// [`set_llm_cost_budget`] — see that function for the re-arm semantics.
117pub fn set_llm_token_budget(max_tokens: Option<u64>) {
118    LLM_TOKEN_BUDGET.with(|b| *b.borrow_mut() = max_tokens);
119}
120
121/// The live per-thread LLM cost ceiling, or `None` when uncapped. Pairs with
122/// [`peek_total_cost`] so a supervisor (or a re-arm acknowledgement) can read
123/// back the ceiling it just set.
124pub fn peek_llm_cost_budget() -> Option<f64> {
125    LLM_BUDGET.with(|b| *b.borrow())
126}
127
128/// The live per-thread LLM token ceiling, or `None` when uncapped. The token
129/// counterpart to [`peek_llm_cost_budget`].
130pub fn peek_llm_token_budget() -> Option<u64> {
131    LLM_TOKEN_BUDGET.with(|b| *b.borrow())
132}
133
134// `Serialize` exists for the typed-options parity test
135// (`orchestration/typed_options_parity.rs`), which compares this struct's
136// serialized default key set against the `LlmBudget` alias in
137// `std/llm/options.harn`.
138#[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/// Resolved pricing for a (provider, model) pair, expressed per 1k tokens.
457/// The `source` discriminates how the rate was found so callers (CLI cost
458/// explanation, economics helpers, `cost_route` summaries) can report it.
459#[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    /// Exact model entry in the catalog (configured `[llm.models.<id>]`).
471    CatalogModel,
472    /// The model's accelerated-serving tier (`serving_tiers[].pricing`), used
473    /// when the provider confirmed it served the request fast.
474    CatalogServingTier,
475    /// Provider-level catalog economics (`[llm.providers.<name>]`).
476    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
489/// Resolve full pricing detail for a (provider, model) pair. Prefers the
490/// exact-id catalog entry, then falls back to provider-level economics.
491/// Returns `None` for unknown pricing — callers must decide whether to
492/// surface that explicitly or coerce to 0.0.
493pub(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
520/// Resolve pricing for a (provider, model) pair, billing at the premium
521/// accelerated-serving tier when `served_fast` is set and the catalog
522/// declares `serving_tiers[].pricing`. Falls back to standard pricing when the
523/// request was served at the standard tier (e.g. a capacity downgrade) or
524/// the fast tier omits explicit rates.
525pub(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
552/// Recover the *authored* decimal value of a catalog rate that was parsed
553/// from a TOML float literal. The pricing catalog
554/// (`crates/harn-vm/src/llm/providers.toml`) writes short, human-authored
555/// decimals like `input_per_mtok = 0.15`; TOML deserializes them to `f64`,
556/// which cannot store `0.15` exactly. Rust's `{}` float formatter emits the
557/// *shortest* decimal string that round-trips to the same `f64` — for these
558/// short literals that is exactly the digits the author wrote (`"0.15"`,
559/// not `0.150000000000000008…`). Parsing that string straight into a
560/// `Decimal` therefore reconstructs the intended exact value without ever
561/// laundering the float's binary rounding error into false precision (which
562/// `Decimal::from_f64_retain` would). The `from_f64_retain` fallback only
563/// fires for non-finite inputs, which the catalog never contains.
564fn 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
571/// Compute the per-call USD cost for a model and token counts as an exact
572/// `Decimal`. Sources each rate directly from the per-MTok catalog literal
573/// via [`authored_rate_decimal`] (never through the derived per-1k rate,
574/// which would round-trip the value through an extra `f64` multiply) and
575/// does all arithmetic in `Decimal`. Division by 1,000,000 is an exact
576/// base-10 rescale, so the result carries no representational error.
577/// Returns `Decimal::ZERO` when the model has no catalog entry.
578pub 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
587/// Calculate cost using catalog model pricing first, then provider catalog
588/// economics when the model has no exact catalog entry. Returns 0.0 when
589/// pricing is unknown (use `pricing_detail_for` to distinguish unknown).
590pub 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
603/// Per-call USD cost for trace attribution, or `None` when the
604/// (provider, model) pair has no catalog pricing. Unlike
605/// [`calculate_cost_for_provider`] (which coerces unknown pricing to
606/// `0.0` for budget arithmetic), this preserves the distinction so a
607/// `cost_usd` span field can honestly report "unpriced" rather than a
608/// misleading zero. Pricing resolution matches `calculate_cost_for_provider`:
609/// catalog model rate first, then provider-level economics.
610pub 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    // Always attribute usage to the active `@step` (if any), even when
678    // the per-call cost is zero — token-only step budgets need the
679    // count regardless of pricing.
680    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        // Return the cost as an exact `decimal` (not a binary `float`): the
736        // value names money, summing it should not drift, and the catalog
737        // rates are recovered exactly. Callers that want a formatted string
738        // pass the result to `llm_format_usd`, which accepts decimals.
739        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        // Accept exact `decimal` amounts (e.g. the result of `llm_cost`).
932        // Formatting rounds to a fixed number of display decimals anyway, so
933        // converting to `f64` for the digit layout is lossless at that
934        // precision; the exact value never feeds back into money math here.
935        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    // Defer rounding to the libc formatter so that values like 81.0 that
988    // arrive as 80.999… don't split into "$80." + "1.0000".
989    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        // 1000*10 + 1000*20 = 30000; /1e6 = 0.03, exactly.
1210        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        // Catalog rates are short literals parsed from TOML into f64; many
1231        // (0.15, 0.8, 0.08) are not exactly representable in binary. The
1232        // recovery must reconstruct the *authored* decimal, never the f64's
1233        // binary-rounding tail that `from_f64_retain` would expose.
1234        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        // Guard the whole point of the helper: it must NOT equal the lossy
1250        // `from_f64_retain` decimal for an inexact literal.
1251        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                // Inexact-in-binary rates, like the real catalog.
1280                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        // 1000 * 0.15 + 500 * 0.60 = 150 + 300 = 450; /1e6 = 0.00045 exactly.
1305        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()); // local has 0/0
1341        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        // Known catalog model: Some(cost) matching the priced arithmetic.
1350        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        // Genuinely unpriced (provider not in the catalog economics table):
1357        // None, not a misleading 0.0. `calculate_cost_for_provider` coerces
1358        // the same case to 0.0, which is exactly the ambiguity this helper
1359        // exists to remove.
1360        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        // 0.00027 * 300_000 produces 80.999… in IEEE-754; the formatter
1384        // must round-then-render rather than splitting the rounded fraction
1385        // back into a separate component (regression: "$80.1.0000").
1386        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        // Opus 4.8 fast mode is 2x standard ($5/$25 -> $10/$50 per MTok).
1397        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        // A model with no fast tier ignores the flag and bills standard.
1405        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        // Simulate accumulation by writing the thread-local directly.
1449        LLM_ACCUMULATED_TOKENS.with(|a| *a.borrow_mut() = 50);
1450
1451        // Nested guard wipes accumulation and installs a tighter cap.
1452        {
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        // Outer scope restored on inner drop.
1459        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        // Install a $1.00 cap and spend $0.60 against it.
1472        let _budget = install_llm_cost_budget(1.0);
1473        LLM_ACCUMULATED_COST.with(|a| *a.borrow_mut() = 0.60);
1474
1475        // Tighten the cap below current spend: the next preflight must trip,
1476        // and the already-accumulated total must be preserved (not reset).
1477        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        // Loosen the cap: spend stays, ceiling rises, room reopens.
1482        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        // Clear the cap entirely.
1487        set_llm_cost_budget(None);
1488        LLM_BUDGET.with(|b| assert_eq!(*b.borrow(), None));
1489
1490        // Negative ceilings clamp to zero (a hard stop), matching `install_*`.
1491        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        // First call within budget — admits.
1523        let first =
1524            accumulate_cost_for_provider("anthropic", "claude-sonnet-4-20250514", 5, 0, false);
1525        assert!(first.is_ok());
1526
1527        // Second call pushes over — raises BudgetExceeded.
1528        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}