diff --git a/CHANGELOG.md b/CHANGELOG.md index b9915c95d3..632108c022 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -195,6 +195,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- The Runtime API test harness builds and serves its router on a thread sized + like the product's runtime workers (`CODEWHALE_MAIN_STACK_BYTES`) instead of + the 2 MiB libtest thread. Two product paths (config load/reload under a + profile, thread lifecycle) marginally overflowed that stack in debug builds + and aborted the whole lib suite with SIGABRT; `runtime_api` now passes with + no `RUST_MIN_STACK` override. - The sandbox read deny-list matches a rule's resolved path as well as its literal spelling. On macOS `/etc` and `/var` are symlinks into `/private`, so a read of `/private/etc/sudoers` walked around the `/etc/sudoers` rule, diff --git a/crates/config/src/provider_templates.rs b/crates/config/src/provider_templates.rs index 430dc94a93..54ba927f62 100644 --- a/crates/config/src/provider_templates.rs +++ b/crates/config/src/provider_templates.rs @@ -35,17 +35,31 @@ pub const SENSENOVA_MODELS: &[&str] = &[SENSENOVA_DEFAULT_MODEL]; pub const AGNES_TEMPLATE_ID: &str = "agnes"; /// Baseten Model APIs — OpenAI Chat Completions, discovered at `/v1/models`. +/// +/// This 2026-08-29 roster is the official offline seed from Baseten's Model +/// APIs overview. The authenticated live catalog replaces this exact provider +/// partition after startup, so upstream additions and removals do not require a +/// Codewhale release. pub const BASETEN_TEMPLATE_ID: &str = "baseten"; pub const BASETEN_BASE_URL: &str = "https://inference.baseten.co/v1"; pub const BASETEN_DEFAULT_MODEL: &str = "deepseek-ai/DeepSeek-V4-Pro"; pub const BASETEN_API_KEY_ENV: &str = "BASETEN_API_KEY"; pub const BASETEN_MODELS: &[&str] = &[ BASETEN_DEFAULT_MODEL, - "deepseek-ai/DeepSeek-V4-Flash-0731", "deepseek-ai/DeepSeek-V4-Pro-0813", + "deepseek-ai/DeepSeek-V4-Flash-0731", + "zai-org/GLM-4.7", "zai-org/GLM-5.2", + "zai-org/GLM-5.2-Fast", + "zai-org/GLM-5.3", "zai-org/GLM-5.3-Flash", + "thinkingmachines/inkling", + "thinkingmachines/inkling-small", + "moonshotai/Kimi-K2.6", "moonshotai/Kimi-K2.7-Code", + "moonshotai/Kimi-K3", + "nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B", + "openai/gpt-oss-120b", ]; /// Groq — OpenAI Chat Completions hosted inference. @@ -505,4 +519,23 @@ mod tests { ); } } + + #[test] + fn baseten_offline_seed_matches_the_current_official_model_api_roster() { + let template = provider_setup_template(BASETEN_TEMPLATE_ID).expect("baseten"); + assert_eq!(template.picker_models(), BASETEN_MODELS); + assert_eq!(BASETEN_MODELS.len(), 15); + for required in [ + "deepseek-ai/DeepSeek-V4-Pro", + "zai-org/GLM-5.3", + "thinkingmachines/inkling", + "moonshotai/Kimi-K3", + "openai/gpt-oss-120b", + ] { + assert!( + BASETEN_MODELS.contains(&required), + "missing current Baseten seed {required}" + ); + } + } } diff --git a/crates/config/src/route/resolver.rs b/crates/config/src/route/resolver.rs index 1a5143386d..3a167ecc60 100644 --- a/crates/config/src/route/resolver.rs +++ b/crates/config/src/route/resolver.rs @@ -145,6 +145,29 @@ impl RouteResolver { /// Returns [`RouteError`] when the model is empty, the provider is invalid, /// or a clearly-foreign model is requested for a strict direct provider. pub fn resolve(&self, req: &RouteRequest) -> Result { + self.resolve_inner(req, false) + } + + /// Resolve with catalog facts authenticated against this exact endpoint. + /// + /// The ordinary [`Self::resolve`] path strips capabilities and pricing when + /// a route uses a custom base URL, because a same-named first-party model is + /// not evidence about an arbitrary proxy. Callers may use this seam only + /// when the injected offering came from the selected provider identity's + /// own endpoint and its base-URL fingerprint matches the request endpoint. + /// All normal routing and protocol validation still applies. + pub fn resolve_with_endpoint_catalog_authority( + &self, + req: &RouteRequest, + ) -> Result { + self.resolve_inner(req, true) + } + + fn resolve_inner( + &self, + req: &RouteRequest, + endpoint_catalog_authoritative: bool, + ) -> Result { // 1. Provider scope from explicit choice only; default otherwise. // The provider is NEVER inferred from a model prefix. let provider_kind = req.explicit_provider.unwrap_or_default(); @@ -235,14 +258,21 @@ impl RouteResolver { selected.endpoint_key = "responses".to_string(); } } - if custom_endpoint { + if custom_endpoint && !endpoint_catalog_authoritative { // Capabilities and pricing belong to the exact provider endpoint // offering that reported them. Reusing a provider enum and a // first-party model id against a custom compatible endpoint does - // not prove that proxy serves the same modality, tool, reasoning, - // or billing contract. Keep the caller's model id and Chat - // pass-through above, but clear every unowned offering fact at the - // authority boundary instead of presenting it as verified. + // not prove that proxy serves the same canonical model, limits, + // modality, tool, reasoning, or billing contract. Keep the + // caller's wire model id, but clear every unowned offering fact at + // the authority boundary instead of presenting it as verified. + // The endpoint_key/protocol stays: it is the provider adapter's + // wire contract (a model-aware roster row or fixed policy), not an + // endpoint-catalog fact, and coercing it to Chat would silently + // change how a Responses- or Messages-bound route speaks. + // Deepseek's custom-endpoint Chat pass-through is handled above. + selected.canonical_model = None; + selected.limits = RouteLimits::default(); selected.capabilities = RouteCapabilities::default(); selected.pricing = PricingSku::UnknownOrStale; } diff --git a/crates/config/src/route/tests.rs b/crates/config/src/route/tests.rs index e251ee7bdc..2e06ed4c1e 100644 --- a/crates/config/src/route/tests.rs +++ b/crates/config/src/route/tests.rs @@ -1663,6 +1663,33 @@ fn custom_endpoint_does_not_inherit_first_party_pricing() { ); } +#[test] +fn exact_endpoint_catalog_may_carry_its_own_pricing_on_a_custom_base_url() { + use super::candidate::PricingSku; + + let request = RouteRequest { + explicit_provider: Some(ProviderKind::Deepseek), + model_selector: Some(LogicalModelRef::from("deepseek-v4-pro")), + saved_provider_model: None, + base_url_override: Some("https://authenticated-catalog.example.test/v1".to_string()), + limit_overrides: Vec::new(), + }; + let out = priced_deepseek_resolver() + .resolve_with_endpoint_catalog_authority(&request) + .expect("exact endpoint-owned catalog route resolves"); + + match out.pricing() { + Some(PricingSku::Token { + input_per_mtok, + output_per_mtok, + }) => { + assert_eq!(*input_per_mtok, Some(0.28)); + assert_eq!(*output_per_mtok, Some(0.42)); + } + other => panic!("expected exact endpoint pricing, got {other:?}"), + } +} + #[test] fn unpriced_offering_stays_unknown() { use super::candidate::PricingSku; diff --git a/crates/tui/CHANGELOG.md b/crates/tui/CHANGELOG.md index ecfd1b3f7e..91c6141f58 100644 --- a/crates/tui/CHANGELOG.md +++ b/crates/tui/CHANGELOG.md @@ -195,6 +195,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- The Runtime API test harness builds and serves its router on a thread sized + like the product's runtime workers (`CODEWHALE_MAIN_STACK_BYTES`) instead of + the 2 MiB libtest thread. Two product paths (config load/reload under a + profile, thread lifecycle) marginally overflowed that stack in debug builds + and aborted the whole lib suite with SIGABRT; `runtime_api` now passes with + no `RUST_MIN_STACK` override. - The sandbox read deny-list matches a rule's resolved path as well as its literal spelling. On macOS `/etc` and `/var` are symlinks into `/private`, so a read of `/private/etc/sudoers` walked around the `/etc/sudoers` rule, diff --git a/crates/tui/src/client.rs b/crates/tui/src/client.rs index 22959ce55a..a739063a32 100644 --- a/crates/tui/src/client.rs +++ b/crates/tui/src/client.rs @@ -246,6 +246,19 @@ pub struct SpeechSynthesisResponse { pub voice: Option, } +/// One decoded provider response from the auxiliary translation path. +/// +/// The immutable route and provider-reported usage travel beside the semantic +/// translation result so callers can account for a successful provider call +/// before rejecting an incomplete, empty, or otherwise unusable translation. +/// `usage == None` is distinct from a transport failure: the provider returned +/// a response, but omitted the receipt needed to price it exactly. +pub(crate) struct TranslationProviderResponse { + pub(crate) translated: Result, + pub(crate) route: crate::cost_status::EffectiveRouteEnvelope, + pub(crate) usage: Option, +} + /// Client for DeepSeek's OpenAI-compatible APIs. #[must_use] pub struct DeepSeekClient { @@ -741,6 +754,8 @@ pub(super) const ERROR_BODY_MAX_BYTES: usize = 64 * 1024; /// the connection and never answers hangs model-list, catalog refresh, and /// health checks forever (ops R4). pub(super) const NON_STREAMING_HTTP_TIMEOUT: Duration = Duration::from_secs(30); +const PROVIDER_CATALOG_MAX_RESPONSE_BYTES: usize = 8 * 1024 * 1024; +const PROVIDER_CATALOG_MAX_ROWS: usize = 10_000; /// Read an error response body with a size limit to prevent unbounded allocation. pub(super) async fn bounded_error_text(response: reqwest::Response, max_bytes: usize) -> String { @@ -758,6 +773,27 @@ pub(super) async fn bounded_error_text(response: reqwest::Response, max_bytes: u String::from_utf8_lossy(&buf).into_owned() } +async fn bounded_provider_catalog_text( + response: reqwest::Response, +) -> Result { + if response + .content_length() + .is_some_and(|length| length > PROVIDER_CATALOG_MAX_RESPONSE_BYTES as u64) + { + return Err(CatalogRefreshError::InvalidResponse); + } + let mut stream = response.bytes_stream(); + let mut body = Vec::new(); + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(|_| CatalogRefreshError::Network)?; + if body.len().saturating_add(chunk.len()) > PROVIDER_CATALOG_MAX_RESPONSE_BYTES { + return Err(CatalogRefreshError::InvalidResponse); + } + body.extend_from_slice(&chunk); + } + String::from_utf8(body).map_err(|_| CatalogRefreshError::InvalidResponse) +} + fn validate_base_url_security(base_url: &str) -> Result<()> { let display_base_url = redact_url_for_display(base_url); if base_url.starts_with("https://") @@ -2186,10 +2222,11 @@ impl DeepSeekClient { ) } - /// Capture the immutable, redacted route envelope for a request immediately - /// before it is dispatched. The wire model is normalized exactly as the - /// transport will normalize it; a provider-returned alias must never replace - /// this billing identity later. + /// Capture the immutable, redacted route envelope at the caller's + /// application-dispatch/admission time. This is not proof of network + /// delivery or provider invoice-time pricing. The wire model is normalized + /// exactly as the transport will normalize it; a provider-returned alias + /// must never replace this billing identity later. #[must_use] pub fn effective_route_envelope( &self, @@ -2198,12 +2235,28 @@ impl DeepSeekClient { ) -> crate::cost_status::EffectiveRouteEnvelope { let model = wire_model_for_provider_route(self.api_provider, &self.base_url, requested_model); + let endpoint_fingerprint = crate::cost_status::endpoint_fingerprint(&self.base_url); + let provider_live_pricing = + u64::try_from(dispatched_at.timestamp()) + .ok() + .and_then(|dispatched_at_unix| { + endpoint_fingerprint.as_deref().and_then(|fingerprint| { + crate::provider_catalog_live::fresh_provider_live_pricing_quote_at( + self.api_provider, + &self.provider_identity, + &model, + fingerprint, + dispatched_at_unix, + ) + }) + }); crate::cost_status::EffectiveRouteEnvelope { provider: self.api_provider, provider_identity: self.provider_identity.clone(), model, billing_surface: self.billing_surface.clone(), - endpoint_fingerprint: crate::cost_status::endpoint_fingerprint(&self.base_url), + endpoint_fingerprint, + provider_live_pricing, billing_mode: self.billing_mode, dispatched_at, } @@ -2293,12 +2346,15 @@ impl DeepSeekClient { /// This is a lightweight translation service — no tool calls, no /// streaming, no conversation history. The dedicated translation agent /// receives the source text and returns only the translated result. - pub async fn translate( + pub(crate) async fn translate_with_usage( &self, text: &str, model: &str, target_language: &str, - ) -> Result { + ) -> Result { + // Freeze pricing before either the remote-control gate or the provider + // permit. A later live-catalog refresh must not reprice this request. + let route = self.effective_route_envelope(model, chrono::Utc::now()); let _inference = self.acquire_remote_control_inference_permit().await; let _permit = self.acquire_provider_request_permit().await; let model = wire_model_for_provider_route(self.api_provider, &self.base_url, model); @@ -2319,11 +2375,25 @@ impl DeepSeekClient { WireDialect::AnthropicMessages => self.handle_anthropic_message(&prepared).await?, WireDialect::ChatCompletions | WireDialect::GoogleCloudCode => unreachable!(), }; - return translation_text_from_response(&response); + let usage = (response.usage != Usage::default()).then_some(response.usage.clone()); + let translated = + if crate::models::is_incomplete_stop_reason(response.stop_reason.as_deref()) { + Err(anyhow::anyhow!( + "translate: provider response incomplete ({})", + crate::models::stop_reason_detail(response.stop_reason.as_deref()) + )) + } else { + translation_text_from_response(&response) + }; + return Ok(TranslationProviderResponse { + translated, + route, + usage, + }); } let url = api_url_with_suffix( - &self.base_url, + self.chat_transport_base_url(), "chat/completions", self.path_suffix.as_deref(), ); @@ -2351,15 +2421,71 @@ impl DeepSeekClient { ); let response = self.send_json_with_retry(&url, &body).await?; + let status = response.status(); + if !status.is_success() { + let raw_error_text = bounded_error_text(response, ERROR_BODY_MAX_BYTES).await; + let error_text = sanitize_http_error_body( + Some(self.api_provider.display_name()), + status.as_u16(), + &raw_error_text, + ); + anyhow::bail!("translate: HTTP {status}: {error_text}"); + } let value: serde_json::Value = response.json().await?; - let translated = value["choices"][0]["message"]["content"] - .as_str() - .ok_or_else(|| anyhow::anyhow!("translate: unexpected API response shape"))? - .trim() - .to_string(); + let usage_reported = value + .get("usage") + .and_then(Value::as_object) + .is_some_and(|usage| { + [ + "input_tokens", + "prompt_tokens", + "output_tokens", + "completion_tokens", + "total_tokens", + ] + .iter() + .any(|field| usage.contains_key(*field)) + }); + let usage = parse_usage(value.get("usage")); + let stop_reason = value["choices"][0]["finish_reason"].as_str(); + let usage = (usage_reported && usage != Usage::default()).then_some(usage); + let translated = if crate::models::is_incomplete_stop_reason(stop_reason) { + Err(anyhow::anyhow!( + "translate: provider response incomplete ({})", + crate::models::stop_reason_detail(stop_reason) + )) + } else { + value["choices"][0]["message"]["content"] + .as_str() + .ok_or_else(|| anyhow::anyhow!("translate: unexpected API response shape")) + .and_then(|translated| { + let translated = translated.trim().to_string(); + if translated.is_empty() { + bail!("translate: provider response did not contain text content"); + } + Ok(translated) + }) + }; + + Ok(TranslationProviderResponse { + translated, + route, + usage, + }) + } - Ok(translated) + /// Compatibility wrapper for callers that only render the translation. + /// Receipt-aware TUI paths use [`Self::translate_with_usage`] directly. + pub async fn translate( + &self, + text: &str, + model: &str, + target_language: &str, + ) -> Result { + self.translate_with_usage(text, model, target_language) + .await? + .translated } /// List available models from the provider. @@ -2396,12 +2522,30 @@ impl DeepSeekClient { /// back to the `ApiProvider` slug for legacy variants without a kind). This /// is the id used as the cache scope and `CatalogOffering.provider`. fn catalog_provider_id(&self) -> String { + if self.api_provider == ApiProvider::Custom { + // This is an ownership key, not a provider-family slug. Exact + // custom identities (including case and built-in-looking names) + // must remain isolated across cache, picker, and runtime layers. + return self.provider_identity.trim().to_string(); + } self.api_provider .kind() .map(|kind| kind.as_str().to_string()) .unwrap_or_else(|| self.api_provider.as_str().to_string()) } + /// Reviewed response schema for a named compatible provider. + /// + /// Schema recognition is intentionally separate from catalog ownership: + /// `base-ten` may use Baseten's `/models` shape, while its exact configured + /// identity remains `base-ten` rather than sharing the `baseten` partition. + fn catalog_setup_template_id(&self) -> Option<&'static str> { + (self.api_provider == ApiProvider::Custom) + .then(|| codewhale_config::provider_setup_template(&self.provider_identity)) + .flatten() + .map(|template| template.id) + } + /// Fetch the provider's live `/models` listing as a secret-free /// [`ProviderCatalogDelta`] (#3385). /// @@ -2442,10 +2586,7 @@ impl DeepSeekClient { }); } - let body = response - .text() - .await - .map_err(|_| CatalogRefreshError::Network)?; + let body = bounded_provider_catalog_text(response).await?; let provider = self.catalog_provider_id(); let fingerprint = base_url_fingerprint(&self.base_url); @@ -2454,7 +2595,7 @@ impl DeepSeekClient { // OpenRouter returns extended capability metadata in its /models // response (#3385). Capture limits, pricing, reasoning, and modalities // from the live API instead of leaving them unknown. - let offerings: Vec = if provider == "openrouter" { + let offerings: Vec = if self.api_provider == ApiProvider::Openrouter { let or_models = parse_openrouter_models_response(&body)?; if or_models.is_empty() { return Err(CatalogRefreshError::EmptyList); @@ -2464,8 +2605,17 @@ impl DeepSeekClient { .map(|item| { openrouter_to_catalog_offering(item, &provider, &fingerprint, fetched_at) }) - .collect() - } else if provider == "telecomjs" { + .collect::, _>>()? + } else if self.catalog_setup_template_id() == Some(codewhale_config::BASETEN_TEMPLATE_ID) { + let baseten_models = parse_baseten_models_response(&body)?; + if baseten_models.is_empty() { + return Err(CatalogRefreshError::EmptyList); + } + baseten_models + .iter() + .map(|item| baseten_to_catalog_offering(item, &provider, &fingerprint, fetched_at)) + .collect::, _>>()? + } else if self.api_provider == ApiProvider::Telecomjs { named_gateway_catalog_offerings_from_body( &body, codewhale_config::ProviderKind::Telecomjs, @@ -2473,7 +2623,7 @@ impl DeepSeekClient { &fingerprint, fetched_at, )? - } else if provider == "edenai" { + } else if self.api_provider == ApiProvider::Edenai { named_gateway_catalog_offerings_from_body( &body, codewhale_config::ProviderKind::Edenai, @@ -2514,6 +2664,10 @@ impl DeepSeekClient { .collect() }; + if offerings.len() > PROVIDER_CATALOG_MAX_ROWS { + return Err(CatalogRefreshError::InvalidResponse); + } + Ok(ProviderCatalogDelta { provider, base_url_fingerprint: fingerprint, @@ -2533,42 +2687,61 @@ impl DeepSeekClient { ) -> CatalogStatus { match self.fetch_catalog_delta().await { Ok(delta) => { + let provider = delta.provider.clone(); + let fingerprint = delta.base_url_fingerprint.clone(); cache.record_success(delta, ttl_secs); - publish_provider_lake_snapshot(cache); + publish_provider_lake_scope(cache, &provider, &fingerprint); CatalogStatus::Fresh } Err(reason) => { - cache.record_failure( - &self.catalog_provider_id(), - &base_url_fingerprint(&self.base_url), - reason, - ); - publish_provider_lake_snapshot(cache); + let provider = self.catalog_provider_id(); + let fingerprint = base_url_fingerprint(&self.base_url); + cache.record_failure(&provider, &fingerprint, reason); + publish_provider_lake_scope(cache, &provider, &fingerprint); CatalogStatus::Failed { reason } } } } /// Best-effort background refresh of the active provider's own `/v1/models` - /// catalog, merging results into the provider lake (#3385). + /// catalog, replacing that provider's exact lake partition (#3385). /// /// Unlike `models_dev_live::spawn_background_refresh` (which fetches the /// cross-provider Models.dev catalog), this calls the provider's own /// `/v1/models` endpoint and merges the results into the existing live - /// snapshot via `provider_lake::merge_live_offerings`, preserving rows - /// from other sources. + /// snapshot via `provider_catalog_live`, preserving other providers while + /// allowing this provider's successful roster to retire removed ids. /// - /// Currently activated for providers whose model list is not covered by the - /// Models.dev catalog (e.g. TelecomJS TokenHub). The refresh is non-fatal: - /// on failure, existing/bundled rows remain available. + /// Activated for model-list authorities that are not satisfied by the + /// cross-provider Models.dev snapshot: OpenRouter, named live gateways, and + /// the reviewed OpenAI-compatible setup templates (including Baseten). + /// The refresh is non-fatal: on failure, persisted prior rows and static + /// seeds remain available with a typed failed receipt. pub fn spawn_active_provider_catalog_refresh(config: &Config) { let provider = config.api_provider(); - // Only refresh for providers that serve their own model list and are - // not already covered by the Models.dev catalog. - if !matches!(provider, ApiProvider::Telecomjs | ApiProvider::Edenai) { + let provider_identity = config.provider_identity_for(provider); + let is_reviewed_compatible_template = provider == ApiProvider::Custom + && codewhale_config::provider_setup_template(&provider_identity) + .is_some_and(|template| template.is_compatible()); + if !matches!( + provider, + ApiProvider::Openrouter | ApiProvider::Telecomjs | ApiProvider::Edenai + ) && !is_reviewed_compatible_template + { return; } + // Invalidate older in-flight fetches before loading any reusable + // scope. Baseten's ticket also clears account-scoped rows because the + // same endpoint can expose a different workspace after a key change. + let refresh_ticket = crate::provider_catalog_live::begin_refresh(&provider_identity); + + // Publish the exact persisted scope immediately so opening `/model` + // never waits on the network and another endpoint's rows cannot leak + // into this route. Account-scoped Baseten rows deliberately do not + // reload from disk until the current credential proves them again. + crate::provider_catalog_live::maybe_load_persisted_cache_for_config(config); + let client = match DeepSeekClient::new(config) { Ok(client) => client, Err(err) => { @@ -2585,7 +2758,18 @@ impl DeepSeekClient { match client.fetch_catalog_delta().await { Ok(delta) => { let count = delta.offerings.len(); - crate::provider_lake::merge_live_offerings(delta.offerings); + if crate::provider_catalog_live::record_success_if_current( + &refresh_ticket, + delta, + ) + .is_none() + { + tracing::debug!( + target: "provider_catalog", + "discarded provider catalog response superseded by a newer refresh" + ); + return; + } tracing::debug!( target: "provider_catalog", offering_count = count, @@ -2593,6 +2777,20 @@ impl DeepSeekClient { ); } Err(err) => { + if crate::provider_catalog_live::record_failure_if_current( + &refresh_ticket, + &client.catalog_provider_id(), + &base_url_fingerprint(&client.base_url), + err, + ) + .is_none() + { + tracing::debug!( + target: "provider_catalog", + "discarded provider catalog failure superseded by a newer refresh" + ); + return; + } tracing::debug!( target: "provider_catalog", error = ?err, @@ -3209,6 +3407,102 @@ struct OpenRouterArchitecture { output_modalities: Option>, } +/// Baseten Model APIs `/v1/models` item. +/// +/// Baseten publishes OpenAI-style ids and per-token prices, while current +/// serving limits and feature fields are additive. Numeric fields accept JSON +/// numbers or numeric strings because both appear in provider catalogs in the +/// wild; malformed or negative known fields reject the refresh so the durable +/// last-known-good snapshot remains authoritative. +#[derive(Debug, Deserialize)] +struct BasetenModelsResponse { + data: Vec, +} + +#[derive(Debug, Deserialize)] +struct BasetenModelItem { + id: String, + #[serde(default)] + context_length: Option, + #[serde(default)] + context_window: Option, + #[serde(default)] + max_output_tokens: Option, + #[serde(default)] + max_completion_tokens: Option, + #[serde(default)] + limits: Option, + #[serde(default)] + top_provider: Option, + #[serde(default)] + pricing: Option, + #[serde(default)] + supported_parameters: Option>, + #[serde(default)] + supported_features: Option>, + #[serde(default)] + features: Option>, + #[serde(default)] + architecture: Option, + #[serde(default)] + input_modalities: Option>, + #[serde(default)] + output_modalities: Option>, + #[serde(default)] + reasoning: Option, + #[serde(default)] + supports_reasoning: Option, + #[serde(default)] + supports_tools: Option, + #[serde(default)] + supports_structured_output: Option, + #[serde(default)] + reasoning_options: Vec, +} + +#[derive(Debug, Deserialize)] +#[serde(untagged)] +enum CatalogNumber { + Number(serde_json::Number), + Text(String), +} + +#[derive(Debug, Default, Deserialize)] +struct BasetenLimits { + #[serde(default)] + context: Option, + #[serde(default)] + output: Option, +} + +#[derive(Debug, Default, Deserialize)] +struct BasetenTopProvider { + #[serde(default)] + context_length: Option, + #[serde(default)] + max_completion_tokens: Option, +} + +#[derive(Debug, Default, Deserialize)] +struct BasetenPricing { + #[serde(default)] + prompt: Option, + #[serde(default)] + completion: Option, + #[serde(default)] + input_cache_read: Option, + #[serde(default)] + input_cache_write: Option, +} + +#[derive(Debug, Default, Deserialize)] +struct BasetenArchitecture { + #[serde(default)] + input_modalities: Option>, + #[serde(default)] + output_modalities: Option>, +} + pub(super) fn parse_models_response(payload: &str) -> Result> { let parsed: ModelsListResponse = serde_json::from_str(payload).context("Failed to parse model list JSON")?; @@ -3342,18 +3636,250 @@ fn parse_openrouter_models_response( Ok(models) } -fn publish_provider_lake_snapshot(cache: &ProviderCatalogCache) { - // Publish fresh *and* stale/prior rows so pickers keep live catalog coverage - // after TTL expiry or a failed refresh (#4139). An empty cache publishes - // nothing: it must not erase a provider-scoped layer populated by another - // refresh path. - let offerings = cache.all_visible_offerings(now_unix()); - if !offerings.is_empty() { - crate::provider_lake::set_live_snapshot( - CatalogSnapshot { offerings }, - crate::provider_lake::LiveSource::PerProvider, - ); +/// Parse Baseten's authenticated Model APIs catalog without inferring facts +/// from an identically named model on another provider. +fn parse_baseten_models_response( + payload: &str, +) -> Result, CatalogRefreshError> { + let parsed: BasetenModelsResponse = + serde_json::from_str(payload).map_err(|_| CatalogRefreshError::InvalidResponse)?; + let mut seen = std::collections::HashSet::new(); + let mut models = Vec::with_capacity(parsed.data.len()); + for mut item in parsed.data { + item.id = item.id.trim().to_string(); + if item.id.is_empty() || !seen.insert(item.id.clone()) { + return Err(CatalogRefreshError::InvalidResponse); + } + models.push(item); } + Ok(models) +} + +fn catalog_number_f64(value: Option<&CatalogNumber>) -> Result, CatalogRefreshError> { + let Some(value) = value else { + return Ok(None); + }; + let parsed = match value { + CatalogNumber::Number(number) => number.as_f64(), + CatalogNumber::Text(text) => text.trim().parse::().ok(), + } + .filter(|number| number.is_finite() && *number >= 0.0) + .ok_or(CatalogRefreshError::InvalidResponse)?; + Ok(Some(parsed)) +} + +fn catalog_number_u64(value: Option<&CatalogNumber>) -> Result, CatalogRefreshError> { + let Some(number) = catalog_number_f64(value)? else { + return Ok(None); + }; + if number.fract() != 0.0 || number > u64::MAX as f64 { + return Err(CatalogRefreshError::InvalidResponse); + } + Ok(Some(number as u64)) +} + +fn checked_per_token_to_per_million(value: f64) -> Result { + let scaled = value * 1_000_000.0; + (scaled.is_finite() + && (0.0..=codewhale_config::pricing::MAX_PLAUSIBLE_PRICE_PER_MILLION).contains(&scaled)) + .then_some(scaled) + .ok_or(CatalogRefreshError::InvalidResponse) +} + +fn catalog_price_per_million( + value: Option<&CatalogNumber>, +) -> Result, CatalogRefreshError> { + catalog_number_f64(value)? + .map(checked_per_token_to_per_million) + .transpose() +} + +fn feature_matches_any(feature: &str, aliases: &[&str]) -> bool { + let normalized = feature.replace('-', "_"); + aliases.iter().any(|alias| normalized == *alias) +} + +fn baseten_features(item: &BasetenModelItem) -> Option> { + let sources = [ + item.supported_parameters.as_ref(), + item.supported_features.as_ref(), + item.features.as_ref(), + ]; + let mut features = Vec::new(); + let mut published = false; + for source in sources.into_iter().flatten() { + published = true; + for feature in source { + let normalized = feature.trim().to_ascii_lowercase(); + if !normalized.is_empty() && !features.contains(&normalized) { + features.push(normalized); + } + } + } + published.then_some(features) +} + +fn baseten_to_catalog_offering( + item: &BasetenModelItem, + provider: &str, + base_url_fingerprint: &str, + fetched_at: u64, +) -> Result { + use codewhale_config::models_dev::{ModelsDevCost, ModelsDevLimit, ModelsDevModalities}; + + let context = catalog_number_u64( + item.top_provider + .as_ref() + .and_then(|provider| provider.context_length.as_ref()) + .or(item.context_length.as_ref()) + .or(item.context_window.as_ref()) + .or_else(|| { + item.limits + .as_ref() + .and_then(|limits| limits.context.as_ref()) + }), + )?; + let output = catalog_number_u64( + item.top_provider + .as_ref() + .and_then(|provider| provider.max_completion_tokens.as_ref()) + .or(item.max_output_tokens.as_ref()) + .or(item.max_completion_tokens.as_ref()) + .or_else(|| { + item.limits + .as_ref() + .and_then(|limits| limits.output.as_ref()) + }), + )?; + let limit = (context.is_some() || output.is_some()).then_some(ModelsDevLimit { + context, + input: context, + output, + }); + + let cost = if let Some(pricing) = item.pricing.as_ref() { + let cost = ModelsDevCost { + input: catalog_price_per_million(pricing.prompt.as_ref())?, + output: catalog_price_per_million(pricing.completion.as_ref())?, + cache_read: catalog_price_per_million(pricing.input_cache_read.as_ref())?, + cache_write: catalog_price_per_million(pricing.input_cache_write.as_ref())?, + }; + if !codewhale_config::pricing::catalog_cost_is_valid(&cost) { + return Err(CatalogRefreshError::InvalidResponse); + } + (cost.input.is_some() + || cost.output.is_some() + || cost.cache_read.is_some() + || cost.cache_write.is_some()) + .then_some(cost) + } else { + None + }; + + let features = baseten_features(item); + let has_feature = |needles: &[&str]| { + features.as_ref().is_some_and(|features| { + features + .iter() + .any(|feature| feature_matches_any(feature, needles)) + }) + }; + let mut input_modalities = item + .architecture + .as_ref() + .and_then(|architecture| architecture.input_modalities.clone()) + .or_else(|| item.input_modalities.clone()); + let mut output_modalities = item + .architecture + .as_ref() + .and_then(|architecture| architecture.output_modalities.clone()) + .or_else(|| item.output_modalities.clone()); + if input_modalities.is_none() { + let supports_vision = has_feature(&["vision", "image", "image_input"]); + let supports_audio = has_feature(&["audio", "audio_input"]); + if supports_vision || supports_audio { + let mut derived = vec!["text".to_string()]; + if supports_vision { + derived.push("image".to_string()); + } + if supports_audio { + derived.push("audio".to_string()); + } + input_modalities = Some(derived); + output_modalities.get_or_insert_with(|| vec!["text".to_string()]); + } + } + let modalities = if input_modalities.is_some() || output_modalities.is_some() { + Some(ModelsDevModalities { + input: input_modalities.unwrap_or_default(), + output: output_modalities.unwrap_or_default(), + }) + } else { + None + }; + let attachment = modalities.as_ref().map(|modalities| { + modalities + .input + .iter() + .any(|modality| !modality.eq_ignore_ascii_case("text") && !modality.trim().is_empty()) + }); + + let feature_support = |needles: &[&str]| { + features.as_ref().map(|features| { + features + .iter() + .any(|feature| feature_matches_any(feature, needles)) + }) + }; + let reasoning = item + .reasoning + .or(item.supports_reasoning) + .or_else(|| feature_support(&["reasoning", "include_reasoning"])); + // Baseten's current Model APIs contract states every catalog model supports + // tool calling and structured outputs. Explicit upstream booleans still + // win if the endpoint publishes a narrower model-specific fact. + // Baseten's Model APIs contract applies these two capabilities to every + // catalog model. `supported_features` is additive and may list only + // model-variable facts such as `reasoning` or `vision`; absence from that + // list is therefore not an explicit false. Only an upstream boolean may + // narrow the universal contract for a specific row. + let tool_call = item.supports_tools.or(Some(true)); + let structured_output = item.supports_structured_output.or(Some(true)); + + Ok(CatalogOffering { + provider: provider.to_string(), + wire_model_id: item.id.clone(), + canonical_model: None, + endpoint_key: "chat".to_string(), + default_for_provider: item + .id + .eq_ignore_ascii_case(codewhale_config::BASETEN_DEFAULT_MODEL), + family: None, + limit, + cost, + modalities, + attachment, + reasoning, + tool_call, + structured_output, + reasoning_options: item.reasoning_options.clone(), + source: CatalogSource::Live { + base_url_fingerprint: base_url_fingerprint.to_string(), + fetched_at, + }, + }) +} + +fn publish_provider_lake_scope(cache: &ProviderCatalogCache, provider: &str, fingerprint: &str) { + // Publish fresh *and* stale/prior rows so pickers keep live catalog coverage + // after TTL expiry or a failed refresh (#4139). Exact replacement is + // essential: a successful smaller roster must remove upstream-retired ids, + // while a failure preserves the rows already stored in this cache scope. + let offerings = cache + .get(provider, fingerprint) + .map(|entry| entry.offerings.clone()) + .unwrap_or_default(); + crate::provider_lake::replace_provider_live_snapshot(provider, CatalogSnapshot { offerings }); } /// Convert an OpenRouter model item into a [`CatalogOffering`] with live-sourced @@ -3363,7 +3889,7 @@ fn openrouter_to_catalog_offering( provider: &str, base_url_fingerprint: &str, fetched_at: u64, -) -> CatalogOffering { +) -> Result { use codewhale_config::models_dev::{ModelsDevCost, ModelsDevLimit, ModelsDevModalities}; let context_length = item @@ -3387,20 +3913,35 @@ fn openrouter_to_catalog_offering( None }; - let cost = item.pricing.as_ref().map(|p| { + let cost = if let Some(p) = item.pricing.as_ref() { // OpenRouter quotes per-token USD strings; ModelsDevCost is per million. - let parse_price = |s: &Option| -> Option { - s.as_ref() - .and_then(|v| v.parse::().ok()) - .map(|price_per_token| price_per_token * 1_000_000.0) + let parse_price = |value: &Option| -> Result, CatalogRefreshError> { + value + .as_ref() + .map(|value| { + let parsed = value + .trim() + .parse::() + .ok() + .filter(|value| value.is_finite() && *value >= 0.0) + .ok_or(CatalogRefreshError::InvalidResponse)?; + checked_per_token_to_per_million(parsed) + }) + .transpose() }; - ModelsDevCost { - input: parse_price(&p.prompt), - output: parse_price(&p.completion), - cache_read: parse_price(&p.input_cache_read), - cache_write: parse_price(&p.input_cache_write), + let cost = ModelsDevCost { + input: parse_price(&p.prompt)?, + output: parse_price(&p.completion)?, + cache_read: parse_price(&p.input_cache_read)?, + cache_write: parse_price(&p.input_cache_write)?, + }; + if !codewhale_config::pricing::catalog_cost_is_valid(&cost) { + return Err(CatalogRefreshError::InvalidResponse); } - }); + Some(cost) + } else { + None + }; let reasoning = item.supported_parameters.as_ref().map(|params| { params @@ -3441,7 +3982,7 @@ fn openrouter_to_catalog_offering( ModelsDevModalities { input, output } }); - CatalogOffering { + Ok(CatalogOffering { provider: provider.to_string(), wire_model_id: item.id.clone(), canonical_model: None, @@ -3460,7 +4001,7 @@ fn openrouter_to_catalog_offering( base_url_fingerprint: base_url_fingerprint.to_string(), fetched_at, }, - } + }) } pub(super) fn system_to_instructions(system: Option) -> Option { @@ -4245,7 +4786,8 @@ mod tests { }]}"#; let items = parse_openrouter_models_response(payload).expect("parses"); - let priced = openrouter_to_catalog_offering(&items[0], "openrouter", "fp", 42); + let priced = openrouter_to_catalog_offering(&items[0], "openrouter", "fp", 42) + .expect("valid priced row"); let cost = priced.cost.as_ref().expect("pricing row"); assert_eq!(cost.input, Some(3.0)); assert_eq!(cost.output, Some(15.0)); @@ -4265,7 +4807,8 @@ mod tests { // A row without a published write rate stays unknown, not zero, and // fails closed for cache-creation turns. - let unwritten = openrouter_to_catalog_offering(&items[1], "openrouter", "fp", 42); + let unwritten = openrouter_to_catalog_offering(&items[1], "openrouter", "fp", 42) + .expect("valid row without cache-write rate"); assert_eq!( unwritten.cost.as_ref().and_then(|cost| cost.cache_write), None @@ -4280,6 +4823,145 @@ mod tests { ); } + #[test] + fn baseten_catalog_maps_provider_stated_prices_limits_and_features() { + // Exact current Baseten shape: pricing is captured from the official + // baseten-switch repository; Model APIs publishes context_length, + // max_completion_tokens, and supported_features including `vision`. + let payload = r#"{"data":[{ + "id":"deepseek-ai/DeepSeek-V4-Pro", + "context_length":"1048576", + "max_completion_tokens":262144, + "pricing":{ + "prompt":0.0000014, + "completion":"0.0000044", + "input_cache_read":0.00000014 + }, + "supported_features":["reasoning","vision"], + "reasoning_options":[{"type":"toggle"}] + }]}"#; + + let items = parse_baseten_models_response(payload).expect("Baseten catalog"); + let offering = + baseten_to_catalog_offering(&items[0], "baseten", "baseten-fp", 42).expect("row"); + assert_eq!(offering.provider, "baseten"); + assert_eq!( + offering.wire_model_id, + codewhale_config::BASETEN_DEFAULT_MODEL + ); + assert!(offering.default_for_provider); + let limit = offering.limit.expect("published limits"); + assert_eq!(limit.context, Some(1_048_576)); + assert_eq!(limit.input, Some(1_048_576)); + assert_eq!(limit.output, Some(262_144)); + let cost = offering.cost.expect("published pricing"); + assert_eq!(cost.input, Some(1.4)); + assert_eq!(cost.output, Some(4.4)); + assert_eq!(cost.cache_read, Some(0.14)); + assert_eq!(cost.cache_write, None); + assert_eq!(offering.reasoning, Some(true)); + assert_eq!(offering.tool_call, Some(true)); + assert_eq!(offering.structured_output, Some(true)); + assert_eq!(offering.attachment, Some(true)); + let modalities = offering.modalities.expect("vision feature modalities"); + assert_eq!(modalities.input, vec!["text", "image"]); + assert_eq!(modalities.output, vec!["text"]); + assert_eq!(offering.reasoning_options, vec![json!({"type":"toggle"})]); + assert!(matches!(offering.source, CatalogSource::Live { .. })); + } + + #[test] + fn baseten_catalog_rejects_duplicate_ids_and_invalid_known_numbers() { + let duplicates = r#"{"data":[{"id":"same/model"},{"id":"same/model"}]}"#; + assert_eq!( + parse_baseten_models_response(duplicates).unwrap_err(), + CatalogRefreshError::InvalidResponse + ); + + let negative = r#"{"data":[{ + "id":"synthetic/model", + "pricing":{"prompt":-0.000001,"completion":0.000002} + }]}"#; + let items = parse_baseten_models_response(negative).expect("shape parses"); + assert_eq!( + baseten_to_catalog_offering(&items[0], "baseten", "fp", 1).unwrap_err(), + CatalogRefreshError::InvalidResponse + ); + } + + #[test] + fn provider_live_price_parsers_reject_present_bad_rates_but_keep_zero_and_omission() { + for invalid in ["not-a-number", "-0.1", "NaN", "inf", "1e308", "0.100001"] { + let openrouter = json!({ + "data": [{ + "id": "synthetic/openrouter-invalid-price", + "pricing": { "prompt": invalid, "completion": "0.000001" } + }] + }) + .to_string(); + let items = parse_openrouter_models_response(&openrouter).expect("OpenRouter shape"); + assert_eq!( + openrouter_to_catalog_offering(&items[0], "openrouter", "fp", 1).unwrap_err(), + CatalogRefreshError::InvalidResponse, + "OpenRouter must reject {invalid:?}" + ); + + let baseten = json!({ + "data": [{ + "id": "synthetic/baseten-invalid-price", + "pricing": { "prompt": invalid, "completion": "0.000001" } + }] + }) + .to_string(); + let items = parse_baseten_models_response(&baseten).expect("Baseten shape"); + assert_eq!( + baseten_to_catalog_offering(&items[0], "baseten", "fp", 1).unwrap_err(), + CatalogRefreshError::InvalidResponse, + "Baseten must reject {invalid:?}" + ); + } + + let openrouter = parse_openrouter_models_response( + r#"{"data":[{"id":"synthetic/openrouter-free","pricing":{"prompt":"0","completion":"0"}}]}"#, + ) + .expect("OpenRouter zero row"); + let openrouter = openrouter_to_catalog_offering(&openrouter[0], "openrouter", "fp", 1) + .expect("explicit zero is a valid published price"); + let cost = openrouter.cost.expect("published zero cost"); + assert_eq!(cost.input, Some(0.0)); + assert_eq!(cost.output, Some(0.0)); + assert_eq!(cost.cache_read, None); + assert_eq!(cost.cache_write, None); + + let baseten = parse_baseten_models_response( + r#"{"data":[{"id":"synthetic/baseten-free","pricing":{"prompt":"0","completion":0}}]}"#, + ) + .expect("Baseten zero row"); + let baseten = baseten_to_catalog_offering(&baseten[0], "baseten", "fp", 1) + .expect("explicit zero is a valid published price"); + let cost = baseten.cost.expect("published zero cost"); + assert_eq!(cost.input, Some(0.0)); + assert_eq!(cost.output, Some(0.0)); + assert_eq!(cost.cache_read, None); + assert_eq!(cost.cache_write, None); + } + + #[test] + fn baseten_feature_names_require_exact_normalized_aliases() { + let payload = r#"{"data":[{ + "id":"synthetic/text-only", + "supported_features":["revision","pre_reasoning_filter"] + }]}"#; + let items = parse_baseten_models_response(payload).expect("Baseten catalog"); + let offering = + baseten_to_catalog_offering(&items[0], "baseten", "fp", 1).expect("valid row"); + assert_eq!(offering.reasoning, Some(false)); + assert_eq!(offering.modalities, None); + assert_eq!(offering.attachment, None); + assert_eq!(offering.tool_call, Some(true)); + assert_eq!(offering.structured_output, Some(true)); + } + fn test_tool(name: &str) -> Tool { Tool { tool_type: None, @@ -7776,6 +8458,112 @@ mod tests { ); } + #[tokio::test] + async fn incomplete_translation_keeps_exact_route_and_usage_before_rejection() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "msg_partial", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "Parcial"}], + // A provider-returned alias must not replace the admitted + // route/model in the frozen cost receipt. + "model": "provider-alias-after-dispatch", + "stop_reason": "max_tokens", + "stop_sequence": null, + "usage": {"input_tokens": 7, "output_tokens": 2} + }))) + .expect(1) + .mount(&server) + .await; + + let client = deepseek_anthropic_client(&server); + let response = client + .translate_with_usage("Hello", "deepseek-chat", "Spanish") + .await + .expect("decoded provider response retains its receipt"); + + assert!( + response.translated.is_err(), + "partial text must be rejected" + ); + let usage = response.usage.expect("provider-reported usage"); + assert_eq!(usage.input_tokens, 7); + assert_eq!(usage.output_tokens, 2); + assert_eq!(response.route.provider, ApiProvider::DeepseekAnthropic); + assert_eq!(response.route.model, "deepseek-chat"); + assert_eq!(response.route.provider_identity, "deepseek-anthropic"); + assert!(response.route.endpoint_fingerprint.is_some()); + } + + #[tokio::test] + async fn chat_translation_without_usage_keeps_unreceipted_success_outcome() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "chatcmpl-no-usage", + "model": "deepseek-chat", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "Hola"}, + "finish_reason": "stop" + }] + }))) + .expect(1) + .mount(&server) + .await; + + let client = deepseek_request_boundary_client("https://api.deepseek.com/v1", server.uri()); + let response = client + .translate_with_usage("Hello", "deepseek-chat", "Spanish") + .await + .expect("provider success must retain its frozen route"); + assert_eq!( + response + .translated + .expect("useful output remains deliverable"), + "Hola" + ); + assert_eq!(response.usage, None, "must not mint a priced-zero receipt"); + assert_eq!(response.route.provider, ApiProvider::Deepseek); + assert_eq!( + response.route.model, + wire_model_for_provider_route( + ApiProvider::Deepseek, + "https://api.deepseek.com/v1", + "deepseek-chat" + ) + ); + } + + #[tokio::test] + async fn chat_translation_http_error_is_not_a_provider_success_outcome() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(429).set_body_json(json!({ + "error": {"message": "rate limited"} + }))) + .mount(&server) + .await; + + let client = deepseek_request_boundary_client("https://api.deepseek.com/v1", server.uri()); + let error = match client + .translate_with_usage("Hello", "deepseek-chat", "Spanish") + .await + { + Ok(_) => panic!("HTTP failure must not become a provider-success receipt"), + Err(error) => error, + }; + let display = error.to_string(); + assert!( + display.to_ascii_lowercase().contains("rate limit"), + "{display}" + ); + assert!(!display.contains("chatcmpl"), "{display}"); + } + #[tokio::test] async fn deepseek_anthropic_health_check_skips_models_probe() { let server = MockServer::start().await; @@ -9822,6 +10610,31 @@ mod tests { .expect("openrouter client") } + fn baseten_client_for(server: &MockServer) -> DeepSeekClient { + baseten_client_for_identity(server, codewhale_config::BASETEN_TEMPLATE_ID) + } + + fn baseten_client_for_identity(server: &MockServer, identity: &str) -> DeepSeekClient { + let _ = rustls::crypto::ring::default_provider().install_default(); + let mut providers = ProvidersConfig::default(); + providers.custom.insert( + identity.to_string(), + ProviderConfig { + kind: Some("openai-compatible".to_string()), + api_key: Some("test-baseten-key".to_string()), + base_url: Some(format!("{}/v1", server.uri())), + model: Some(codewhale_config::BASETEN_DEFAULT_MODEL.to_string()), + ..ProviderConfig::default() + }, + ); + DeepSeekClient::new(&Config { + provider: Some(identity.to_string()), + providers: Some(providers), + ..Config::default() + }) + .expect("Baseten client") + } + fn opencode_go_client_for(server: &MockServer) -> DeepSeekClient { let _ = rustls::crypto::ring::default_provider().install_default(); DeepSeekClient::new(&Config { @@ -9896,6 +10709,76 @@ mod tests { .expect("mocked /models success should verify"); } + #[tokio::test] + async fn baseten_live_catalog_keeps_exact_identity_auth_and_metadata() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/models")) + .and(header("authorization", "Bearer test-baseten-key")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "data": [{ + "id": codewhale_config::BASETEN_DEFAULT_MODEL, + "context_length": 1_048_576, + "max_completion_tokens": 262_144, + "pricing": { + "prompt": 0.0000014, + "completion": 0.0000044, + "input_cache_read": 0.00000014 + }, + "supported_features": ["reasoning", "tools", "structured_outputs", "vision"] + }] + }))) + .mount(&server) + .await; + + let client = baseten_client_for(&server); + assert_eq!(client.catalog_provider_id(), "baseten"); + let delta = client.fetch_catalog_delta().await.expect("Baseten delta"); + assert_eq!(delta.provider, "baseten"); + assert_eq!(delta.offerings.len(), 1); + let offering = &delta.offerings[0]; + assert_eq!( + offering.wire_model_id, + codewhale_config::BASETEN_DEFAULT_MODEL + ); + assert_eq!( + offering.limit.as_ref().and_then(|limit| limit.context), + Some(1_048_576) + ); + assert_eq!( + offering.cost.as_ref().and_then(|cost| cost.input), + Some(1.4) + ); + assert_eq!(offering.tool_call, Some(true)); + assert_eq!(offering.structured_output, Some(true)); + + let alias = baseten_client_for_identity(&server, "base-ten"); + assert_eq!(alias.catalog_provider_id(), "base-ten"); + assert_eq!( + alias.catalog_setup_template_id(), + Some(codewhale_config::BASETEN_TEMPLATE_ID) + ); + let alias_delta = alias + .fetch_catalog_delta() + .await + .expect("Baseten alias delta"); + assert_eq!(alias_delta.provider, "base-ten"); + assert!( + alias_delta + .offerings + .iter() + .all(|row| row.provider == "base-ten"), + "schema aliases must preserve exact catalog ownership" + ); + assert_eq!(offering.attachment, Some(true)); + assert!( + offering + .modalities + .as_ref() + .is_some_and(|modalities| modalities.input.iter().any(|value| value == "image")) + ); + } + #[tokio::test] async fn verify_provider_api_key_returns_status_and_unicode_body_without_panic() { let server = MockServer::start().await; @@ -10173,6 +11056,39 @@ mod tests { ); } + #[tokio::test] + async fn fetch_catalog_delta_rejects_oversized_bodies_and_rosters() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/models")) + .respond_with(ResponseTemplate::new(200).set_body_raw( + "x".repeat(PROVIDER_CATALOG_MAX_RESPONSE_BYTES + 1), + "application/json", + )) + .mount(&server) + .await; + assert_eq!( + openrouter_client_for(&server) + .fetch_catalog_delta() + .await + .expect_err("oversized body"), + CatalogRefreshError::InvalidResponse + ); + + let server = MockServer::start().await; + let rows: Vec<_> = (0..=PROVIDER_CATALOG_MAX_ROWS) + .map(|index| json!({"id": format!("synthetic-model-{index}")})) + .collect(); + mount_models_json(&server, 200, json!({"data": rows})).await; + assert_eq!( + openrouter_client_for(&server) + .fetch_catalog_delta() + .await + .expect_err("oversized roster"), + CatalogRefreshError::InvalidResponse + ); + } + #[tokio::test] async fn refresh_catalog_cache_records_success_then_preserves_rows_on_failure() { // First refresh succeeds and caches live rows. @@ -10224,6 +11140,111 @@ mod tests { ); } + #[tokio::test] + async fn invalid_live_prices_fail_refresh_and_preserve_each_provider_last_known_good() { + let openrouter_server = MockServer::start().await; + mount_models_json( + &openrouter_server, + 200, + json!({"data": [{ + "id": "synthetic/openrouter-priced", + "pricing": {"prompt": "0.000001", "completion": "0.000002"} + }]}), + ) + .await; + let openrouter = openrouter_client_for(&openrouter_server); + let mut openrouter_cache = ProviderCatalogCache::new(); + assert_eq!( + openrouter + .refresh_catalog_cache(&mut openrouter_cache, 3_600) + .await, + CatalogStatus::Fresh + ); + let openrouter_fp = base_url_fingerprint(&openrouter_server.uri()); + let openrouter_lkg = openrouter_cache + .get("openrouter", &openrouter_fp) + .expect("OpenRouter LKG") + .offerings + .clone(); + + openrouter_server.reset().await; + mount_models_json( + &openrouter_server, + 200, + json!({"data": [{ + "id": "synthetic/openrouter-priced", + "pricing": {"prompt": "1e308", "completion": "0.000002"} + }]}), + ) + .await; + assert!(matches!( + openrouter + .refresh_catalog_cache(&mut openrouter_cache, 3_600) + .await, + CatalogStatus::Failed { + reason: CatalogRefreshError::InvalidResponse + } + )); + assert_eq!( + openrouter_cache + .get("openrouter", &openrouter_fp) + .expect("preserved OpenRouter LKG") + .offerings, + openrouter_lkg + ); + + let baseten_server = MockServer::start().await; + mount_models_json( + &baseten_server, + 200, + json!({"data": [{ + "id": "synthetic/baseten-priced", + "pricing": {"prompt": "0.000001", "completion": "0.000002"} + }]}), + ) + .await; + let baseten = baseten_client_for(&baseten_server); + let mut baseten_cache = ProviderCatalogCache::new(); + assert_eq!( + baseten + .refresh_catalog_cache(&mut baseten_cache, 3_600) + .await, + CatalogStatus::Fresh + ); + let baseten_fp = base_url_fingerprint(&format!("{}/v1", baseten_server.uri())); + let baseten_lkg = baseten_cache + .get(codewhale_config::BASETEN_TEMPLATE_ID, &baseten_fp) + .expect("Baseten LKG") + .offerings + .clone(); + + baseten_server.reset().await; + mount_models_json( + &baseten_server, + 200, + json!({"data": [{ + "id": "synthetic/baseten-priced", + "pricing": {"prompt": "0.100001", "completion": "0.000002"} + }]}), + ) + .await; + assert!(matches!( + baseten + .refresh_catalog_cache(&mut baseten_cache, 3_600) + .await, + CatalogStatus::Failed { + reason: CatalogRefreshError::InvalidResponse + } + )); + assert_eq!( + baseten_cache + .get(codewhale_config::BASETEN_TEMPLATE_ID, &baseten_fp) + .expect("preserved Baseten LKG") + .offerings, + baseten_lkg + ); + } + #[tokio::test] async fn live_catalog_is_scoped_by_base_url_fingerprint() { // Same provider, two different base URLs -> two distinct cache scopes. diff --git a/crates/tui/src/core/engine.rs b/crates/tui/src/core/engine.rs index 9e87270d5c..c848cfa556 100644 --- a/crates/tui/src/core/engine.rs +++ b/crates/tui/src/core/engine.rs @@ -1870,6 +1870,8 @@ impl Engine { .tx_event .send(Event::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status, error, tool_catalog: None, @@ -2427,6 +2429,7 @@ impl Engine { self.current_mode, route, self.config.compaction.clone(), + crate::cost_status::RuntimeUsageBatch::default(), self.config.goal_objective.clone(), self.config.goal_token_budget, self.config.goal_status, @@ -2488,6 +2491,7 @@ impl Engine { mode, route, compaction, + initial_routed_usage, goal_objective, goal_token_budget, goal_status, @@ -2510,6 +2514,7 @@ impl Engine { mode, *route, *compaction, + *initial_routed_usage, goal_objective, goal_token_budget, goal_status, @@ -2614,6 +2619,7 @@ impl Engine { self.current_mode, route, self.config.compaction.clone(), + crate::cost_status::RuntimeUsageBatch::default(), goal_snapshot.objective, goal_snapshot.token_budget, GoalStatus::Active, @@ -3155,6 +3161,7 @@ impl Engine { mode, route, self.config.compaction.clone(), + crate::cost_status::RuntimeUsageBatch::default(), self.config.goal_objective.clone(), self.config.goal_token_budget, self.config.goal_status, @@ -3803,6 +3810,7 @@ impl Engine { self.current_mode, route, self.config.compaction.clone(), + crate::cost_status::RuntimeUsageBatch::default(), self.config.goal_objective.clone(), self.config.goal_token_budget, self.config.goal_status, @@ -3925,6 +3933,8 @@ impl Engine { .tx_event .send(Event::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Failed, error: Some(message.clone()), tool_catalog: None, @@ -4570,6 +4580,7 @@ impl Engine { mode: AppMode, route: ResolvedRuntimeRoute, compaction: CompactionConfig, + initial_routed_usage: crate::cost_status::RuntimeUsageBatch, goal_objective: Option, goal_token_budget: Option, goal_status: GoalStatus, @@ -4590,6 +4601,7 @@ impl Engine { let mut goal_objective = goal_objective; let mut goal_token_budget = goal_token_budget; let mut goal_status = goal_status; + let initial_usage_owner = compaction.runtime_cost_owner.clone(); // A literal natural-language `/goal` declaration is control-plane // intent, not a suggestion that each provider may acknowledge or @@ -4657,6 +4669,12 @@ impl Engine { let dispatched_product = crate::route_billing::capture_product(&route.config, effective_provider); if let Err(err) = self.install_resolved_runtime_route(route) { + let cost_scope = crate::cost_status::scope_token(); + crate::cost_status::report_runtime_usage_batch( + cost_scope, + initial_usage_owner.as_deref(), + &initial_routed_usage, + ); let _ = self .tx_event .send(Event::error(ErrorEnvelope::fatal_auth(format!( @@ -4756,8 +4774,8 @@ impl Engine { model: model.clone(), auto_model, receipt: route_receipt, - // A start is not a dispatch. The billing envelope is attached - // below, on the route held for the wire boundary only. + // A start is not an application dispatch. The billing envelope is + // attached below, then stamped at the pre-permit admission boundary. billing: None, // The classification receipt, by contrast, is frozen here at the // client-freeze boundary and is readable from `TurnStarted` on. @@ -4778,10 +4796,14 @@ impl Engine { ) .map(str::to_string), endpoint_fingerprint: route_base_url.and_then(crate::cost_status::endpoint_fingerprint), + // A live rate is not evidence at turn creation. `turn_loop` + // freezes it from the exact fresh cache scope at CodeWhale's + // pre-permit application-dispatch boundary. + provider_live_pricing: None, // Classified from this turn's own frozen receipt, not from a // second ambient `for_route` read. Both halves of the route then // answer from the same captured endpoint + credential product, so - // the envelope stamped on the wire and the receipt carried on + // the application-dispatch envelope and the receipt carried on // `TurnRoute` cannot disagree about how this turn bills. billing_mode: crate::route_billing::for_dispatched_receipt( crate::route_billing::DispatchedReceipt { @@ -4792,8 +4814,8 @@ impl Engine { }, ) .into(), - // Provisional. Replaced with the true wire-boundary instant - // when `run_turn` emits `Event::RouteDispatched`. + // Provisional. Replaced with the pre-permit application-dispatch + // instant when `run_turn` emits `Event::RouteDispatched`. dispatched_at: turn_started_at, }; turn.pending_route = Some(TurnRoute { @@ -4813,30 +4835,57 @@ impl Engine { }) .await; + // Auto's classifier completed before this parent turn was admitted. + // Bind its exact routed records to the now-accepted turn: total tokens + // and model-call telemetry include the work, while parent_route_usage + // remains untouched so the parent quote can never price it. + turn.add_routed_usages( + initial_routed_usage + .records + .iter() + .map(|record| &record.usage.usage), + ); + // Exact missing-usage records are persisted route-aware by the runtime + // sink. TurnComplete carries only any count whose exact route record + // was truncated, otherwise the terminal merge would count the same + // provider response twice and misclassify subscription/local calls. + let residual_dropped_records = initial_routed_usage.dropped_records.saturating_sub( + u64::try_from(initial_routed_usage.drop_records.len()).unwrap_or(u64::MAX), + ); + turn.add_routed_usage_dropped_records(residual_dropped_records); + let initial_cost_scope = crate::cost_status::scope_token(); + for record in &initial_routed_usage.records { + crate::cost_status::report_effective_route_for_runtime( + initial_cost_scope, + initial_usage_owner.as_deref(), + &record.source_id, + &record.usage.route, + &record.usage.usage, + ); + let _ = self + .tx_event + .send(Event::RoutedTurnUsage { + usage: record.usage.usage.clone(), + duration_ms: 0, + first_token_ms: None, + request_ms: None, + }) + .await; + } + for record in &initial_routed_usage.drop_records { + crate::cost_status::report_unreceipted_provider_success( + initial_cost_scope, + initial_usage_owner.as_deref(), + &record.source_id, + &record.route, + ); + } + // Apply the host-resolved route budget before building the request. // The model, limits, and compaction policy arrive in one operation so // no provider request can observe a partially updated route. self.active_route_limits = route_limits; self.config.compaction = compaction; - // Headless/runtime hosts supply their durable turn owner. Interactive - // turns historically supplied none, leaving a detached child with - // only the soon-to-be-sealed mailbox. Give this turn an owner whose - // sink folds into the existing session cost pool; cloned child leases - // keep it live beyond TurnComplete without reopening the mailbox. - let interactive_runtime_cost_owner = if self.config.terminal_chrome_enabled - && self.config.compaction.runtime_cost_owner.is_none() - { - let owner = format!("interactive:{}:{}", self.session.id, turn.id); - crate::cost_status::register_interactive_runtime_usage_sink( - &owner, - crate::cost_status::scope_token(), - ); - self.config.compaction.runtime_cost_owner = Some(owner.clone()); - Some(owner) - } else { - None - }; - // Snapshot the workspace BEFORE we touch a single tool. Run the git // work on the blocking pool so the async runtime stays responsive; // failure is non-fatal (the helper logs at WARN). @@ -4887,6 +4936,8 @@ impl Engine { .tx_event .send(Event::TurnComplete { usage: turn.usage.clone(), + parent_route_usage: turn.parent_route_usage.clone(), + routed_usage_dropped_records: turn.routed_usage_dropped_records, status: TurnOutcomeStatus::Failed, error: Some(message.clone()), tool_catalog: None, @@ -4905,6 +4956,27 @@ impl Engine { return outcome; } + // Headless/runtime hosts supply their durable turn owner. Interactive + // turns historically supplied none, leaving a detached child with + // only the soon-to-be-sealed mailbox. Install this turn-local sink + // only after every pre-dispatch failure return, and retire/clear it at + // settlement so the next turn always receives a fresh owner. + let interactive_runtime_cost_owner = if self.config.terminal_chrome_enabled + && self.config.compaction.runtime_cost_owner.is_none() + { + let owner = format!("interactive:{}:{}", self.session.id, turn.id); + crate::cost_status::register_persistent_interactive_runtime_usage_sink( + &owner, + crate::cost_status::scope_token(), + &self.session.id, + &turn.id, + ); + self.config.compaction.runtime_cost_owner = Some(owner.clone()); + Some(owner) + } else { + None + }; + let previous_goal_objective = self.config.goal_objective.clone(); let previous_goal_token_budget = self.config.goal_token_budget; let previous_goal_status = self.config.goal_status; @@ -5102,8 +5174,25 @@ impl Engine { barrier.cancel_and_flush().await; } } + // The advisor is dispatched after TurnComplete, but its usage still + // belongs to this originating turn. Acquire the owner lease before an + // interactive owner is marked terminal so a late provider response + // retains its exact sink instead of falling into a later session. + let advisor_usage_context = (self.config.advisor_config.enabled + && status == TurnOutcomeStatus::Completed + && self.deepseek_client.is_some()) + .then(|| { + crate::tools::subagent::advisor::AdvisorUsageContext::capture( + self.config.compaction.runtime_cost_owner.as_deref(), + ) + }); if let Some(owner) = interactive_runtime_cost_owner.as_deref() { crate::cost_status::finish_runtime_usage_owner(owner); + // This owner is turn-local. Leaving it in the reusable engine + // config makes the next interactive turn skip registration and + // route background usage into a retired sink/journal. Host-owned + // runtime turn ids never enter this branch and remain untouched. + self.config.compaction.runtime_cost_owner = None; } // Emit turn complete event — after all post-turn bookkeeping so @@ -5116,6 +5205,8 @@ impl Engine { .tx_event .send(Event::TurnComplete { usage: turn.usage, + parent_route_usage: turn.parent_route_usage, + routed_usage_dropped_records: turn.routed_usage_dropped_records, status, error: error.clone(), tool_catalog: tool_catalog_for_event, @@ -5161,6 +5252,7 @@ impl Engine { if self.config.advisor_config.enabled && matches!(status, TurnOutcomeStatus::Completed) && let Some(client) = self.deepseek_client.clone() + && let Some(usage_context) = advisor_usage_context { // Lazily create the shared emission guard on first use. let guard = self @@ -5174,6 +5266,10 @@ impl Engine { let advisor_messages: Vec = self.session.messages.to_vec(); let advisor_config = self.config.advisor_config.clone(); + // This clone is frozen before the detached task starts and keeps + // every configured provider route available for an explicit + // cross-provider advisor model without consulting later UI state. + let advisor_route_config = self.api_config.clone(); let advisor_model = self.session.model.clone(); let advisor_tx = self.tx_event.clone(); let advisor_turn_id = turn.id.clone(); @@ -5187,7 +5283,9 @@ impl Engine { advisor_messages, advisor_config, client, + advisor_route_config, advisor_model, + usage_context, guard, advisor_tx, ) @@ -5253,6 +5351,8 @@ impl Engine { .tx_event .send(Event::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Interrupted, error: None, tool_catalog: None, @@ -5296,6 +5396,8 @@ impl Engine { .tx_event .send(Event::TurnComplete { usage: zero_usage, + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Failed, error: Some(message), tool_catalog: None, @@ -5338,6 +5440,8 @@ impl Engine { .tx_event .send(Event::TurnComplete { usage: zero_usage, + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Interrupted, error: None, tool_catalog: None, @@ -5365,6 +5469,8 @@ impl Engine { .tx_event .send(Event::TurnComplete { usage: zero_usage, + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Interrupted, error: None, tool_catalog: None, @@ -5429,6 +5535,8 @@ impl Engine { .tx_event .send(Event::TurnComplete { usage: zero_usage, + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: turn_status, error: turn_error, tool_catalog: None, @@ -5454,6 +5562,8 @@ impl Engine { .tx_event .send(Event::TurnComplete { usage: zero_usage, + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Failed, error: Some(message), tool_catalog: None, @@ -5511,6 +5621,8 @@ impl Engine { .tx_event .send(Event::TurnComplete { usage: zero_usage, + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status, error, tool_catalog: None, diff --git a/crates/tui/src/core/engine/preview/tests.rs b/crates/tui/src/core/engine/preview/tests.rs index 3d61944b9d..8c037fab3a 100644 --- a/crates/tui/src/core/engine/preview/tests.rs +++ b/crates/tui/src/core/engine/preview/tests.rs @@ -963,6 +963,7 @@ async fn assert_preview_matches_first_wire_body( AppMode::Agent, production_route, compaction, + crate::cost_status::RuntimeUsageBatch::default(), goal_objective, None, goal_status, @@ -1927,6 +1928,7 @@ async fn provider_reported_usage_is_unavailable_until_a_response_reports_it() { AppMode::Agent, production_route, compaction, + crate::cost_status::RuntimeUsageBatch::default(), None, None, GoalStatus::Active, diff --git a/crates/tui/src/core/engine/tests.rs b/crates/tui/src/core/engine/tests.rs index 0bea822d27..9913275a70 100644 --- a/crates/tui/src/core/engine/tests.rs +++ b/crates/tui/src/core/engine/tests.rs @@ -659,6 +659,7 @@ async fn exact_turn_snapshot_restores_custom_endpoint_and_turn_receipt_after_bui .expect("resolve exact custom route"), ), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -1029,6 +1030,7 @@ async fn goal_continuation_preserves_goal_and_resolves_updated_authoritative_rou mode: AppMode::Agent, route: resolved_route_for_test(&config, "local-model"), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: Some("keep going".to_string()), goal_token_budget: Some(50_000), goal_status: crate::tools::goal::GoalStatus::Active, @@ -1307,6 +1309,7 @@ async fn saturated_mailbox_does_not_deadlock_goal_continuation_self_dispatch() { mode: AppMode::Agent, route: resolved_route_for_test(&config, "local-model"), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: Some("survive a saturated mailbox".to_string()), goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -1434,6 +1437,7 @@ async fn queued_ordinary_turn_does_not_multiply_engine_goal_continuations() { mode: AppMode::Agent, route: resolved_route_for_test(&config, "local-model"), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: Some("coalesce queued goal turns".to_string()), goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -2626,6 +2630,7 @@ async fn cross_turn_token_budget_exhaustion_does_not_pause_goal() { mode: AppMode::Agent, route: resolved_route_for_test(&config, crate::config::DEFAULT_TEXT_MODEL), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: Some("finish within budget".to_string()), goal_token_budget: Some(10), goal_status: crate::tools::goal::GoalStatus::Active, @@ -3082,6 +3087,7 @@ async fn explicit_natural_goal_activates_before_provider_request() { mode: AppMode::Agent, route: resolved_route_for_test(&config, "local-model"), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -3765,6 +3771,7 @@ async fn host_managed_engine_does_not_self_dispatch_goal_continuation() { mode: AppMode::Agent, route: resolved_route_for_test(&config, "local-model"), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: Some("keep going".to_string()), goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -3889,6 +3896,7 @@ async fn host_managed_engine_defers_idle_subagent_completion_to_explicit_turn() mode: AppMode::Agent, route: resolved_route_for_test(&config, "local-model"), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -4094,14 +4102,14 @@ fn subagent_mailbox_keeps_lifecycle_events_reliable() { &MailboxMessage::TokenUsage { agent_id: "agent_a".to_string(), source_id: "response-a".to_string(), - route: crate::cost_status::EffectiveRouteEnvelope::capture( + route: Box::new(crate::cost_status::EffectiveRouteEnvelope::capture( None, ApiProvider::Deepseek, "deepseek", "model", Some(ApiProvider::Deepseek.default_base_url()), chrono::Utc::now(), - ), + )), usage: Usage::default(), } )); @@ -4188,14 +4196,14 @@ fn subagent_mailbox_never_samples_lifecycle_or_usage_events() { &MailboxMessage::TokenUsage { agent_id: "agent_a".to_string(), source_id: "response-a".to_string(), - route: crate::cost_status::EffectiveRouteEnvelope::capture( + route: Box::new(crate::cost_status::EffectiveRouteEnvelope::capture( None, ApiProvider::Deepseek, "deepseek", "model", Some(ApiProvider::Deepseek.default_base_url()), chrono::Utc::now(), - ), + )), usage: Usage::default(), }, start, @@ -5311,6 +5319,7 @@ fn active_goal_message_op( mode: AppMode::Agent, route: resolved_route_for_test(config, "local-model"), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: Some(objective.to_string()), goal_token_budget: token_budget, goal_status: crate::tools::goal::GoalStatus::Active, @@ -5347,6 +5356,7 @@ fn external_user_message_op(content: &str, mode: AppMode, config: &Config) -> Op mode, route: resolved_route_for_test(config, crate::config::DEFAULT_TEXT_MODEL), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -5372,6 +5382,7 @@ fn auto_review_message_op(content: &str, config: &Config) -> Op { mode: AppMode::Agent, route: resolved_route_for_test(config, crate::config::DEFAULT_TEXT_MODEL), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -5648,8 +5659,11 @@ async fn normal_repl_kernel_persists_across_user_turns() { assert_eq!(first_turn.usage.output_tokens, 11); let child_usage_event = { let mut events = handle.rx_event.write().await; + // Kernel child calls carry their own routed cost receipt, so their + // per-call telemetry arrives as `RoutedTurnUsage` rather than the + // parent-route `TurnUsage` receipt. std::iter::from_fn(|| events.try_recv().ok()).find_map(|event| match event { - Event::TurnUsage { usage, .. } + Event::RoutedTurnUsage { usage, .. } if usage.input_tokens == 7 && usage.output_tokens == 11 => { Some(usage) @@ -6003,6 +6017,131 @@ fn deterministic_engine_config(workspace: &Path) -> EngineConfig { } } +#[tokio::test] +async fn initial_routed_usage_is_total_only_emitted_once_and_keeps_parent_route_separate() { + use crate::llm_client::mock::{MockLlmClient, canned}; + + let _cost_scope = crate::cost_status::test_scope(); + let workspace = tempdir().expect("tempdir"); + let parent_usage = Usage { + input_tokens: 11, + output_tokens: 3, + ..Usage::default() + }; + let classifier_usage = Usage { + input_tokens: 7, + output_tokens: 5, + ..Usage::default() + }; + let mock = std::sync::Arc::new(MockLlmClient::new(vec![vec![ + canned::message_start("parent-response"), + canned::text_block_start(0), + canned::text_delta(0, "done"), + canned::block_stop(0), + canned::message_delta("end_turn", Some(parent_usage.clone())), + canned::message_stop(), + ]])); + let client: crate::core::model_client::SharedModelClient = mock; + let api_config = Config::default(); + let (engine, handle) = Engine::new_with_model_client( + deterministic_engine_config(workspace.path()), + &api_config, + client, + ); + let task = tokio::spawn(engine.run()); + + let mut op = external_user_message_op("account for classifier", AppMode::Agent, &api_config); + let Op::SendMessage { + initial_routed_usage, + .. + } = &mut op + else { + unreachable!("external_user_message_op always builds SendMessage"); + }; + let mut missing_usage_route = crate::cost_status::EffectiveRouteEnvelope::capture( + None, + ApiProvider::Openai, + "openai", + "classifier-model", + Some(ApiProvider::Openai.default_base_url()), + chrono::Utc::now(), + ); + missing_usage_route.billing_mode = crate::cost_status::RouteBillingMode::Metered; + **initial_routed_usage = crate::cost_status::RuntimeUsageBatch { + records: vec![crate::cost_status::RuntimeUsageRecord { + source_id: "auto-router:engine-fixture".to_string(), + usage: crate::cost_status::EffectiveRouteUsage { + route: crate::cost_status::EffectiveRouteEnvelope::capture( + None, + ApiProvider::Openai, + "openai", + "classifier-model", + Some(ApiProvider::Openai.default_base_url()), + chrono::Utc::now(), + ), + usage: classifier_usage.clone(), + }, + }], + drop_records: vec![crate::cost_status::RuntimeUsageDropRecord { + source_id: "auto-router:engine-fixture:missing-usage".to_string(), + route: missing_usage_route, + }], + // One exact route-aware missing receipt plus two residual gaps whose + // route identity was truncated upstream. + dropped_records: 3, + }; + handle.send(op).await.expect("send routed-usage turn"); + + let mut routed_events = Vec::new(); + let mut rx = handle.rx_event.write().await; + let (total_usage, terminal_parent_usage, dropped_records) = loop { + let event = tokio::time::timeout(model_turn_event_timeout(), rx.recv()) + .await + .expect("timed out waiting for routed-usage turn") + .expect("engine event stream closed"); + match event { + Event::RoutedTurnUsage { usage, .. } => routed_events.push(usage), + Event::TurnComplete { + usage, + parent_route_usage, + routed_usage_dropped_records, + status, + error, + .. + } => { + assert_eq!(status, TurnOutcomeStatus::Completed, "{error:?}"); + break (usage, parent_route_usage, routed_usage_dropped_records); + } + _ => {} + } + }; + drop(rx); + + assert_eq!(routed_events, vec![classifier_usage.clone()]); + assert_eq!(terminal_parent_usage, parent_usage); + assert_eq!(total_usage.input_tokens, 18); + assert_eq!(total_usage.output_tokens, 8); + assert_eq!(dropped_records, 2); + let initial_cost = crate::cost_status::drain(); + assert!( + initial_cost.usage_source_fingerprints.contains( + &crate::cost_status::usage_source_fingerprint( + "auto-router:engine-fixture:missing-usage" + ) + ), + "exact missing-usage response identity was not settled" + ); + assert!( + initial_cost + .unpriced_reasons + .contains("provider_success_missing_usage"), + "metered missing-usage route was not marked incomplete" + ); + + handle.send(Op::Shutdown).await.expect("shutdown engine"); + task.await.expect("engine task"); +} + #[tokio::test] async fn isolated_runtime_chat_provider_request_contains_no_host_context_or_tools() { use crate::llm_client::mock::{MockLlmClient, canned}; @@ -7259,7 +7398,12 @@ async fn collect_guardian_journey_with_receipts( "duplicate tool result" ); } - Event::TurnUsage { usage, .. } => usage_events.push(usage), + // Guardian consults carry their own routed receipt; both the + // parent-route and routed per-call telemetry count as reaching + // the cost UI. + Event::TurnUsage { usage, .. } | Event::RoutedTurnUsage { usage, .. } => { + usage_events.push(usage); + } Event::ToolGateDecision { tool_id, gate, @@ -11038,6 +11182,7 @@ async fn operate_model_shell_uses_normal_approval_and_workspace_sandbox() { mode: AppMode::Operate, route: resolved_route_for_test(&api_config, crate::config::DEFAULT_TEXT_MODEL), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -11179,6 +11324,7 @@ async fn full_access_subagent_handoff_keeps_model_shell_free_of_approval_prompts mode: AppMode::Agent, route: resolved_route_for_test(&api_config, crate::config::DEFAULT_TEXT_MODEL), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -11315,6 +11461,7 @@ async fn assert_full_access_model_tool_batch_is_blocked( mode: AppMode::Agent, route: resolved_route_for_test(&api_config, crate::config::DEFAULT_TEXT_MODEL), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -11520,6 +11667,7 @@ async fn assert_full_access_model_tool_batch_runs( mode: AppMode::Agent, route: resolved_route_for_test(&api_config, crate::config::DEFAULT_TEXT_MODEL), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -11796,6 +11944,7 @@ async fn auto_review_auto_resolves_hallucinated_question_without_prompting() { mode: AppMode::Agent, route: resolved_route_for_test(&api_config, crate::config::DEFAULT_TEXT_MODEL), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -11982,6 +12131,7 @@ async fn full_access_permission_allow_cannot_bypass_background_catastrophic_floo mode: AppMode::Agent, route: resolved_route_for_test(&api_config, crate::config::DEFAULT_TEXT_MODEL), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -12122,6 +12272,7 @@ async fn yolo_mode_does_not_prompt_for_background_shell() { mode: AppMode::Yolo, route: resolved_route_for_test(&api_config, crate::config::DEFAULT_TEXT_MODEL), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -12258,6 +12409,7 @@ async fn yolo_mode_executes_publish_like_shell_without_prompt() { mode: AppMode::Yolo, route: resolved_route_for_test(&api_config, crate::config::DEFAULT_TEXT_MODEL), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -12398,6 +12550,7 @@ async fn yolo_mode_does_not_prompt_for_mcp_action() { mode: AppMode::Yolo, route: resolved_route_for_test(&api_config, crate::config::DEFAULT_TEXT_MODEL), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -18756,6 +18909,7 @@ async fn run_headless_turn_with_flaky_network( mode: AppMode::Agent, route: resolved_route_for_test(&config, crate::config::DEFAULT_TEXT_MODEL), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -18881,6 +19035,7 @@ async fn terminal_output_limit_followed_by_stream_error_is_charged_and_not_retri mode: AppMode::Agent, route: resolved_route_for_test(&config, crate::config::DEFAULT_TEXT_MODEL), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -18985,6 +19140,7 @@ async fn midstream_error_frame_stops_the_stream_and_drops_trailing_deltas() { mode: AppMode::Agent, route: resolved_route_for_test(&config, crate::config::DEFAULT_TEXT_MODEL), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -19232,6 +19388,7 @@ async fn run_interactive_turn_with_flaky_network( mode: AppMode::Agent, route: resolved_route_for_test(&config, crate::config::DEFAULT_TEXT_MODEL), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -19460,6 +19617,7 @@ async fn interactive_thinking_only_drop_preserves_nothing_and_never_claims_it_di mode: AppMode::Agent, route: resolved_route_for_test(&config, crate::config::DEFAULT_TEXT_MODEL), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -19683,6 +19841,7 @@ async fn run_reasoning_only_turn( mode: AppMode::Agent, route: resolved_route_for_test(&config, crate::config::DEFAULT_TEXT_MODEL), compaction: Box::new(CompactionConfig::default()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, diff --git a/crates/tui/src/core/engine/turn_loop.rs b/crates/tui/src/core/engine/turn_loop.rs index 2e3faf1021..cf1278c22d 100644 --- a/crates/tui/src/core/engine/turn_loop.rs +++ b/crates/tui/src/core/engine/turn_loop.rs @@ -519,15 +519,24 @@ impl Engine { context.tool_name ))) .await; + let cost_scope = crate::cost_status::scope_token(); + let review_route = client.effective_route_envelope(client.model(), chrono::Utc::now()); let started = Instant::now(); let review = super::reviewer::consult_reviewer(client, &context_text, &self.cancel_token).await; if let Some(usage) = &review.usage { turn.add_usage(usage); + crate::cost_status::report_effective_route_for_runtime( + cost_scope, + self.config.compaction.runtime_cost_owner.as_deref(), + &format!("auto-review:{}:{tool_id}", turn.id), + &review_route, + usage, + ); if usage_has_reported_data(usage) { let _ = self .tx_event - .send(Event::TurnUsage { + .send(Event::RoutedTurnUsage { usage: usage.clone(), duration_ms: u64::try_from(started.elapsed().as_millis()) .unwrap_or(u64::MAX), @@ -536,6 +545,12 @@ impl Engine { }) .await; } + } else if matches!( + &review.outcome, + super::reviewer::ReviewerOutcome::Unavailable { reason } + if reason == "the reviewer timed out" || reason == "the reviewer request failed" + ) { + turn.add_routed_usage_dropped_records(1); } let decision = review.outcome.audit_decision(); let risk = review.outcome.audit_risk(); @@ -1354,7 +1369,26 @@ impl Engine { .await; if let Some(mut route) = turn.pending_route.take() { if let Some(billing) = route.billing.as_mut() { - billing.dispatched_at = chrono::Utc::now(); + // Freeze the exact provider-live row at CodeWhale's + // pre-permit application-dispatch boundary. This is an + // admission contract, not provider invoice-time evidence; + // a later cancellation/preparation failure has no usage + // and therefore contributes no usage cost. + let dispatched_at = chrono::Utc::now(); + billing.dispatched_at = dispatched_at; + billing.provider_live_pricing = u64::try_from(dispatched_at.timestamp()) + .ok() + .and_then(|dispatched_at_unix| { + billing.endpoint_fingerprint.as_deref().and_then(|fingerprint| { + crate::provider_catalog_live::fresh_provider_live_pricing_quote_at( + route.provider, + &route.provider_identity, + &route.model, + fingerprint, + dispatched_at_unix, + ) + }) + }); } let _ = self .tx_event @@ -1972,8 +2006,7 @@ impl Engine { 1, ) }); - let bridge_usage_handle = - bridge.as_ref().map(crate::rlm::RlmBridge::usage_handle); + let repl_cost_scope = crate::cost_status::scope_token(); let repl_started = Instant::now(); let mut final_result: Option = None; @@ -2111,14 +2144,18 @@ impl Engine { // into the parent turn exactly once, including failures // after a partial fan-out, so `/cost`, goals, and the // final receipt cannot undercount the working kernel. - if let Some(usage_handle) = bridge_usage_handle { - let child_usage = usage_handle.lock().await.clone(); - turn.add_usage(&child_usage); - if usage_has_reported_data(&child_usage) { + if let Some(bridge) = bridge.as_ref() { + let snapshot = bridge.usage_snapshot().await; + turn.add_usage(&snapshot.usage); + let residual_dropped_records = snapshot.dropped_records.saturating_sub( + u64::try_from(snapshot.drop_records.len()).unwrap_or(u64::MAX), + ); + turn.add_routed_usage_dropped_records(residual_dropped_records); + if usage_has_reported_data(&snapshot.usage) { let _ = self .tx_event - .send(Event::TurnUsage { - usage: child_usage, + .send(Event::RoutedTurnUsage { + usage: snapshot.usage.clone(), duration_ms: u64::try_from(repl_started.elapsed().as_millis()) .unwrap_or(u64::MAX), first_token_ms: None, @@ -2126,6 +2163,23 @@ impl Engine { }) .await; } + for record in snapshot.records { + crate::cost_status::report_effective_route_for_runtime( + repl_cost_scope, + self.config.compaction.runtime_cost_owner.as_deref(), + &record.source_id, + &record.usage.route, + &record.usage.usage, + ); + } + for record in snapshot.drop_records { + crate::cost_status::report_unreceipted_provider_success( + repl_cost_scope, + self.config.compaction.runtime_cost_owner.as_deref(), + &record.source_id, + &record.route, + ); + } } if let Some(final_val) = final_result { @@ -2418,6 +2472,7 @@ impl Engine { self.process_tool_results( outcomes, + turn, &mut tool_catalog, &mut active_tool_names, &hook_contexts, @@ -3918,6 +3973,7 @@ impl Engine { async fn process_tool_results( &mut self, outcomes: Vec>, + turn: &mut TurnContext, tool_catalog: &mut Vec, active_tool_names: &mut std::collections::HashSet, hook_contexts: &std::collections::HashMap, @@ -3935,12 +3991,43 @@ impl Engine { let tool_input = outcome.input.clone(); let tool_name_for_ws = outcome.name.clone(); let terminal_status = outcome.terminal.status; + let routed_duration_ms = + u64::try_from(outcome.started_at.elapsed().as_millis()).unwrap_or(u64::MAX); let result = outcome.terminal.into_legacy_result(); if matches!(outcome.name.as_str(), "create_goal" | "update_goal") { goal_tool_ran = true; } match result { Ok(output) => { + let routed_usage = if let Some(metadata) = output.metadata.as_ref() + && let Some(batch) = + crate::cost_status::child_usage_records_from_metadata(metadata) + { + let residual_dropped_records = batch.dropped_records.saturating_sub( + u64::try_from(batch.drop_records.len()).unwrap_or(u64::MAX), + ); + turn.add_routed_usage_dropped_records(residual_dropped_records); + turn.add_routed_usages( + batch.records.iter().map(|record| &record.usage.usage), + ) + } else if let Some(metadata) = output.metadata.as_ref() + && let Some(usage) = crate::cost_status::child_usage_from_metadata(metadata) + { + turn.add_routed_usages(std::iter::once(&usage)) + } else { + Usage::default() + }; + if usage_has_reported_data(&routed_usage) { + let _ = self + .tx_event + .send(Event::RoutedTurnUsage { + usage: routed_usage, + duration_ms: routed_duration_ms, + first_token_ms: None, + request_ms: None, + }) + .await; + } let mut tool_surface_changed = super::tool_catalog::activate_result_dependencies( tool_catalog, diff --git a/crates/tui/src/core/events.rs b/crates/tui/src/core/events.rs index ec6a5a9d55..ef79aba35e 100644 --- a/crates/tui/src/core/events.rs +++ b/crates/tui/src/core/events.rs @@ -48,13 +48,15 @@ pub struct TurnRoute { /// `None` when no concrete client was installed (injected-client engines, /// or a client that failed to construct). pub receipt: Option, - /// Billing evidence for the request that was actually put on the wire. + /// Billing evidence for a request admitted to application dispatch. /// /// `None` at `TurnStarted`: a lifecycle start is not a dispatch, and a - /// route that has not been sent has no billing time, no metering surface, - /// and no endpoint to attest. Populated exactly once, at the wire - /// boundary, and delivered on `RouteDispatched`. Consumers that price a - /// turn must treat `None` as *unknown*, never as a zero-cost turn. + /// route that has not reached admission has no billing time, no metering + /// surface, and no endpoint to attest. Populated exactly once at the + /// pre-permit application-dispatch boundary and delivered on + /// `RouteDispatched`. This does not attest network delivery or a provider + /// invoice-time rate. Consumers that price a turn must treat `None` as + /// *unknown*, never as a zero-cost turn. pub billing: Option, /// Endpoint this turn's client was frozen against, verbatim. /// @@ -92,14 +94,16 @@ pub struct TurnRoute { /// it bill* — a [`crate::route_billing::DispatchedReceipt`]. They must be /// readable from `TurnStarted` onward so a child turn arriving mid-flight /// can be billed against the parent's frozen route. -/// - This envelope is stamped at the **wire** boundary and answers *what was -/// actually put on the wire, when*. A planned-but-unsent route has no -/// metering surface and no dispatch instant, so it must be structurally -/// absent rather than defaulted. +/// - This envelope is stamped at the **pre-permit application-dispatch** +/// boundary and answers *what CodeWhale admitted for provider execution, +/// when*. It does not claim network delivery or provider invoice-time +/// pricing. A merely planned route has no metering surface or dispatch +/// instant, so it must be structurally absent rather than defaulted. #[derive(Debug, Clone, PartialEq, Eq)] pub struct RouteBillingEnvelope { pub billing_surface: Option, pub endpoint_fingerprint: Option, + pub provider_live_pricing: Option, pub billing_mode: crate::cost_status::RouteBillingMode, pub dispatched_at: DateTime, } @@ -117,6 +121,7 @@ impl TurnRoute { model: self.model.clone(), billing_surface: billing.billing_surface.clone(), endpoint_fingerprint: billing.endpoint_fingerprint.clone(), + provider_live_pricing: billing.provider_live_pricing.clone(), billing_mode: billing.billing_mode, dispatched_at: billing.dispatched_at, }) @@ -218,7 +223,7 @@ pub enum Event { turn_id: String, created_at: DateTime, /// Legacy/non-model hosts may still attach a route at start. Model - /// turns emit it separately at the real provider dispatch boundary. + /// turns emit it separately at the application dispatch boundary. route: Option, }, @@ -228,13 +233,23 @@ pub enum Event { snapshot: crate::tool_inspection::ToolInspectionSnapshot, }, - /// Immutable billing route captured immediately before the first provider - /// request, after snapshots and other potentially slow pre-dispatch work. + /// Immutable billing route captured at CodeWhale's pre-permit application + /// dispatch boundary, after request preparation. This is admission-time + /// evidence, not proof of network delivery or provider invoice-time rates. RouteDispatched { turn_id: String, route: TurnRoute }, /// The turn is complete (no more tool calls) TurnComplete { + /// Total usage for session/goal/token metrics, including programmatic + /// child calls performed inline during this turn. usage: Usage, + /// Usage served by the parent turn's frozen route only. Consumers + /// price this under the parent quote and price routed children from + /// their own receipts, avoiding double billing without subtraction. + parent_route_usage: Usage, + /// Provider calls whose execution/usage could not be receipted. + /// Non-zero makes cost coverage explicitly incomplete. + routed_usage_dropped_records: u64, status: TurnOutcomeStatus, error: Option, /// Tool catalog sent with this turn's model request. @@ -266,6 +281,17 @@ pub enum Event { request_ms: Option, }, + /// Usage telemetry for a programmatic provider call whose cost is carried + /// by its own routed receipt rather than the active parent route. TUI + /// consumers fold this into model-call metrics only; `TurnComplete.usage` + /// remains the authoritative total-token reconciliation. + RoutedTurnUsage { + usage: Usage, + duration_ms: u64, + first_token_ms: Option, + request_ms: Option, + }, + /// Runtime goal state changed inside the engine, usually from model-visible /// `create_goal` or `update_goal` tool calls. GoalUpdated { snapshot: GoalSnapshot }, diff --git a/crates/tui/src/core/ops.rs b/crates/tui/src/core/ops.rs index f4264491f8..4ded3b0229 100644 --- a/crates/tui/src/core/ops.rs +++ b/crates/tui/src/core/ops.rs @@ -117,6 +117,11 @@ pub enum Op { /// Compaction policy derived from the same provider route. Carrying it /// atomically avoids a model/limit mismatch before `SendMessage`. compaction: Box, + /// Auxiliary provider calls completed while planning this exact turn + /// (currently Auto's classifier), bounded and paired with their own + /// immutable routes. The engine folds their tokens into total usage + /// only; they never enter the parent route's billing aggregate. + initial_routed_usage: Box, goal_objective: Option, goal_token_budget: Option, goal_status: GoalStatus, diff --git a/crates/tui/src/core/turn.rs b/crates/tui/src/core/turn.rs index f650f6280f..246f8d0998 100644 --- a/crates/tui/src/core/turn.rs +++ b/crates/tui/src/core/turn.rs @@ -44,6 +44,16 @@ pub struct TurnContext { /// Usage for this turn pub usage: Usage, + /// Subset of `usage` served by the parent turn's frozen route. Programmatic + /// reviewer/RLM calls remain in the total above but are billed only from + /// their own routed receipts. + pub parent_route_usage: Usage, + + /// Provider calls whose usage became ambiguous after dispatch (for + /// example an RLM timeout). A non-zero value makes cost coverage + /// explicitly incomplete instead of inventing a zero-usage receipt. + pub routed_usage_dropped_records: u64, + /// Input tokens reported for the most recent parent-route model request. /// This is deliberately separate from `usage`, which accumulates every /// parent step and programmatic child call for billing. @@ -73,6 +83,8 @@ impl TurnContext { output_tokens: 0, ..Usage::default() }, + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, latest_parent_input_tokens: None, compaction_refusal_notified: false, pending_route: None, @@ -110,33 +122,7 @@ impl TurnContext { /// Add usage from an API response pub fn add_usage(&mut self, usage: &Usage) { - self.usage.input_tokens = self.usage.input_tokens.saturating_add(usage.input_tokens); - self.usage.output_tokens = self.usage.output_tokens.saturating_add(usage.output_tokens); - self.usage.prompt_cache_hit_tokens = add_optional_usage( - self.usage.prompt_cache_hit_tokens, - usage.prompt_cache_hit_tokens, - ); - self.usage.prompt_cache_miss_tokens = add_optional_usage( - self.usage.prompt_cache_miss_tokens, - usage.prompt_cache_miss_tokens, - ); - self.usage.prompt_cache_write_tokens = add_optional_usage( - self.usage.prompt_cache_write_tokens, - usage.prompt_cache_write_tokens, - ); - self.usage.reasoning_tokens = - add_optional_usage(self.usage.reasoning_tokens, usage.reasoning_tokens); - self.usage.reasoning_replay_tokens = add_optional_usage( - self.usage.reasoning_replay_tokens, - usage.reasoning_replay_tokens, - ); - if let Some(delta) = usage.server_tool_use.as_ref() { - let total = self.usage.server_tool_use.get_or_insert_default(); - total.code_execution_requests = - add_optional_usage(total.code_execution_requests, delta.code_execution_requests); - total.tool_search_requests = - add_optional_usage(total.tool_search_requests, delta.tool_search_requests); - } + add_usage_to(&mut self.usage, usage); } /// Record one parent-route response for both billing and live-context @@ -145,6 +131,53 @@ impl TurnContext { pub fn add_parent_usage(&mut self, usage: &Usage) { self.latest_parent_input_tokens = (usage.input_tokens > 0).then_some(usage.input_tokens); self.add_usage(usage); + add_usage_to(&mut self.parent_route_usage, usage); + } + + pub fn add_routed_usage_dropped_records(&mut self, dropped_records: u64) { + self.routed_usage_dropped_records = self + .routed_usage_dropped_records + .saturating_add(dropped_records); + } + + /// Add programmatic child-call usage to the authoritative total and + /// return the same batch aggregate for telemetry emission. + pub fn add_routed_usages<'a>(&mut self, usages: impl IntoIterator) -> Usage { + let mut aggregate = Usage::default(); + for usage in usages { + self.add_usage(usage); + add_usage_to(&mut aggregate, usage); + } + aggregate + } +} + +fn add_usage_to(total: &mut Usage, delta: &Usage) { + total.input_tokens = total.input_tokens.saturating_add(delta.input_tokens); + total.output_tokens = total.output_tokens.saturating_add(delta.output_tokens); + total.prompt_cache_hit_tokens = + add_optional_usage(total.prompt_cache_hit_tokens, delta.prompt_cache_hit_tokens); + total.prompt_cache_miss_tokens = add_optional_usage( + total.prompt_cache_miss_tokens, + delta.prompt_cache_miss_tokens, + ); + total.prompt_cache_write_tokens = add_optional_usage( + total.prompt_cache_write_tokens, + delta.prompt_cache_write_tokens, + ); + total.reasoning_tokens = add_optional_usage(total.reasoning_tokens, delta.reasoning_tokens); + total.reasoning_replay_tokens = + add_optional_usage(total.reasoning_replay_tokens, delta.reasoning_replay_tokens); + if let Some(delta) = delta.server_tool_use.as_ref() { + let server_total = total.server_tool_use.get_or_insert_default(); + server_total.code_execution_requests = add_optional_usage( + server_total.code_execution_requests, + delta.code_execution_requests, + ); + server_total.tool_search_requests = add_optional_usage( + server_total.tool_search_requests, + delta.tool_search_requests, + ); } } diff --git a/crates/tui/src/cost_status.rs b/crates/tui/src/cost_status.rs index d3a37ef65a..cbe7335717 100644 --- a/crates/tui/src/cost_status.rs +++ b/crates/tui/src/cost_status.rs @@ -96,6 +96,15 @@ pub struct EffectiveRouteEnvelope { pub model: String, pub billing_surface: Option, pub endpoint_fingerprint: Option, + /// Frozen provider-live rates captured from the exact fresh catalog scope + /// at CodeWhale's pre-permit application-dispatch boundary. Legacy + /// receipts omit this and therefore cannot meter a reviewed custom route + /// retroactively. + #[serde( + default, + deserialize_with = "crate::provider_catalog_live::deserialize_optional_provider_live_pricing" + )] + pub provider_live_pricing: Option, #[serde(default)] pub billing_mode: RouteBillingMode, pub dispatched_at: DateTime, @@ -109,12 +118,13 @@ impl serde::Serialize for EffectiveRouteEnvelope { use serde::ser::SerializeStruct as _; let route = self.sanitized_for_persistence(); - let mut state = serializer.serialize_struct("EffectiveRouteEnvelope", 7)?; + let mut state = serializer.serialize_struct("EffectiveRouteEnvelope", 8)?; state.serialize_field("provider", &route.provider)?; state.serialize_field("provider_identity", &route.provider_identity)?; state.serialize_field("model", &route.model)?; state.serialize_field("billing_surface", &route.billing_surface)?; state.serialize_field("endpoint_fingerprint", &route.endpoint_fingerprint)?; + state.serialize_field("provider_live_pricing", &route.provider_live_pricing)?; state.serialize_field("billing_mode", &route.billing_mode)?; state.serialize_field("dispatched_at", &route.dispatched_at)?; state.end() @@ -170,6 +180,21 @@ impl EffectiveRouteEnvelope { || crate::route_billing::for_endpoint_without_config(provider, base_url), |config| crate::route_billing::for_route(config, provider), ); + let endpoint_fingerprint = base_url.and_then(endpoint_fingerprint); + let provider_live_pricing = + u64::try_from(dispatched_at.timestamp()) + .ok() + .and_then(|dispatched_at_unix| { + endpoint_fingerprint.as_deref().and_then(|fingerprint| { + crate::provider_catalog_live::fresh_provider_live_pricing_quote_at( + provider, + &provider_identity, + &model, + fingerprint, + dispatched_at_unix, + ) + }) + }); Self { provider, provider_identity: sanitize_persisted_route_label(&provider_identity), @@ -178,7 +203,8 @@ impl EffectiveRouteEnvelope { config, provider, base_url, ) .map(str::to_string), - endpoint_fingerprint: base_url.and_then(endpoint_fingerprint), + endpoint_fingerprint, + provider_live_pricing, billing_mode: billing.into(), dispatched_at, } @@ -186,22 +212,29 @@ impl EffectiveRouteEnvelope { #[must_use] pub fn audit(&self, usage: &Usage) -> TurnCostAudit { + let reviewed_custom_metered = crate::pricing::reviewed_custom_route_is_metered( + self.provider, + Some(&self.provider_identity), + self.endpoint_fingerprint.as_deref(), + ); match self.billing_mode { RouteBillingMode::Subscription | RouteBillingMode::Local => { return TurnCostAudit::unpriced(crate::pricing::UnpricedReason::NotMoneyMetered); } - RouteBillingMode::Unknown => { + RouteBillingMode::Unknown if !reviewed_custom_metered => { return TurnCostAudit::unpriced( crate::pricing::UnpricedReason::UnknownBillingBasis, ); } - RouteBillingMode::Metered => {} + RouteBillingMode::Metered | RouteBillingMode::Unknown => {} } - crate::pricing::audit_turn_cost_for_route_on_endpoint_at( + crate::pricing::audit_turn_cost_for_route_on_endpoint_for_identity_at( self.provider, + Some(&self.provider_identity), &self.model, self.billing_surface.as_deref(), self.endpoint_fingerprint.as_deref(), + self.provider_live_pricing.as_ref(), usage, self.dispatched_at, ) @@ -243,6 +276,27 @@ impl EffectiveRouteEnvelope { && fingerprint.bytes().all(|byte| byte.is_ascii_hexdigit())) .then(|| fingerprint.to_ascii_lowercase()) }); + let quote_is_valid = route + .provider_live_pricing + .as_ref() + .zip(route.endpoint_fingerprint.as_deref()) + .and_then(|(quote, fingerprint)| { + u64::try_from(route.dispatched_at.timestamp()) + .ok() + .and_then(|dispatched_at_unix| { + quote.pricing_for_route( + route.provider, + &route.provider_identity, + &route.model, + fingerprint, + dispatched_at_unix, + ) + }) + }) + .is_some(); + if !quote_is_valid { + route.provider_live_pricing = None; + } route } } @@ -292,6 +346,10 @@ pub fn child_usage_metadata_fields( "child_endpoint_fingerprint".into(), serde_json::json!(route.endpoint_fingerprint), ); + fields.insert( + "child_provider_live_pricing".into(), + serde_json::json!(route.provider_live_pricing), + ); fields.insert( "child_billing_mode".into(), serde_json::json!(route.billing_mode), @@ -347,6 +405,163 @@ pub fn attach_child_usage_metadata( } } +/// Maximum number of distinct routed-usage segments accepted from one tool +/// result. RLM reserves against the same bound before dispatch, so a valid +/// producer never has to discard a provider receipt after doing the work. +pub const MAX_CHILD_USAGE_RECORDS: usize = 64; + +const CHILD_USAGE_RECORDS_KEY: &str = "child_usage_records"; +const CHILD_USAGE_DROP_RECORDS_KEY: &str = "child_usage_drop_records"; +const CHILD_USAGE_DROPPED_RECORDS_KEY: &str = "child_usage_dropped_records"; + +/// Attach a bounded batch of routed child usage to tool metadata. +/// +/// The source identity is reduced to a one-way fingerprint before metadata +/// can enter a transcript. Routes pass through their persistence sanitizer, +/// so neither a raw response id nor an endpoint/credential can hitch a ride. +/// New consumers prefer this batch over the legacy single `child_*` fields. +/// Attach a bounded batch containing both exact usage receipts and exact +/// provider-success/missing-usage route receipts. +pub fn attach_child_usage_batch_metadata( + metadata: &mut serde_json::Value, + batch: &RuntimeUsageBatch, +) { + let Some(object) = metadata.as_object_mut() else { + return; + }; + let retained_records = batch + .records + .iter() + .take(MAX_CHILD_USAGE_RECORDS) + .map(|record| { + serde_json::json!({ + "source_id": format!( + "routed:{}", + usage_source_fingerprint(&record.source_id) + ), + "route": record.usage.route.sanitized_for_persistence(), + "usage": record.usage.usage, + }) + }) + .collect::>(); + let remaining = MAX_CHILD_USAGE_RECORDS.saturating_sub(retained_records.len()); + let retained_drops = batch + .drop_records + .iter() + .take(remaining) + .map(|record| { + serde_json::json!({ + "source_id": format!( + "routed:{}", + usage_source_fingerprint(&record.source_id) + ), + "route": record.route.sanitized_for_persistence(), + }) + }) + .collect::>(); + object.insert( + CHILD_USAGE_RECORDS_KEY.into(), + serde_json::json!(retained_records), + ); + object.insert( + CHILD_USAGE_DROP_RECORDS_KEY.into(), + serde_json::json!(retained_drops), + ); + let usage_overflow = batch.records.len().saturating_sub(MAX_CHILD_USAGE_RECORDS); + let dropped_records = batch + .dropped_records + .max(u64::try_from(batch.drop_records.len()).unwrap_or(u64::MAX)) + .saturating_add(u64::try_from(usage_overflow).unwrap_or(u64::MAX)); + if dropped_records > 0 { + object.insert( + CHILD_USAGE_DROPPED_RECORDS_KEY.into(), + serde_json::json!(dropped_records), + ); + } else { + object.remove(CHILD_USAGE_DROPPED_RECORDS_KEY); + } +} + +/// Parse the preferred routed child-usage batch. +/// +/// `None` means the batch key was absent and callers may use the legacy +/// single-record parser. Once the key is present, malformed/overflow entries +/// are represented by `dropped_records` instead of falling back and risking a +/// partial subtotal being presented as complete. +#[must_use] +pub fn child_usage_records_from_metadata( + metadata: &serde_json::Value, +) -> Option { + let value = metadata.get(CHILD_USAGE_RECORDS_KEY)?; + let drop_values = metadata + .get(CHILD_USAGE_DROP_RECORDS_KEY) + .and_then(serde_json::Value::as_array) + .cloned() + .unwrap_or_default(); + let declared_dropped = metadata + .get(CHILD_USAGE_DROPPED_RECORDS_KEY) + .and_then(serde_json::Value::as_u64) + .unwrap_or(0); + let Some(values) = value.as_array() else { + return Some(RuntimeUsageBatch { + records: Vec::new(), + drop_records: Vec::new(), + dropped_records: declared_dropped.saturating_add(1), + }); + }; + + let overflow = values.len().saturating_sub(MAX_CHILD_USAGE_RECORDS); + let mut batch = RuntimeUsageBatch { + records: Vec::with_capacity(values.len().min(MAX_CHILD_USAGE_RECORDS)), + drop_records: Vec::with_capacity(drop_values.len().min(MAX_CHILD_USAGE_RECORDS)), + dropped_records: declared_dropped + .max(u64::try_from(drop_values.len()).unwrap_or(u64::MAX)) + .saturating_add(u64::try_from(overflow).unwrap_or(u64::MAX)), + }; + for value in values.iter().take(MAX_CHILD_USAGE_RECORDS) { + let parsed = (|| { + let source_id = value.get("source_id")?.as_str()?; + let route = + serde_json::from_value::(value.get("route")?.clone()) + .ok()? + .sanitized_for_persistence(); + let usage = serde_json::from_value::(value.get("usage")?.clone()).ok()?; + Some(RuntimeUsageRecord { + // Treat metadata as an untrusted persistence boundary. A + // stable hash preserves idempotence without retaining the + // producer's raw identifier. + source_id: format!("metadata:{}", usage_source_fingerprint(source_id)), + usage: EffectiveRouteUsage { route, usage }, + }) + })(); + if let Some(record) = parsed { + batch.records.push(record); + } else { + batch.dropped_records = batch.dropped_records.saturating_add(1); + } + } + let remaining = MAX_CHILD_USAGE_RECORDS.saturating_sub(batch.records.len()); + for value in drop_values.iter().take(remaining) { + let parsed = (|| { + let source_id = value.get("source_id")?.as_str()?; + let route = + serde_json::from_value::(value.get("route")?.clone()) + .ok()? + .sanitized_for_persistence(); + Some(RuntimeUsageDropRecord { + source_id: format!("metadata:{}", usage_source_fingerprint(source_id)), + route, + }) + })(); + if let Some(record) = parsed { + batch.drop_records.push(record); + } else if declared_dropped < u64::try_from(drop_values.len()).unwrap_or(u64::MAX) { + batch.dropped_records = batch.dropped_records.saturating_add(1); + } + } + Some(batch) +} + /// Rehydrate the immutable route envelope emitted with child usage. Legacy or /// incomplete metadata becomes an explicitly unknown route and never borrows /// mutable parent-session facts. @@ -389,6 +604,10 @@ pub fn child_route_envelope_from_metadata( .get("child_endpoint_fingerprint") .and_then(serde_json::Value::as_str) .map(str::to_string), + provider_live_pricing: metadata + .get("child_provider_live_pricing") + .cloned() + .and_then(|value| serde_json::from_value(value).ok()), billing_mode: billing_mode .filter(|_| complete) .unwrap_or(RouteBillingMode::Unknown), @@ -507,7 +726,10 @@ const MAX_RUNTIME_USAGE_RECORDS_PER_OWNER: usize = 64; #[derive(Default)] struct OwnerRuntimeUsageJournal { records: VecDeque, + drop_records: VecDeque, dropped_records: u64, + dropped_source_fingerprints: HashSet, + dropped_fingerprint_overflowed: bool, } type RuntimeUsageJournal = HashMap; @@ -518,6 +740,10 @@ type RuntimeUsageJournal = HashMap; #[derive(Debug, Clone, Default, PartialEq, Eq)] pub struct RuntimeUsageBatch { pub records: Vec, + /// Exact provider-success calls whose usage payload was absent. The + /// bounded records retain route billing truth; `dropped_records` remains + /// the authoritative total and may exceed this vector after overflow. + pub drop_records: Vec, pub dropped_records: u64, } @@ -529,10 +755,22 @@ pub struct RuntimeUsageRecord { pub usage: EffectiveRouteUsage, } +/// One provider-success response that omitted usage metadata. +/// +/// The frozen route is required to distinguish money-metered calls from +/// subscription/local calls without consulting mutable completion-time config. +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub struct RuntimeUsageDropRecord { + pub source_id: String, + pub route: EffectiveRouteEnvelope, +} + pub(crate) type RuntimeUsageSink = Arc bool + Send + Sync>; +pub(crate) type RuntimeUsageDropSink = Arc bool + Send + Sync>; struct RuntimeUsageSinkEntry { sink: RuntimeUsageSink, + dropped_sink: Option, leases: usize, terminal: bool, } @@ -665,10 +903,87 @@ fn record_runtime_usage( }); } +fn record_runtime_usage_drop(owner: &str, source_id: &str, route: &EffectiveRouteEnvelope) { + let owner = owner.trim(); + if owner.is_empty() { + return; + } + let fingerprint = usage_source_fingerprint(source_id); + let sink = with_runtime_usage_sinks(|sinks| { + sinks + .get(owner) + .and_then(|entry| entry.dropped_sink.as_ref().map(Arc::clone)) + }); + let record = RuntimeUsageDropRecord { + source_id: source_id.to_string(), + route: route.sanitized_for_persistence(), + }; + if sink.is_some_and(|sink| sink(record.clone())) { + return; + } + with_runtime_usage_journal_mut(|journal| { + let owner_journal = journal.entry(owner.to_string()).or_default(); + if owner_journal + .dropped_source_fingerprints + .contains(&fingerprint) + { + return; + } + if owner_journal.dropped_source_fingerprints.len() < MAX_RUNTIME_USAGE_RECORDS_PER_OWNER { + owner_journal + .dropped_source_fingerprints + .insert(fingerprint); + owner_journal.drop_records.push_back(record); + owner_journal.dropped_records = owner_journal.dropped_records.saturating_add(1); + } else if !owner_journal.dropped_fingerprint_overflowed { + // Preserve a bounded fail-closed overflow marker. Once the exact + // identity ledger is full, further unknown ids share this one + // marker so replays cannot grow the count without bound. + owner_journal.dropped_fingerprint_overflowed = true; + owner_journal.dropped_records = owner_journal.dropped_records.saturating_add(1); + } + }); +} + +fn record_runtime_usage_drop_count(owner: &str, source_id: &str, count: u64) { + let owner = owner.trim(); + if owner.is_empty() || count == 0 { + return; + } + let fingerprint = usage_source_fingerprint(source_id); + with_runtime_usage_journal_mut(|journal| { + let owner_journal = journal.entry(owner.to_string()).or_default(); + if owner_journal + .dropped_source_fingerprints + .contains(&fingerprint) + { + return; + } + if owner_journal.dropped_source_fingerprints.len() < MAX_RUNTIME_USAGE_RECORDS_PER_OWNER { + owner_journal + .dropped_source_fingerprints + .insert(fingerprint); + owner_journal.dropped_records = owner_journal.dropped_records.saturating_add(count); + } else if !owner_journal.dropped_fingerprint_overflowed { + owner_journal.dropped_fingerprint_overflowed = true; + owner_journal.dropped_records = owner_journal.dropped_records.saturating_add(1); + } + }); +} + /// Install a synchronous durability sink for one active runtime turn. /// Compaction calls invoke this before they return to the engine, so a process /// crash cannot erase already-reported usage from an in-memory journal. +#[cfg(test)] pub(crate) fn register_runtime_usage_sink(owner: &str, sink: RuntimeUsageSink) { + register_runtime_usage_sink_with_drop(owner, sink, None); +} + +pub(crate) fn register_runtime_usage_sink_with_drop( + owner: &str, + sink: RuntimeUsageSink, + dropped_sink: Option, +) { let owner = owner.trim(); if owner.is_empty() { return; @@ -678,6 +993,7 @@ pub(crate) fn register_runtime_usage_sink(owner: &str, sink: RuntimeUsageSink) { owner.to_string(), RuntimeUsageSinkEntry { sink, + dropped_sink, leases: 0, terminal: false, }, @@ -697,13 +1013,80 @@ pub(crate) fn usage_source_fingerprint(source_id: &str) -> String { /// child may report after the parent mailbox has sealed; its owner lease keeps /// this sink alive, while the captured scope prevents a later session from /// inheriting the spend. +#[cfg(test)] pub(crate) fn register_interactive_runtime_usage_sink(owner: &str, scope: CostScopeToken) { - register_runtime_usage_sink( + register_runtime_usage_sink_with_drop( + owner, + Arc::new(move |record| record_interactive_runtime_usage(scope, record)), + Some(Arc::new(move |record| { + record_interactive_runtime_usage_drop(scope, record) + })), + ); +} + +/// Install an interactive sink whose stale-scope fallback is an origin-session +/// sidecar. `/new` may close the foreground pool while a detached provider call +/// is still running; the sidecar keeps that exact response with the old saved +/// session instead of either dropping it or contaminating the new one. +pub(crate) fn register_persistent_interactive_runtime_usage_sink( + owner: &str, + scope: CostScopeToken, + session_id: &str, + turn_id: &str, +) { + let Ok(manager) = crate::session_manager::SessionManager::default_location() else { + register_runtime_usage_sink_with_drop( + owner, + Arc::new(move |record| record_interactive_runtime_usage(scope, record)), + Some(Arc::new(move |record| { + record_interactive_runtime_usage_drop(scope, record) + })), + ); + return; + }; + register_persistent_interactive_runtime_usage_sink_at( + owner, + scope, + session_id, + turn_id, + manager.sessions_dir().to_path_buf(), + ); +} + +fn register_persistent_interactive_runtime_usage_sink_at( + owner: &str, + scope: CostScopeToken, + session_id: &str, + turn_id: &str, + sessions_dir: std::path::PathBuf, +) { + let usage_session_id = session_id.to_string(); + let usage_turn_id = turn_id.to_string(); + let drop_session_id = usage_session_id.clone(); + let drop_turn_id = usage_turn_id.clone(); + let usage_sessions_dir = sessions_dir.clone(); + register_runtime_usage_sink_with_drop( owner, Arc::new(move |record| { - record_interactive_runtime_usage(scope, record); - true + if record_interactive_runtime_usage(scope, record.clone()) { + return true; + } + crate::session_manager::SessionManager::new(usage_sessions_dir.clone()) + .and_then(|manager| { + manager.persist_late_runtime_usage(&usage_session_id, &usage_turn_id, &record) + }) + .unwrap_or(false) }), + Some(Arc::new(move |record| { + if record_interactive_runtime_usage_drop(scope, record.clone()) { + return true; + } + crate::session_manager::SessionManager::new(sessions_dir.clone()) + .and_then(|manager| { + manager.persist_late_runtime_drop(&drop_session_id, &drop_turn_id, &record) + }) + .unwrap_or(false) + })), ); } @@ -794,6 +1177,7 @@ pub fn take_runtime_usage(owner: &str) -> RuntimeUsageBatch { .remove(owner) .map_or_else(RuntimeUsageBatch::default, |entry| RuntimeUsageBatch { records: entry.records.into_iter().collect(), + drop_records: entry.drop_records.into_iter().collect(), dropped_records: entry.dropped_records, }) }) @@ -825,7 +1209,7 @@ pub fn close_current_scope() -> PendingBackgroundCost { pub(crate) fn restore_usage_source_fingerprints(fingerprints: impl IntoIterator) { with_pending_state_mut(|state| { state.seen_usage_source_fingerprints = fingerprints.into_iter().collect(); - }); + }) } /// Whether this session generation already accepted a provider response. @@ -1208,10 +1592,246 @@ pub fn report_effective_route_for_runtime( if let Some(owner) = runtime_owner { record_runtime_usage(owner, source_id, route, usage); } else { - report_effective_route(scope, route, usage); + record_interactive_runtime_usage( + scope, + RuntimeUsageRecord { + source_id: source_id.to_string(), + usage: EffectiveRouteUsage { + route: route.sanitized_for_persistence(), + usage: usage.clone(), + }, + }, + ); + } +} + +/// Report an interactive auxiliary response against its immutable origin. +/// A stale foreground scope is not an error: it means `/new` or session load +/// already moved on, so the exact receipt is appended to the old session's +/// durable sidecar instead of being redirected to the active session. +pub(crate) fn report_effective_route_for_interactive_origin( + scope: CostScopeToken, + session_id: &str, + turn_id: &str, + source_id: &str, + route: &EffectiveRouteEnvelope, + usage: &Usage, +) { + let record = RuntimeUsageRecord { + source_id: source_id.to_string(), + usage: EffectiveRouteUsage { + route: route.sanitized_for_persistence(), + usage: usage.clone(), + }, + }; + if record_interactive_runtime_usage(scope, record.clone()) { + return; + } + let persisted = + crate::session_manager::SessionManager::default_location().is_ok_and(|manager| { + manager + .persist_late_runtime_usage(session_id, turn_id, &record) + .unwrap_or(false) + }); + if !persisted { + tracing::warn!("late interactive usage could not be persisted for its origin session"); + } +} + +#[cfg(test)] +fn report_effective_route_for_interactive_origin_with_manager( + scope: CostScopeToken, + session_id: &str, + turn_id: &str, + source_id: &str, + route: &EffectiveRouteEnvelope, + usage: &Usage, + manager: &crate::session_manager::SessionManager, +) -> bool { + let record = RuntimeUsageRecord { + source_id: source_id.to_string(), + usage: EffectiveRouteUsage { + route: route.sanitized_for_persistence(), + usage: usage.clone(), + }, + }; + if record_interactive_runtime_usage(scope, record.clone()) { + return true; + } + manager + .persist_late_runtime_usage(session_id, turn_id, &record) + .unwrap_or(false) +} + +pub(crate) fn report_unreceipted_for_interactive_origin( + scope: CostScopeToken, + session_id: &str, + turn_id: &str, + source_id: &str, + route: &EffectiveRouteEnvelope, +) { + let record = RuntimeUsageDropRecord { + source_id: source_id.to_string(), + route: route.sanitized_for_persistence(), + }; + if record_interactive_runtime_usage_drop(scope, record.clone()) { + return; + } + let persisted = + crate::session_manager::SessionManager::default_location().is_ok_and(|manager| { + manager + .persist_late_runtime_drop(session_id, turn_id, &record) + .unwrap_or(false) + }); + if !persisted { + tracing::warn!( + "late interactive missing-usage receipt could not be persisted for its origin session" + ); + } +} + +#[cfg(test)] +fn report_unreceipted_for_interactive_origin_with_manager( + scope: CostScopeToken, + session_id: &str, + turn_id: &str, + source_id: &str, + route: &EffectiveRouteEnvelope, + manager: &crate::session_manager::SessionManager, +) -> bool { + let record = RuntimeUsageDropRecord { + source_id: source_id.to_string(), + route: route.sanitized_for_persistence(), + }; + if record_interactive_runtime_usage_drop(scope, record.clone()) { + return true; + } + manager + .persist_late_runtime_drop(session_id, turn_id, &record) + .unwrap_or(false) +} + +/// Record one provider-success response whose usage payload was absent. +/// +/// Callers must supply the same fixed-length, non-secret source identity they +/// would use for a normal routed usage receipt. Runtime owners persist one +/// bounded dropped-coverage marker; ownerless/interactive calls add one +/// unpriced coverage turn to the captured session scope. Replays are +/// idempotent, and a stale scope cannot contaminate a later session. +pub(crate) fn report_unreceipted_provider_success( + scope: CostScopeToken, + runtime_owner: Option<&str>, + source_id: &str, + route: &EffectiveRouteEnvelope, +) { + if let Some(owner) = runtime_owner { + record_runtime_usage_drop(owner, source_id, route); + } else { + record_interactive_runtime_usage_drop( + scope, + RuntimeUsageDropRecord { + source_id: source_id.to_string(), + route: route.sanitized_for_persistence(), + }, + ); + } +} + +/// Settle one bounded routed-usage batch without repricing or losing exact +/// missing-usage route evidence. Replaying the same batch is idempotent by the +/// stable per-response source ids. Any residual count whose exact record was +/// truncated remains an explicit fail-closed coverage gap. +pub(crate) fn report_runtime_usage_batch( + scope: CostScopeToken, + runtime_owner: Option<&str>, + batch: &RuntimeUsageBatch, +) { + for record in &batch.records { + report_effective_route_for_runtime( + scope, + runtime_owner, + &record.source_id, + &record.usage.route, + &record.usage.usage, + ); + } + for record in &batch.drop_records { + report_unreceipted_provider_success(scope, runtime_owner, &record.source_id, &record.route); + } + + let residual = batch + .dropped_records + .saturating_sub(u64::try_from(batch.drop_records.len()).unwrap_or(u64::MAX)); + if residual == 0 { + return; + } + let mut identities = batch + .records + .iter() + .map(|record| usage_source_fingerprint(&record.source_id)) + .chain( + batch + .drop_records + .iter() + .map(|record| usage_source_fingerprint(&record.source_id)), + ) + .take(MAX_RUNTIME_USAGE_RECORDS_PER_OWNER) + .collect::>(); + identities.sort_unstable(); + let residual_source = format!( + "runtime-usage-batch-residual:{}", + usage_source_fingerprint(&format!( + "{}:{}:{}:{}", + batch.records.len(), + batch.drop_records.len(), + batch.dropped_records, + identities.join(":") + )) + ); + if let Some(owner) = runtime_owner { + record_runtime_usage_drop_count(owner, &residual_source, residual); + } else { + record_interactive_runtime_usage_drop_count(scope, &residual_source, residual); } } +#[must_use] +pub(crate) fn background_cost_for_runtime_usage( + record: &RuntimeUsageRecord, +) -> PendingBackgroundCost { + let mut pending = PendingBackgroundCost::default(); + let fingerprint = usage_source_fingerprint(&record.source_id); + let audit = record.usage.route.audit(&record.usage.usage); + let receipt = record.usage.route.receipt(&audit); + pending.usage_source_fingerprints.insert(fingerprint); + fold_audit_into_pending(&mut pending, receipt, &audit, &record.usage.usage); + pending +} + +#[must_use] +pub(crate) fn background_cost_for_runtime_drop( + record: &RuntimeUsageDropRecord, +) -> PendingBackgroundCost { + let mut pending = PendingBackgroundCost::default(); + pending + .usage_source_fingerprints + .insert(usage_source_fingerprint(&record.source_id)); + if !matches!( + record.route.billing_mode, + RouteBillingMode::Subscription | RouteBillingMode::Local + ) { + pending.unpriced_turns = 1; + pending.cny_unpriced_turns = 1; + pending + .unpriced_reasons + .insert("provider_success_missing_usage"); + pending + .cny_unpriced_reasons + .insert("provider_success_missing_usage"); + } + pending +} + /// Fold one already-computed audit into the pending pool. fn record(scope: CostScopeToken, route_receipt: String, audit: &TurnCostAudit, usage: &Usage) { with_pending_state_mut(|state| { @@ -1222,23 +1842,92 @@ fn record(scope: CostScopeToken, route_receipt: String, audit: &TurnCostAudit, u }); } -fn record_interactive_runtime_usage(scope: CostScopeToken, record: RuntimeUsageRecord) { +fn record_interactive_runtime_usage(scope: CostScopeToken, record: RuntimeUsageRecord) -> bool { with_pending_state_mut(|state| { if state.generation != scope.0 { - return; + return false; } let fingerprint = usage_source_fingerprint(&record.source_id); if !state .seen_usage_source_fingerprints .insert(fingerprint.clone()) { - return; + return true; } let audit = record.usage.route.audit(&record.usage.usage); let receipt = record.usage.route.receipt(&audit); state.pending.usage_source_fingerprints.insert(fingerprint); fold_audit_into_pending(&mut state.pending, receipt, &audit, &record.usage.usage); - }); + true + }) +} + +fn record_interactive_runtime_usage_drop( + scope: CostScopeToken, + record: RuntimeUsageDropRecord, +) -> bool { + with_pending_state_mut(|state| { + if state.generation != scope.0 { + return false; + } + let fingerprint = usage_source_fingerprint(&record.source_id); + if !state + .seen_usage_source_fingerprints + .insert(fingerprint.clone()) + { + return true; + } + state.pending.usage_source_fingerprints.insert(fingerprint); + if matches!( + record.route.billing_mode, + RouteBillingMode::Subscription | RouteBillingMode::Local + ) { + return true; + } + state.pending.unpriced_turns = state.pending.unpriced_turns.saturating_add(1); + state.pending.cny_unpriced_turns = state.pending.cny_unpriced_turns.saturating_add(1); + state + .pending + .unpriced_reasons + .insert("provider_success_missing_usage"); + state + .pending + .cny_unpriced_reasons + .insert("provider_success_missing_usage"); + true + }) +} + +fn record_interactive_runtime_usage_drop_count( + scope: CostScopeToken, + source_id: &str, + count: u64, +) -> bool { + with_pending_state_mut(|state| { + if state.generation != scope.0 { + return false; + } + let fingerprint = usage_source_fingerprint(source_id); + if !state + .seen_usage_source_fingerprints + .insert(fingerprint.clone()) + { + return true; + } + state.pending.usage_source_fingerprints.insert(fingerprint); + let count = u32::try_from(count).unwrap_or(u32::MAX); + state.pending.unpriced_turns = state.pending.unpriced_turns.saturating_add(count); + state.pending.cny_unpriced_turns = state.pending.cny_unpriced_turns.saturating_add(count); + state + .pending + .unpriced_reasons + .insert("routed_usage_receipt_missing"); + state + .pending + .cny_unpriced_reasons + .insert("routed_usage_receipt_missing"); + true + }) } fn fold_audit_into_pending( @@ -1340,6 +2029,102 @@ pub(crate) fn test_scope() -> TestCostScope { mod tests { use super::*; + struct ProviderCatalogTestReset; + + impl Drop for ProviderCatalogTestReset { + fn drop(&mut self) { + crate::provider_catalog_live::reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + } + } + + fn priced_provider_delta( + provider: &str, + model: &str, + fingerprint: &str, + fetched_at: u64, + ) -> codewhale_config::catalog::ProviderCatalogDelta { + priced_provider_delta_with_rates(provider, model, fingerprint, fetched_at, 1.25, 5.0) + } + + fn priced_provider_delta_with_rates( + provider: &str, + model: &str, + fingerprint: &str, + fetched_at: u64, + input: f64, + output: f64, + ) -> codewhale_config::catalog::ProviderCatalogDelta { + codewhale_config::catalog::ProviderCatalogDelta { + provider: provider.to_string(), + base_url_fingerprint: fingerprint.to_string(), + fetched_at, + offerings: vec![codewhale_config::catalog::CatalogOffering { + provider: provider.to_string(), + wire_model_id: model.to_string(), + endpoint_key: "chat".to_string(), + cost: Some(codewhale_config::models_dev::ModelsDevCost { + input: Some(input), + output: Some(output), + cache_read: Some(0.25), + cache_write: None, + }), + ..Default::default() + }], + } + } + + fn custom_usage_envelope( + identity: &str, + model: &str, + fingerprint: &str, + billing_mode: RouteBillingMode, + dispatched_at: DateTime, + ) -> EffectiveRouteEnvelope { + provider_live_usage_envelope( + ApiProvider::Custom, + identity, + model, + fingerprint, + Some(crate::pricing::UNCLASSIFIED_BILLING_SURFACE), + billing_mode, + dispatched_at, + ) + } + + fn provider_live_usage_envelope( + provider: ApiProvider, + identity: &str, + model: &str, + fingerprint: &str, + billing_surface: Option<&str>, + billing_mode: RouteBillingMode, + dispatched_at: DateTime, + ) -> EffectiveRouteEnvelope { + let provider_live_pricing = + u64::try_from(dispatched_at.timestamp()) + .ok() + .and_then(|dispatched_at_unix| { + crate::provider_catalog_live::fresh_provider_live_pricing_quote_at( + provider, + identity, + model, + fingerprint, + dispatched_at_unix, + ) + }); + EffectiveRouteEnvelope { + provider, + provider_identity: identity.to_string(), + model: model.to_string(), + billing_surface: billing_surface.map(str::to_string), + endpoint_fingerprint: Some(fingerprint.to_string()), + provider_live_pricing, + billing_mode, + dispatched_at, + } + } + fn small_usage() -> Usage { Usage { input_tokens: 1_000, @@ -1348,6 +2133,680 @@ mod tests { } } + #[test] + fn baseten_usage_prices_only_the_reviewed_identity_on_the_official_endpoint() { + let _env = crate::test_support::lock_test_env(); + let _live = crate::provider_lake::lock_live_snapshot(); + let home = tempfile::tempdir().expect("test home"); + let _home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", home.path()); + let _reset = ProviderCatalogTestReset; + crate::provider_catalog_live::reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + + let now = Utc::now(); + let fetched_at = u64::try_from(now.timestamp()).expect("nonnegative timestamp"); + let model = "synthetic-baseten-priced-model"; + let fingerprint = + codewhale_config::catalog::base_url_fingerprint(codewhale_config::BASETEN_BASE_URL); + crate::provider_catalog_live::record_success(priced_provider_delta( + codewhale_config::BASETEN_TEMPLATE_ID, + model, + &fingerprint, + fetched_at, + )); + let usage = Usage { + input_tokens: 1_000_000, + ..Usage::default() + }; + + let exact = custom_usage_envelope( + codewhale_config::BASETEN_TEMPLATE_ID, + model, + &fingerprint, + RouteBillingMode::Unknown, + now, + ) + .audit(&usage); + assert!(exact.is_priced(), "{exact:?}"); + assert_eq!( + exact.provenance, + Some(codewhale_config::pricing::PricingProvenance::ProviderLive) + ); + assert_eq!(exact.estimate.expect("priced").usd, 1.25); + + // A reviewed schema alias remains a distinct custom ownership scope. + // It becomes billable only after that exact identity refreshed its own + // catalog; it cannot borrow the canonical `baseten` partition above. + let alias = "base-ten"; + crate::provider_catalog_live::record_success(priced_provider_delta( + alias, + model, + &fingerprint, + fetched_at, + )); + let alias_audit = + custom_usage_envelope(alias, model, &fingerprint, RouteBillingMode::Unknown, now) + .audit(&usage); + assert!(alias_audit.is_priced(), "{alias_audit:?}"); + assert_eq!( + alias_audit.provenance, + Some(codewhale_config::pricing::PricingProvenance::ProviderLive) + ); + assert_eq!(alias_audit.estimate.expect("priced").usd, 1.25); + + let generic = custom_usage_envelope( + "custom-lab", + model, + &fingerprint, + RouteBillingMode::Metered, + now, + ) + .audit(&usage); + assert!(!generic.is_priced(), "{generic:?}"); + assert_eq!( + generic.unpriced_reason, + Some(crate::pricing::UnpricedReason::UnknownBillingBasis) + ); + + let wrong_fingerprint = + codewhale_config::catalog::base_url_fingerprint("https://proxy.example/v1"); + let wrong_endpoint = custom_usage_envelope( + codewhale_config::BASETEN_TEMPLATE_ID, + model, + &wrong_fingerprint, + RouteBillingMode::Metered, + now, + ) + .audit(&usage); + assert!(!wrong_endpoint.is_priced(), "{wrong_endpoint:?}"); + assert_eq!( + wrong_endpoint.unpriced_reason, + Some(crate::pricing::UnpricedReason::UnknownBillingBasis) + ); + } + + #[test] + fn baseten_usage_rejects_unknown_stale_and_failed_live_catalogs() { + let _env = crate::test_support::lock_test_env(); + let _live = crate::provider_lake::lock_live_snapshot(); + let home = tempfile::tempdir().expect("test home"); + let _home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", home.path()); + let _reset = ProviderCatalogTestReset; + crate::provider_catalog_live::reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + + let now = Utc::now(); + let now_unix = u64::try_from(now.timestamp()).expect("nonnegative timestamp"); + let model = "synthetic-baseten-status-model"; + let fingerprint = + codewhale_config::catalog::base_url_fingerprint(codewhale_config::BASETEN_BASE_URL); + let unknown_route = custom_usage_envelope( + codewhale_config::BASETEN_TEMPLATE_ID, + model, + &fingerprint, + RouteBillingMode::Unknown, + now, + ); + assert!(unknown_route.provider_live_pricing.is_none()); + let usage = Usage { + input_tokens: 1_000_000, + ..Usage::default() + }; + + // A same-model price owned by another custom partition cannot price a + // Baseten receipt whose exact catalog was never refreshed. + crate::provider_catalog_live::record_success(priced_provider_delta( + "other-custom", + model, + &fingerprint, + now_unix, + )); + let unknown = unknown_route.audit(&usage); + assert!(!unknown.is_priced(), "{unknown:?}"); + assert_eq!( + unknown.unpriced_reason, + Some(crate::pricing::UnpricedReason::UnverifiedLivePricing) + ); + + let stale_at = now_unix + .saturating_sub(crate::provider_catalog_live::DEFAULT_PROVIDER_CATALOG_TTL_SECS) + .saturating_sub(1); + crate::provider_catalog_live::record_success(priced_provider_delta( + codewhale_config::BASETEN_TEMPLATE_ID, + model, + &fingerprint, + stale_at, + )); + let stale_route = custom_usage_envelope( + codewhale_config::BASETEN_TEMPLATE_ID, + model, + &fingerprint, + RouteBillingMode::Unknown, + now, + ); + assert!(stale_route.provider_live_pricing.is_none()); + let stale = stale_route.audit(&usage); + assert!(!stale.is_priced(), "{stale:?}"); + assert_eq!( + stale.unpriced_reason, + Some(crate::pricing::UnpricedReason::UnverifiedLivePricing) + ); + + crate::provider_catalog_live::record_success(priced_provider_delta( + codewhale_config::BASETEN_TEMPLATE_ID, + model, + &fingerprint, + now_unix, + )); + crate::provider_catalog_live::record_failure( + codewhale_config::BASETEN_TEMPLATE_ID, + &fingerprint, + codewhale_config::catalog::CatalogRefreshError::Network, + ); + let failed_route = custom_usage_envelope( + codewhale_config::BASETEN_TEMPLATE_ID, + model, + &fingerprint, + RouteBillingMode::Unknown, + now, + ); + assert!(failed_route.provider_live_pricing.is_none()); + let failed = failed_route.audit(&usage); + assert!(!failed.is_priced(), "{failed:?}"); + assert_eq!( + failed.unpriced_reason, + Some(crate::pricing::UnpricedReason::UnverifiedLivePricing) + ); + } + + #[test] + fn reviewed_provider_live_quotes_survive_same_second_refresh_and_key_state_changes() { + let _env = crate::test_support::lock_test_env(); + let _live = crate::provider_lake::lock_live_snapshot(); + let home = tempfile::tempdir().expect("test home"); + let _home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", home.path()); + let _reset = ProviderCatalogTestReset; + crate::provider_catalog_live::reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + + let now = Utc::now(); + let fetched_at = u64::try_from(now.timestamp()).expect("nonnegative timestamp"); + let cases = [ + ( + ApiProvider::Openrouter, + ApiProvider::Openrouter.as_str(), + "synthetic-openrouter-frozen-price", + codewhale_config::catalog::base_url_fingerprint( + crate::config::DEFAULT_OPENROUTER_BASE_URL, + ), + crate::pricing::AGGREGATOR_BILLING_SURFACE, + RouteBillingMode::Metered, + ), + ( + ApiProvider::Custom, + codewhale_config::BASETEN_TEMPLATE_ID, + "synthetic-baseten-frozen-price", + codewhale_config::catalog::base_url_fingerprint(codewhale_config::BASETEN_BASE_URL), + crate::pricing::UNCLASSIFIED_BILLING_SURFACE, + RouteBillingMode::Unknown, + ), + ]; + let usage = Usage { + input_tokens: 1_000_000, + ..Usage::default() + }; + + for (provider, identity, model, fingerprint, surface, mode) in cases { + crate::provider_catalog_live::record_success(priced_provider_delta_with_rates( + identity, + model, + &fingerprint, + fetched_at, + 1.25, + 5.0, + )); + let first = provider_live_usage_envelope( + provider, + identity, + model, + &fingerprint, + Some(surface), + mode, + now, + ); + let first_quote = first + .provider_live_pricing + .as_ref() + .expect("fresh exact scope freezes a quote"); + + // A second refresh in the same Unix second must still be a distinct + // catalog revision and must not retroactively change `first`. + crate::provider_catalog_live::record_success(priced_provider_delta_with_rates( + identity, + model, + &fingerprint, + fetched_at, + 9.5, + 19.0, + )); + let second = provider_live_usage_envelope( + provider, + identity, + model, + &fingerprint, + Some(surface), + mode, + now, + ); + let second_quote = second + .provider_live_pricing + .as_ref() + .expect("replacement fresh scope freezes a quote"); + assert_ne!( + first_quote.catalog_revision, second_quote.catalog_revision, + "same-second price changes need distinct revisions" + ); + + crate::provider_catalog_live::record_failure( + identity, + &fingerprint, + codewhale_config::catalog::CatalogRefreshError::Unauthorized, + ); + if provider == ApiProvider::Custom { + // Baseten's same URL can represent another account after a key + // switch. Starting that refresh clears the mutable old scope. + let _new_key_refresh = crate::provider_catalog_live::begin_refresh(identity); + } + + let first_audit = first.audit(&usage); + let second_audit = second.audit(&usage); + assert_eq!(first_audit.estimate.expect("first quote priced").usd, 1.25); + assert_eq!(second_audit.estimate.expect("second quote priced").usd, 9.5); + + let after_mutation = provider_live_usage_envelope( + provider, + identity, + model, + &fingerprint, + Some(surface), + mode, + now, + ); + assert!( + after_mutation.provider_live_pricing.is_none(), + "failed or cleared mutable state cannot mint a new quote" + ); + } + } + + #[test] + fn legacy_no_quote_receipts_cannot_be_retro_priced_by_a_later_refresh() { + let _env = crate::test_support::lock_test_env(); + let _live = crate::provider_lake::lock_live_snapshot(); + let home = tempfile::tempdir().expect("test home"); + let _home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", home.path()); + let _reset = ProviderCatalogTestReset; + crate::provider_catalog_live::reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + + let now = Utc::now(); + let fetched_at = u64::try_from(now.timestamp()).expect("nonnegative timestamp"); + let routes = [ + provider_live_usage_envelope( + ApiProvider::Openrouter, + ApiProvider::Openrouter.as_str(), + "synthetic-openrouter-legacy", + &codewhale_config::catalog::base_url_fingerprint( + crate::config::DEFAULT_OPENROUTER_BASE_URL, + ), + Some(crate::pricing::AGGREGATOR_BILLING_SURFACE), + RouteBillingMode::Metered, + now, + ), + custom_usage_envelope( + codewhale_config::BASETEN_TEMPLATE_ID, + "synthetic-baseten-legacy", + &codewhale_config::catalog::base_url_fingerprint( + codewhale_config::BASETEN_BASE_URL, + ), + RouteBillingMode::Unknown, + now, + ), + ]; + assert!( + routes + .iter() + .all(|route| route.provider_live_pricing.is_none()) + ); + + for route in &routes { + crate::provider_catalog_live::record_success(priced_provider_delta( + &route.provider_identity, + &route.model, + route.endpoint_fingerprint.as_deref().expect("fingerprint"), + fetched_at, + )); + let audit = route.audit(&Usage { + input_tokens: 1_000_000, + ..Usage::default() + }); + assert_eq!( + audit.unpriced_reason, + Some(if route.provider == ApiProvider::Openrouter { + crate::pricing::UnpricedReason::NoPricingRow + } else { + crate::pricing::UnpricedReason::UnverifiedLivePricing + }), + "a completion-time refresh must not price {route:?}" + ); + } + } + + #[test] + fn openrouter_offline_bundled_price_is_immutable_after_dispatch() { + let _env = crate::test_support::lock_test_env(); + let _live = crate::provider_lake::lock_live_snapshot(); + let home = tempfile::tempdir().expect("test home"); + let _home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", home.path()); + let _reset = ProviderCatalogTestReset; + crate::provider_catalog_live::reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + + let dispatched_at = Utc::now(); + let fetched_at = u64::try_from(dispatched_at.timestamp()).expect("timestamp"); + let model = "qwen/qwen3.8-flash"; + let fingerprint = codewhale_config::catalog::base_url_fingerprint( + crate::config::DEFAULT_OPENROUTER_BASE_URL, + ); + let route = provider_live_usage_envelope( + ApiProvider::Openrouter, + ApiProvider::Openrouter.as_str(), + model, + &fingerprint, + Some(crate::pricing::AGGREGATOR_BILLING_SURFACE), + RouteBillingMode::Metered, + dispatched_at, + ); + assert!(route.provider_live_pricing.is_none()); + + let usage = Usage { + input_tokens: 1_000_000, + ..Usage::default() + }; + let offline = route.audit(&usage); + assert_eq!( + offline.estimate.expect("bundled OpenRouter price").usd, + 0.16 + ); + assert_eq!( + offline.provenance, + Some(codewhale_config::pricing::PricingProvenance::ModelsDevBundled) + ); + + // A later mutable refresh cannot change a turn that had no quote at + // the application-dispatch boundary. + crate::provider_catalog_live::record_success(priced_provider_delta_with_rates( + ApiProvider::Openrouter.as_str(), + model, + &fingerprint, + fetched_at, + 19.0, + 29.0, + )); + let after_refresh = route.audit(&usage); + assert_eq!(after_refresh, offline); + + // Admission without provider usage does not create a charge. + let no_usage = route.audit(&Usage::default()); + let no_usage_estimate = no_usage.estimate.expect("known zero usage is priced"); + assert_eq!(no_usage_estimate.usd, 0.0); + assert_eq!(no_usage_estimate.cny, 0.0); + } + + #[test] + fn provider_live_quotes_reject_future_prices_and_every_route_binding_mismatch() { + let _env = crate::test_support::lock_test_env(); + let _live = crate::provider_lake::lock_live_snapshot(); + let home = tempfile::tempdir().expect("test home"); + let _home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", home.path()); + let _reset = ProviderCatalogTestReset; + crate::provider_catalog_live::reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + + let dispatched_at = Utc::now(); + let dispatch_unix = u64::try_from(dispatched_at.timestamp()).expect("timestamp"); + let future_at = dispatched_at + chrono::Duration::seconds(1); + let future_unix = dispatch_unix.saturating_add(1); + let cases = [ + ( + ApiProvider::Openrouter, + ApiProvider::Openrouter.as_str(), + "synthetic-openrouter-future", + codewhale_config::catalog::base_url_fingerprint( + crate::config::DEFAULT_OPENROUTER_BASE_URL, + ), + crate::pricing::AGGREGATOR_BILLING_SURFACE, + RouteBillingMode::Metered, + ), + ( + ApiProvider::Custom, + codewhale_config::BASETEN_TEMPLATE_ID, + "synthetic-baseten-future", + codewhale_config::catalog::base_url_fingerprint(codewhale_config::BASETEN_BASE_URL), + crate::pricing::UNCLASSIFIED_BILLING_SURFACE, + RouteBillingMode::Unknown, + ), + ]; + let usage = Usage { + input_tokens: 1_000_000, + ..Usage::default() + }; + + for (provider, identity, model, fingerprint, surface, mode) in cases { + crate::provider_catalog_live::record_success(priced_provider_delta( + identity, + model, + &fingerprint, + future_unix, + )); + let no_future_quote = provider_live_usage_envelope( + provider, + identity, + model, + &fingerprint, + Some(surface), + mode, + dispatched_at, + ); + assert!(no_future_quote.provider_live_pricing.is_none()); + assert_eq!( + no_future_quote.audit(&usage).unpriced_reason, + Some(if provider == ApiProvider::Openrouter { + crate::pricing::UnpricedReason::NoPricingRow + } else { + crate::pricing::UnpricedReason::UnverifiedLivePricing + }) + ); + + let captured = provider_live_usage_envelope( + provider, + identity, + model, + &fingerprint, + Some(surface), + mode, + future_at, + ); + assert!(captured.provider_live_pricing.is_some()); + + let mut future_relative_to_dispatch = captured.clone(); + future_relative_to_dispatch.dispatched_at = dispatched_at; + assert_eq!( + future_relative_to_dispatch.audit(&usage).unpriced_reason, + Some(crate::pricing::UnpricedReason::UnverifiedLivePricing) + ); + let persisted = serde_json::to_value(&future_relative_to_dispatch) + .expect("invalid future quote serializes only as absent"); + assert!(persisted["provider_live_pricing"].is_null()); + + let mut wrong_model = captured.clone(); + wrong_model.model.push_str("-other"); + assert_eq!( + wrong_model.audit(&usage).unpriced_reason, + Some(crate::pricing::UnpricedReason::UnverifiedLivePricing) + ); + + let mut wrong_identity = captured.clone(); + wrong_identity.provider_identity.push_str("-other"); + assert_eq!( + wrong_identity.audit(&usage).unpriced_reason, + Some(if provider == ApiProvider::Custom { + crate::pricing::UnpricedReason::UnknownBillingBasis + } else { + crate::pricing::UnpricedReason::UnverifiedLivePricing + }) + ); + + let mut wrong_endpoint = captured; + wrong_endpoint.endpoint_fingerprint = Some( + codewhale_config::catalog::base_url_fingerprint("https://proxy.example/v1"), + ); + assert_eq!( + wrong_endpoint.audit(&usage).unpriced_reason, + Some(if provider == ApiProvider::Custom { + crate::pricing::UnpricedReason::UnknownBillingBasis + } else { + crate::pricing::UnpricedReason::UnverifiedLivePricing + }) + ); + } + } + + #[test] + fn provider_live_quote_serialization_is_secret_free_and_legacy_compatible() { + let _env = crate::test_support::lock_test_env(); + let _live = crate::provider_lake::lock_live_snapshot(); + let home = tempfile::tempdir().expect("test home"); + let _home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", home.path()); + let _reset = ProviderCatalogTestReset; + crate::provider_catalog_live::reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + + let now = Utc::now(); + let fetched_at = u64::try_from(now.timestamp()).expect("nonnegative timestamp"); + let model = "synthetic-baseten-serialized-quote"; + let fingerprint = + codewhale_config::catalog::base_url_fingerprint(codewhale_config::BASETEN_BASE_URL); + crate::provider_catalog_live::record_success(priced_provider_delta( + codewhale_config::BASETEN_TEMPLATE_ID, + model, + &fingerprint, + fetched_at, + )); + let route = custom_usage_envelope( + codewhale_config::BASETEN_TEMPLATE_ID, + model, + &fingerprint, + RouteBillingMode::Unknown, + now, + ); + assert!(route.provider_live_pricing.is_some()); + + let serialized = serde_json::to_string(&route).expect("serialize frozen route"); + assert!(serialized.contains("provider_live_pricing")); + assert!(serialized.contains("catalog_revision")); + assert!(serialized.contains("input_per_million")); + for secret in [codewhale_config::BASETEN_BASE_URL, "api_key", "Bearer "] { + // The assertion message must not itself become a logging sink for + // the credential fragment it checks for — name the check, not the + // secret. + assert!( + !serialized.contains(secret), + "frozen route serialization leaked a credential fragment" + ); + } + + let mut child = serde_json::json!({}); + attach_child_usage_metadata(&mut child, &route, &Usage::default()); + let child_route = child_route_envelope_from_metadata(&child).expect("child route"); + assert_eq!(child_route, route.sanitized_for_persistence()); + + let mut legacy: serde_json::Value = + serde_json::from_str(&serialized).expect("route JSON value"); + legacy + .as_object_mut() + .expect("route object") + .remove("provider_live_pricing"); + let legacy: EffectiveRouteEnvelope = + serde_json::from_value(legacy).expect("legacy route remains readable"); + assert!(legacy.provider_live_pricing.is_none()); + let audit = legacy.audit(&Usage { + input_tokens: 1_000_000, + ..Usage::default() + }); + assert_eq!( + audit.unpriced_reason, + Some(crate::pricing::UnpricedReason::UnverifiedLivePricing) + ); + + let mut wrong_model = route.clone(); + wrong_model.model.push_str("-other"); + assert_eq!( + wrong_model.audit(&Usage::default()).unpriced_reason, + Some(crate::pricing::UnpricedReason::UnverifiedLivePricing) + ); + } + + #[test] + fn routed_child_batch_is_preferred_bounded_and_sanitized() { + let route = deepseek_envelope(); + let records = vec![ + RuntimeUsageRecord { + source_id: "raw-provider-response-id-one".to_string(), + usage: EffectiveRouteUsage { + route: route.clone(), + usage: Usage { + input_tokens: 11, + ..Usage::default() + }, + }, + }, + RuntimeUsageRecord { + source_id: "raw-provider-response-id-two".to_string(), + usage: EffectiveRouteUsage { + route: route.clone(), + usage: Usage { + output_tokens: 7, + ..Usage::default() + }, + }, + }, + ]; + let mut metadata = serde_json::json!({}); + attach_child_usage_metadata(&mut metadata, &route, &Usage::default()); + attach_child_usage_batch_metadata( + &mut metadata, + &RuntimeUsageBatch { + records, + drop_records: Vec::new(), + dropped_records: 0, + }, + ); + + let serialized = serde_json::to_string(&metadata).expect("batch metadata"); + assert!(!serialized.contains("raw-provider-response-id")); + let batch = child_usage_records_from_metadata(&metadata).expect("preferred batch"); + assert_eq!(batch.records.len(), 2); + assert_eq!(batch.records[0].usage.usage.input_tokens, 11); + assert_eq!(batch.records[1].usage.usage.output_tokens, 7); + assert_eq!(batch.dropped_records, 0); + + metadata[CHILD_USAGE_RECORDS_KEY] = serde_json::json!([{"bad": true}]); + let malformed = child_usage_records_from_metadata(&metadata).expect("batch key wins"); + assert!(malformed.records.is_empty()); + assert_eq!(malformed.dropped_records, 1); + } + fn deepseek() -> BackgroundRoute<'static> { BackgroundRoute::new(ApiProvider::Deepseek, "deepseek-v4-flash") .with_base_url(Some(crate::config::DEFAULT_DEEPSEEK_BASE_URL)) @@ -1520,6 +2979,7 @@ mod tests { model: "/Volumes/private/checkpoints/model.gguf".to_string(), billing_surface: None, endpoint_fingerprint: None, + provider_live_pricing: None, billing_mode: RouteBillingMode::Metered, dispatched_at: Utc::now(), }; @@ -1586,6 +3046,7 @@ mod tests { "https://alice:password@example.test/v1?token=secret#fragment".to_string(), ), endpoint_fingerprint: Some("../.ssh/provider_key".to_string()), + provider_live_pricing: None, billing_mode: RouteBillingMode::Metered, dispatched_at: Utc::now(), }; @@ -1663,6 +3124,118 @@ mod tests { assert_eq!(drain().priced_turns, 1); } + #[test] + fn detached_advisor_and_translation_receipts_survive_new_exactly_once() { + let _g = test_scope(); + let tmp = tempfile::tempdir().expect("tempdir"); + let manager = crate::session_manager::SessionManager::new(tmp.path().join("sessions")) + .expect("session manager"); + let old_session_id = "origin-session"; + let new_session_id = "replacement-session"; + for session_id in [old_session_id, new_session_id] { + let session = crate::session_manager::create_saved_session_with_id_and_mode( + session_id.to_string(), + &[], + "deepseek-v4-flash", + tmp.path(), + 0, + None, + Some("agent"), + ); + manager.save_session(&session).expect("save session"); + } + + let origin_scope = scope_token(); + let owner = "interactive:origin-session:origin-turn"; + register_persistent_interactive_runtime_usage_sink_at( + owner, + origin_scope, + old_session_id, + "origin-turn", + manager.sessions_dir().to_path_buf(), + ); + let advisor_lease = acquire_runtime_usage_lease(owner).expect("advisor owner lease"); + finish_runtime_usage_owner(owner); + + // `/new` closes the old foreground generation while the detached + // advisor and translation requests are still in flight. + assert!(close_current_scope().is_empty()); + let route = deepseek_envelope(); + let usage = Usage { + input_tokens: 17, + output_tokens: 5, + ..Usage::default() + }; + for _ in 0..2 { + report_effective_route_for_runtime( + origin_scope, + Some(owner), + "advisor:origin-turn:response", + &route, + &usage, + ); + report_unreceipted_provider_success( + origin_scope, + Some(owner), + "advisor:origin-turn:missing-usage", + &route, + ); + assert!(report_effective_route_for_interactive_origin_with_manager( + origin_scope, + old_session_id, + "origin-turn", + "translation:origin-turn:assistant", + &route, + &usage, + &manager, + )); + assert!(report_unreceipted_for_interactive_origin_with_manager( + origin_scope, + old_session_id, + "origin-turn", + "translation:origin-turn:thinking-missing-usage", + &route, + &manager, + )); + } + drop(advisor_lease); + + let fallback = take_runtime_usage(owner); + assert!(fallback.records.is_empty()); + assert!(fallback.drop_records.is_empty()); + assert_eq!(fallback.dropped_records, 0); + assert!(drain().is_empty(), "late receipts polluted the new scope"); + + let old = manager + .load_session_snapshot(old_session_id) + .expect("load origin session"); + assert_eq!(old.metadata.total_tokens, 44); + assert_eq!(old.metadata.cost.priced_turns, 2); + assert_eq!(old.metadata.cost.unpriced_turns, 2); + assert_eq!(old.metadata.cost.cny_unpriced_turns, 2); + assert_eq!(old.metadata.cost.usage_source_fingerprints.len(), 4); + + let replay = manager + .load_session_snapshot(old_session_id) + .expect("replay origin session"); + assert_eq!(replay.metadata.total_tokens, 44); + assert_eq!(replay.metadata.cost.usage_source_fingerprints.len(), 4); + + let replacement = manager + .load_session_snapshot(new_session_id) + .expect("load replacement session"); + assert_eq!(replacement.metadata.total_tokens, 0); + assert_eq!(replacement.metadata.cost.priced_turns, 0); + assert_eq!(replacement.metadata.cost.unpriced_turns, 0); + assert!( + replacement + .metadata + .cost + .usage_source_fingerprints + .is_empty() + ); + } + #[test] fn report_counts_unknown_models_as_missing_spend_not_as_free() { let _g = test_scope(); diff --git a/crates/tui/src/exec_agent.rs b/crates/tui/src/exec_agent.rs index d40d51c213..3b7489c80e 100644 --- a/crates/tui/src/exec_agent.rs +++ b/crates/tui/src/exec_agent.rs @@ -397,6 +397,7 @@ pub(crate) async fn run_exec_agent( mode, route: Box::new(validated_route.into_resolved()), compaction: Box::new(compaction.clone()), + initial_routed_usage: Box::default(), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, diff --git a/crates/tui/src/lib.rs b/crates/tui/src/lib.rs index 50c8f2fee2..0dccca53be 100644 --- a/crates/tui/src/lib.rs +++ b/crates/tui/src/lib.rs @@ -92,6 +92,7 @@ mod project_context; mod project_context_cache; mod prompt_zones; mod prompts; +mod provider_catalog_live; mod provider_lake; mod provider_readiness; mod purge; diff --git a/crates/tui/src/model_catalog.rs b/crates/tui/src/model_catalog.rs index 99f5d538e2..b62d14b1d7 100644 --- a/crates/tui/src/model_catalog.rs +++ b/crates/tui/src/model_catalog.rs @@ -1,9 +1,13 @@ -//! Offline model metadata catalog (#3072). +//! Legacy offline model metadata compatibility catalog (#3072). //! //! This module adds a secret-free metadata layer in front of the legacy model //! tables. It is intentionally conservative: startup reads a local cache plus a //! bundled snapshot, never performs a network refresh, and only overrides a //! legacy fact when the active catalog entry actually carries that field. +//! The unscoped `catalog/openrouter.json` file cannot safely own a live provider +//! roster; new provider/base-URL-scoped refreshes and runtime consumers use +//! `provider_catalog_live` + `provider_lake`. Keep this reader only until the +//! remaining `models`, `pricing`, and safe-label compatibility callers migrate. use std::collections::BTreeMap; use std::path::PathBuf; diff --git a/crates/tui/src/model_routing.rs b/crates/tui/src/model_routing.rs index 45a78dd356..80e1a9b23b 100644 --- a/crates/tui/src/model_routing.rs +++ b/crates/tui/src/model_routing.rs @@ -8,8 +8,11 @@ use std::time::Duration; use anyhow::Result; use serde::{Deserialize, Serialize}; -use crate::client::DeepSeekClient; +use crate::client::{DeepSeekClient, WireDialect}; use crate::config::{ApiProvider, Config, normalize_model_name_for_provider}; +use crate::cost_status::{ + EffectiveRouteEnvelope, EffectiveRouteUsage, RuntimeUsageDropRecord, RuntimeUsageRecord, +}; use crate::llm_client::LlmClient; use crate::model_inventory::ModelInventory; use crate::models::Role; @@ -514,6 +517,21 @@ pub(crate) struct AutoRouteSelection { /// Present for Auto decisions; explicit inventory lookups intentionally do /// not pretend to be Auto routing receipts. pub(crate) receipt: Option, + /// Provider calls made to choose this route. These are deliberately kept + /// separate from the selected parent route: a classifier may run on a + /// different provider/model/quote, so pricing it under the eventual turn + /// would double-charge the parent and lose the classifier's real route. + /// + /// Auto currently admits at most one classifier request per selection. + pub(crate) routed_usage: Vec, + /// Exact frozen routes for admitted classifier calls whose provider + /// response omitted usage. The count below remains authoritative and may + /// exceed this bounded vector after overflow. + pub(crate) routed_usage_drop_records: Vec, + /// Classifier requests admitted to dispatch whose response usage could not + /// be recovered (timeout/transport failure). Consumers must surface this + /// as incomplete coverage rather than silently treating it as zero spend. + pub(crate) routed_usage_dropped_records: u64, } fn extract_first_json_object(raw: &str) -> Option<&str> { @@ -605,6 +623,17 @@ struct InventoryAutoRouteRecommendation { reasoning_effort: Option, } +/// One provider-backed classifier attempt. A provider-success response always +/// reaches this shape before its content is interpreted, so invalid JSON and +/// provider-declared incomplete output retain their exact routed usage. +#[derive(Debug, Clone, PartialEq, Eq)] +struct InventoryAutoRouteAttempt { + recommendation: Option, + routed_usage: Vec, + routed_usage_drop_records: Vec, + routed_usage_dropped_records: u64, +} + pub(crate) async fn resolve_auto_route_with_inventory( config: &Config, latest_request: &str, @@ -679,8 +708,11 @@ pub(crate) async fn resolve_auto_route_with_inventory_for_session_and_cache_poli ) .await { - Ok(Some(recommendation)) => auto_route_from_classifier(&inventory, recommendation), - Ok(None) | Err(_) => auto_route_classifier_fallback(heuristic, &inventory), + Ok(attempt) => auto_route_from_classifier_attempt(heuristic, &inventory, attempt), + // Client construction/preparation failed before a provider request was + // admitted. There is no provider usage to invent and no dropped + // response receipt to claim. + Err(_) => auto_route_classifier_fallback(heuristic, &inventory), }; Ok(normalize_auto_route_selection_for_config(config, selection)) } @@ -714,6 +746,9 @@ pub(crate) fn resolve_explicit_route_with_inventory( }), source: AutoRouteSource::Heuristic, receipt: None, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, }); } @@ -739,6 +774,9 @@ pub(crate) fn resolve_explicit_route_with_inventory( }), source: AutoRouteSource::Heuristic, receipt: None, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, }) } @@ -794,6 +832,9 @@ fn auto_route_from_inventory_heuristic( model, reasoning_effort: Some(crate::auto_reasoning::select(false, latest_request)), source: AutoRouteSource::Heuristic, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, }; }; // Use the candidates' cheap/big info for complexity-based routing. @@ -829,6 +870,9 @@ fn auto_route_from_inventory_heuristic( model: decision.model, reasoning_effort: Some(crate::auto_reasoning::select(false, latest_request)), source: AutoRouteSource::Heuristic, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, } } @@ -860,9 +904,33 @@ fn auto_route_from_classifier( model: recommendation.model, reasoning_effort: recommendation.reasoning_effort, source: AutoRouteSource::FlashRouter, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, } } +fn auto_route_from_classifier_attempt( + heuristic: AutoRouteSelection, + inventory: &ModelInventory, + attempt: InventoryAutoRouteAttempt, +) -> AutoRouteSelection { + let InventoryAutoRouteAttempt { + recommendation, + routed_usage, + routed_usage_drop_records, + routed_usage_dropped_records, + } = attempt; + let mut selection = recommendation.map_or_else( + || auto_route_classifier_fallback(heuristic, inventory), + |recommendation| auto_route_from_classifier(inventory, recommendation), + ); + selection.routed_usage = routed_usage; + selection.routed_usage_drop_records = routed_usage_drop_records; + selection.routed_usage_dropped_records = routed_usage_dropped_records; + selection +} + fn auto_route_classifier_fallback( mut heuristic: AutoRouteSelection, inventory: &ModelInventory, @@ -970,6 +1038,135 @@ fn auto_route_pair( AutoRoutePair { strong, fast } } +fn auto_route_usage_has_reported_data(usage: &crate::models::Usage) -> bool { + usage.input_tokens > 0 + || usage.output_tokens > 0 + || usage.prompt_cache_hit_tokens.is_some() + || usage.prompt_cache_miss_tokens.is_some() + || usage.prompt_cache_write_tokens.is_some() + || usage.reasoning_tokens.is_some() + || usage.reasoning_replay_tokens.is_some() + || usage.server_tool_use.is_some() +} + +/// Stable, persistence-safe identity for one classifier response. The raw +/// provider response id is hashed with the frozen dispatch route and instant; +/// neither it nor any custom route label crosses into telemetry/persistence. +fn auto_route_usage_source_id(route: &EffectiveRouteEnvelope, response_id: &str) -> String { + use sha2::{Digest as _, Sha256}; + + let mut digest = Sha256::new(); + let dispatched_at = route.dispatched_at.to_rfc3339(); + for part in [ + b"codewhale:auto-route-classifier:v1".as_slice(), + route.provider.as_str().as_bytes(), + route.provider_identity.as_bytes(), + route.model.as_bytes(), + route + .endpoint_fingerprint + .as_deref() + .unwrap_or_default() + .as_bytes(), + dispatched_at.as_bytes(), + response_id.as_bytes(), + ] { + digest.update((part.len() as u64).to_le_bytes()); + digest.update(part); + } + format!( + "auto-router:{}", + crate::hashing::hex_bytes(digest.finalize()) + ) +} + +fn auto_route_attempt_from_response( + request_route: EffectiveRouteEnvelope, + response: &MessageResponse, + inventory: &ModelInventory, +) -> InventoryAutoRouteAttempt { + // All-zero usage cannot price a routed segment. The dispatch caller owns + // the stronger cache/provenance context and must explicitly classify this + // as either a proven cache replay or missing provider billing evidence. + let routed_usage = auto_route_usage_has_reported_data(&response.usage) + .then(|| RuntimeUsageRecord { + source_id: auto_route_usage_source_id(&request_route, &response.id), + usage: EffectiveRouteUsage { + route: request_route.sanitized_for_persistence(), + usage: response.usage.clone(), + }, + }) + .into_iter() + .collect(); + let recommendation = + (!crate::models::is_incomplete_stop_reason(response.stop_reason.as_deref())) + .then(|| { + parse_inventory_auto_route_recommendation( + &message_response_text(response), + inventory, + ) + }) + .flatten(); + InventoryAutoRouteAttempt { + recommendation, + routed_usage, + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, + } +} + +fn auto_route_attempt_from_provider_response( + request_route: EffectiveRouteEnvelope, + response: &MessageResponse, + inventory: &ModelInventory, +) -> InventoryAutoRouteAttempt { + let drop_route = request_route.sanitized_for_persistence(); + let mut attempt = auto_route_attempt_from_response(request_route, response, inventory); + // This classifier request is currently not cacheable (temperature is + // provider-default, not the deterministic Some(0.0) cache contract), so a + // decoded all-zero response is missing provider billing evidence even when + // the caller permits response-cache use. Do not silently reinterpret the + // policy boolean as cache-hit provenance. + if attempt.routed_usage.is_empty() { + attempt.routed_usage_drop_records = vec![RuntimeUsageDropRecord { + source_id: auto_route_usage_source_id( + &drop_route, + &format!("missing-usage:{}", response.id), + ), + route: drop_route, + }]; + attempt.routed_usage_dropped_records = 1; + } + attempt +} + +fn auto_route_attempt_with_dropped_response( + request_route: EffectiveRouteEnvelope, +) -> InventoryAutoRouteAttempt { + let request_route = request_route.sanitized_for_persistence(); + InventoryAutoRouteAttempt { + recommendation: None, + routed_usage: Vec::new(), + routed_usage_drop_records: vec![RuntimeUsageDropRecord { + source_id: auto_route_usage_source_id(&request_route, "transport-error"), + route: request_route, + }], + routed_usage_dropped_records: 1, + } +} + +/// Prove that the deterministic request seam accepts this classifier request +/// before capturing a quote or entering any provider permit/network path. +/// Blocking Cloud Code is rejected here too: that dialect is stream-only, so +/// `create_message` would otherwise fail locally after the apparent dispatch +/// boundary and incorrectly look like missing provider usage. +fn preflight_auto_route_request(client: &DeepSeekClient, request: &MessageRequest) -> Result<()> { + let prepared = client.prepare_outbound_request(request.clone(), false)?; + if prepared.dialect == WireDialect::GoogleCloudCode { + anyhow::bail!("auto-route classifier requires a blocking-capable provider route"); + } + Ok(()) +} + #[allow(clippy::too_many_arguments)] async fn auto_route_inventory_recommendation( config: &Config, @@ -980,7 +1177,7 @@ async fn auto_route_inventory_recommendation( selected_model_mode: &str, selected_thinking_mode: &str, allow_response_cache: bool, -) -> Result> { +) -> Result { let mut router_config = config.clone(); // The classifier runs on the inventory's router route: the explicit // [auto.router] route when configured, else the DeepSeek flash default. @@ -997,8 +1194,7 @@ async fn auto_route_inventory_recommendation( selected_model_mode, selected_thinking_mode, ); - let request_route = - client.effective_route_envelope(&inventory.router_model, chrono::Utc::now()); + let max_tokens = client.effective_max_output_tokens(&inventory.router_model); let request = MessageRequest { model: inventory.router_model.to_string(), messages: vec![Message { @@ -1008,7 +1204,7 @@ async fn auto_route_inventory_recommendation( cache_control: None, }], }], - max_tokens: client.effective_max_output_tokens(&request_route.model), + max_tokens, system: Some(SystemPrompt::Text(router_system)), tools: None, tool_choice: None, @@ -1025,27 +1221,49 @@ async fn auto_route_inventory_recommendation( top_p: None, }; + // Freeze pricing at the last application seam before the provider future + // starts. Prompt shaping above may be slow and may overlap a catalog + // refresh; completion-time mutable catalog state must never reprice this + // already-admitted classifier request. + preflight_auto_route_request(&client, &request)?; + let request_route = + client.effective_route_envelope(&inventory.router_model, chrono::Utc::now()); let response = if allow_response_cache { tokio::time::timeout( Duration::from_secs(inventory.router_timeout_secs), client.create_message(request), ) - .await?? + .await } else { tokio::time::timeout( Duration::from_secs(inventory.router_timeout_secs), client.create_message_without_response_cache(request), ) - .await?? + .await }; - if crate::models::is_incomplete_stop_reason(response.stop_reason.as_deref()) { - anyhow::bail!( - "auto-route classifier response incomplete: provider stop reason `{}`", - crate::models::stop_reason_detail(response.stop_reason.as_deref()) - ); - } - Ok(parse_inventory_auto_route_recommendation( - &message_response_text(&response), + let response = match response { + Ok(Ok(response)) => response, + // The request crossed Codewhale's dispatch boundary, but no exact + // provider usage came back. Preserve the fallback while explicitly + // failing cost coverage closed. + Ok(Err(_)) => return Ok(auto_route_attempt_with_dropped_response(request_route)), + // The local deadline cancels the future and can fire while the request + // is still waiting on an application/provider permit. With no response + // evidence we must not invent a provider call or a missing-usage + // receipt. Transport errors returned by the client remain the + // conservative explicit-dropped path above. + Err(_) => { + return Ok(InventoryAutoRouteAttempt { + recommendation: None, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, + }); + } + }; + Ok(auto_route_attempt_from_provider_response( + request_route, + &response, inventory, )) } @@ -1217,6 +1435,297 @@ fn truncate_for_auto_router(text: &str, max_chars: usize) -> String { mod tests { use super::*; + struct ProviderCatalogReset; + + impl Drop for ProviderCatalogReset { + fn drop(&mut self) { + crate::provider_catalog_live::reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + } + } + + fn priced_openrouter_delta( + model: &str, + fingerprint: &str, + fetched_at: u64, + input: f64, + output: f64, + ) -> codewhale_config::catalog::ProviderCatalogDelta { + use codewhale_config::catalog::{CatalogOffering, CatalogSource, ProviderCatalogDelta}; + + ProviderCatalogDelta { + provider: ApiProvider::Openrouter.as_str().to_string(), + base_url_fingerprint: fingerprint.to_string(), + fetched_at, + offerings: vec![CatalogOffering { + provider: ApiProvider::Openrouter.as_str().to_string(), + wire_model_id: model.to_string(), + endpoint_key: "chat".to_string(), + source: CatalogSource::Live { + base_url_fingerprint: fingerprint.to_string(), + fetched_at, + }, + cost: Some(codewhale_config::models_dev::ModelsDevCost { + input: Some(input), + output: Some(output), + cache_read: Some(input / 2.0), + cache_write: None, + }), + ..CatalogOffering::default() + }], + } + } + + fn classifier_response( + id: &str, + text: &str, + stop_reason: &str, + usage: crate::models::Usage, + ) -> MessageResponse { + MessageResponse { + id: id.to_string(), + r#type: "message".to_string(), + role: "assistant".to_string(), + content: vec![ContentBlock::Text { + text: text.to_string(), + cache_control: None, + }], + model: "router-response-alias-must-not-price".to_string(), + stop_reason: Some(stop_reason.to_string()), + stop_sequence: None, + container: None, + usage, + } + } + + #[test] + fn classifier_semantic_fallbacks_keep_exact_quotes_and_replay_once() { + let _env_lock = crate::test_support::lock_test_env(); + let _live = crate::provider_lake::lock_live_snapshot(); + let home = tempfile::tempdir().expect("test home"); + let _home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", home.path()); + let _reset = ProviderCatalogReset; + crate::provider_catalog_live::reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + + let model = "synthetic/openrouter-auto-classifier"; + let config = Config { + provider: Some("openrouter".to_string()), + providers: Some(crate::config::ProvidersConfig { + openrouter: crate::config::ProviderConfig { + api_key: Some("test-openrouter-key".to_string()), + base_url: Some(crate::config::DEFAULT_OPENROUTER_BASE_URL.to_string()), + model: Some(model.to_string()), + ..Default::default() + }, + ..Default::default() + }), + auto: Some(crate::config::AutoConfig { + cost_saving: None, + cross_provider: None, + router: Some(crate::config::AutoRouterConfig { + provider: Some("openrouter".to_string()), + model: Some(model.to_string()), + thinking: Some("off".to_string()), + timeout_secs: None, + }), + }), + ..Default::default() + }; + let inventory = ModelInventory::from_config(&config); + assert!( + inventory + .candidate(ApiProvider::Openrouter, model) + .is_some() + ); + let client = DeepSeekClient::new(&config).expect("OpenRouter classifier client"); + let fingerprint = codewhale_config::catalog::base_url_fingerprint( + crate::config::DEFAULT_OPENROUTER_BASE_URL, + ); + let first_at = chrono::Utc::now(); + let fetched_at = u64::try_from(first_at.timestamp()).expect("nonnegative timestamp"); + + crate::provider_catalog_live::record_success(priced_openrouter_delta( + model, + &fingerprint, + fetched_at, + 1.0, + 4.0, + )); + let first_route = client.effective_route_envelope(model, first_at); + let valid = auto_route_attempt_from_response( + first_route, + &classifier_response( + "same-provider-response-id", + &format!(r#"{{"provider":"openrouter","model":"{model}","thinking":"off"}}"#), + "stop", + crate::models::Usage { + input_tokens: 10, + output_tokens: 2, + prompt_cache_hit_tokens: Some(3), + ..Default::default() + }, + ), + &inventory, + ); + + // Replace the live row in the same Unix second. The first routed + // record must keep its old immutable revision and the new attempt must + // freeze a distinct one at its own dispatch boundary. + crate::provider_catalog_live::record_success(priced_openrouter_delta( + model, + &fingerprint, + fetched_at, + 9.0, + 19.0, + )); + let second_at = first_at + chrono::Duration::nanoseconds(1); + let invalid = auto_route_attempt_from_response( + client.effective_route_envelope(model, second_at), + &classifier_response( + "same-provider-response-id", + "not valid route json", + "stop", + crate::models::Usage { + input_tokens: 11, + output_tokens: 3, + ..Default::default() + }, + ), + &inventory, + ); + let incomplete = auto_route_attempt_from_response( + client.effective_route_envelope(model, second_at + chrono::Duration::nanoseconds(1)), + &classifier_response( + "same-provider-response-id", + &format!(r#"{{"provider":"openrouter","model":"{model}"}}"#), + "length", + crate::models::Usage { + input_tokens: 12, + output_tokens: 4, + ..Default::default() + }, + ), + &inventory, + ); + let missing_usage_route = + client.effective_route_envelope(model, second_at + chrono::Duration::nanoseconds(2)); + let missing_usage = auto_route_attempt_from_provider_response( + missing_usage_route.clone(), + &classifier_response( + "missing-usage-response-id", + "not valid route json", + "stop", + crate::models::Usage::default(), + ), + &inventory, + ); + assert!(missing_usage.routed_usage.is_empty()); + assert_eq!(missing_usage.routed_usage_dropped_records, 1); + assert_eq!(missing_usage.routed_usage_drop_records.len(), 1); + assert_eq!( + missing_usage.routed_usage_drop_records[0].route, + missing_usage_route.sanitized_for_persistence() + ); + assert!( + missing_usage.routed_usage_drop_records[0] + .source_id + .starts_with("auto-router:") + ); + assert!( + !missing_usage.routed_usage_drop_records[0] + .source_id + .contains("missing-usage-response-id") + ); + + let transport = auto_route_attempt_with_dropped_response( + client.effective_route_envelope(model, second_at + chrono::Duration::nanoseconds(3)), + ); + assert_eq!(transport.routed_usage_dropped_records, 1); + assert_eq!(transport.routed_usage_drop_records.len(), 1); + assert!(transport.routed_usage.is_empty()); + + let heuristic = auto_route_from_inventory_heuristic(&config, "quick status", &inventory); + let valid = auto_route_from_classifier_attempt(heuristic.clone(), &inventory, valid); + let invalid = auto_route_from_classifier_attempt(heuristic.clone(), &inventory, invalid); + let incomplete = auto_route_from_classifier_attempt(heuristic, &inventory, incomplete); + assert_eq!(valid.source, AutoRouteSource::FlashRouter); + for fallback in [&invalid, &incomplete] { + assert_eq!(fallback.source, AutoRouteSource::Heuristic); + assert!(matches!( + fallback.receipt.as_ref().map(|receipt| receipt.reason), + Some(AutoRouteReason::ClassifierFallback(_)) + )); + assert_eq!(fallback.routed_usage.len(), 1); + assert_eq!(fallback.routed_usage_dropped_records, 0); + } + assert_eq!(valid.routed_usage.len(), 1); + assert_eq!(valid.routed_usage[0].usage.usage.input_tokens, 10); + assert_eq!( + valid.routed_usage[0].usage.usage.prompt_cache_hit_tokens, + Some(3) + ); + + let first_quote = valid.routed_usage[0] + .usage + .route + .provider_live_pricing + .as_ref() + .expect("first exact quote"); + let second_quote = invalid.routed_usage[0] + .usage + .route + .provider_live_pricing + .as_ref() + .expect("replacement exact quote"); + assert_ne!(first_quote.catalog_revision, second_quote.catalog_revision); + assert_eq!(first_quote.input_per_million.as_deref(), Some("1")); + assert_eq!(second_quote.input_per_million.as_deref(), Some("9")); + + let records = valid + .routed_usage + .iter() + .chain(&invalid.routed_usage) + .chain(&incomplete.routed_usage) + .cloned() + .collect::>(); + assert_eq!(records.len(), 3); + assert!(records.iter().all(|record| { + record.source_id.starts_with("auto-router:") + && record.source_id.len() == "auto-router:".len() + 64 + && !record.source_id.contains("same-provider-response-id") + })); + + // Exercise the canonical sink exactly as selection consumers do: + // replaying any record cannot add parent-route spend or a second + // routed segment, while distinct dispatches remain distinct. + let _cost_scope = crate::cost_status::test_scope(); + let owner = "auto-router-selection-test-owner"; + crate::cost_status::register_interactive_runtime_usage_sink( + owner, + crate::cost_status::scope_token(), + ); + let lease = crate::cost_status::acquire_runtime_usage_lease(owner) + .expect("runtime usage owner lease"); + for record in &records { + for _ in 0..2 { + crate::cost_status::report_effective_route_for_runtime( + crate::cost_status::scope_token(), + Some(lease.owner()), + &record.source_id, + &record.usage.route, + &record.usage.usage, + ); + } + } + crate::cost_status::finish_runtime_usage_owner(owner); + drop(lease); + let pending = crate::cost_status::drain(); + assert_eq!(pending.usage_source_fingerprints.len(), records.len()); + assert_eq!(pending.priced_turns, records.len() as u32); + assert_eq!(pending.unpriced_turns, 0); + } + #[test] fn auto_model_reasoning_keeps_model_and_thinking_choices_independent() { assert_eq!( diff --git a/crates/tui/src/pricing.rs b/crates/tui/src/pricing.rs index b13d696c99..40afc3049d 100644 --- a/crates/tui/src/pricing.rs +++ b/crates/tui/src/pricing.rs @@ -1344,6 +1344,32 @@ pub(crate) fn audit_turn_cost_for_provider_on_endpoint_at( endpoint_fingerprint: Option<&str>, usage: &Usage, recorded_at: DateTime, +) -> TurnCostAudit { + audit_turn_cost_for_provider_on_endpoint_for_identity_at( + provider, + None, + model, + endpoint_fingerprint, + usage, + recorded_at, + ) +} + +/// Identity-aware provider audit for named compatible routes. +/// +/// `ApiProvider::Custom` is only a transport family, so it is never sufficient +/// pricing provenance on its own. Baseten is the first reviewed compatible +/// provider whose authenticated live catalog can price actual usage; every +/// other custom identity stays unknown until it receives an equivalent +/// provider/endpoint contract. +#[must_use] +fn audit_turn_cost_for_provider_on_endpoint_for_identity_at( + provider: ApiProvider, + provider_identity: Option<&str>, + model: &str, + endpoint_fingerprint: Option<&str>, + usage: &Usage, + recorded_at: DateTime, ) -> TurnCostAudit { if !usage_cache_partition_is_consistent(usage) { return TurnCostAudit::unpriced(UnpricedReason::InconsistentUsage); @@ -1351,6 +1377,12 @@ pub(crate) fn audit_turn_cost_for_provider_on_endpoint_at( if provider == ApiProvider::OpenaiCodex { return TurnCostAudit::unpriced(UnpricedReason::NotMoneyMetered); } + if provider == ApiProvider::Custom { + // A transport family plus current mutable catalog state is not a + // billing receipt. Reviewed custom routes are priced only by the + // frozen dispatch quote handled in the route-audit path below. + return TurnCostAudit::unpriced(UnpricedReason::UnknownBillingBasis); + } if route_requires_billing_surface(provider, model) { return TurnCostAudit::unpriced(UnpricedReason::AmbiguousBillingSurface); } @@ -1426,6 +1458,7 @@ pub(crate) fn audit_turn_cost_for_provider_on_endpoint_at( let mut live_defect = None; let offering = match verified_catalog_offering( provider, + provider_identity, &catalog_model, endpoint_fingerprint, recorded_at, @@ -1439,6 +1472,9 @@ pub(crate) fn audit_turn_cost_for_provider_on_endpoint_at( live_defect = Some(defect); None } + VerifiedOffering::FutureEffective => { + return TurnCostAudit::unpriced(UnpricedReason::UnverifiedLivePricing); + } VerifiedOffering::Absent => None, }; @@ -1518,6 +1554,9 @@ enum VerifiedOffering { }, /// The live row could not be verified and no bundled row exists. Unusable(LivePricingDefect), + /// The row claims it was fetched after this turn was dispatched. Clock + /// saturation must never turn a future price into an age-zero price. + FutureEffective, /// No catalog row for this provider/model at all. Absent, } @@ -1532,12 +1571,16 @@ enum VerifiedOffering { /// billing against a rate whose endpoint scope is unproven. fn verified_catalog_offering( provider: ApiProvider, + provider_identity: Option<&str>, catalog_model: &str, endpoint_fingerprint: Option<&str>, recorded_at: DateTime, ) -> VerifiedOffering { - let Some(offering) = crate::provider_lake::catalog_offering_for_model(provider, catalog_model) - else { + let Some(offering) = crate::provider_lake::catalog_offering_for_model_identity( + provider, + provider_identity, + catalog_model, + ) else { return VerifiedOffering::Absent; }; // Models.dev is a capabilities catalog. A live overlay from that fetch @@ -1562,6 +1605,14 @@ fn verified_catalog_offering( // `recorded_at` is the turn's own clock, which is the right reference for // "was this price current when the turn happened". let now_unix = u64::try_from(recorded_at.timestamp()).ok(); + if pricing.provenance == PricingProvenance::ProviderLive + && pricing + .effective_at + .zip(now_unix) + .is_some_and(|(effective_at, dispatched_at)| effective_at > dispatched_at) + { + return VerifiedOffering::FutureEffective; + } let Some(defect) = pricing.live_pricing_defect(endpoint_fingerprint, now_unix, LIVE_PRICING_MAX_AGE_SECS) else { @@ -1649,6 +1700,34 @@ pub(crate) fn audit_turn_cost_for_route_on_endpoint_at( usage: &Usage, recorded_at: DateTime, ) -> TurnCostAudit { + audit_turn_cost_for_route_on_endpoint_for_identity_at( + provider, + None, + model, + billing_surface, + endpoint_fingerprint, + None, + usage, + recorded_at, + ) +} + +/// Identity-aware route audit for an immutable dispatch receipt. +#[must_use] +pub(crate) fn audit_turn_cost_for_route_on_endpoint_for_identity_at( + provider: ApiProvider, + provider_identity: Option<&str>, + model: &str, + billing_surface: Option<&str>, + endpoint_fingerprint: Option<&str>, + provider_live_pricing: Option<&crate::provider_catalog_live::ProviderLivePricingQuote>, + usage: &Usage, + recorded_at: DateTime, +) -> TurnCostAudit { + let reviewed_custom_metered = + reviewed_custom_route_is_metered(provider, provider_identity, endpoint_fingerprint); + let reviewed_provider_live = + reviewed_provider_live_route_is_metered(provider, provider_identity, endpoint_fingerprint); // An explicitly recorded surface is evidence. Exact non-metered surfaces // override provider guesses; an explicit unknown/unrecognized surface must // fail closed and may never fall through to a familiar model's hand row. @@ -1656,7 +1735,7 @@ pub(crate) fn audit_turn_cost_for_route_on_endpoint_at( EndpointMetering::ExactSubscription | EndpointMetering::LocalNoBill => { return TurnCostAudit::unpriced(UnpricedReason::NotMoneyMetered); } - EndpointMetering::Unknown if billing_surface.is_some() => { + EndpointMetering::Unknown if billing_surface.is_some() && !reviewed_custom_metered => { return TurnCostAudit::unpriced(UnpricedReason::UnknownBillingBasis); } EndpointMetering::Unknown | EndpointMetering::Money => {} @@ -1700,8 +1779,79 @@ pub(crate) fn audit_turn_cost_for_route_on_endpoint_at( if billing_surface.is_none() { return TurnCostAudit::unpriced(UnpricedReason::UnestablishedEndpoint); } - audit_turn_cost_for_provider_on_endpoint_at( + if reviewed_provider_live { + let Some(provider_identity) = provider_identity.map(str::trim).filter(|id| !id.is_empty()) + else { + return TurnCostAudit::unpriced(UnpricedReason::UnknownBillingBasis); + }; + let Some(endpoint_fingerprint) = endpoint_fingerprint else { + return TurnCostAudit::unpriced(UnpricedReason::UnknownBillingBasis); + }; + let Some(dispatched_at_unix) = u64::try_from(recorded_at.timestamp()).ok() else { + return TurnCostAudit::unpriced(UnpricedReason::UnverifiedLivePricing); + }; + let pricing = match provider_live_pricing { + Some(quote) => { + let Some(pricing) = quote.pricing_for_route( + provider, + provider_identity, + model, + endpoint_fingerprint, + dispatched_at_unix, + ) else { + return TurnCostAudit::unpriced(UnpricedReason::UnverifiedLivePricing); + }; + pricing + } + None if provider == ApiProvider::Openrouter => { + // An offline/startup OpenRouter dispatch has no mutable live + // quote to freeze. Audit it only against the immutable bundled + // snapshot (and provider-owned hand rows, if one is added), so + // a refresh that lands after dispatch cannot retro-price it. + return audit_openrouter_immutable_pricing(model, usage, recorded_at); + } + None => { + // Baseten has no reviewed immutable price card. Its compatible + // custom route therefore requires the exact frozen live quote. + return TurnCostAudit::unpriced(UnpricedReason::UnverifiedLivePricing); + } + }; + let classes = token_usage_for_pricing(usage); + let unpriced_classes = pricing.unpriced_used_classes(&classes); + if !unpriced_classes.is_empty() { + return TurnCostAudit::missing_classes(pricing.provenance, unpriced_classes); + } + let Some(amount) = pricing.estimate_cost(&classes) else { + return TurnCostAudit::unpriced(UnpricedReason::InvalidPricingRow); + }; + let (estimate, usd_priced, cny_priced) = match pricing.currency { + Currency::Usd => (CostEstimate::usd_only(amount), true, false), + Currency::Cny => ( + CostEstimate { + usd: 0.0, + cny: amount, + }, + false, + true, + ), + Currency::Other(_) => { + return TurnCostAudit::unpriced(UnpricedReason::UnsupportedCurrency); + } + }; + return TurnCostAudit::priced(estimate, pricing.provenance, usd_priced, cny_priced); + } + if provider == ApiProvider::Openrouter && provider_identity.is_some() { + // A persisted built-in OpenRouter receipt that is missing the exact + // official identity/endpoint binding (or its frozen quote) must not + // fall through to the mutable process-wide provider lake. + return TurnCostAudit::unpriced(UnpricedReason::UnverifiedLivePricing); + } + if provider == ApiProvider::Custom { + return TurnCostAudit::unpriced(UnpricedReason::UnknownBillingBasis); + } + audit_turn_cost_for_provider_on_endpoint_for_identity_at( provider, + provider_identity, model, endpoint_fingerprint, usage, @@ -1709,6 +1859,119 @@ pub(crate) fn audit_turn_cost_for_route_on_endpoint_at( ) } +/// Price an exact official OpenRouter route without consulting mutable live +/// catalog state. This is the no-quote application-dispatch fallback used when +/// CodeWhale starts offline or the provider refresh has not completed yet. +fn audit_openrouter_immutable_pricing( + model: &str, + usage: &Usage, + recorded_at: DateTime, +) -> TurnCostAudit { + let Some(canonical_model) = canonical_model_id_for_provider(ApiProvider::Openrouter, model) + else { + return TurnCostAudit::unpriced(UnpricedReason::NoPricingRow); + }; + let classes = token_usage_for_pricing(usage); + if let Some(offering) = crate::provider_lake::bundled_catalog_offering_for_model( + ApiProvider::Openrouter, + &canonical_model, + ) { + if let Some(audit) = invalid_catalog_pricing_audit(&offering) { + return audit; + } + if let Some(pricing) = effective_offering_pricing( + ApiProvider::Openrouter, + &canonical_model, + &offering, + &classes, + ) { + let unpriced_classes = pricing.unpriced_used_classes(&classes); + if !unpriced_classes.is_empty() { + return TurnCostAudit::missing_classes(pricing.provenance, unpriced_classes); + } + let Some(estimate) = catalog_cost_estimate_for_route( + ApiProvider::Openrouter, + &canonical_model, + &offering, + usage, + ) else { + return TurnCostAudit::unpriced(UnpricedReason::UnsupportedCurrency); + }; + let (usd_priced, cny_priced) = match pricing.currency { + Currency::Usd => (true, false), + Currency::Cny => (false, true), + Currency::Other(_) => { + return TurnCostAudit::unpriced(UnpricedReason::UnsupportedCurrency); + } + }; + return TurnCostAudit::priced(estimate, pricing.provenance, usd_priced, cny_priced); + } + } + + hand_priced_audit( + provider_owned_hand_pricing_at(ApiProvider::Openrouter, &canonical_model, recorded_at), + usage, + ) +} + +/// Whether a named compatible route has a reviewed per-token billing contract. +/// +/// Baseten is accepted only through its setup-template identity (including the +/// aliases that resolve to that canonical template) and the fingerprint of its +/// documented Model APIs endpoint. A generic custom table, a Baseten-like name, +/// or a Baseten identity pointed at another host cannot become metered merely by +/// publishing a priced `/models` row. +#[must_use] +pub(crate) fn reviewed_custom_route_is_metered( + provider: ApiProvider, + provider_identity: Option<&str>, + endpoint_fingerprint: Option<&str>, +) -> bool { + if provider != ApiProvider::Custom { + return false; + } + let is_baseten = provider_identity + .and_then(codewhale_config::provider_setup_template) + .is_some_and(|template| template.id == codewhale_config::BASETEN_TEMPLATE_ID); + if !is_baseten { + return false; + } + endpoint_fingerprint.is_some_and(|fingerprint| { + fingerprint + == codewhale_config::catalog::base_url_fingerprint(codewhale_config::BASETEN_BASE_URL) + }) +} + +/// Exact routes whose mutable provider-live rates must be frozen at the +/// pre-permit application-dispatch boundary. +/// +/// OpenRouter is accepted only as the built-in identity on its official API; +/// a custom table shadowing that name or an endpoint override is a different +/// billing contract. Baseten-compatible custom identities follow the reviewed +/// setup template but retain their exact, case-sensitive cache ownership. +#[must_use] +fn reviewed_provider_live_route_is_metered( + provider: ApiProvider, + provider_identity: Option<&str>, + endpoint_fingerprint: Option<&str>, +) -> bool { + match provider { + ApiProvider::Openrouter => { + provider_identity.map(str::trim) == Some(ApiProvider::Openrouter.as_str()) + && endpoint_fingerprint.is_some_and(|fingerprint| { + fingerprint + == codewhale_config::catalog::base_url_fingerprint( + crate::config::DEFAULT_OPENROUTER_BASE_URL, + ) + }) + } + ApiProvider::Custom => { + reviewed_custom_route_is_metered(provider, provider_identity, endpoint_fingerprint) + } + _ => false, + } +} + /// Audit a turn against the route's billing presentation. /// /// The three non-metered presentations are **not** interchangeable, and @@ -4630,6 +4893,62 @@ mod tests { ); } + #[test] + fn future_effective_provider_live_rate_is_not_treated_as_age_zero() { + let _live = crate::provider_lake::lock_live_snapshot(); + crate::provider_lake::clear_live_snapshot(); + let dispatched_at = Utc::now(); + let future_fetched_at = u64::try_from(dispatched_at.timestamp()) + .expect("timestamp") + .saturating_add(1); + let fingerprint = codewhale_config::catalog::base_url_fingerprint( + crate::config::DEFAULT_FIREWORKS_BASE_URL, + ); + crate::provider_lake::set_live_snapshot( + codewhale_config::catalog::CatalogSnapshot { + offerings: vec![codewhale_config::catalog::CatalogOffering { + provider: "fireworks".to_string(), + wire_model_id: "accounts/fireworks/models/future-price-only".to_string(), + endpoint_key: "chat".to_string(), + cost: Some(codewhale_config::models_dev::ModelsDevCost { + input: Some(9.0), + output: Some(18.0), + cache_read: Some(1.0), + cache_write: None, + }), + source: codewhale_config::catalog::CatalogSource::Live { + base_url_fingerprint: fingerprint.clone(), + fetched_at: future_fetched_at, + }, + ..Default::default() + }], + }, + crate::provider_lake::LiveSource::PerProvider, + ); + + let audit = audit_turn_cost_for_route_on_endpoint_at( + ApiProvider::Fireworks, + "accounts/fireworks/models/future-price-only", + billing_surface_for_route( + ApiProvider::Fireworks, + Some(crate::config::DEFAULT_FIREWORKS_BASE_URL), + ), + Some(&fingerprint), + &million_input_usage(), + dispatched_at, + ); + crate::provider_lake::clear_live_snapshot(); + + assert!( + !audit.is_priced(), + "future price must fail closed: {audit:?}" + ); + assert_eq!( + audit.unpriced_reason, + Some(UnpricedReason::UnverifiedLivePricing) + ); + } + #[test] fn models_dev_live_overlay_does_not_replace_bundled_catalog_rates() { let _live = crate::provider_lake::lock_live_snapshot(); diff --git a/crates/tui/src/provider_catalog_live.rs b/crates/tui/src/provider_catalog_live.rs new file mode 100644 index 0000000000..6cb9f821b7 --- /dev/null +++ b/crates/tui/src/provider_catalog_live.rs @@ -0,0 +1,1624 @@ +//! Durable, secret-free per-provider `/models` catalog cache. +//! +//! This is the persistence owner for [`codewhale_config::catalog::ProviderCatalogCache`]. +//! It replaces the previous process-only provider refresh path: successful +//! refreshes replace one exact `(provider identity, base URL fingerprint)` +//! partition, failures retain that partition's prior rows, and startup loads +//! only the active route's exact partition. Credentials authorize the fetch in +//! `client`; they never enter this module or its disk envelope. +//! +//! This deliberately does not import the legacy `model_catalog` cache at +//! `catalog/openrouter.json`: that file has no provider/base-URL scope, so +//! treating it as a provider-owned roster could leak stale facts across custom +//! endpoints. `model_catalog` remains a read-only compatibility fallback for +//! older model-metadata consumers while provider-lake/runtime consumers migrate; +//! `catalog/provider-catalogs.json` is the sole writer-owned live roster store. + +use std::collections::BTreeMap; +use std::fs::{self, OpenOptions}; +use std::io::Read as _; +use std::path::{Path, PathBuf}; +use std::sync::{LazyLock, RwLock}; + +use anyhow::{Context, Result}; +use codewhale_config::catalog::now_unix; +use codewhale_config::catalog::{ + CatalogRefreshError, CatalogSnapshot, CatalogStatus, ProviderCatalogCache, + ProviderCatalogDelta, base_url_fingerprint, +}; +use codewhale_config::persistence::atomic_write_json; +use codewhale_config::pricing::{Currency, OfferingPricing, PricingProvenance}; +use serde::{Deserialize, Serialize}; + +use crate::config::{ApiProvider, Config}; + +const CACHE_SCHEMA_VERSION: u32 = 1; +const CACHE_FILE: &str = "provider-catalogs.json"; +const MAX_CACHE_BYTES: u64 = 32 * 1024 * 1024; +const MAX_CACHE_SCOPES: usize = 64; +const MAX_CACHE_ROWS: usize = 50_000; + +#[derive(Debug, Clone, Copy)] +struct CachePersistenceLimits { + max_bytes: u64, + max_scopes: usize, + max_rows: usize, +} + +const CACHE_PERSISTENCE_LIMITS: CachePersistenceLimits = CachePersistenceLimits { + max_bytes: MAX_CACHE_BYTES, + max_scopes: MAX_CACHE_SCOPES, + max_rows: MAX_CACHE_ROWS, +}; + +/// Provider-owned catalogs are refreshed daily. Past-TTL rows remain visible +/// with an explicit stale receipt until a successful replacement arrives. +pub const DEFAULT_PROVIDER_CATALOG_TTL_SECS: u64 = 24 * 60 * 60; + +static CACHE: LazyLock> = + LazyLock::new(|| RwLock::new(ProviderCatalogCache::new())); +static REFRESH_GENERATIONS: LazyLock>> = + LazyLock::new(|| RwLock::new(BTreeMap::new())); + +#[derive(Debug, Clone)] +pub struct ProviderCatalogRefreshTicket { + provider: String, + generation: u64, +} + +/// Immutable, secret-free provider-live rate evidence captured at dispatch. +/// +/// Rates are stored as canonical decimal strings rather than `f64` so route +/// receipts retain exact equality and stable JSON. `catalog_revision` binds +/// every identity, scope, timestamp, currency, provenance, and rate field; it +/// therefore changes even when two refreshes land in the same Unix second. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderLivePricingQuote { + pub(crate) provider: ApiProvider, + pub(crate) provider_identity: String, + pub(crate) wire_model: String, + pub(crate) endpoint_fingerprint: String, + pub(crate) catalog_fetched_at: u64, + pub(crate) catalog_revision: String, + pub(crate) currency: Currency, + pub(crate) provenance: PricingProvenance, + pub(crate) input_per_million: Option, + pub(crate) output_per_million: Option, + pub(crate) cache_read_per_million: Option, + pub(crate) cache_write_per_million: Option, +} + +#[derive(Serialize, Deserialize)] +struct ProviderLivePricingQuoteWire { + provider: ApiProvider, + provider_identity: String, + wire_model: String, + endpoint_fingerprint: String, + catalog_fetched_at: u64, + catalog_revision: String, + currency: Currency, + provenance: PricingProvenance, + #[serde(default, skip_serializing_if = "Option::is_none")] + input_per_million: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + output_per_million: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + cache_read_per_million: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + cache_write_per_million: Option, +} + +impl From<&ProviderLivePricingQuote> for ProviderLivePricingQuoteWire { + fn from(quote: &ProviderLivePricingQuote) -> Self { + Self { + provider: quote.provider, + provider_identity: quote.provider_identity.clone(), + wire_model: quote.wire_model.clone(), + endpoint_fingerprint: quote.endpoint_fingerprint.clone(), + catalog_fetched_at: quote.catalog_fetched_at, + catalog_revision: quote.catalog_revision.clone(), + currency: quote.currency.clone(), + provenance: quote.provenance.clone(), + input_per_million: quote.input_per_million.clone(), + output_per_million: quote.output_per_million.clone(), + cache_read_per_million: quote.cache_read_per_million.clone(), + cache_write_per_million: quote.cache_write_per_million.clone(), + } + } +} + +impl Serialize for ProviderLivePricingQuote { + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + if !self.is_structurally_valid() { + return serializer.serialize_none(); + } + ProviderLivePricingQuoteWire::from(self).serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for ProviderLivePricingQuote { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + let wire = ProviderLivePricingQuoteWire::deserialize(deserializer)?; + let quote = Self { + provider: wire.provider, + provider_identity: wire.provider_identity, + wire_model: wire.wire_model, + endpoint_fingerprint: wire.endpoint_fingerprint, + catalog_fetched_at: wire.catalog_fetched_at, + catalog_revision: wire.catalog_revision, + currency: wire.currency, + provenance: wire.provenance, + input_per_million: wire.input_per_million, + output_per_million: wire.output_per_million, + cache_read_per_million: wire.cache_read_per_million, + cache_write_per_million: wire.cache_write_per_million, + }; + quote + .is_structurally_valid() + .then_some(quote) + .ok_or_else(|| serde::de::Error::custom("invalid provider-live pricing quote")) + } +} + +pub(crate) fn deserialize_optional_provider_live_pricing<'de, D>( + deserializer: D, +) -> std::result::Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + let value = Option::::deserialize(deserializer)?; + Ok(value.and_then(|value| serde_json::from_value(value).ok())) +} + +impl ProviderLivePricingQuote { + fn is_structurally_valid(&self) -> bool { + self.pricing_for_route( + self.provider, + &self.provider_identity, + &self.wire_model, + &self.endpoint_fingerprint, + self.catalog_fetched_at, + ) + .is_some() + } + fn canonical_rate(rate: Option) -> Option { + rate.map(|rate| rate.to_string()) + } + + fn revision_for( + provider: ApiProvider, + provider_identity: &str, + wire_model: &str, + endpoint_fingerprint: &str, + catalog_fetched_at: u64, + currency: &Currency, + provenance: &PricingProvenance, + input_per_million: &Option, + output_per_million: &Option, + cache_read_per_million: &Option, + cache_write_per_million: &Option, + ) -> Option { + let payload = serde_json::to_vec(&( + "codewhale-provider-live-pricing-quote-v1", + provider, + provider_identity, + wire_model, + endpoint_fingerprint, + catalog_fetched_at, + currency, + provenance, + input_per_million, + output_per_million, + cache_read_per_million, + cache_write_per_million, + )) + .ok()?; + Some(format!("sha256:{}", crate::hashing::sha256_hex(payload))) + } + + fn from_pricing( + provider: ApiProvider, + provider_identity: &str, + wire_model: &str, + endpoint_fingerprint: &str, + catalog_fetched_at: u64, + pricing: &OfferingPricing, + ) -> Option { + let provider_identity = provider_identity.trim(); + let wire_model = wire_model.trim(); + if crate::cost_status::sanitize_persisted_route_label(provider_identity) + != provider_identity + || crate::cost_status::sanitize_persisted_route_label(wire_model) != wire_model + { + return None; + } + let input_per_million = Self::canonical_rate(pricing.input_per_million); + let output_per_million = Self::canonical_rate(pricing.output_per_million); + let cache_read_per_million = Self::canonical_rate(pricing.cache_read_per_million); + let cache_write_per_million = Self::canonical_rate(pricing.cache_write_per_million); + let catalog_revision = Self::revision_for( + provider, + provider_identity, + wire_model, + endpoint_fingerprint, + catalog_fetched_at, + &pricing.currency, + &pricing.provenance, + &input_per_million, + &output_per_million, + &cache_read_per_million, + &cache_write_per_million, + )?; + Some(Self { + provider, + provider_identity: provider_identity.to_string(), + wire_model: wire_model.to_string(), + endpoint_fingerprint: endpoint_fingerprint.to_string(), + catalog_fetched_at, + catalog_revision, + currency: pricing.currency.clone(), + provenance: pricing.provenance.clone(), + input_per_million, + output_per_million, + cache_read_per_million, + cache_write_per_million, + }) + } + + fn parse_rate(rate: &Option) -> Option> { + let Some(rate) = rate else { + return Some(None); + }; + let parsed = rate.parse::().ok()?; + (parsed.is_finite() && parsed >= 0.0 && parsed.to_string() == *rate).then_some(Some(parsed)) + } + + /// Rehydrate the frozen row only when every receipt binding is intact. + /// This is deliberately cache-free: a refresh after dispatch cannot alter + /// an earlier turn, while malformed or legacy receipts fail closed. + pub(crate) fn pricing_for_route( + &self, + provider: ApiProvider, + provider_identity: &str, + wire_model: &str, + endpoint_fingerprint: &str, + dispatched_at_unix: u64, + ) -> Option { + let provider_identity = provider_identity.trim(); + let wire_model = wire_model.trim(); + if self.provider != provider + || crate::cost_status::sanitize_persisted_route_label(&self.provider_identity) + != self.provider_identity + || crate::cost_status::sanitize_persisted_route_label(&self.wire_model) + != self.wire_model + || self.endpoint_fingerprint.len() != 64 + || !self + .endpoint_fingerprint + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) + || self.provider_identity != provider_identity + || self.wire_model != wire_model + || self.endpoint_fingerprint != endpoint_fingerprint + || self.catalog_fetched_at > dispatched_at_unix + || dispatched_at_unix.saturating_sub(self.catalog_fetched_at) + >= DEFAULT_PROVIDER_CATALOG_TTL_SECS + || self.currency != Currency::Usd + || self.provenance != PricingProvenance::ProviderLive + || !reviewed_provider_live_scope(provider, provider_identity, endpoint_fingerprint) + { + return None; + } + let input_per_million = Self::parse_rate(&self.input_per_million)?; + let output_per_million = Self::parse_rate(&self.output_per_million)?; + let cache_read_per_million = Self::parse_rate(&self.cache_read_per_million)?; + let cache_write_per_million = Self::parse_rate(&self.cache_write_per_million)?; + let cost = codewhale_config::models_dev::ModelsDevCost { + input: input_per_million, + output: output_per_million, + cache_read: cache_read_per_million, + cache_write: cache_write_per_million, + }; + if !codewhale_config::pricing::catalog_cost_is_valid(&cost) { + return None; + } + // A reviewed per-token route needs both ordinary request classes. Cache + // classes remain optional and fail closed later if a turn used them. + if cost.input.is_none() || cost.output.is_none() { + return None; + } + let expected_revision = Self::revision_for( + self.provider, + &self.provider_identity, + &self.wire_model, + &self.endpoint_fingerprint, + self.catalog_fetched_at, + &self.currency, + &self.provenance, + &self.input_per_million, + &self.output_per_million, + &self.cache_read_per_million, + &self.cache_write_per_million, + )?; + if self.catalog_revision != expected_revision { + return None; + } + Some(OfferingPricing { + provider: self.provider_identity.clone(), + wire_model_id: self.wire_model.clone(), + canonical_model: None, + currency: self.currency.clone(), + input_per_million: cost.input, + output_per_million: cost.output, + cache_read_per_million: cost.cache_read, + cache_write_per_million: cost.cache_write, + provenance: self.provenance.clone(), + effective_at: Some(self.catalog_fetched_at), + endpoint_fingerprint: Some(self.endpoint_fingerprint.clone()), + }) + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct PersistedProviderCatalogs { + schema_version: u32, + cache: ProviderCatalogCache, +} + +#[derive(Serialize)] +struct PersistedProviderCatalogsRef<'a> { + schema_version: u32, + cache: &'a ProviderCatalogCache, +} + +/// Resolve the cache under Codewhale's catalog state directory. +/// +/// Unguarded tests are confined to the TUI test root, matching the Models.dev +/// cache contract, so they never inspect a developer's real provider catalog. +#[must_use] +pub fn cache_path() -> Option { + #[cfg(test)] + { + if !crate::test_support::guarded_environment_provides_state_paths() { + return Some( + crate::test_support::unsealed_test_state_root() + .join("catalog") + .join(CACHE_FILE), + ); + } + } + codewhale_config::resolve_state_dir("catalog") + .ok() + .map(|dir| dir.join(CACHE_FILE)) +} + +fn canonical_provider_scope(provider: &str) -> String { + // Despite the historical name, this is the exact configured ownership + // scope. Never collapse a custom table that happens to resemble a built-in + // or setup-template alias. + crate::provider_lake::catalog_partition_key(provider) +} + +fn is_account_scoped_provider(provider: &str) -> bool { + codewhale_config::provider_setup_template(provider) + .is_some_and(|template| template.id == codewhale_config::BASETEN_TEMPLATE_ID) +} + +fn cache_lock_path(path: &Path) -> PathBuf { + let mut name = path + .file_name() + .map(|name| name.to_os_string()) + .unwrap_or_else(|| CACHE_FILE.into()); + name.push(".lock"); + path.with_file_name(name) +} + +fn open_cache_lock(path: &Path) -> Result { + let parent = path + .parent() + .context("provider catalog lock path has no parent")?; + fs::create_dir_all(parent) + .with_context(|| format!("create provider catalog directory {}", parent.display()))?; + let mut options = OpenOptions::new(); + options.read(true).write(true).create(true).truncate(false); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt as _; + options + .mode(0o600) + .custom_flags(libc::O_NOFOLLOW | libc::O_CLOEXEC); + } + #[cfg(windows)] + { + use std::os::windows::fs::OpenOptionsExt as _; + options.custom_flags(0x0020_0000); // FILE_FLAG_OPEN_REPARSE_POINT + } + let file = options + .open(path) + .with_context(|| format!("open provider catalog lock {}", path.display()))?; + let metadata = file + .metadata() + .with_context(|| format!("inspect provider catalog lock {}", path.display()))?; + anyhow::ensure!( + metadata.is_file(), + "provider catalog lock {} must be a regular file", + path.display() + ); + #[cfg(unix)] + { + use std::os::unix::fs::MetadataExt as _; + anyhow::ensure!( + metadata.nlink() == 1, + "provider catalog lock {} must not be hard linked", + path.display() + ); + } + #[cfg(windows)] + { + use std::os::windows::fs::MetadataExt as _; + const FILE_ATTRIBUTE_REPARSE_POINT: u32 = 0x0000_0400; + anyhow::ensure!( + metadata.file_attributes() & FILE_ATTRIBUTE_REPARSE_POINT == 0, + "provider catalog lock {} must not be a reparse point", + path.display() + ); + } + Ok(file) +} + +fn load_from_disk_unlocked_with_limit(path: &Path, max_bytes: u64) -> Option { + let file = fs::File::open(path).ok()?; + if file.metadata().ok()?.len() > max_bytes { + tracing::debug!( + target: "provider_catalog", + path = %path.display(), + max_bytes, + "provider catalog cache exceeds read limit" + ); + return None; + } + // Re-check through `take`: the file can grow after metadata is sampled. + let mut body = Vec::new(); + file.take(max_bytes.saturating_add(1)) + .read_to_end(&mut body) + .ok()?; + if body.len() as u64 > max_bytes { + return None; + } + let persisted: PersistedProviderCatalogs = serde_json::from_slice(&body).ok()?; + if persisted.schema_version != CACHE_SCHEMA_VERSION { + return None; + } + let mut cache = persisted.cache; + // Older builds could durably cache Baseten's account-scoped roster. Scrub + // those entries on every load so upgrading cannot attach one workspace's + // catalog to a different credential. + cache + .entries + .retain(|_, entry| !is_account_scoped_provider(&entry.provider)); + Some(cache) +} + +fn load_from_disk_unlocked(path: &Path) -> Option { + load_from_disk_unlocked_with_limit(path, MAX_CACHE_BYTES) +} + +fn load_from_disk() -> Option { + let path = cache_path()?; + if !path.is_file() { + return None; + } + let lock_file = open_cache_lock(&cache_lock_path(&path)).ok()?; + let lock = fd_lock::RwLock::new(lock_file); + let _guard = lock.read().ok()?; + load_from_disk_unlocked(&path) +} + +fn merge_durable_scope( + mut durable_cache: ProviderCatalogCache, + process_cache: &ProviderCatalogCache, + provider: &str, + fingerprint: &str, +) -> ProviderCatalogCache { + durable_cache + .entries + .retain(|_, entry| !is_account_scoped_provider(&entry.provider)); + if !is_account_scoped_provider(provider) + && let Some(entry) = process_cache.get(provider, fingerprint).cloned() + { + durable_cache.entries.insert( + ProviderCatalogCache::cache_key(provider, fingerprint), + entry, + ); + } + durable_cache +} + +fn persisted_envelope_len(cache: &ProviderCatalogCache) -> Result { + let envelope = PersistedProviderCatalogsRef { + schema_version: CACHE_SCHEMA_VERSION, + cache, + }; + let mut body = serde_json::to_vec_pretty(&envelope) + .context("serialize provider catalog cache for bounded persistence")?; + body.push(b'\n'); + u64::try_from(body.len()).context("provider catalog cache length exceeds u64") +} + +fn cached_row_count(cache: &ProviderCatalogCache) -> usize { + cache.entries.values().fold(0usize, |total, entry| { + total.saturating_add(entry.offerings.len()) + }) +} + +/// Compact a durable cache without ever truncating one provider roster. +/// +/// The exact scope being written is protected: if that scope alone fits, older +/// failed/stale scopes are evicted whole until the envelope is bounded. If the +/// protected scope alone does not fit, persistence is refused and the prior +/// atomic file remains intact. This avoids both self-bricking the 32 MiB read +/// limit and turning a partial provider roster into false authoritative truth. +fn bounded_cache_for_persistence( + mut cache: ProviderCatalogCache, + protected_scope: Option<(&str, &str)>, + now: u64, + limits: CachePersistenceLimits, +) -> Result { + cache + .entries + .retain(|_, entry| !is_account_scoped_provider(&entry.provider)); + + let protected_key = protected_scope + .filter(|(provider, _)| !is_account_scoped_provider(provider)) + .map(|(provider, fingerprint)| ProviderCatalogCache::cache_key(provider, fingerprint)); + + if let Some(key) = protected_key.as_deref() + && let Some(entry) = cache.entries.get(key).cloned() + { + let mut protected_only = ProviderCatalogCache::new(); + protected_only.entries.insert(key.to_string(), entry); + anyhow::ensure!( + protected_only.entries.len() <= limits.max_scopes.min(MAX_CACHE_SCOPES) + && cached_row_count(&protected_only) <= limits.max_rows + && persisted_envelope_len(&protected_only)? <= limits.max_bytes, + "provider catalog scope {key:?} exceeds bounded persistence limits" + ); + } + + // Rank once while the cache/file locks are held. An older implementation + // reserialized and rescanned the entire envelope for every eviction, which + // made a valid sub-32-MiB file with many tiny scopes quadratic to compact. + let mut eviction_keys = cache + .entries + .iter() + .filter(|(key, _)| protected_key.as_deref() != Some(key.as_str())) + .map(|(key, entry)| { + let health_rank = if matches!(entry.status, CatalogStatus::Failed { .. }) { + 0u8 + } else if entry.is_stale(now) || matches!(entry.status, CatalogStatus::Stale { .. }) { + 1u8 + } else { + 2u8 + }; + (health_rank, entry.fetched_at, key.clone()) + }) + .collect::>(); + eviction_keys.sort(); + let eviction_keys = eviction_keys + .into_iter() + .map(|(_, _, key)| key) + .collect::>(); + let mut eviction_index = 0usize; + let mut rows = cached_row_count(&cache); + let max_scopes = limits.max_scopes.min(MAX_CACHE_SCOPES); + + let mut evict_next = |cache: &mut ProviderCatalogCache| -> Result { + let key = eviction_keys + .get(eviction_index) + .context("provider catalog envelope cannot fit even after whole-scope compaction")?; + eviction_index = eviction_index.saturating_add(1); + let entry = cache + .entries + .remove(key) + .context("provider catalog eviction candidate disappeared")?; + Ok(entry.offerings.len()) + }; + + // First enforce the cheap cardinality limits in bulk. Only after at most 64 + // scopes remain do we serialize to enforce the exact on-disk byte limit. + while cache.entries.len() > max_scopes || rows > limits.max_rows { + rows = rows.saturating_sub(evict_next(&mut cache)?); + } + while persisted_envelope_len(&cache)? > limits.max_bytes { + let _removed_rows = evict_next(&mut cache)?; + } + + Ok(cache) +} + +fn write_bounded_cache( + path: &Path, + cache: ProviderCatalogCache, + protected_scope: Option<(&str, &str)>, + limits: CachePersistenceLimits, +) -> Result<()> { + let cache = bounded_cache_for_persistence(cache, protected_scope, now_unix(), limits)?; + let envelope = PersistedProviderCatalogs { + schema_version: CACHE_SCHEMA_VERSION, + cache, + }; + anyhow::ensure!( + persisted_envelope_len(&envelope.cache)? <= limits.max_bytes, + "bounded provider catalog cache exceeds its write limit" + ); + atomic_write_json(path, &envelope) +} + +fn persist_scope(cache: &ProviderCatalogCache, provider: &str, fingerprint: &str) { + let Some(path) = cache_path() else { + return; + }; + let provider = canonical_provider_scope(provider); + let result = (|| -> Result<()> { + let lock_file = open_cache_lock(&cache_lock_path(&path))?; + let mut lock = fd_lock::RwLock::new(lock_file); + let _guard = lock + .write() + .with_context(|| format!("write-lock provider catalog cache {}", path.display()))?; + // Merge only the exact scope this process just changed into the latest + // disk snapshot. A stale long-running TUI therefore cannot erase a + // different scope written by the Runtime API (or vice versa). + let durable_cache = merge_durable_scope( + load_from_disk_unlocked(&path).unwrap_or_default(), + cache, + &provider, + fingerprint, + ); + write_bounded_cache( + &path, + durable_cache, + Some((&provider, fingerprint)), + CACHE_PERSISTENCE_LIMITS, + ) + .with_context(|| format!("atomically write provider catalog {}", path.display())) + })(); + if let Err(error) = result { + tracing::debug!( + target: "provider_catalog", + error = %error, + "provider catalog cache write failed" + ); + } +} + +/// Persist a failure without letting a stale process replace newer rows from +/// another Codewhale process for the same exact scope. +/// +/// The ordinary scoped merge is sufficient for successes because the response +/// being committed is the new roster. A failure is different: its process may +/// have started with an older last-known-good entry. Re-read the durable exact +/// scope while holding the cross-process write lock, prefer it when it is at +/// least as recent, then change only the status before writing. Thus a failed +/// refresh can preserve the newest roster without resurrecting its own stale +/// snapshot over another process's success. +fn persist_failure_scope( + cache: &mut ProviderCatalogCache, + provider: &str, + fingerprint: &str, + reason: CatalogRefreshError, +) { + let Some(path) = cache_path() else { + return; + }; + let provider = canonical_provider_scope(provider); + let result = (|| -> Result<()> { + let lock_file = open_cache_lock(&cache_lock_path(&path))?; + let mut lock = fd_lock::RwLock::new(lock_file); + let _guard = lock + .write() + .with_context(|| format!("write-lock provider catalog cache {}", path.display()))?; + let durable_cache = load_from_disk_unlocked(&path).unwrap_or_default(); + + if !is_account_scoped_provider(&provider) + && let Some(durable_entry) = durable_cache.get(&provider, fingerprint).cloned() + { + let durable_is_newer = cache + .get(&provider, fingerprint) + .is_none_or(|local| durable_entry.fetched_at >= local.fetched_at); + if durable_is_newer { + cache.entries.insert( + ProviderCatalogCache::cache_key(&provider, fingerprint), + durable_entry, + ); + cache.record_failure(&provider, fingerprint, reason); + } + } + + let durable_cache = merge_durable_scope(durable_cache, cache, &provider, fingerprint); + write_bounded_cache( + &path, + durable_cache, + Some((&provider, fingerprint)), + CACHE_PERSISTENCE_LIMITS, + ) + .with_context(|| format!("atomically write provider catalog {}", path.display())) + })(); + if let Err(error) = result { + tracing::debug!( + target: "provider_catalog", + error = %error, + "provider catalog failure receipt write failed" + ); + } +} + +fn publish_exact_scope(cache: &ProviderCatalogCache, provider: &str, fingerprint: &str) -> usize { + let provider = canonical_provider_scope(provider); + let provider_kind = if codewhale_config::provider_setup_template(&provider) + .is_some_and(|template| template.is_compatible()) + { + ApiProvider::Custom + } else { + ApiProvider::parse(&provider).unwrap_or(ApiProvider::Custom) + }; + publish_exact_scope_for_identity(cache, provider_kind, &provider, fingerprint) +} + +fn publish_exact_scope_for_identity( + cache: &ProviderCatalogCache, + provider_kind: ApiProvider, + provider_identity: &str, + fingerprint: &str, +) -> usize { + let provider = canonical_provider_scope(provider_identity); + let offerings = cache + .get(&provider, fingerprint) + .map(|entry| entry.offerings.clone()) + .unwrap_or_default(); + let count = offerings.len(); + crate::provider_lake::replace_provider_live_snapshot_for_identity( + provider_kind, + &provider, + CatalogSnapshot { offerings }, + ); + count +} + +/// Load and publish only the active route's exact provider/base-URL scope. +/// +/// A cache created for another custom endpoint or for an old endpoint override +/// is retained on disk but cannot leak into the active picker. +pub fn maybe_load_persisted_cache_for_config(config: &Config) -> usize { + let provider = config.api_provider(); + let provider_identity = canonical_provider_scope(&config.provider_identity_for(provider)); + let fingerprint = base_url_fingerprint(&config.deepseek_base_url()); + if is_account_scoped_provider(&provider_identity) { + forget_account_scoped_provider(&provider_identity); + return 0; + } + if let Ok(mut guard) = CACHE.write() + && let Some(loaded) = load_from_disk() + { + // Keep session-only scopes that cannot exist on disk, while allowing a + // newer durable scope from another Codewhale process to refresh this + // process. Every in-process writer takes CACHE before the file lock, so + // this read/merge cannot overwrite a concurrent local refresh. + for (key, entry) in loaded.entries { + let should_replace = guard + .entries + .get(&key) + .is_none_or(|current| entry.fetched_at >= current.fetched_at); + if should_replace { + guard.entries.insert(key, entry); + } + } + } + CACHE + .read() + .map(|guard| { + publish_exact_scope_for_identity(&guard, provider, &provider_identity, &fingerprint) + }) + .unwrap_or(0) +} + +fn forget_account_scoped_provider(provider: &str) { + let provider = canonical_provider_scope(provider); + if let Ok(mut cache) = CACHE.write() { + cache + .entries + .retain(|_, entry| canonical_provider_scope(&entry.provider) != provider); + } + crate::provider_lake::replace_provider_live_snapshot_for_identity( + ApiProvider::Custom, + &provider, + CatalogSnapshot::default(), + ); +} + +/// Begin a provider refresh and invalidate older in-flight results. +/// +/// Baseten additionally drops its prior in-memory roster because the same URL +/// can expose a different workspace catalog after an API-key change and the +/// endpoint supplies no safe account identifier for cache reuse. +pub fn begin_refresh(provider: &str) -> ProviderCatalogRefreshTicket { + let provider = canonical_provider_scope(provider); + let generation = if let Ok(mut generations) = REFRESH_GENERATIONS.write() { + let generation = generations.entry(provider.clone()).or_default(); + *generation = generation.saturating_add(1); + *generation + } else { + 0 + }; + if is_account_scoped_provider(&provider) { + forget_account_scoped_provider(&provider); + } + ProviderCatalogRefreshTicket { + provider, + generation, + } +} + +fn with_current_ticket( + ticket: &ProviderCatalogRefreshTicket, + provider: &str, + operation: impl FnOnce() -> T, +) -> Option { + let provider = canonical_provider_scope(provider); + if ticket.provider != provider { + return None; + } + let generations = REFRESH_GENERATIONS.read().ok()?; + if generations.get(&ticket.provider).copied() != Some(ticket.generation) { + return None; + } + // Keep the generation read guard alive through publication. A newer + // `begin_refresh` needs the write lock, so it cannot slip between the + // current-ticket check and this operation's cache/lake update. + let result = operation(); + drop(generations); + Some(result) +} + +/// Record a successful refresh only if no newer refresh superseded it. +pub fn record_success_if_current( + ticket: &ProviderCatalogRefreshTicket, + delta: ProviderCatalogDelta, +) -> Option { + let provider = canonical_provider_scope(&delta.provider); + with_current_ticket(ticket, &provider, || record_success(delta)) +} + +/// Record a failed refresh only if no newer refresh superseded it. +pub fn record_failure_if_current( + ticket: &ProviderCatalogRefreshTicket, + provider: &str, + fingerprint: &str, + reason: CatalogRefreshError, +) -> Option { + let provider = canonical_provider_scope(provider); + with_current_ticket(ticket, &provider, || { + record_failure(&provider, fingerprint, reason) + }) +} + +/// Current freshness receipt for one exact provider/base-URL scope. +/// +/// Runtime route resolution uses this independently from picker visibility: +/// stale or failed rows may remain selectable as an explicit fallback, but +/// their limits, capabilities, and prices are not treated as current endpoint +/// facts during execution. +pub fn status_for_scope(provider: &str, base_url: &str) -> CatalogStatus { + let fingerprint = base_url_fingerprint(base_url); + status_for_fingerprint(provider, &fingerprint) +} + +/// Current freshness receipt when the caller already owns the endpoint +/// fingerprint (for example, an immutable usage-pricing receipt). +pub(crate) fn status_for_fingerprint(provider: &str, fingerprint: &str) -> CatalogStatus { + let provider = canonical_provider_scope(provider); + CACHE + .read() + .map(|cache| cache.status(&provider, fingerprint, now_unix())) + .unwrap_or(CatalogStatus::Unknown) +} + +/// Freeze the exact reviewed provider-live rate row fresh at CodeWhale's +/// pre-permit application-dispatch boundary. +/// +/// Status, scope, model, source, and rates are all read beneath one `CACHE` +/// read guard. The returned value owns every fact needed by later auditing, so +/// completion-time code never re-opens mutable catalog or provider-lake state. +fn reviewed_provider_live_scope( + provider: ApiProvider, + provider_identity: &str, + endpoint_fingerprint: &str, +) -> bool { + match provider { + ApiProvider::Openrouter => { + provider_identity == ApiProvider::Openrouter.as_str() + && endpoint_fingerprint + == base_url_fingerprint(crate::config::DEFAULT_OPENROUTER_BASE_URL) + } + ApiProvider::Custom => { + codewhale_config::provider_setup_template(provider_identity) + .is_some_and(|template| template.id == codewhale_config::BASETEN_TEMPLATE_ID) + && endpoint_fingerprint == base_url_fingerprint(codewhale_config::BASETEN_BASE_URL) + } + _ => false, + } +} + +#[must_use] +pub(crate) fn fresh_provider_live_pricing_quote_at( + provider: ApiProvider, + provider_identity: &str, + wire_model: &str, + endpoint_fingerprint: &str, + dispatched_at_unix: u64, +) -> Option { + let provider_identity = canonical_provider_scope(provider_identity); + let wire_model = wire_model.trim(); + let endpoint_fingerprint = endpoint_fingerprint.trim(); + if provider_identity.is_empty() + || wire_model.is_empty() + || !reviewed_provider_live_scope(provider, &provider_identity, endpoint_fingerprint) + { + return None; + } + + let cache = CACHE.read().ok()?; + if cache.status(&provider_identity, endpoint_fingerprint, dispatched_at_unix) + != CatalogStatus::Fresh + { + return None; + } + let entry = cache.get(&provider_identity, endpoint_fingerprint)?; + if entry.provider.trim() != provider_identity + || entry.base_url_fingerprint.trim() != endpoint_fingerprint + || entry.fetched_at > dispatched_at_unix + { + return None; + } + let offering = entry.offerings.iter().find(|offering| { + offering.provider.trim() == provider_identity && offering.wire_model_id.trim() == wire_model + })?; + let pricing = OfferingPricing::from_catalog_offering(offering)?; + if pricing.provider.trim() != provider_identity + || pricing.wire_model_id.trim() != wire_model + || pricing.currency != Currency::Usd + || pricing.provenance != PricingProvenance::ProviderLive + || pricing.effective_at != Some(entry.fetched_at) + || pricing.endpoint_fingerprint.as_deref() != Some(endpoint_fingerprint) + || pricing.input_per_million.is_none() + || pricing.output_per_million.is_none() + { + return None; + } + ProviderLivePricingQuote::from_pricing( + provider, + &provider_identity, + wire_model, + endpoint_fingerprint, + entry.fetched_at, + &pricing, + ) +} + +/// Record and atomically persist a successful provider refresh. +/// +/// `ProviderCatalogCache::record_success` replaces the exact scope, so models +/// removed upstream disappear instead of accumulating forever. +pub fn record_success(mut delta: ProviderCatalogDelta) -> CatalogStatus { + let provider = canonical_provider_scope(&delta.provider); + delta.provider.clone_from(&provider); + delta.offerings.retain_mut(|row| { + if canonical_provider_scope(&row.provider) != provider { + return false; + } + row.provider.clone_from(&provider); + true + }); + let fingerprint = delta.base_url_fingerprint.clone(); + let Ok(mut guard) = CACHE.write() else { + return CatalogStatus::Unknown; + }; + guard.record_success(delta, DEFAULT_PROVIDER_CATALOG_TTL_SECS); + persist_scope(&guard, &provider, &fingerprint); + publish_exact_scope(&guard, &provider, &fingerprint); + CatalogStatus::Fresh +} + +/// Record a typed failure while preserving and republishing prior rows for the +/// exact route scope. +pub fn record_failure( + provider: &str, + fingerprint: &str, + reason: CatalogRefreshError, +) -> CatalogStatus { + let provider = canonical_provider_scope(provider); + let Ok(mut guard) = CACHE.write() else { + return CatalogStatus::Failed { reason }; + }; + guard.record_failure(&provider, fingerprint, reason); + persist_failure_scope(&mut guard, &provider, fingerprint, reason); + publish_exact_scope(&guard, &provider, fingerprint); + CatalogStatus::Failed { reason } +} + +#[cfg(test)] +pub(crate) fn reset_cache_for_test() { + if let Ok(mut cache) = CACHE.write() { + *cache = ProviderCatalogCache::new(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::{ApiProvider, ProviderConfig, ProvidersConfig}; + use crate::test_support::{EnvVarGuard, lock_test_env}; + use codewhale_config::catalog::{CatalogOffering, CatalogSource}; + + fn delta(provider: &str, fingerprint: &str, ids: &[&str]) -> ProviderCatalogDelta { + delta_at(provider, fingerprint, ids, now_unix()) + } + + fn delta_at( + provider: &str, + fingerprint: &str, + ids: &[&str], + fetched_at: u64, + ) -> ProviderCatalogDelta { + ProviderCatalogDelta { + provider: provider.to_string(), + base_url_fingerprint: fingerprint.to_string(), + fetched_at, + offerings: ids + .iter() + .map(|id| CatalogOffering { + provider: provider.to_string(), + wire_model_id: (*id).to_string(), + endpoint_key: "chat".to_string(), + source: CatalogSource::Live { + base_url_fingerprint: fingerprint.to_string(), + fetched_at, + }, + ..CatalogOffering::default() + }) + .collect(), + } + } + + #[test] + fn success_replaces_scope_and_failure_preserves_last_rows_on_disk() { + let _env = lock_test_env(); + let _live = crate::provider_lake::lock_live_snapshot(); + let home = tempfile::tempdir().expect("home"); + let _home = EnvVarGuard::set("CODEWHALE_HOME", home.path()); + if let Ok(mut cache) = CACHE.write() { + *cache = ProviderCatalogCache::new(); + } + + assert_eq!( + record_success(delta("openrouter", "fp", &["old"])), + CatalogStatus::Fresh + ); + assert_eq!( + record_success(delta("openrouter", "fp", &["new"])), + CatalogStatus::Fresh + ); + assert!(matches!( + record_failure("openrouter", "fp", CatalogRefreshError::RateLimited), + CatalogStatus::Failed { + reason: CatalogRefreshError::RateLimited + } + )); + + let loaded = load_from_disk().expect("persisted cache"); + let entry = loaded.get("openrouter", "fp").expect("OpenRouter scope"); + assert_eq!(entry.offerings.len(), 1); + assert_eq!(entry.offerings[0].wire_model_id, "new"); + assert!(matches!(entry.status, CatalogStatus::Failed { .. })); + } + + #[test] + fn baseten_workspace_roster_is_session_only_and_clears_before_reauthentication() { + let _env = lock_test_env(); + let _live = crate::provider_lake::lock_live_snapshot(); + let home = tempfile::tempdir().expect("home"); + let _home = EnvVarGuard::set("CODEWHALE_HOME", home.path()); + reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + + let base_url = codewhale_config::BASETEN_BASE_URL; + let fingerprint = base_url_fingerprint(base_url); + record_success(delta( + codewhale_config::BASETEN_TEMPLATE_ID, + &fingerprint, + &["workspace-a-only-model"], + )); + assert!( + crate::provider_lake::all_catalog_models_for_provider_identity( + ApiProvider::Custom, + Some(codewhale_config::BASETEN_TEMPLATE_ID), + ) + .contains(&"workspace-a-only-model".to_string()) + ); + assert!( + load_from_disk().is_none_or(|cache| cache + .get(codewhale_config::BASETEN_TEMPLATE_ID, &fingerprint) + .is_none()), + "an account-scoped Baseten roster must never be durable without a safe account id" + ); + + let mut custom = std::collections::HashMap::new(); + custom.insert( + codewhale_config::BASETEN_TEMPLATE_ID.to_string(), + ProviderConfig { + kind: Some("openai-compatible".to_string()), + base_url: Some(base_url.to_string()), + model: Some(codewhale_config::BASETEN_DEFAULT_MODEL.to_string()), + ..ProviderConfig::default() + }, + ); + let config = Config { + provider: Some(codewhale_config::BASETEN_TEMPLATE_ID.to_string()), + providers: Some(ProvidersConfig { + custom, + ..ProvidersConfig::default() + }), + ..Config::default() + }; + assert_eq!(maybe_load_persisted_cache_for_config(&config), 0); + assert!(matches!( + status_for_scope(codewhale_config::BASETEN_TEMPLATE_ID, base_url), + CatalogStatus::Unknown + )); + assert!( + !crate::provider_lake::all_catalog_models_for_provider_identity( + ApiProvider::Custom, + Some(codewhale_config::BASETEN_TEMPLATE_ID), + ) + .contains(&"workspace-a-only-model".to_string()), + "a new credential attempt must not see the previous workspace roster" + ); + + reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + } + + #[test] + fn superseded_refresh_ticket_cannot_publish_a_late_response() { + let _env = lock_test_env(); + let _live = crate::provider_lake::lock_live_snapshot(); + let home = tempfile::tempdir().expect("home"); + let _home = EnvVarGuard::set("CODEWHALE_HOME", home.path()); + reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + + let old = begin_refresh("openrouter"); + let current = begin_refresh("openrouter"); + assert!( + record_success_if_current(&old, delta("openrouter", "fp", &["late-old-model"])) + .is_none() + ); + assert!( + record_success_if_current(¤t, delta("openrouter", "fp", &["current-model"]),) + .is_some() + ); + assert_eq!( + CACHE + .read() + .expect("cache") + .get("openrouter", "fp") + .expect("current scope") + .offerings[0] + .wire_model_id, + "current-model" + ); + + reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + } + + #[test] + fn current_ticket_holds_generation_gate_through_publication() { + let ticket = begin_refresh("generation-barrier-provider"); + let entered = std::sync::Arc::new(std::sync::Barrier::new(2)); + let release = std::sync::Arc::new(std::sync::Barrier::new(2)); + let publish_entered = std::sync::Arc::clone(&entered); + let publish_release = std::sync::Arc::clone(&release); + let publisher = std::thread::spawn(move || { + with_current_ticket(&ticket, "generation-barrier-provider", || { + publish_entered.wait(); + publish_release.wait(); + }) + }); + entered.wait(); + + let (started_tx, started_rx) = std::sync::mpsc::channel(); + let (finished_tx, finished_rx) = std::sync::mpsc::channel(); + let newer = std::thread::spawn(move || { + started_tx.send(()).expect("signal refresh start"); + let next = begin_refresh("generation-barrier-provider"); + finished_tx.send(next).expect("signal refresh finish"); + }); + started_rx.recv().expect("new refresh thread started"); + assert!( + finished_rx + .recv_timeout(std::time::Duration::from_millis(50)) + .is_err(), + "a newer generation must wait until the accepted result finishes publication" + ); + + release.wait(); + assert!(publisher.join().expect("publisher thread").is_some()); + assert!( + finished_rx + .recv_timeout(std::time::Duration::from_secs(1)) + .is_ok() + ); + newer.join().expect("newer refresh thread"); + } + + #[test] + fn stale_process_snapshots_merge_exact_scopes_under_file_lock() { + let _env = lock_test_env(); + let home = tempfile::tempdir().expect("home"); + let _home = EnvVarGuard::set("CODEWHALE_HOME", home.path()); + + let mut process_a = ProviderCatalogCache::new(); + process_a.record_success(delta("CustomA", "fp-a", &["upper-model"]), 60); + persist_scope(&process_a, "CustomA", "fp-a"); + + // Simulate another process that started before A wrote and therefore + // has an empty/stale in-memory snapshot. Its scoped write must merge A + // from disk rather than replacing the whole envelope. + let mut process_b = ProviderCatalogCache::new(); + process_b.record_success(delta("customa", "fp-b", &["lower-model"]), 60); + persist_scope(&process_b, "customa", "fp-b"); + + let loaded = load_from_disk().expect("merged durable cache"); + assert_eq!( + loaded + .get("CustomA", "fp-a") + .expect("case-sensitive upper scope") + .offerings[0] + .wire_model_id, + "upper-model" + ); + assert_eq!( + loaded + .get("customa", "fp-b") + .expect("case-sensitive lower scope") + .offerings[0] + .wire_model_id, + "lower-model" + ); + } + + #[test] + fn stale_process_failure_preserves_newer_durable_rows_for_the_same_scope() { + let _env = lock_test_env(); + let _live = crate::provider_lake::lock_live_snapshot(); + let home = tempfile::tempdir().expect("home"); + let _home = EnvVarGuard::set("CODEWHALE_HOME", home.path()); + reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + + // Process B began with this old roster and still holds it in memory. + let mut stale_process = ProviderCatalogCache::new(); + stale_process.record_success(delta_at("openrouter", "fp", &["old-model"], 1), 60); + persist_scope(&stale_process, "openrouter", "fp"); + *CACHE.write().expect("cache") = stale_process; + + // Process A completes a newer successful refresh for the same scope. + let mut newer_process = ProviderCatalogCache::new(); + newer_process.record_success(delta_at("openrouter", "fp", &["new-model"], 2), 60); + persist_scope(&newer_process, "openrouter", "fp"); + + // B then fails. Its failure status is current, but its old rows are + // not: the transaction must retain A's newer durable roster. + assert!(matches!( + record_failure("openrouter", "fp", CatalogRefreshError::Network), + CatalogStatus::Failed { + reason: CatalogRefreshError::Network + } + )); + let in_memory = CACHE.read().expect("cache"); + let entry = in_memory.get("openrouter", "fp").expect("failed scope"); + assert_eq!(entry.offerings[0].wire_model_id, "new-model"); + assert!(matches!(entry.status, CatalogStatus::Failed { .. })); + drop(in_memory); + + let durable = load_from_disk().expect("durable cache"); + let entry = durable + .get("openrouter", "fp") + .expect("durable failed scope"); + assert_eq!(entry.offerings[0].wire_model_id, "new-model"); + assert!(matches!(entry.status, CatalogStatus::Failed { .. })); + + reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + } + + #[test] + fn oversized_cache_file_is_rejected_before_allocation() { + let _env = lock_test_env(); + let home = tempfile::tempdir().expect("home"); + let _home = EnvVarGuard::set("CODEWHALE_HOME", home.path()); + let path = cache_path().expect("cache path"); + fs::create_dir_all(path.parent().expect("catalog directory")).expect("catalog directory"); + fs::File::create(&path) + .and_then(|file| file.set_len(MAX_CACHE_BYTES + 1)) + .expect("sparse oversized cache"); + assert!(load_from_disk().is_none()); + } + + #[test] + fn bounded_persistence_evicts_failed_then_stale_scopes_and_keeps_exact_owner() { + let mut cache = ProviderCatalogCache::new(); + cache.record_success(delta_at("failed", "fp", &["failed-model"], 10), 1_000); + cache.record_failure("failed", "fp", CatalogRefreshError::Network); + cache.record_success(delta_at("stale", "fp", &["stale-model"], 20), 1); + cache.record_success(delta_at("fresh", "fp", &["fresh-model"], 30), 1_000); + cache.record_success(delta_at("protected", "fp", &["protected-model"], 40), 1_000); + + let compacted = bounded_cache_for_persistence( + cache, + Some(("protected", "fp")), + 100, + CachePersistenceLimits { + max_bytes: u64::MAX, + max_scopes: 2, + max_rows: 100, + }, + ) + .expect("bounded cache"); + + assert!(compacted.get("protected", "fp").is_some()); + assert!(compacted.get("fresh", "fp").is_some()); + assert!(compacted.get("failed", "fp").is_none()); + assert!(compacted.get("stale", "fp").is_none()); + } + + #[test] + fn bounded_persistence_evicts_whole_scopes_and_refuses_an_oversized_owner() { + let mut cache = ProviderCatalogCache::new(); + cache.record_success(delta_at("protected", "fp", &["one", "two"], 40), 1_000); + cache.record_success(delta_at("other", "fp", &["other"], 30), 1_000); + let limits = CachePersistenceLimits { + max_bytes: u64::MAX, + max_scopes: 10, + max_rows: 2, + }; + + let compacted = + bounded_cache_for_persistence(cache.clone(), Some(("protected", "fp")), 50, limits) + .expect("other scope can be evicted whole"); + assert_eq!( + compacted + .get("protected", "fp") + .expect("protected roster") + .offerings + .len(), + 2 + ); + assert!(compacted.get("other", "fp").is_none()); + + let mut oversized = cache; + oversized.record_success( + delta_at("protected", "fp", &["one", "two", "three"], 50), + 1_000, + ); + assert!( + bounded_cache_for_persistence(oversized, Some(("protected", "fp")), 50, limits,) + .is_err(), + "a provider roster must be refused, never partially persisted" + ); + } + + #[test] + fn bounded_cache_write_matches_read_limit_and_round_trips_after_compaction() { + let directory = tempfile::tempdir().expect("cache directory"); + let path = directory.path().join(CACHE_FILE); + let mut protected_only = ProviderCatalogCache::new(); + protected_only.record_success(delta_at("protected", "fp", &["protected-model"], 40), 1_000); + let exact_bytes = persisted_envelope_len(&protected_only).expect("encoded length"); + let limits = CachePersistenceLimits { + max_bytes: exact_bytes, + max_scopes: 10, + max_rows: 10, + }; + + let mut combined = protected_only; + combined.record_success( + delta_at( + "evicted", + "fp", + &["this-entire-scope-does-not-fit-the-byte-bound"], + 30, + ), + 1_000, + ); + write_bounded_cache(&path, combined, Some(("protected", "fp")), limits) + .expect("bounded disk write"); + + assert!(fs::metadata(&path).expect("cache metadata").len() <= exact_bytes); + let loaded = load_from_disk_unlocked_with_limit(&path, exact_bytes) + .expect("bounded cache must remain readable under the same cap"); + assert!(loaded.get("protected", "fp").is_some()); + assert!(loaded.get("evicted", "fp").is_none()); + } + + #[test] + fn baseten_alias_roster_is_session_only_and_keeps_exact_ownership() { + let _env = lock_test_env(); + let _live = crate::provider_lake::lock_live_snapshot(); + let home = tempfile::tempdir().expect("home"); + let _home = EnvVarGuard::set("CODEWHALE_HOME", home.path()); + reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + + let alias = "base-ten"; + let fingerprint = base_url_fingerprint(codewhale_config::BASETEN_BASE_URL); + record_success(delta(alias, &fingerprint, &["alias-workspace-model"])); + + assert!( + crate::provider_lake::all_catalog_models_for_provider_identity( + ApiProvider::Custom, + Some(alias), + ) + .contains(&"alias-workspace-model".to_string()) + ); + assert!( + !crate::provider_lake::all_catalog_models_for_provider_identity( + ApiProvider::Custom, + Some(codewhale_config::BASETEN_TEMPLATE_ID), + ) + .contains(&"alias-workspace-model".to_string()), + "a reviewed schema alias must not collapse distinct exact table ownership" + ); + assert!( + load_from_disk().is_none_or(|cache| cache.get(alias, &fingerprint).is_none()), + "every Baseten schema alias must remain session-only" + ); + + reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + } + + #[test] + fn different_base_url_fingerprints_do_not_share_rows() { + let mut cache = ProviderCatalogCache::new(); + cache.record_success(delta("baseten", "one", &["model-one"]), 60); + cache.record_success(delta("baseten", "two", &["model-two"]), 60); + assert_eq!( + cache.get("baseten", "one").unwrap().offerings[0].wire_model_id, + "model-one" + ); + assert_eq!( + cache.get("baseten", "two").unwrap().offerings[0].wire_model_id, + "model-two" + ); + } + + #[test] + fn missing_cache_for_changed_base_url_clears_the_previous_provider_partition() { + let _live = crate::provider_lake::lock_live_snapshot(); + crate::provider_lake::clear_live_snapshot(); + let mut cache = ProviderCatalogCache::new(); + cache.record_success(delta("baseten", "old-fp", &["old-endpoint-model"]), 60); + + assert_eq!(publish_exact_scope(&cache, "baseten", "old-fp"), 1); + assert_eq!( + crate::provider_lake::all_catalog_models_for_provider_identity( + crate::config::ApiProvider::Custom, + Some("baseten"), + ), + vec!["old-endpoint-model".to_string()] + ); + + assert_eq!(publish_exact_scope(&cache, "baseten", "new-fp"), 0); + let after_switch = crate::provider_lake::all_catalog_models_for_provider_identity( + crate::config::ApiProvider::Custom, + Some("baseten"), + ); + assert!( + !after_switch.contains(&"old-endpoint-model".to_string()), + "rows from the old Baseten endpoint must not survive a fingerprint change" + ); + assert!( + after_switch.contains(&codewhale_config::BASETEN_DEFAULT_MODEL.to_string()), + "the exact provider should fall back to its offline seed" + ); + crate::provider_lake::clear_live_snapshot(); + } + + #[test] + fn disk_reload_rehydrates_and_exposes_six_hundred_openrouter_models() { + let _env = lock_test_env(); + let _live = crate::provider_lake::lock_live_snapshot(); + let home = tempfile::tempdir().expect("home"); + let _home = EnvVarGuard::set("CODEWHALE_HOME", home.path()); + crate::provider_lake::clear_live_snapshot(); + if let Ok(mut cache) = CACHE.write() { + *cache = ProviderCatalogCache::new(); + } + + let config = Config { + provider: Some("openrouter".to_string()), + providers: Some(ProvidersConfig { + openrouter: ProviderConfig { + base_url: Some("https://synthetic.openrouter.invalid/api/v1".to_string()), + ..ProviderConfig::default() + }, + ..ProvidersConfig::default() + }), + ..Config::default() + }; + let provider = config.provider_identity_for(config.api_provider()); + let fingerprint = base_url_fingerprint(&config.deepseek_base_url()); + let fetched_at = now_unix(); + let ids: Vec = (0..600) + .map(|index| format!("synthetic/openrouter-model-{index:03}")) + .collect(); + let status = record_success(ProviderCatalogDelta { + provider: provider.clone(), + base_url_fingerprint: fingerprint, + fetched_at, + offerings: ids + .iter() + .map(|id| CatalogOffering { + provider: provider.clone(), + wire_model_id: id.clone(), + endpoint_key: "chat".to_string(), + source: CatalogSource::Live { + base_url_fingerprint: base_url_fingerprint(&config.deepseek_base_url()), + fetched_at, + }, + ..CatalogOffering::default() + }) + .collect(), + }); + assert_eq!(status, CatalogStatus::Fresh); + assert!(cache_path().is_some_and(|path| path.is_file())); + assert_eq!( + crate::provider_lake::all_catalog_models_for_provider(ApiProvider::Openrouter), + ids, + "the string compatibility publisher must retain built-in OpenRouter ownership" + ); + assert!( + crate::provider_lake::all_catalog_models_for_provider_identity( + ApiProvider::Custom, + Some("openrouter"), + ) + .is_empty(), + "built-in OpenRouter rows must not enter the custom namespace" + ); + + // Simulate a new process: remove both in-memory owners, then republish + // only through the durable startup load path. + if let Ok(mut cache) = CACHE.write() { + *cache = ProviderCatalogCache::new(); + } + crate::provider_lake::clear_live_snapshot(); + + assert_eq!(maybe_load_persisted_cache_for_config(&config), 600); + let visible = + crate::provider_lake::all_catalog_models_for_provider(ApiProvider::Openrouter); + assert_eq!(visible.len(), 600); + assert_eq!(visible.first(), ids.first()); + assert_eq!(visible.last(), ids.last()); + + if let Ok(mut cache) = CACHE.write() { + *cache = ProviderCatalogCache::new(); + } + crate::provider_lake::clear_live_snapshot(); + } +} diff --git a/crates/tui/src/provider_lake.rs b/crates/tui/src/provider_lake.rs index 596fa8e257..f1d33e128a 100644 --- a/crates/tui/src/provider_lake.rs +++ b/crates/tui/src/provider_lake.rs @@ -11,11 +11,16 @@ //! does not represent (and for unbundled gateways until the live catalog covers //! them). +use std::borrow::Cow; use std::collections::BTreeMap; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, RwLock}; -use codewhale_config::catalog::{CatalogOffering, CatalogSnapshot, bundled_catalog_offerings}; +use codewhale_config::catalog::{ + CatalogOffering, CatalogSnapshot, CatalogSource, CatalogStatus, base_url_fingerprint, + bundled_catalog_offerings, +}; +use codewhale_config::route::{ProviderModelOffering, RouteResolver, bundled_offerings}; use crate::codex_model_cache; use crate::config::{ @@ -54,50 +59,59 @@ static LIVE_SNAPSHOT: RwLock = RwLock::new(LiveSnapshotP #[derive(Default)] struct LiveSnapshotPartitions { models_dev: Option, - per_provider: BTreeMap, + per_provider: BTreeMap, } -impl LiveSnapshotPartitions { - /// Collect all live rows from every partition into a single flat snapshot. - fn flattened(&self) -> Option { - if self.models_dev.is_none() && self.per_provider.is_empty() { - return None; - } +/// Internal ownership key for one provider-owned live roster. +/// +/// Catalog rows intentionally keep their public provider string for receipts and +/// cache compatibility. The storage key carries the route kind separately so an +/// exact custom table named `openai` cannot overwrite, suppress, or borrow the +/// built-in OpenAI partition. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +enum LivePartitionOwner { + BuiltIn(String), + Custom(String), +} - // Merge by (provider, wire_model_id); provider-scoped rows win on - // collision because they came from that gateway's own live endpoint. - let mut merged: BTreeMap<(String, String), CatalogOffering> = BTreeMap::new(); - if let Some(models_dev) = &self.models_dev { - for row in &models_dev.offerings { - merged.insert( - (row.provider.clone(), row.wire_model_id.clone()), - row.clone(), - ); - } +impl LivePartitionOwner { + fn identity(&self) -> &str { + match self { + Self::BuiltIn(identity) | Self::Custom(identity) => identity, } - for provider_snapshot in self.per_provider.values() { - for row in &provider_snapshot.offerings { - merged.insert( - (row.provider.clone(), row.wire_model_id.clone()), - row.clone(), - ); - } - } - Some(CatalogSnapshot { - offerings: merged.into_values().collect(), - }) } } +fn live_partition_owner_for_route( + provider: ApiProvider, + provider_identity: Option<&str>, +) -> LivePartitionOwner { + let identity = catalog_provider_id_for_identity(provider, provider_identity); + if provider == ApiProvider::Custom { + LivePartitionOwner::Custom(catalog_partition_key(identity.as_ref())) + } else { + LivePartitionOwner::BuiltIn(catalog_partition_key(identity.as_ref())) + } +} + +fn inferred_live_partition_owner(provider: &str) -> LivePartitionOwner { + let identity = catalog_partition_key(provider); + ApiProvider::parse(&identity).map_or_else( + || LivePartitionOwner::Custom(identity), + |provider| { + LivePartitionOwner::BuiltIn(catalog_partition_key(catalog_provider_id(provider))) + }, + ) +} + fn offerings_by_provider( offerings: Vec, -) -> BTreeMap> { +) -> BTreeMap> { let mut grouped = BTreeMap::new(); - for offering in offerings { - grouped - .entry(offering.provider.trim().to_ascii_lowercase()) - .or_insert_with(Vec::new) - .push(offering); + for mut offering in offerings { + let owner = inferred_live_partition_owner(&offering.provider); + offering.provider = owner.identity().to_string(); + grouped.entry(owner).or_insert_with(Vec::new).push(offering); } grouped } @@ -114,6 +128,27 @@ static LIVE_GENERATION: AtomicU64 = AtomicU64::new(0); /// only changes when the live snapshot changes, so cache it. static MERGED_CACHE: RwLock)>> = RwLock::new(None); +/// Generation/freshness-scoped route resolvers for provider-owned catalogs. +/// Picker calls read the merged snapshot directly; execution projects that +/// snapshot into the immutable `RouteResolver` seam and must not rebuild a +/// 600+ row OpenRouter catalog for every route candidate. +static RUNTIME_RESOLVER_CACHE: RwLock> = + RwLock::new(BTreeMap::new()); + +#[derive(Clone)] +struct RuntimeResolverCacheEntry { + generation: u64, + status_is_fresh: bool, + endpoint_catalog_authoritative: bool, + resolver: RouteResolver, +} + +#[derive(Clone)] +pub(crate) struct RuntimeCatalogResolver { + pub(crate) resolver: RouteResolver, + pub(crate) endpoint_catalog_authoritative: bool, +} + fn bundled_snapshot() -> &'static CatalogSnapshot { BUNDLED_SNAPSHOT.get_or_init(|| CatalogSnapshot { offerings: bundled_catalog_offerings(), @@ -187,6 +222,66 @@ pub fn set_live_snapshot(snapshot: CatalogSnapshot, source: LiveSource) { } } +/// Replace one exact provider-owned live partition, including with no rows. +/// +/// The generic [`set_live_snapshot`] derives partitions from rows, so an empty +/// snapshot cannot say which previous partition should disappear. Endpoint- +/// scoped persistent caches need that distinction: switching Baseten to a new +/// base URL with no matching cache must remove the old URL's Baseten rows +/// immediately instead of presenting them as if they belonged to the new host. +pub fn replace_provider_live_snapshot(provider: &str, snapshot: CatalogSnapshot) { + let provider = provider.trim(); + if provider.is_empty() { + return; + } + let owner = inferred_live_partition_owner(provider); + replace_provider_live_snapshot_for_owner(owner, snapshot); +} + +/// Replace one provider-owned partition with an explicit route-kind boundary. +/// +/// Callers that know the concrete route must use this form. The legacy +/// string-only wrapper above remains for built-in publishers and older generic +/// tests, where a built-in-looking string necessarily denotes the built-in. +pub(crate) fn replace_provider_live_snapshot_for_identity( + provider: ApiProvider, + provider_identity: &str, + snapshot: CatalogSnapshot, +) { + let owner = live_partition_owner_for_route(provider, Some(provider_identity)); + if owner.identity().is_empty() { + return; + } + replace_provider_live_snapshot_for_owner(owner, snapshot); +} + +fn replace_provider_live_snapshot_for_owner(owner: LivePartitionOwner, snapshot: CatalogSnapshot) { + let provider_key = owner.identity().to_string(); + let mut snapshot = if matches!(&owner, LivePartitionOwner::Custom(_)) { + snapshot + } else { + apply_provider_model_cutlines(snapshot) + }; + snapshot.offerings.retain_mut(|row| { + if catalog_partition_key(&row.provider) != provider_key { + return false; + } + row.provider.clone_from(&provider_key); + true + }); + + if let Ok(mut guard) = LIVE_SNAPSHOT.write() { + let previous = guard.per_provider.remove(&owner); + let next = (!snapshot.offerings.is_empty()).then_some(snapshot); + if let Some(next) = next.clone() { + guard.per_provider.insert(owner, next); + } + if previous != next { + LIVE_GENERATION.fetch_add(1, Ordering::SeqCst); + } + } +} + /// Clear all live snapshots (both Models.dev and per-provider partitions). /// Used by tests and shutdown paths that need a full reset. #[allow(dead_code)] @@ -239,6 +334,7 @@ pub fn merge_live_offerings(new_offerings: Vec) { #[must_use] pub fn live_catalog_origin(provider: ApiProvider, wire_model_id: &str) -> Option { let catalog_id = catalog_provider_id(provider); + let owner = LivePartitionOwner::BuiltIn(catalog_partition_key(catalog_id)); let needle = wire_model_id.trim(); if needle.is_empty() { return None; @@ -252,8 +348,8 @@ pub fn live_catalog_origin(provider: ApiProvider, wire_model_id: &str) -> Option }; if guard .per_provider - .values() - .any(|snap| snap.offerings.iter().any(matches)) + .get(&owner) + .is_some_and(|snap| snap.offerings.iter().any(matches)) { return Some(LiveSource::PerProvider); } @@ -276,11 +372,14 @@ pub(crate) fn lock_live_snapshot() -> std::sync::MutexGuard<'static, ()> { .unwrap_or_else(|poisoned| poisoned.into_inner()) } -/// The merged catalog snapshot: live rows override bundled rows on -/// `(provider, wire_model_id)` identity (#4188). When no live snapshot is -/// present, this is just the offline bundled snapshot. Per-provider live rows -/// override Models.dev live rows on collision (gateway-specific wins over -/// cross-provider). +/// The merged catalog snapshot: Models.dev rows override bundled rows on +/// `(provider, wire_model_id)` identity (#4188). A provider-owned live +/// partition is authoritative for that provider's complete roster, so it +/// suppresses both bundled and Models.dev rows for the provider rather than +/// merely overlaying matching ids. This is what lets a successful +/// `/v1/models` refresh remove models retired upstream. Failed refreshes retain +/// the last successful provider partition; clearing a partition restores the +/// offline/cross-provider fallbacks. /// /// Memoized: the merge is recomputed only after a live-layer mutation bumps /// `LIVE_GENERATION`; every other call returns the cached `Arc` (the picker @@ -305,30 +404,64 @@ fn merged_snapshot() -> Arc { /// Uncached merge (see [`merged_snapshot`] for the caching seam). fn compute_merged_snapshot() -> CatalogSnapshot { - let live = LIVE_SNAPSHOT - .read() - .ok() - .and_then(|guard| guard.flattened()); - let merged = match live { - None => bundled_snapshot().clone(), - Some(live) => { - let mut merged: BTreeMap<(String, String), CatalogOffering> = BTreeMap::new(); - for row in &bundled_snapshot().offerings { - merged.insert( - (row.provider.clone(), row.wire_model_id.clone()), - row.clone(), - ); - } - for row in &live.offerings { + let Ok(live) = LIVE_SNAPSHOT.read() else { + return apply_provider_model_cutlines(bundled_snapshot().clone()); + }; + if live.models_dev.is_none() && live.per_provider.is_empty() { + return apply_provider_model_cutlines(bundled_snapshot().clone()); + } + + let authoritative_providers: std::collections::BTreeSet<&str> = live + .per_provider + .keys() + .filter_map(|owner| match owner { + LivePartitionOwner::BuiltIn(identity) => Some(identity.as_str()), + LivePartitionOwner::Custom(_) => None, + }) + .collect(); + let is_authoritative = |provider: &str| { + let key = catalog_partition_key(provider); + authoritative_providers.contains(key.as_str()) + }; + let mut merged: BTreeMap<(String, String), CatalogOffering> = BTreeMap::new(); + for row in &bundled_snapshot().offerings { + if !is_authoritative(&row.provider) { + merged.insert( + (row.provider.clone(), row.wire_model_id.clone()), + row.clone(), + ); + } + } + if let Some(models_dev) = &live.models_dev { + for row in &models_dev.offerings { + if !is_authoritative(&row.provider) { merged.insert( (row.provider.clone(), row.wire_model_id.clone()), row.clone(), ); } - CatalogSnapshot { - offerings: merged.into_values().collect(), + } + } + for provider_snapshot in live + .per_provider + .iter() + .filter_map(|(owner, snapshot)| match owner { + LivePartitionOwner::BuiltIn(_) => Some(snapshot), + LivePartitionOwner::Custom(identity) if ApiProvider::parse(identity).is_none() => { + Some(snapshot) } + LivePartitionOwner::Custom(_) => None, + }) + { + for row in &provider_snapshot.offerings { + merged.insert( + (row.provider.clone(), row.wire_model_id.clone()), + row.clone(), + ); } + } + let merged = CatalogSnapshot { + offerings: merged.into_values().collect(), }; apply_provider_model_cutlines(merged) } @@ -342,6 +475,250 @@ fn catalog_provider_id(provider: ApiProvider) -> &'static str { } } +/// Exact partition key for one provider-owned catalog. +/// +/// Publishers of built-in catalogs already emit their canonical provider id. +/// Custom table identities are ownership boundaries and therefore remain +/// case-sensitive even when their spelling resembles a built-in provider or a +/// reviewed setup-template alias: `[providers.openai]` may intentionally shadow +/// the built-in, and `CustomA` / `customa` may be different hosts. +pub(crate) fn catalog_partition_key(provider: &str) -> String { + provider.trim().to_string() +} + +/// Resolve the catalog partition for a concrete route. +/// +/// `ApiProvider::Custom` is only the wire family. Named compatible providers +/// such as Baseten own independent catalogs and must keep their exact config +/// identity instead of collapsing into a shared `custom` bucket. +fn catalog_provider_id_for_identity<'a>( + provider: ApiProvider, + provider_identity: Option<&'a str>, +) -> Cow<'a, str> { + if provider == ApiProvider::Custom + && let Some(identity) = provider_identity.map(str::trim).filter(|id| !id.is_empty()) + { + return Cow::Owned(catalog_partition_key(identity)); + } + Cow::Borrowed(catalog_provider_id(provider)) +} + +fn offering_key(offering: &ProviderModelOffering) -> (String, String) { + ( + offering.provider.as_str().trim().to_ascii_lowercase(), + offering.wire_model_id.as_str().to_string(), + ) +} + +fn row_matches_endpoint_fingerprint(row: &CatalogOffering, fingerprint: &str) -> bool { + matches!( + &row.source, + CatalogSource::Live { + base_url_fingerprint, + .. + } if base_url_fingerprint == fingerprint + ) +} + +/// Build or reuse the runtime resolver for an exact provider identity. +/// +/// Only a fresh provider-owned partition whose source fingerprint matches the +/// selected endpoint can carry live limits, capabilities, and pricing into an +/// executable route. Stale, failed, unknown, or wrong-endpoint partitions stay +/// visible to the picker but are removed from this resolver and replaced by the +/// ordinary Models.dev/bundled fallback. Named compatible providers such as +/// Baseten are remapped from their exact catalog identity to the resolver's +/// `custom` transport scope only after this check. +pub(crate) fn runtime_catalog_resolver_for_identity( + provider: ApiProvider, + provider_identity: Option<&str>, + base_url: &str, + status: CatalogStatus, +) -> RuntimeCatalogResolver { + let catalog_id = catalog_provider_id_for_identity(provider, provider_identity); + let catalog_key = catalog_partition_key(catalog_id.as_ref()); + let fingerprint = base_url_fingerprint(base_url); + let status_is_fresh = matches!(status, CatalogStatus::Fresh); + let generation = LIVE_GENERATION.load(Ordering::SeqCst); + let cache_key = format!( + "{}\u{1f}{}\u{1f}{}", + provider.as_str(), + catalog_key, + fingerprint + ); + + if let Ok(cache) = RUNTIME_RESOLVER_CACHE.read() + && let Some(cached) = cache.get(&cache_key) + && cached.generation == generation + && cached.status_is_fresh == status_is_fresh + { + return RuntimeCatalogResolver { + resolver: cached.resolver.clone(), + endpoint_catalog_authoritative: cached.endpoint_catalog_authoritative, + }; + } + + let mut snapshot = (*merged_snapshot()).clone(); + let partition_owner = live_partition_owner_for_route(provider, provider_identity); + let ( + endpoint_catalog_authoritative, + authoritative_providers, + models_dev_fallback, + exact_custom_rows, + ) = if let Ok(live) = LIVE_SNAPSHOT.read() { + let exact_partition = live.per_provider.get(&partition_owner); + let exact_matches = status_is_fresh + && exact_partition.is_some_and(|partition| { + !partition.offerings.is_empty() + && partition.offerings.iter().all(|row| { + catalog_partition_key(&row.provider) == catalog_key + && row_matches_endpoint_fingerprint(row, &fingerprint) + }) + }); + let mut authoritative: std::collections::BTreeSet = live + .per_provider + .keys() + .filter_map(|owner| match owner { + LivePartitionOwner::BuiltIn(identity) => Some(identity.clone()), + LivePartitionOwner::Custom(_) => None, + }) + .collect(); + if !exact_matches && let LivePartitionOwner::BuiltIn(identity) = &partition_owner { + authoritative.remove(identity); + } + let fallback = if provider == ApiProvider::Custom { + Vec::new() + } else { + live.models_dev + .as_ref() + .map(|models_dev| { + models_dev + .offerings + .iter() + .filter(|row| catalog_partition_key(&row.provider) == catalog_key) + .cloned() + .collect::>() + }) + .unwrap_or_default() + }; + let custom_rows = if provider == ApiProvider::Custom && exact_matches { + exact_partition + .map(|partition| partition.offerings.clone()) + .unwrap_or_default() + } else { + Vec::new() + }; + (exact_matches, authoritative, fallback, custom_rows) + } else { + ( + false, + std::collections::BTreeSet::new(), + Vec::new(), + Vec::new(), + ) + }; + + if provider != ApiProvider::Custom && !endpoint_catalog_authoritative { + let mut restored: BTreeMap<(String, String), CatalogOffering> = snapshot + .offerings + .into_iter() + .filter(|row| catalog_partition_key(&row.provider) != catalog_key) + .map(|row| { + ( + ( + catalog_partition_key(&row.provider), + row.wire_model_id.clone(), + ), + row, + ) + }) + .collect(); + for row in bundled_snapshot() + .offerings + .iter() + .filter(|row| catalog_partition_key(&row.provider) == catalog_key) + .cloned() + .chain(models_dev_fallback) + { + restored.insert( + ( + catalog_partition_key(&row.provider), + row.wire_model_id.clone(), + ), + row, + ); + } + snapshot.offerings = restored.into_values().collect(); + } + + let mut route_offerings: BTreeMap<(String, String), ProviderModelOffering> = snapshot + .to_offerings() + .into_iter() + .map(|offering| (offering_key(&offering), offering)) + .collect(); + + // Restore the curated transport seam only where a provider-owned roster is + // not authoritative. Curated endpoint/protocol facts still win over the + // bundled Models.dev asset, matching `RouteResolver::new()`. + for offering in bundled_offerings() { + if !authoritative_providers.contains(offering.provider.as_str()) { + route_offerings.insert(offering_key(&offering), offering); + } + } + + if provider == ApiProvider::Custom { + for row in &exact_custom_rows { + let mut mapped = row.clone(); + mapped.provider = ApiProvider::Custom.as_str().to_string(); + let offering = mapped.to_offering(); + route_offerings.insert(offering_key(&offering), offering); + } + } + + let resolver = RouteResolver::from_offerings(route_offerings.into_values().collect()); + if let Ok(mut cache) = RUNTIME_RESOLVER_CACHE.write() { + cache.insert( + cache_key, + RuntimeResolverCacheEntry { + generation, + status_is_fresh, + endpoint_catalog_authoritative, + resolver: resolver.clone(), + }, + ); + } + RuntimeCatalogResolver { + resolver, + endpoint_catalog_authoritative, + } +} + +fn offerings_for_provider_identity<'a>( + snapshot: &'a CatalogSnapshot, + provider_id: &str, +) -> Vec<&'a CatalogOffering> { + let provider_key = catalog_partition_key(provider_id); + snapshot + .offerings + .iter() + .filter(|row| catalog_partition_key(&row.provider) == provider_key) + .collect() +} + +fn exact_custom_offerings(provider_identity: &str) -> Vec { + let provider_identity = provider_identity.trim(); + if provider_identity.is_empty() { + return Vec::new(); + } + let owner = LivePartitionOwner::Custom(catalog_partition_key(provider_identity)); + LIVE_SNAPSHOT + .read() + .ok() + .and_then(|live| live.per_provider.get(&owner).cloned()) + .map(|snapshot| snapshot.offerings) + .unwrap_or_default() +} + fn push_unique_model(models: &mut Vec, model: &str) { let model = model.trim(); if model.is_empty() { @@ -381,6 +758,19 @@ fn catalog_models_from_offerings<'a>( /// local providers (and gateways not yet in the offline seed) keep defaults. #[must_use] pub fn all_catalog_models_for_provider(provider: ApiProvider) -> Vec { + all_catalog_models_for_provider_identity(provider, None) +} + +/// Catalog-backed model ids for one exact provider route. +/// +/// Built-in providers retain their canonical ids. Named compatible custom +/// routes use `provider_identity`, so Baseten's live `/v1/models` rows and its +/// offline setup-template seeds remain isolated from every other custom host. +#[must_use] +pub fn all_catalog_models_for_provider_identity( + provider: ApiProvider, + provider_identity: Option<&str>, +) -> Vec { // ChatGPT OAuth availability is account-scoped. A generic OpenAI or // Models.dev catalog is not evidence that a model can be routed through // the Codex backend, so this provider owns a separate secret-free source. @@ -388,9 +778,26 @@ pub fn all_catalog_models_for_provider(provider: ApiProvider) -> Vec { return codex_model_cache::model_roster().model_ids(); } - let catalog_id = catalog_provider_id(provider); + let catalog_id = catalog_provider_id_for_identity(provider, provider_identity); + let custom_offerings = + (provider == ApiProvider::Custom).then(|| exact_custom_offerings(catalog_id.as_ref())); let merged = merged_snapshot(); - let mut models = catalog_models_from_offerings(merged.offerings_for_provider(catalog_id)); + let mut models = match custom_offerings.as_ref() { + Some(rows) => catalog_models_from_offerings(rows.iter()), + None => catalog_models_from_offerings(offerings_for_provider_identity( + &merged, + catalog_id.as_ref(), + )), + }; + if models.is_empty() + && provider == ApiProvider::Custom + && let Some(template) = codewhale_config::provider_setup_template(catalog_id.as_ref()) + && template.is_compatible() + { + for model in template.picker_models() { + push_unique_model(&mut models, model); + } + } if models.is_empty() { for model in model_completion_names_for_provider(provider) { push_unique_model(&mut models, model); @@ -409,17 +816,31 @@ pub fn all_catalog_models_for_provider(provider: ApiProvider) -> Vec { pub fn catalog_offering_for_model( provider: ApiProvider, wire_model_id: &str, +) -> Option { + catalog_offering_for_model_identity(provider, None, wire_model_id) +} + +/// Look up a merged-catalog offering for one exact provider route. +#[must_use] +pub fn catalog_offering_for_model_identity( + provider: ApiProvider, + provider_identity: Option<&str>, + wire_model_id: &str, ) -> Option { if provider == ApiProvider::OpenaiCodex { return None; } - let catalog_id = catalog_provider_id(provider); + let catalog_id = catalog_provider_id_for_identity(provider, provider_identity); let needle = wire_model_id.trim(); if needle.is_empty() { return None; } - merged_snapshot() - .offerings_for_provider(catalog_id) + if provider == ApiProvider::Custom { + return exact_custom_offerings(catalog_id.as_ref()) + .into_iter() + .find(|row| row.wire_model_id.eq_ignore_ascii_case(needle)); + } + offerings_for_provider_identity(&merged_snapshot(), catalog_id.as_ref()) .into_iter() .find(|row| row.wire_model_id.eq_ignore_ascii_case(needle)) .cloned() @@ -676,6 +1097,234 @@ mod tests { assert_eq!(after_clear, bundled); } + #[test] + fn provider_owned_roster_replaces_bundled_and_models_dev_rows() { + let _live = lock_live_snapshot(); + clear_live_snapshot(); + let bundled = all_catalog_models_for_provider(ApiProvider::Openrouter); + assert!( + !bundled.is_empty(), + "OpenRouter must have an offline fallback roster" + ); + + set_live_snapshot( + CatalogSnapshot { + offerings: vec![CatalogOffering { + provider: "openrouter".to_string(), + wire_model_id: "models-dev-only-openrouter-model".to_string(), + endpoint_key: "chat".to_string(), + ..Default::default() + }], + }, + LiveSource::ModelsDev, + ); + set_live_snapshot( + CatalogSnapshot { + offerings: vec![CatalogOffering { + provider: "openrouter".to_string(), + wire_model_id: "provider-owned-openrouter-model".to_string(), + endpoint_key: "chat".to_string(), + ..Default::default() + }], + }, + LiveSource::PerProvider, + ); + + assert_eq!( + all_catalog_models_for_provider(ApiProvider::Openrouter), + vec!["provider-owned-openrouter-model".to_string()], + "a successful provider-owned refresh must remove stale bundled and Models.dev ids" + ); + + replace_provider_live_snapshot("openrouter", CatalogSnapshot::default()); + let restored_cross_provider = all_catalog_models_for_provider(ApiProvider::Openrouter); + assert!( + restored_cross_provider.contains(&"models-dev-only-openrouter-model".to_string()), + "clearing the exact partition must restore the cross-provider fallback" + ); + assert!( + restored_cross_provider + .iter() + .any(|model| bundled.contains(model)), + "clearing the exact partition must restore bundled fallbacks" + ); + + clear_live_snapshot(); + assert_eq!( + all_catalog_models_for_provider(ApiProvider::Openrouter), + bundled + ); + } + + #[test] + fn named_custom_catalogs_keep_exact_identity_and_baseten_offline_seeds() { + let _live = lock_live_snapshot(); + clear_live_snapshot(); + + let offline = all_catalog_models_for_provider_identity( + ApiProvider::Custom, + Some(codewhale_config::BASETEN_TEMPLATE_ID), + ); + assert_eq!( + offline, + codewhale_config::BASETEN_MODELS + .iter() + .map(|model| (*model).to_string()) + .collect::>() + ); + assert!( + !all_catalog_models_for_provider_identity( + ApiProvider::Custom, + Some("another-compatible-host"), + ) + .iter() + .any(|model| offline.contains(model)), + "Baseten seeds must not leak into another custom provider" + ); + + set_live_snapshot( + CatalogSnapshot { + offerings: vec![CatalogOffering { + provider: "baseten".to_string(), + wire_model_id: "synthetic-live-baseten-model".to_string(), + endpoint_key: "chat".to_string(), + source: CatalogSource::Live { + base_url_fingerprint: "baseten-fp".to_string(), + fetched_at: 42, + }, + ..Default::default() + }], + }, + LiveSource::PerProvider, + ); + + assert_eq!( + all_catalog_models_for_provider_identity(ApiProvider::Custom, Some("baseten")), + vec!["synthetic-live-baseten-model".to_string()] + ); + let case_distinct = + all_catalog_models_for_provider_identity(ApiProvider::Custom, Some("BASETEN")); + assert_eq!( + case_distinct, + codewhale_config::BASETEN_MODELS + .iter() + .map(|model| (*model).to_string()) + .collect::>(), + "template schema aliases may share offline seeds, but not another exact table's live roster" + ); + assert!( + catalog_offering_for_model_identity( + ApiProvider::Custom, + Some("baseten"), + "synthetic-live-baseten-model", + ) + .is_some() + ); + assert!( + catalog_offering_for_model(ApiProvider::Custom, "synthetic-live-baseten-model",) + .is_none(), + "the generic custom bucket must not see Baseten rows" + ); + + clear_live_snapshot(); + } + + #[test] + fn case_colliding_and_builtin_named_custom_catalogs_stay_isolated() { + let _live = lock_live_snapshot(); + clear_live_snapshot(); + + for (provider, model) in [("CustomA", "upper-model"), ("customa", "lower-model")] { + replace_provider_live_snapshot( + provider, + CatalogSnapshot { + offerings: vec![CatalogOffering { + provider: provider.to_string(), + wire_model_id: model.to_string(), + endpoint_key: "chat".to_string(), + ..Default::default() + }], + }, + ); + } + + assert_eq!( + all_catalog_models_for_provider_identity(ApiProvider::Custom, Some("CustomA")), + vec!["upper-model".to_string()] + ); + assert_eq!( + all_catalog_models_for_provider_identity(ApiProvider::Custom, Some("customa")), + vec!["lower-model".to_string()] + ); + let built_in_openai = all_catalog_models_for_provider(ApiProvider::Openai); + assert!(!built_in_openai.is_empty()); + assert!( + all_catalog_models_for_provider_identity(ApiProvider::Custom, Some("openai")) + .is_empty(), + "a custom table named openai must not borrow the first-class OpenAI template" + ); + for model in &built_in_openai { + assert!( + catalog_offering_for_model_identity(ApiProvider::Custom, Some("openai"), model) + .is_none(), + "an exact custom table named openai must not inherit built-in model {model}" + ); + } + + let custom_model = "custom-openai-only-model"; + replace_provider_live_snapshot_for_identity( + ApiProvider::Custom, + "openai", + CatalogSnapshot { + offerings: vec![CatalogOffering { + provider: "openai".to_string(), + wire_model_id: custom_model.to_string(), + endpoint_key: "chat".to_string(), + ..Default::default() + }], + }, + ); + assert_eq!( + all_catalog_models_for_provider_identity(ApiProvider::Custom, Some("openai")), + vec![custom_model.to_string()], + "the exact custom table must retrieve its own built-in-looking roster" + ); + assert_eq!( + all_catalog_models_for_provider(ApiProvider::Openai), + built_in_openai, + "publishing custom openai must not replace or suppress built-in OpenAI" + ); + assert!( + catalog_offering_for_model(ApiProvider::Openai, custom_model).is_none(), + "the built-in OpenAI route must not see the custom table's row" + ); + + let built_in_live_model = "built-in-openai-only-model"; + replace_provider_live_snapshot_for_identity( + ApiProvider::Openai, + "openai", + CatalogSnapshot { + offerings: vec![CatalogOffering { + provider: "openai".to_string(), + wire_model_id: built_in_live_model.to_string(), + endpoint_key: "chat".to_string(), + ..Default::default() + }], + }, + ); + assert_eq!( + all_catalog_models_for_provider(ApiProvider::Openai), + vec![built_in_live_model.to_string()] + ); + assert_eq!( + all_catalog_models_for_provider_identity(ApiProvider::Custom, Some("openai")), + vec![custom_model.to_string()], + "publishing built-in OpenAI must not replace the custom table's roster" + ); + + clear_live_snapshot(); + } + #[test] fn live_catalog_origin_prefers_per_provider_over_models_dev() { let _live = lock_live_snapshot(); diff --git a/crates/tui/src/remote_control.rs b/crates/tui/src/remote_control.rs index 0aced1218e..a10037bf71 100644 --- a/crates/tui/src/remote_control.rs +++ b/crates/tui/src/remote_control.rs @@ -7491,6 +7491,8 @@ mod tests { fn turn_complete_event() -> EngineEvent { EngineEvent::TurnComplete { usage: crate::models::Usage::default(), + parent_route_usage: crate::models::Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -7977,6 +7979,8 @@ mod tests { }; controller.observe_engine_event(&EngineEvent::TurnComplete { usage: usage.clone(), + parent_route_usage: usage.clone(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, diff --git a/crates/tui/src/rlm/bridge.rs b/crates/tui/src/rlm/bridge.rs index 7fd69da5d9..dba61c5cac 100644 --- a/crates/tui/src/rlm/bridge.rs +++ b/crates/tui/src/rlm/bridge.rs @@ -18,6 +18,7 @@ use std::{future::Future, pin::Pin}; use anyhow::Result; use futures_util::future::join_all; use tokio::sync::Mutex; +use uuid::Uuid; use crate::llm_client::LlmClient; use crate::models::Role; @@ -28,6 +29,183 @@ use crate::models::{ use crate::repl::runtime::{BatchResp, RpcDispatcher, RpcRequest, RpcResponse, SingleResp}; use crate::utils::spawn_supervised; +/// One pre-dispatch reservation in the shared routed-usage ledger. +#[derive(Debug, Clone, Copy)] +pub(crate) struct RlmUsageReservation { + index: usize, +} + +#[derive(Debug, Default)] +struct RlmUsageState { + ledger_id: String, + usage: Usage, + records: Vec>, + drop_records: Vec, + dropped_records: u64, +} + +#[derive(Debug)] +struct RlmUsageSlot { + record: crate::cost_status::RuntimeUsageRecord, + completed: bool, +} + +/// Shared, bounded provider-call ledger for one complete RLM tree. +/// +/// Every root, child, batch member, and recursive call reserves one slot +/// before invoking a provider. A distinct call is never coalesced merely +/// because it used the same route: its dispatch instant and frozen quote are +/// independent accounting evidence. Sharing one accumulator across recursion +/// makes the bound global instead of allowing every nested bridge to reset it. +#[derive(Debug, Clone)] +pub(crate) struct RlmUsageAccumulator { + state: Arc>, +} + +/// Atomic snapshot returned after all RPC work for a round has settled. +#[derive(Debug, Clone, Default)] +pub(crate) struct RlmUsageSnapshot { + pub usage: Usage, + pub records: Vec, + /// Exact frozen routes for provider-success responses that did not carry + /// authoritative usage. Keeping these separate prevents a missing payload + /// from becoming a priced zero-usage receipt. + pub drop_records: Vec, + /// Calls whose execution/usage became ambiguous (for example a timeout). + /// They are never represented as authoritative zero-usage responses. + pub dropped_records: u64, +} + +impl RlmUsageAccumulator { + #[must_use] + pub(crate) fn new() -> Self { + Self { + state: Arc::new(Mutex::new(RlmUsageState { + ledger_id: Uuid::new_v4().simple().to_string(), + ..RlmUsageState::default() + })), + } + } + + /// Reserve durable accounting capacity before a provider request. + /// Definite transport failure cancels the slot; ambiguous cancellation is + /// explicit incomplete coverage. Reaching the cap rejects before any + /// unreceipted provider work can occur. + pub(crate) async fn reserve( + &self, + route: crate::cost_status::EffectiveRouteEnvelope, + ) -> std::result::Result { + let mut state = self.state.lock().await; + if state.records.len() == crate::cost_status::MAX_CHILD_USAGE_RECORDS { + return Err(format!( + "RLM provider-call receipt limit reached ({}); request rejected before dispatch", + crate::cost_status::MAX_CHILD_USAGE_RECORDS + )); + } + let index = state.records.len(); + let source_id = format!("rlm:{}:request:{index}", state.ledger_id); + state.records.push(Some(RlmUsageSlot { + record: crate::cost_status::RuntimeUsageRecord { + source_id, + usage: crate::cost_status::EffectiveRouteUsage { + route: route.sanitized_for_persistence(), + usage: Usage::default(), + }, + }, + completed: false, + })); + Ok(RlmUsageReservation { index }) + } + + /// Attach a provider's reported usage to its already-reserved exact route. + pub(crate) async fn complete(&self, reservation: RlmUsageReservation, usage: &Usage) { + let mut state = self.state.lock().await; + let completed = if let Some(Some(slot)) = state.records.get_mut(reservation.index) + && !slot.completed + { + super::add_usage_with_prompt_cache(&mut slot.record.usage.usage, usage); + slot.completed = true; + true + } else { + false + }; + if completed { + super::add_usage_with_prompt_cache(&mut state.usage, usage); + } + } + + /// Settle a decoded provider-success response without inventing usage. + /// `MessageResponse::usage == Usage::default()` is also what adapters + /// produce when the provider omitted the payload, so it is not proof of a + /// genuine zero-token request. + pub(crate) async fn settle_provider_success( + &self, + reservation: RlmUsageReservation, + usage: &Usage, + ) { + if usage == &Usage::default() { + self.cancel(reservation, true).await; + } else { + self.complete(reservation, usage).await; + } + } + + /// Remove a reservation that never produced provider-reported usage. + /// Ambiguous execution increments explicit incomplete coverage instead of + /// being persisted as a priced-zero response. + pub(crate) async fn cancel(&self, reservation: RlmUsageReservation, coverage_unknown: bool) { + let mut state = self.state.lock().await; + let cancelled = state.records.get_mut(reservation.index).and_then(|slot| { + if slot.as_ref().is_some_and(|slot| !slot.completed) { + slot.take() + } else { + None + } + }); + if let Some(slot) = cancelled + && coverage_unknown + { + state + .drop_records + .push(crate::cost_status::RuntimeUsageDropRecord { + source_id: slot.record.source_id, + route: slot.record.usage.route, + }); + state.dropped_records = state.dropped_records.saturating_add(1); + } + } + + pub(crate) async fn snapshot(&self) -> RlmUsageSnapshot { + let state = self.state.lock().await; + let pending = state + .records + .iter() + .flatten() + .filter(|slot| !slot.completed) + .map(|slot| crate::cost_status::RuntimeUsageDropRecord { + source_id: slot.record.source_id.clone(), + route: slot.record.usage.route.clone(), + }) + .collect::>(); + let mut drop_records = state.drop_records.clone(); + drop_records.extend(pending.iter().cloned()); + RlmUsageSnapshot { + usage: state.usage.clone(), + records: state + .records + .iter() + .flatten() + .filter(|slot| slot.completed) + .map(|slot| slot.record.clone()) + .collect(), + drop_records, + dropped_records: state + .dropped_records + .saturating_add(u64::try_from(pending.len()).unwrap_or(u64::MAX)), + } + } +} + /// Object-safe runtime-model adapter for a working kernel. /// /// The normal turn loop owns a `SharedModelClient`, while the original RLM @@ -124,7 +302,7 @@ pub struct RlmBridge { /// Recursion budget remaining for `Rlm` / `RlmBatch` requests. When /// zero, those requests fall back to plain `Llm` completions. depth_remaining: u32, - usage: Arc>, + usage: RlmUsageAccumulator, } impl RlmBridge { @@ -132,17 +310,31 @@ impl RlmBridge { client: Arc, child_model: String, depth_remaining: u32, + ) -> Self { + Self::with_usage_accumulator( + client, + child_model, + depth_remaining, + RlmUsageAccumulator::new(), + ) + } + + pub(crate) fn with_usage_accumulator( + client: Arc, + child_model: String, + depth_remaining: u32, + usage: RlmUsageAccumulator, ) -> Self { Self { client, child_model, depth_remaining, - usage: Arc::new(Mutex::new(Usage::default())), + usage, } } - pub fn usage_handle(&self) -> Arc> { - Arc::clone(&self.usage) + pub(crate) async fn usage_snapshot(&self) -> RlmUsageSnapshot { + self.usage.snapshot().await } async fn dispatch_llm( @@ -155,6 +347,15 @@ impl RlmBridge { let request_route = self .client .effective_route_envelope(&self.child_model, chrono::Utc::now()); + let reservation = match self.usage.reserve(request_route.clone()).await { + Ok(reservation) => reservation, + Err(error) => { + return SingleResp { + text: String::new(), + error: Some(error), + }; + } + }; let route_max_tokens = self .client .effective_max_output_tokens(&request_route.model); @@ -191,12 +392,14 @@ impl RlmBridge { match tokio::time::timeout(Duration::from_secs(CHILD_TIMEOUT_SECS), fut).await { Ok(Ok(r)) => r, Ok(Err(e)) => { + self.usage.cancel(reservation, false).await; return SingleResp { text: String::new(), error: Some(format!("llm_query failed: {e}")), }; } Err(_) => { + self.usage.cancel(reservation, true).await; return SingleResp { text: String::new(), error: Some(format!("llm_query timed out after {CHILD_TIMEOUT_SECS}s")), @@ -204,10 +407,12 @@ impl RlmBridge { } }; - { - let mut u = self.usage.lock().await; - super::add_usage_with_prompt_cache(&mut u, &response.usage); - } + // Incomplete output is rejected below, but it is still a successful + // provider response and therefore billed. Complete the reserved route + // before inspecting the stop reason. + self.usage + .settle_provider_success(reservation, &response.usage) + .await; if is_incomplete_stop_reason(response.stop_reason.as_deref()) { return SingleResp { @@ -279,7 +484,7 @@ impl RlmBridge { // Recursive call. The dyn-erasure on `run_rlm_turn_inner` breaks // the `bridge → turn → bridge` opaque-future cycle. - let result = super::turn::run_rlm_turn_inner( + let result = super::turn::run_rlm_turn_inner_with_usage( Arc::clone(&self.client), child_model.clone(), prompt, @@ -287,16 +492,12 @@ impl RlmBridge { child_model, tx, self.depth_remaining.saturating_sub(1), + self.usage.clone(), ) .await; drain.abort(); - { - let mut u = self.usage.lock().await; - super::add_usage_with_prompt_cache(&mut u, &result.usage); - } - SingleResp { text: result.answer, error: result.error, @@ -518,9 +719,87 @@ mod tests { Some(SystemPrompt::Text("child system".to_string())) ); - let usage = bridge.usage.lock().await; - assert_eq!(usage.input_tokens, 7); - assert_eq!(usage.output_tokens, 11); + let snapshot = bridge.usage_snapshot().await; + assert_eq!(snapshot.usage.input_tokens, 7); + assert_eq!(snapshot.usage.output_tokens, 11); + assert_eq!(snapshot.records.len(), 1); + assert_eq!(snapshot.records[0].usage.usage, snapshot.usage); + assert!(snapshot.drop_records.is_empty()); + assert_eq!(snapshot.dropped_records, 0); + } + + #[tokio::test] + async fn llm_dispatch_keeps_semantic_success_but_marks_missing_usage_once() { + let mock = Arc::new(MockLlmClient::new(Vec::new())); + mock.push_message_response(mock_response_with_usage( + "usable child answer", + Usage::default(), + )); + let bridge = bridge_for(Arc::clone(&mock), 1); + + let response = bridge + .dispatch(RpcRequest::Llm { + prompt: "child prompt".to_string(), + model: None, + max_tokens: None, + system: None, + }) + .await; + + let RpcResponse::Single(response) = response else { + panic!("expected single response"); + }; + assert_eq!(response.text, "usable child answer"); + assert!(response.error.is_none()); + + let first = bridge.usage_snapshot().await; + let replay = bridge.usage_snapshot().await; + assert_eq!(first.usage, Usage::default()); + assert!(first.records.is_empty()); + assert_eq!(first.drop_records.len(), 1); + assert_eq!(first.dropped_records, 1); + assert_eq!(replay.drop_records, first.drop_records); + assert_eq!(replay.dropped_records, 1); + assert_eq!(first.drop_records[0].route.model, "child-model"); + assert!(first.drop_records[0].source_id.starts_with("rlm:")); + } + + #[tokio::test] + async fn repeated_reservation_settlement_cannot_duplicate_usage_or_missing_coverage() { + let mock = Arc::new(MockLlmClient::new(Vec::new())); + let bridge = bridge_for(Arc::clone(&mock), 1); + let route = RlmLlmClient::effective_route_envelope( + mock.as_ref(), + "child-model", + chrono::Utc::now(), + ); + + let usage_reservation = bridge + .usage + .reserve(route.clone()) + .await + .expect("usage reservation"); + let reported = Usage { + input_tokens: 3, + output_tokens: 5, + ..Usage::default() + }; + bridge.usage.complete(usage_reservation, &reported).await; + bridge.usage.complete(usage_reservation, &reported).await; + + let missing_reservation = bridge + .usage + .reserve(route) + .await + .expect("missing reservation"); + bridge.usage.cancel(missing_reservation, true).await; + bridge.usage.cancel(missing_reservation, true).await; + + let snapshot = bridge.usage_snapshot().await; + assert_eq!(snapshot.usage, reported); + assert_eq!(snapshot.records.len(), 1); + assert_eq!(snapshot.drop_records.len(), 1); + assert_eq!(snapshot.dropped_records, 1); } #[tokio::test] @@ -555,7 +834,7 @@ mod tests { other => panic!("expected single response, got {other:?}"), } - let usage = bridge.usage.lock().await; + let usage = bridge.usage_snapshot().await.usage; assert_eq!(usage.input_tokens, 1000); assert_eq!(usage.output_tokens, 100); assert_eq!(usage.prompt_cache_hit_tokens, Some(800)); @@ -601,7 +880,10 @@ mod tests { other => panic!("expected single response, got {other:?}"), } - assert_eq!(*bridge.usage.lock().await, usage); + let snapshot = bridge.usage_snapshot().await; + assert_eq!(snapshot.usage, usage); + assert_eq!(snapshot.records.len(), 1); + assert_eq!(snapshot.records[0].usage.usage, usage); assert_eq!(mock.call_count(), 1, "truncation must not retry"); } @@ -643,9 +925,77 @@ mod tests { .all(|request| request.model == "child-model") ); - let usage = bridge.usage.lock().await; - assert_eq!(usage.input_tokens, 9); - assert_eq!(usage.output_tokens, 12); + let snapshot = bridge.usage_snapshot().await; + assert_eq!(snapshot.usage.input_tokens, 9); + assert_eq!(snapshot.usage.output_tokens, 12); + assert_eq!(snapshot.records.len(), 3); + assert_ne!( + snapshot.records[0].source_id, snapshot.records[1].source_id, + "distinct provider calls must keep distinct stable identities" + ); + } + + #[tokio::test] + async fn shared_accumulator_rejects_the_first_unreceipted_request_before_provider_work() { + let mock = Arc::new(MockLlmClient::new(Vec::new())); + let client: Arc = mock.clone(); + let usage = RlmUsageAccumulator::new(); + let bridge = RlmBridge::with_usage_accumulator( + Arc::clone(&client), + "child-model".to_string(), + 1, + usage.clone(), + ); + let nested_bridge = + RlmBridge::with_usage_accumulator(client, "child-model".to_string(), 1, usage); + let route = RlmLlmClient::effective_route_envelope( + mock.as_ref(), + "child-model", + chrono::Utc::now(), + ); + for _ in 0..crate::cost_status::MAX_CHILD_USAGE_RECORDS { + let reservation = bridge + .usage + .reserve(route.clone()) + .await + .expect("receipt slot below cap"); + bridge + .usage + .complete( + reservation, + &Usage { + input_tokens: 1, + ..Usage::default() + }, + ) + .await; + } + + let response = nested_bridge + .dispatch(RpcRequest::Llm { + prompt: "must not reach provider".to_string(), + model: None, + max_tokens: None, + system: None, + }) + .await; + let RpcResponse::Single(response) = response else { + panic!("expected single response"); + }; + assert!( + response + .error + .as_deref() + .is_some_and(|error| error.contains("rejected before dispatch")) + ); + assert_eq!(mock.call_count(), 0); + let snapshot = bridge.usage_snapshot().await; + assert_eq!( + snapshot.records.len(), + crate::cost_status::MAX_CHILD_USAGE_RECORDS + ); + assert_eq!(snapshot.dropped_records, 0); + assert!(snapshot.drop_records.is_empty()); } #[tokio::test] @@ -669,7 +1019,7 @@ mod tests { other => panic!("expected single response, got {other:?}"), } - let usage = bridge.usage.lock().await; + let usage = bridge.usage_snapshot().await.usage; assert_eq!(usage.input_tokens, 3); assert_eq!(usage.output_tokens, 5); diff --git a/crates/tui/src/rlm/turn.rs b/crates/tui/src/rlm/turn.rs index ebef384746..c26aef0bee 100644 --- a/crates/tui/src/rlm/turn.rs +++ b/crates/tui/src/rlm/turn.rs @@ -16,7 +16,7 @@ use crate::models::{ }; use crate::repl::PythonRuntime; -use super::bridge::{RlmBridge, RlmLlmClient}; +use super::bridge::{RlmBridge, RlmLlmClient, RlmUsageAccumulator}; use super::prompt::rlm_system_prompt; use crate::models::Role; @@ -79,6 +79,13 @@ pub struct RlmTurnResult { pub duration: Duration, pub error: Option, pub usage: Usage, + /// One exact frozen route/quote receipt per admitted provider request. + /// Distinct calls are never coalesced, even when they share a route. + pub routed_usage: Vec, + /// Exact routes for provider-success responses that omitted authoritative + /// usage metadata. + pub routed_usage_drop_records: Vec, + pub routed_usage_dropped_records: u64, pub termination: RlmTermination, /// Per-round trace. Empty when the loop never reached the REPL. pub trace: Vec, @@ -145,7 +152,7 @@ pub(crate) fn run_rlm_turn_inner( tx_event: mpsc::Sender, max_depth: u32, ) -> std::pin::Pin + Send>> { - Box::pin(run_rlm_turn_impl( + run_rlm_turn_inner_with_usage( client, model, prompt, @@ -153,7 +160,41 @@ pub(crate) fn run_rlm_turn_inner( child_model, tx_event, max_depth, - )) + RlmUsageAccumulator::new(), + ) +} + +/// Recursive entry point that keeps one pre-dispatch receipt bound across the +/// entire nested/batched RLM tree. +pub(crate) fn run_rlm_turn_inner_with_usage( + client: Arc, + model: String, + prompt: String, + root_prompt: Option, + child_model: String, + tx_event: mpsc::Sender, + max_depth: u32, + usage: RlmUsageAccumulator, +) -> std::pin::Pin + Send>> { + Box::pin(async move { + let mut result = run_rlm_turn_impl( + client, + model, + prompt, + root_prompt, + child_model, + tx_event, + max_depth, + usage.clone(), + ) + .await; + let snapshot = usage.snapshot().await; + result.usage = snapshot.usage; + result.routed_usage = snapshot.records; + result.routed_usage_drop_records = snapshot.drop_records; + result.routed_usage_dropped_records = snapshot.dropped_records; + result + }) } /// RLM turns are long-running background-style work. Do not kill the whole @@ -175,6 +216,7 @@ async fn run_rlm_turn_impl( child_model: String, tx_event: mpsc::Sender, max_depth: u32, + routed_usage: RlmUsageAccumulator, ) -> RlmTurnResult { let start = Instant::now(); let mut total_usage = Usage::default(); @@ -193,6 +235,9 @@ async fn run_rlm_turn_impl( duration: start.elapsed(), error: Some(format!("rlm: failed to stage context: {e}")), usage: total_usage, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, termination: RlmTermination::Error, trace, total_rpcs, @@ -211,6 +256,9 @@ async fn run_rlm_turn_impl( duration: start.elapsed(), error: Some(format!("rlm: failed to spawn REPL: {e}")), usage: total_usage, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, termination: RlmTermination::Error, trace, total_rpcs, @@ -219,8 +267,12 @@ async fn run_rlm_turn_impl( }; // 3. Build the bridge that services llm_query / rlm_query RPCs. - let bridge = RlmBridge::new(Arc::clone(&client), child_model.clone(), max_depth); - let usage_handle = bridge.usage_handle(); + let bridge = RlmBridge::with_usage_accumulator( + Arc::clone(&client), + child_model.clone(), + max_depth, + routed_usage.clone(), + ); let _ = tx_event .send(Event::status(format!( @@ -254,6 +306,9 @@ async fn run_rlm_turn_impl( duration: start.elapsed(), error: Some(format!("RLM turn timed out after {}s", timeout.as_secs())), usage: total_usage, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, termination: RlmTermination::Error, trace: trace.clone(), total_rpcs, @@ -270,6 +325,24 @@ async fn run_rlm_turn_impl( // 4a. Root LLM generates code from metadata-only context. let request_route = client.effective_route_envelope(&model, chrono::Utc::now()); + let reservation = match routed_usage.reserve(request_route.clone()).await { + Ok(reservation) => reservation, + Err(error) => { + break 'turn RlmTurnResult { + answer: String::new(), + iterations: iteration, + duration: start.elapsed(), + error: Some(error), + usage: total_usage, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, + termination: RlmTermination::Error, + trace: trace.clone(), + total_rpcs, + }; + } + }; let request = build_root_request( &model, &messages, @@ -280,12 +353,16 @@ async fn run_rlm_turn_impl( let response = match client.create_message_boxed(request).await { Ok(r) => r, Err(e) => { + routed_usage.cancel(reservation, false).await; break 'turn RlmTurnResult { answer: String::new(), iterations: iteration + 1, duration: start.elapsed(), error: Some(format!("Root LLM call failed: {e}")), usage: total_usage, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, termination: RlmTermination::Error, trace: trace.clone(), total_rpcs, @@ -293,6 +370,11 @@ async fn run_rlm_turn_impl( } }; + // Preserve billed usage even when the response is incomplete and + // its partial FINAL/REPL output is rejected below. + routed_usage + .settle_provider_success(reservation, &response.usage) + .await; super::add_usage_with_prompt_cache(&mut total_usage, &response.usage); if is_incomplete_stop_reason(response.stop_reason.as_deref()) { @@ -305,6 +387,9 @@ async fn run_rlm_turn_impl( "RLM root model response incomplete: provider stop reason `{reason}`; partial FINAL/REPL output was not accepted." )), usage: total_usage, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, termination: RlmTermination::Error, trace: trace.clone(), total_rpcs, @@ -331,6 +416,9 @@ async fn run_rlm_turn_impl( duration: start.elapsed(), error: None, usage: total_usage, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, termination: RlmTermination::NoCode, trace: trace.clone(), total_rpcs, @@ -369,6 +457,9 @@ async fn run_rlm_turn_impl( duration: start.elapsed(), error: None, usage: total_usage, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, termination: RlmTermination::Final, trace: trace.clone(), total_rpcs, @@ -393,6 +484,9 @@ async fn run_rlm_turn_impl( "RLM: model failed to emit ```repl after {MAX_CONSECUTIVE_NO_CODE} consecutive rounds" )), usage: total_usage, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, termination: RlmTermination::NoCode, trace: trace.clone(), total_rpcs, @@ -441,6 +535,9 @@ async fn run_rlm_turn_impl( duration: start.elapsed(), error: Some(format!("REPL execution failed: {e}")), usage: total_usage, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, termination: RlmTermination::Error, trace: trace.clone(), total_rpcs, @@ -484,6 +581,9 @@ async fn run_rlm_turn_impl( duration: start.elapsed(), error: None, usage: total_usage, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, termination: RlmTermination::Final, trace: trace.clone(), total_rpcs, @@ -534,6 +634,9 @@ async fn run_rlm_turn_impl( "RLM: {MAX_CONSECUTIVE_NO_CODE} consecutive empty REPL rounds" )), usage: total_usage, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, termination: RlmTermination::NoCode, trace: trace.clone(), total_rpcs, @@ -608,24 +711,17 @@ async fn run_rlm_turn_impl( "RLM loop exhausted after {MAX_RLM_ITERATIONS} iterations without FINAL" )), usage: total_usage, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_dropped_records: 0, termination: RlmTermination::Exhausted, trace: trace.clone(), total_rpcs, } }; - // Fold bridge usage (children + nested sub_rlm) into totals. - let bridge_usage = usage_handle.lock().await; - let mut final_usage = result.usage.clone(); - super::add_usage_with_prompt_cache(&mut final_usage, &bridge_usage); - drop(bridge_usage); - repl.shutdown().await; - - RlmTurnResult { - usage: final_usage, - ..result - } + result } // --------------------------------------------------------------------------- @@ -956,6 +1052,9 @@ mod tests { assert!(error.contains("incomplete"), "{error}"); assert!(error.contains("max_tokens"), "{error}"); assert_eq!(result.usage, usage, "billed usage must still be charged"); + assert_eq!(result.routed_usage.len(), 1); + assert_eq!(result.routed_usage[0].usage.usage, usage); + assert_eq!(result.routed_usage_dropped_records, 0); assert_eq!(mock.call_count(), 1, "truncation must not retry"); assert!( !marker.exists(), @@ -963,6 +1062,48 @@ mod tests { ); } + #[tokio::test] + async fn root_provider_success_without_usage_retains_exact_missing_receipt() { + let mock = Arc::new(MockLlmClient::new(Vec::new())); + mock.push_message_response(MessageResponse { + id: "mock_missing_rlm_usage".to_string(), + r#type: "message".to_string(), + role: "assistant".to_string(), + content: vec![ContentBlock::Text { + text: "partial".to_string(), + cache_control: None, + }], + model: "mock-model".to_string(), + stop_reason: Some("max_tokens".to_string()), + stop_sequence: None, + container: None, + usage: Usage::default(), + }); + let client: Arc = mock; + let (tx, _rx) = mpsc::channel(8); + + let result = run_rlm_turn_inner( + client, + "root-model".to_string(), + "long context".to_string(), + None, + "child-model".to_string(), + tx, + 0, + ) + .await; + + assert_eq!(result.termination, RlmTermination::Error); + assert_eq!(result.usage, Usage::default()); + assert!(result.routed_usage.is_empty()); + assert_eq!(result.routed_usage_drop_records.len(), 1); + assert_eq!(result.routed_usage_dropped_records, 1); + assert_eq!( + result.routed_usage_drop_records[0].route.model, + "root-model" + ); + } + #[test] fn extract_repl_code_finds_simple_block() { let text = "Here:\n```repl\nprint('hi')\n```\nEnd."; diff --git a/crates/tui/src/route_runtime.rs b/crates/tui/src/route_runtime.rs index a49ceffa30..045c4e9231 100644 --- a/crates/tui/src/route_runtime.rs +++ b/crates/tui/src/route_runtime.rs @@ -425,6 +425,28 @@ pub(crate) fn resolve_route_candidate_with_context_metadata( base_url_override: Option, context_window_override: Option, provider_reported_context: Option, +) -> Result { + resolve_route_candidate_with_catalog_resolver( + provider, + model_selector, + saved_provider_model, + base_url_override, + context_window_override, + provider_reported_context, + &RouteResolver::new(), + false, + ) +} + +fn resolve_route_candidate_with_catalog_resolver( + provider: ApiProvider, + model_selector: Option<&str>, + saved_provider_model: Option<&str>, + base_url_override: Option, + context_window_override: Option, + provider_reported_context: Option, + resolver: &RouteResolver, + endpoint_catalog_authoritative: bool, ) -> Result { let effective_base_url = base_url_override .as_deref() @@ -432,7 +454,6 @@ pub(crate) fn resolve_route_candidate_with_context_metadata( if let Some(model) = model_selector.or(saved_provider_model) { validate_kimi_code_api_model_id(provider, effective_base_url, model)?; } - let resolver = RouteResolver::new(); let base_request = RouteRequest { explicit_provider: provider.kind(), model_selector: model_selector.map(|model| LogicalModelRef::from(model.to_string())), @@ -447,9 +468,14 @@ pub(crate) fn resolve_route_candidate_with_context_metadata( // requested through `RouteRequest::limit_overrides` on a second pass; the // resolver applies them BEFORE minting the final candidate and records // their provenance on it. - let resolved = resolver - .resolve(&base_request) - .map_err(|err| err.to_string())?; + let resolve = |request: &RouteRequest| { + if endpoint_catalog_authoritative { + resolver.resolve_with_endpoint_catalog_authority(request) + } else { + resolver.resolve(request) + } + }; + let resolved = resolve(&base_request).map_err(|err| err.to_string())?; let plan = plan_limit_overrides( provider, &resolved, @@ -459,12 +485,11 @@ pub(crate) fn resolve_route_candidate_with_context_metadata( let candidate = if plan.overrides.is_empty() { resolved } else { - resolver - .resolve(&RouteRequest { - limit_overrides: plan.overrides, - ..base_request - }) - .map_err(|err| err.to_string())? + resolve(&RouteRequest { + limit_overrides: plan.overrides, + ..base_request + }) + .map_err(|err| err.to_string())? }; Ok(RouteCandidateResolution { candidate, @@ -666,14 +691,39 @@ pub(crate) fn resolve_runtime_route_for_identity( .then(|| model_roster().preferred_model_id().map(str::to_string)) .flatten(); let model_selector = model_selector.or(roster_preferred.as_deref()); - let resolution = resolve_route_candidate_with_context_metadata( - provider, - model_selector, - saved_provider_model, - Some(route_config.deepseek_base_url()), - route_config.context_window_for_provider_config(provider), - None, - )?; + let base_url = route_config.deepseek_base_url(); + let uses_provider_owned_catalog = provider == ApiProvider::Openrouter + || (provider == ApiProvider::Custom + && codewhale_config::provider_setup_template(&identity.key) + .is_some_and(|template| template.id == codewhale_config::BASETEN_TEMPLATE_ID)); + let resolution = if uses_provider_owned_catalog { + let status = crate::provider_catalog_live::status_for_scope(&identity.key, &base_url); + let catalog = crate::provider_lake::runtime_catalog_resolver_for_identity( + provider, + Some(&identity.key), + &base_url, + status, + ); + resolve_route_candidate_with_catalog_resolver( + provider, + model_selector, + saved_provider_model, + Some(base_url), + route_config.context_window_for_provider_config(provider), + None, + &catalog.resolver, + catalog.endpoint_catalog_authoritative, + )? + } else { + resolve_route_candidate_with_context_metadata( + provider, + model_selector, + saved_provider_model, + Some(base_url), + route_config.context_window_for_provider_config(provider), + None, + )? + }; let candidate = resolution.candidate; let model = candidate.wire_model_id().as_str().to_string(); set_model_for_route(&mut route_config, provider, &model); @@ -1518,6 +1568,231 @@ mod tests { assert_eq!(custom.wire_model_id().as_str(), "deepseek-v4-pro"); } + fn live_catalog_offering( + provider: &str, + model: &str, + base_url: &str, + ) -> codewhale_config::catalog::CatalogOffering { + codewhale_config::catalog::CatalogOffering { + provider: provider.to_string(), + wire_model_id: model.to_string(), + endpoint_key: "chat".to_string(), + default_for_provider: true, + limit: Some(codewhale_config::models_dev::ModelsDevLimit { + context: Some(654_321), + input: Some(600_000), + output: Some(54_321), + }), + cost: Some(codewhale_config::models_dev::ModelsDevCost { + input: Some(1.25), + output: Some(3.5), + cache_read: None, + cache_write: None, + }), + modalities: Some(codewhale_config::models_dev::ModelsDevModalities { + input: vec!["text".to_string(), "image".to_string()], + output: vec!["text".to_string()], + }), + attachment: Some(true), + reasoning: Some(true), + tool_call: Some(true), + structured_output: Some(true), + source: codewhale_config::catalog::CatalogSource::Live { + base_url_fingerprint: codewhale_config::catalog::base_url_fingerprint(base_url), + fetched_at: codewhale_config::catalog::now_unix(), + }, + ..Default::default() + } + } + + fn assert_live_catalog_route_facts(route: &ResolvedRuntimeRoute) { + use codewhale_config::route::{CapabilityState, PricingSku}; + + assert_eq!(route.candidate.limits().context_tokens, Some(654_321)); + assert_eq!(route.candidate.limits().input_tokens, Some(600_000)); + assert_eq!(route.candidate.limits().output_tokens, Some(54_321)); + assert_eq!(route.context_window.tokens, 654_321); + assert_eq!(route.context_window.source, ContextWindowSource::Catalog); + let capabilities = route.candidate.capabilities(); + assert_eq!(capabilities.attachments, CapabilityState::Supported); + assert_eq!(capabilities.image_input, CapabilityState::Supported); + assert_eq!(capabilities.reasoning, CapabilityState::Supported); + assert_eq!(capabilities.native_tool_calls, CapabilityState::Supported); + assert_eq!(capabilities.structured_output, CapabilityState::Supported); + match route.candidate.pricing() { + Some(PricingSku::Token { + input_per_mtok, + output_per_mtok, + }) => { + assert_eq!(*input_per_mtok, Some(1.25)); + assert_eq!(*output_per_mtok, Some(3.5)); + } + other => panic!("expected provider-live token pricing, got {other:?}"), + } + } + + #[test] + fn live_only_openrouter_model_facts_reach_runtime_and_fail_closed_on_refresh_error() { + use codewhale_config::catalog::{CatalogRefreshError, ProviderCatalogDelta}; + use codewhale_config::route::{CapabilityState, PricingSku}; + + let _env = crate::test_support::lock_test_env(); + let _live = crate::provider_lake::lock_live_snapshot(); + let home = tempfile::tempdir().expect("home"); + let _home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", home.path()); + crate::provider_catalog_live::reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + + let base_url = "https://synthetic.openrouter.invalid/api/v1"; + let model = "synthetic/live-only-openrouter-model"; + let config = Config { + provider: Some("openrouter".to_string()), + providers: Some(ProvidersConfig { + openrouter: ProviderConfig { + base_url: Some(base_url.to_string()), + model: Some(model.to_string()), + ..Default::default() + }, + ..Default::default() + }), + ..Default::default() + }; + let fingerprint = codewhale_config::catalog::base_url_fingerprint(base_url); + crate::provider_catalog_live::record_success(ProviderCatalogDelta { + provider: "openrouter".to_string(), + base_url_fingerprint: fingerprint.clone(), + fetched_at: codewhale_config::catalog::now_unix(), + offerings: vec![live_catalog_offering("openrouter", model, base_url)], + }); + + let route = resolve_runtime_route(&config, ApiProvider::Openrouter, Some(model)) + .expect("live-only OpenRouter route resolves"); + assert_eq!(route.model, model); + assert_live_catalog_route_facts(&route); + + crate::provider_catalog_live::record_failure( + "openrouter", + &fingerprint, + CatalogRefreshError::Network, + ); + let failed = resolve_runtime_route(&config, ApiProvider::Openrouter, Some(model)) + .expect("wire id remains routable after a failed refresh"); + assert!(!failed.candidate.limits().has_known_limit()); + assert_eq!( + failed.candidate.capabilities().image_input, + CapabilityState::Unknown + ); + assert!(matches!( + failed.candidate.pricing(), + Some(PricingSku::UnknownOrStale) + )); + + crate::provider_catalog_live::reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + } + + #[test] + fn named_baseten_live_facts_reach_exact_custom_runtime_without_leaking() { + use codewhale_config::catalog::ProviderCatalogDelta; + + let _env = crate::test_support::lock_test_env(); + let _live = crate::provider_lake::lock_live_snapshot(); + let home = tempfile::tempdir().expect("home"); + let _home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", home.path()); + crate::provider_catalog_live::reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + + let base_url = codewhale_config::BASETEN_BASE_URL; + let model = "synthetic-live-baseten-model"; + let mut custom = std::collections::HashMap::new(); + custom.insert( + codewhale_config::BASETEN_TEMPLATE_ID.to_string(), + ProviderConfig { + kind: Some("openai-compatible".to_string()), + base_url: Some(base_url.to_string()), + model: Some(model.to_string()), + ..Default::default() + }, + ); + let config = Config { + provider: Some(codewhale_config::BASETEN_TEMPLATE_ID.to_string()), + providers: Some(ProvidersConfig { + custom, + ..Default::default() + }), + ..Default::default() + }; + crate::provider_catalog_live::record_success(ProviderCatalogDelta { + provider: codewhale_config::BASETEN_TEMPLATE_ID.to_string(), + base_url_fingerprint: codewhale_config::catalog::base_url_fingerprint(base_url), + fetched_at: codewhale_config::catalog::now_unix(), + offerings: vec![live_catalog_offering( + codewhale_config::BASETEN_TEMPLATE_ID, + model, + base_url, + )], + }); + + let route = resolve_runtime_route(&config, ApiProvider::Custom, Some(model)) + .expect("named Baseten route resolves"); + assert_eq!(route.identity.key, codewhale_config::BASETEN_TEMPLATE_ID); + assert_eq!(route.model, model); + assert_live_catalog_route_facts(&route); + + let alias_identity = "base-ten"; + let alias_model = "synthetic-alias-baseten-model"; + let mut alias_custom = std::collections::HashMap::new(); + alias_custom.insert( + alias_identity.to_string(), + ProviderConfig { + kind: Some("openai-compatible".to_string()), + base_url: Some(base_url.to_string()), + model: Some(alias_model.to_string()), + ..Default::default() + }, + ); + let alias_config = Config { + provider: Some(alias_identity.to_string()), + providers: Some(ProvidersConfig { + custom: alias_custom, + ..Default::default() + }), + ..Default::default() + }; + crate::provider_catalog_live::record_success(ProviderCatalogDelta { + provider: alias_identity.to_string(), + base_url_fingerprint: codewhale_config::catalog::base_url_fingerprint(base_url), + fetched_at: codewhale_config::catalog::now_unix(), + offerings: vec![live_catalog_offering(alias_identity, alias_model, base_url)], + }); + let alias_route = + resolve_runtime_route(&alias_config, ApiProvider::Custom, Some(alias_model)) + .expect("Baseten schema alias route resolves"); + assert_eq!(alias_route.identity.key, alias_identity); + assert_live_catalog_route_facts(&alias_route); + assert!( + crate::provider_lake::catalog_offering_for_model_identity( + ApiProvider::Custom, + Some(codewhale_config::BASETEN_TEMPLATE_ID), + alias_model, + ) + .is_none(), + "a Baseten schema alias must not share another exact table's live roster" + ); + + let unrelated = custom_config("https://other-compatible.invalid/v1", model); + let unrelated_route = resolve_runtime_route(&unrelated, ApiProvider::Custom, Some(model)) + .expect("another compatible provider remains routable"); + assert!(!unrelated_route.candidate.limits().has_known_limit()); + assert_eq!( + unrelated_route.candidate.capabilities(), + codewhale_config::route::RouteCapabilities::default() + ); + + crate::provider_catalog_live::reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + } + fn custom_config(base_url: &str, model: &str) -> Config { let mut custom = std::collections::HashMap::new(); custom.insert( diff --git a/crates/tui/src/runtime_api.rs b/crates/tui/src/runtime_api.rs index 23b9425738..1aef358646 100644 --- a/crates/tui/src/runtime_api.rs +++ b/crates/tui/src/runtime_api.rs @@ -19,6 +19,8 @@ use axum::response::sse::{Event as SseEvent, KeepAlive, Sse}; use axum::response::{IntoResponse, Response}; use axum::routing::{get, post}; use axum::{Json, Router}; +use base64::Engine as _; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; use chrono::Utc; use codewhale_protocol::agent_mail::{ AgentMailDeliveryMode, AgentMailEnvelope, AgentMailMessageId, AgentMailSendRequest, @@ -5463,6 +5465,137 @@ struct ProviderModelEntry { struct ProviderModelsResponse { provider: String, models: Vec, + total: usize, + #[serde(rename = "nextCursor", skip_serializing_if = "Option::is_none")] + next_cursor: Option, +} + +const DEFAULT_PROVIDER_MODELS_PAGE_SIZE: usize = 100; +const MAX_PROVIDER_MODELS_PAGE_SIZE: usize = 250; +const MAX_PROVIDER_MODELS_CATALOG_SIZE: usize = 10_000; +const PROVIDER_MODELS_CURSOR_VERSION: u8 = 1; +const MAX_PROVIDER_MODELS_CURSOR_BYTES: usize = 1_024; +const MAX_PROVIDER_MODELS_FILTER_CHARS: usize = 128; + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct ProviderModelsCursor { + version: u8, + provider: String, + filter: String, + after: String, +} + +fn normalized_provider_model_filter(filter: Option<&str>) -> Result { + let filter = filter.unwrap_or_default().trim(); + if filter.chars().count() > MAX_PROVIDER_MODELS_FILTER_CHARS { + return Err(ApiError::bad_request(format!( + "Provider model filter exceeds {MAX_PROVIDER_MODELS_FILTER_CHARS} characters" + ))); + } + Ok(filter.to_lowercase()) +} + +fn encode_provider_models_cursor(cursor: &ProviderModelsCursor) -> Result { + let bytes = serde_json::to_vec(cursor) + .map_err(|error| ApiError::internal(format!("Could not encode model cursor: {error}")))?; + Ok(URL_SAFE_NO_PAD.encode(bytes)) +} + +fn decode_provider_models_cursor(value: &str) -> Result { + if value.is_empty() || value.len() > MAX_PROVIDER_MODELS_CURSOR_BYTES { + return Err(ApiError::bad_request("Invalid provider model cursor")); + } + let bytes = URL_SAFE_NO_PAD + .decode(value) + .map_err(|_| ApiError::bad_request("Invalid provider model cursor"))?; + if bytes.len() > MAX_PROVIDER_MODELS_CURSOR_BYTES { + return Err(ApiError::bad_request("Invalid provider model cursor")); + } + let cursor: ProviderModelsCursor = serde_json::from_slice(&bytes) + .map_err(|_| ApiError::bad_request("Invalid provider model cursor"))?; + if cursor.version != PROVIDER_MODELS_CURSOR_VERSION + || cursor.provider.is_empty() + || cursor.after.is_empty() + || cursor.after.len() > 256 + { + return Err(ApiError::bad_request("Invalid provider model cursor")); + } + Ok(cursor) +} + +fn paginate_provider_models( + provider: &str, + mut models: Vec, + params: &ListProviderModelsParams, +) -> Result { + let filter = normalized_provider_model_filter(params.filter.as_deref())?; + let limit = params.limit.unwrap_or(DEFAULT_PROVIDER_MODELS_PAGE_SIZE); + if limit == 0 || limit > MAX_PROVIDER_MODELS_PAGE_SIZE { + return Err(ApiError::bad_request(format!( + "Provider model page limit must be between 1 and {MAX_PROVIDER_MODELS_PAGE_SIZE}" + ))); + } + + models.sort_by(|left, right| { + left.id + .to_lowercase() + .cmp(&right.id.to_lowercase()) + .then_with(|| left.id.cmp(&right.id)) + }); + models.dedup_by(|left, right| left.id.eq_ignore_ascii_case(&right.id)); + if models.len() > MAX_PROVIDER_MODELS_CATALOG_SIZE { + return Err(ApiError::internal(format!( + "Provider model catalog exceeds the safe {MAX_PROVIDER_MODELS_CATALOG_SIZE}-row limit" + ))); + } + if !filter.is_empty() { + models.retain(|entry| entry.id.to_lowercase().contains(&filter)); + } + + let start = if let Some(encoded) = params.cursor.as_deref() { + let cursor = decode_provider_models_cursor(encoded)?; + if cursor.provider != provider || cursor.filter != filter { + return Err(ApiError::bad_request( + "Provider model cursor does not match this provider and filter", + )); + } + models + .iter() + .position(|entry| entry.id == cursor.after) + .map(|index| index.saturating_add(1)) + .ok_or_else(|| { + ApiError::bad_request("Provider model cursor is stale; restart from the first page") + })? + } else { + 0 + }; + let total = models.len(); + let end = start.saturating_add(limit).min(total); + let page = models + .get(start..end) + .ok_or_else(|| ApiError::bad_request("Provider model cursor is outside the catalog"))? + .to_vec(); + let next_cursor = if end < total { + let after = page + .last() + .map(|entry| entry.id.clone()) + .ok_or_else(|| ApiError::internal("Provider model page made no progress"))?; + Some(encode_provider_models_cursor(&ProviderModelsCursor { + version: PROVIDER_MODELS_CURSOR_VERSION, + provider: provider.to_string(), + filter, + after, + })?) + } else { + None + }; + + Ok(ProviderModelsResponse { + provider: provider.to_string(), + models: page, + total, + next_cursor, + }) } fn push_unique_model(models: &mut Vec, model: &str) { @@ -5585,6 +5718,21 @@ pub(crate) fn runtime_chat_route_id_is_safe(value: &str) -> bool { .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'-')) } +fn runtime_chat_safe_models(mut models: Vec) -> Result, String> { + models.retain(|model| runtime_chat_model_id_is_safe(model)); + models.sort(); + models.dedup(); + if models.len() > MAX_PROVIDER_MODELS_CATALOG_SIZE { + return Err(format!( + "The active Runtime provider catalog exceeds the safe {MAX_PROVIDER_MODELS_CATALOG_SIZE}-model relay limit." + )); + } + if models.is_empty() { + return Err("The active Runtime provider has no safe model catalog.".to_string()); + } + Ok(models) +} + /// Build the deliberately narrow provider projection used by the account-owned /// Runtime Chat relay. This is the same active-route truth exposed by the /// authenticated native `/v1/runtime/info`, `/v1/providers`, and @@ -5622,14 +5770,7 @@ pub(crate) fn runtime_chat_relay_catalog( } }; - let mut models = provider_models_for_api(config, provider, provider); - models.retain(|model| runtime_chat_model_id_is_safe(model)); - models.sort(); - models.dedup(); - models.truncate(256); - if models.is_empty() { - return Err("The active Runtime provider has no safe model catalog.".to_string()); - } + let models = runtime_chat_safe_models(provider_models_for_api(config, provider, provider))?; let requested_default = provider_default_model_for_api(config, provider, provider); let default_model = models .iter() @@ -5712,18 +5853,21 @@ async fn list_providers( #[derive(Debug, Deserialize)] struct ListProviderModelsParams { - /// Optional filter: when provided, models whose id contains this - /// substring (case-insensitive) are returned. Currently informational — - /// the catalog is small enough to filter client-side. + /// Optional case-insensitive substring filter applied before pagination. #[serde(default)] - #[allow(dead_code)] filter: Option, + /// Opaque continuation cursor returned as `nextCursor` by the prior page. + #[serde(default)] + cursor: Option, + /// Page size. The bounded default is 100 and the maximum is 250. + #[serde(default)] + limit: Option, } async fn list_provider_models( State(state): State, Path(id): Path, - _params: Query, + Query(params): Query, ) -> Result, ApiError> { let config = state.config.read().clone(); let active_provider = config.api_provider(); @@ -5743,10 +5887,7 @@ async fn list_provider_models( id, }) .collect(); - Ok(Json(ProviderModelsResponse { - provider: api_provider.as_str().to_string(), - models, - })) + paginate_provider_models(api_provider.as_str(), models, ¶ms).map(Json) } /// Request body for `POST /v1/providers/{id}/switch`. diff --git a/crates/tui/src/runtime_api/tests.rs b/crates/tui/src/runtime_api/tests.rs index 610c586a4a..c254bf3fea 100644 --- a/crates/tui/src/runtime_api/tests.rs +++ b/crates/tui/src/runtime_api/tests.rs @@ -441,15 +441,18 @@ fn messages_from_thread_detail_batches_tool_results() { ended_at: Some(now), duration_ms: Some(0), usage: None, + effective_route_usage: None, permission_posture: Some("ask".to_string()), effective_provider: None, effective_provider_id: None, effective_billing_surface: None, effective_endpoint_fingerprint: None, + effective_provider_live_pricing: None, effective_billing_mode: None, effective_dispatched_at: None, effective_model: None, routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), routed_usage_source_ids: Vec::new(), routed_usage_dropped_records: 0, error: None, @@ -919,6 +922,155 @@ async fn spawn_test_server_with_root_token_mobile_workspace_and_overrides( SharedRuntimeThreadManager, tokio::task::JoinHandle<()>, )>, +> { + let (setup_tx, setup_rx) = oneshot::channel(); + // If this test sealed the process environment (`lock_test_env`), the + // server thread must join that scope: its `Config::load` reads env through + // `with_test_env_lock`, which would otherwise block on the mutex the test + // holds while awaiting this very setup — a deadlock. + let env_ticket = crate::test_support::env_scope_ticket(); + spawn_product_stack_server( + root, + sessions_dir, + runtime_token, + mobile_enabled, + workspace, + overrides, + env_ticket, + setup_tx, + ); + let Some((addr, runtime_threads, shutdown_tx)) = setup_rx + .await + .context("runtime-api test server thread ended during setup")?? + else { + return Ok(None); + }; + // Owns the shutdown side for as long as the test keeps the handle alive: + // aborting it (or the test runtime dropping it) closes the listener and + // ends the server thread. + let handle = tokio::spawn(async move { + let _shutdown = shutdown_tx; + std::future::pending::<()>().await + }); + Ok(Some((addr, runtime_threads, handle))) +} + +/// What the server thread hands back once the router is bound: everything the +/// test body needs, plus the shutdown side of the server. +type TestServerSetup = (SocketAddr, SharedRuntimeThreadManager, oneshot::Sender<()>); + +/// Builds and serves the Runtime API router where the product builds and +/// serves it: on a thread with the product's `CODEWHALE_MAIN_STACK_BYTES` +/// stack. +/// +/// `#[tokio::test]` drives its current-thread runtime on the 2 MiB libtest +/// thread, so a harness that built its state and spawned its server onto that +/// runtime ran every product path — config load and reload (the serde +/// `toml::de::visit_map` frames for the full `Config`), manager construction, +/// thread lifecycle, streaming — on a stack the product never gives it +/// (`lib.rs` sizes the runtime workers with `CODEWHALE_MAIN_STACK_BYTES`). +/// Config load under a profile and the thread-lifecycle path marginally +/// overflowed 2 MiB in debug builds (`has overflowed its stack`, SIGABRT for +/// the whole lib suite), which CI masked with `RUST_MIN_STACK`. Running setup +/// *and* serving on one product-sized thread removes the class: the libtest +/// thread keeps only the test body and its HTTP client, and no product frame +/// depth can overflow it. +/// +/// Setup results come back through `setup_tx`; the caller wraps the shutdown +/// sender in the `JoinHandle` the call sites expect. Nothing here runs on the +/// test's runtime, so the server thread must outlive setup: it serves until +/// the shutdown sender is dropped. +/// +/// `env_ticket` adopts the thread into the calling test's sealed env scope +/// (`test_env_lock::join_env_scope`), so its `Config::load` env reads see the +/// test's environment instead of blocking on the mutex the test holds while +/// awaiting setup. `None` when the caller sealed nothing. +fn spawn_product_stack_server( + root: PathBuf, + sessions_dir: PathBuf, + runtime_token: Option, + mobile_enabled: bool, + workspace: PathBuf, + overrides: TestServerOverrides, + env_ticket: Option, + setup_tx: oneshot::Sender>>, +) { + std::thread::Builder::new() + .name("runtime-api-test-server".to_string()) + .stack_size(crate::CODEWHALE_MAIN_STACK_BYTES) + .spawn(move || { + // Adopted for the thread's lifetime; the scope's generation check + // refuses enrollment once the sealing test has ended. + let _membership = crate::test_support::join_env_scope(env_ticket); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("runtime-api test server runtime"); + runtime.block_on(async move { + match build_test_server( + root, + sessions_dir, + runtime_token, + mobile_enabled, + workspace, + overrides, + ) + .await + { + Ok(Some((listener, app, addr, runtime_threads))) => { + let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); + if setup_tx + .send(Ok(Some((addr, runtime_threads, shutdown_tx)))) + .is_err() + { + // The test gave up waiting; do not serve. + return; + } + listener + .set_nonblocking(true) + .expect("nonblocking test listener"); + let listener = + TcpListener::from_std(listener).expect("register test listener"); + tokio::select! { + _ = async { + let _ = axum::serve( + listener, + app.into_make_service_with_connect_info::(), + ) + .await; + } => {} + _ = shutdown_rx => {} + } + } + Ok(None) => { + let _ = setup_tx.send(Ok(None)); + } + Err(err) => { + let _ = setup_tx.send(Err(err)); + } + } + }); + }) + .expect("spawn runtime-api test server thread"); +} + +/// The whole harness body — config load, managers, router build — formerly +/// inline in `spawn_test_server_with_root_token_mobile_workspace_and_overrides`. +/// Runs on the product-stack server thread (see `spawn_product_stack_server`). +async fn build_test_server( + root: PathBuf, + sessions_dir: PathBuf, + runtime_token: Option, + mobile_enabled: bool, + workspace: PathBuf, + overrides: TestServerOverrides, +) -> Result< + Option<( + std::net::TcpListener, + axum::Router, + SocketAddr, + SharedRuntimeThreadManager, + )>, > { let _ = rustls::crypto::ring::default_provider().install_default(); fs::create_dir_all(&sessions_dir)?; @@ -934,8 +1086,6 @@ async fn spawn_test_server_with_root_token_mobile_workspace_and_overrides( ..Config::default() } }; - config.mcp_config_path = Some(root.join("mcp.json").to_string_lossy().to_string()); - config.mcp_config_path = Some(root.join("mcp.json").to_string_lossy().to_string()); let manager = TaskManager::start_with_executor( TaskManagerConfig { @@ -966,7 +1116,9 @@ async fn spawn_test_server_with_root_token_mobile_workspace_and_overrides( let sub_agent_manager = overrides .sub_agent_manager .unwrap_or_else(|| runtime_api_sub_agent_manager(&workspace, 2)); - let listener = match TcpListener::bind("127.0.0.1:0").await { + // A std listener: the server thread registers it with its own runtime + // after setup is reported (see `spawn_product_stack_server`). + let listener = match std::net::TcpListener::bind("127.0.0.1:0") { Ok(listener) => listener, Err(err) if err.kind() == std::io::ErrorKind::PermissionDenied => return Ok(None), Err(err) => return Err(err.into()), @@ -1002,14 +1154,7 @@ async fn spawn_test_server_with_root_token_mobile_workspace_and_overrides( compat_stream_test_hook: overrides.compat_stream_test_hook, }; let app = build_router(state); - let handle = tokio::spawn(async move { - let _ = axum::serve( - listener, - app.into_make_service_with_connect_info::(), - ) - .await; - }); - Ok(Some((addr, runtime_threads, handle))) + Ok(Some((listener, app, addr, runtime_threads))) } async fn spawn_test_server() -> Result< @@ -2514,6 +2659,12 @@ async fn compatibility_stream_closes_losslessly_across_replay_live_handoff() -> output_tokens: 1, ..Usage::default() }, + parent_route_usage: Usage { + input_tokens: 3, + output_tokens: 1, + ..Usage::default() + }, + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -2762,6 +2913,8 @@ async fn compatibility_stream_exposes_and_resolves_user_input_without_answer_ech .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -3125,6 +3278,12 @@ async fn thread_endpoints_expose_lifecycle_contract() -> Result<()> { output_tokens: 5, ..Usage::default() }, + parent_route_usage: Usage { + input_tokens: 10, + output_tokens: 5, + ..Usage::default() + }, + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -3140,6 +3299,8 @@ async fn thread_endpoints_expose_lifecycle_contract() -> Result<()> { output_tokens: 0, ..Usage::default() }, + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -3277,6 +3438,8 @@ async fn turn_endpoint_operation_key_returns_original_and_conflicts_on_mismatch( let _ = tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -3385,6 +3548,12 @@ async fn events_endpoint_respects_since_seq_cursor() -> Result<()> { output_tokens: 3, ..Usage::default() }, + parent_route_usage: Usage { + input_tokens: 5, + output_tokens: 3, + ..Usage::default() + }, + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -3604,6 +3773,12 @@ async fn steer_and_interrupt_endpoints_work_on_active_turn() -> Result<()> { output_tokens: 1, ..Usage::default() }, + parent_route_usage: Usage { + input_tokens: 2, + output_tokens: 1, + ..Usage::default() + }, + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -3961,6 +4136,12 @@ async fn stream_endpoint_remains_backward_compatible() -> Result<()> { output_tokens: 2, ..Usage::default() }, + parent_route_usage: Usage { + input_tokens: 4, + output_tokens: 2, + ..Usage::default() + }, + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -4650,6 +4831,12 @@ async fn session_create_from_thread_rejects_active_turn() -> Result<()> { output_tokens: 1, ..Usage::default() }, + parent_route_usage: Usage { + input_tokens: 2, + output_tokens: 1, + ..Usage::default() + }, + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -4793,7 +4980,10 @@ async fn session_summary_route_projects_rows_and_honours_archive_filters() -> Re let client = crate::tls::reqwest_client(); let active: Vec = client - .get(format!("http://{addr}/v1/sessions/summary")) + .get(format!( + "http://127.0.0.1:{}/v1/sessions/summary", + addr.port() + )) .send() .await? .error_for_status()? @@ -4823,7 +5013,8 @@ async fn session_summary_route_projects_rows_and_honours_archive_filters() -> Re let archived: Vec = client .get(format!( - "http://{addr}/v1/sessions/summary?archived_only=true" + "http://127.0.0.1:{}/v1/sessions/summary?archived_only=true", + addr.port() )) .send() .await? @@ -4854,7 +5045,10 @@ async fn session_patch_route_renames_archives_and_reports_real_changes() -> Resu let client = crate::tls::reqwest_client(); let patched: serde_json::Value = client - .patch(format!("http://{addr}/v1/sessions/sess-patch")) + .patch(format!( + "http://127.0.0.1:{}/v1/sessions/sess-patch", + addr.port() + )) .json(&json!({ "title": "After", "archived": true })) .send() .await? @@ -4874,7 +5068,10 @@ async fn session_patch_route_renames_archives_and_reports_real_changes() -> Resu // A re-patch to the same state changes nothing, and says so. let repeat: serde_json::Value = client - .patch(format!("http://{addr}/v1/sessions/sess-patch")) + .patch(format!( + "http://127.0.0.1:{}/v1/sessions/sess-patch", + addr.port() + )) .json(&json!({ "archived": true })) .send() .await? @@ -4891,7 +5088,10 @@ async fn session_patch_route_renames_archives_and_reports_real_changes() -> Resu // An empty body is a client error, not a silent no-op. let empty = client - .patch(format!("http://{addr}/v1/sessions/sess-patch")) + .patch(format!( + "http://127.0.0.1:{}/v1/sessions/sess-patch", + addr.port() + )) .json(&json!({})) .send() .await?; @@ -4899,7 +5099,10 @@ async fn session_patch_route_renames_archives_and_reports_real_changes() -> Resu // A blank title is rejected with the reason, not accepted. let blank = client - .patch(format!("http://{addr}/v1/sessions/sess-patch")) + .patch(format!( + "http://127.0.0.1:{}/v1/sessions/sess-patch", + addr.port() + )) .json(&json!({ "title": " " })) .send() .await?; @@ -4926,7 +5129,10 @@ async fn session_patch_route_refuses_a_live_session_with_a_conflict() -> Result< crate::session_manager::set_live_session(Some("sess-live")); let conflict = client - .patch(format!("http://{addr}/v1/sessions/sess-live")) + .patch(format!( + "http://127.0.0.1:{}/v1/sessions/sess-live", + addr.port() + )) .json(&json!({ "title": "Renamed from the dashboard" })) .send() .await?; @@ -4934,7 +5140,10 @@ async fn session_patch_route_refuses_a_live_session_with_a_conflict() -> Result< crate::session_manager::set_live_session(None); let allowed = client - .patch(format!("http://{addr}/v1/sessions/sess-live")) + .patch(format!( + "http://127.0.0.1:{}/v1/sessions/sess-live", + addr.port() + )) .json(&json!({ "title": "Renamed from the dashboard" })) .send() .await?; @@ -4957,7 +5166,8 @@ async fn session_detail_route_serves_a_bounded_redacted_peek_on_request() -> Res let peek: serde_json::Value = client .get(format!( - "http://{addr}/v1/sessions/sess-peek?peek=true&entries=12" + "http://127.0.0.1:{}/v1/sessions/sess-peek?peek=true&entries=12", + addr.port() )) .send() .await? @@ -4979,7 +5189,10 @@ async fn session_detail_route_serves_a_bounded_redacted_peek_on_request() -> Res } let detail: serde_json::Value = client - .get(format!("http://{addr}/v1/sessions/sess-peek")) + .get(format!( + "http://127.0.0.1:{}/v1/sessions/sess-peek", + addr.port() + )) .send() .await? .error_for_status()? @@ -5463,11 +5676,13 @@ async fn session_save_merges_thread_cost_split_and_records_coverage() -> Result< output_tokens: 1_000, ..Usage::default() }), + effective_route_usage: None, permission_posture: None, effective_provider: None, effective_provider_id: None, effective_billing_surface: None, effective_endpoint_fingerprint: None, + effective_provider_live_pricing: None, effective_billing_mode: None, effective_dispatched_at: None, effective_model: None, @@ -5478,6 +5693,7 @@ async fn session_save_merges_thread_cost_split_and_records_coverage() -> Result< model: "deepseek-v4-flash".to_string(), billing_surface: Some(crate::pricing::FIRST_PARTY_PAYG_BILLING_SURFACE.to_string()), endpoint_fingerprint: None, + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Metered, dispatched_at: now, }, @@ -5487,6 +5703,7 @@ async fn session_save_merges_thread_cost_split_and_records_coverage() -> Result< ..Usage::default() }, }], + routed_usage_drop_records: Vec::new(), routed_usage_source_ids: Vec::new(), routed_usage_dropped_records: 0, error: None, @@ -5648,11 +5865,13 @@ async fn session_save_persists_parent_cny_unpriced_reasons_without_double_count( output_tokens: 1_000, ..Usage::default() }), + effective_route_usage: None, permission_posture: None, effective_provider: None, effective_provider_id: None, effective_billing_surface: None, effective_endpoint_fingerprint: None, + effective_provider_live_pricing: None, effective_billing_mode: None, effective_dispatched_at: None, effective_model: None, @@ -5663,6 +5882,7 @@ async fn session_save_persists_parent_cny_unpriced_reasons_without_double_count( model: "deepseek-v4-flash".to_string(), billing_surface: Some(crate::pricing::FIRST_PARTY_PAYG_BILLING_SURFACE.to_string()), endpoint_fingerprint: None, + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Metered, dispatched_at: now, }, @@ -5672,6 +5892,7 @@ async fn session_save_persists_parent_cny_unpriced_reasons_without_double_count( ..Usage::default() }, }], + routed_usage_drop_records: Vec::new(), routed_usage_source_ids: Vec::new(), routed_usage_dropped_records: 0, error: None, @@ -6416,15 +6637,18 @@ fn seed_summary_search_transcript( ended_at: Some(created_at), duration_ms: Some(0), usage: None, + effective_route_usage: None, permission_posture: None, effective_provider: None, effective_provider_id: None, effective_billing_surface: None, effective_endpoint_fingerprint: None, + effective_provider_live_pricing: None, effective_billing_mode: None, effective_dispatched_at: None, effective_model: None, routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), routed_usage_source_ids: Vec::new(), routed_usage_dropped_records: 0, error: None, @@ -7290,6 +7514,120 @@ async fn get_provider_models( .expect("GET /v1/providers/{id}/models should return valid JSON") } +#[test] +fn provider_model_catalog_paginates_all_six_hundred_rows_without_truncation() { + let models = (0..600) + .rev() + .map(|index| ProviderModelEntry { + id: format!("openrouter/model-{index:03}"), + image_input: codewhale_config::route::CapabilityState::Unknown, + }) + .collect::>(); + let mut cursor = None; + let mut observed = Vec::new(); + let mut page_count = 0usize; + + loop { + let response = paginate_provider_models( + "openrouter", + models.clone(), + &ListProviderModelsParams { + filter: None, + cursor, + limit: Some(MAX_PROVIDER_MODELS_PAGE_SIZE), + }, + ) + .expect("page should be valid"); + page_count += 1; + assert_eq!(response.provider, "openrouter"); + assert_eq!(response.total, 600); + assert!(!response.models.is_empty()); + assert!(response.models.len() <= MAX_PROVIDER_MODELS_PAGE_SIZE); + observed.extend(response.models.into_iter().map(|entry| entry.id)); + cursor = response.next_cursor; + if cursor.is_none() { + break; + } + } + + assert_eq!(page_count, 3); + assert_eq!(observed.len(), 600); + assert_eq!( + observed.first().map(String::as_str), + Some("openrouter/model-000") + ); + assert_eq!( + observed.last().map(String::as_str), + Some("openrouter/model-599") + ); + let unique = observed.iter().collect::>(); + assert_eq!(unique.len(), 600); +} + +#[test] +fn provider_model_catalog_applies_filter_before_cursor_and_rejects_cross_scope_replay() { + let models = ["Alpha-One", "alpha-two", "beta"] + .into_iter() + .map(|id| ProviderModelEntry { + id: id.to_string(), + image_input: codewhale_config::route::CapabilityState::Unknown, + }) + .collect::>(); + let first = paginate_provider_models( + "openrouter", + models.clone(), + &ListProviderModelsParams { + filter: Some(" ALPHA ".to_string()), + cursor: None, + limit: Some(1), + }, + ) + .expect("filtered first page should be valid"); + assert_eq!(first.total, 2); + assert_eq!(first.models[0].id, "Alpha-One"); + let cursor = first.next_cursor.expect("a second filtered row remains"); + + let second = paginate_provider_models( + "openrouter", + models.clone(), + &ListProviderModelsParams { + filter: Some("alpha".to_string()), + cursor: Some(cursor.clone()), + limit: Some(1), + }, + ) + .expect("matching filter should continue"); + assert_eq!(second.models[0].id, "alpha-two"); + assert!(second.next_cursor.is_none()); + + let replay = paginate_provider_models( + "baseten", + models, + &ListProviderModelsParams { + filter: Some("alpha".to_string()), + cursor: Some(cursor), + limit: Some(1), + }, + ); + assert!(replay.is_err(), "a cursor cannot cross provider ownership"); +} + +#[test] +fn runtime_chat_relay_projection_keeps_six_hundred_safe_models_without_silent_cutoff() { + let models = (0..600) + .map(|index| format!("openrouter/model-{index:03}")) + .collect::>(); + let projected = runtime_chat_safe_models(models.clone()).expect("600 rows are safely bounded"); + assert_eq!(projected, models); + + let oversized = (0..=MAX_PROVIDER_MODELS_CATALOG_SIZE) + .map(|index| format!("provider/model-{index:05}")) + .collect::>(); + let error = runtime_chat_safe_models(oversized).expect_err("oversized relay must fail loudly"); + assert!(error.contains("safe")); + assert!(error.contains(&MAX_PROVIDER_MODELS_CATALOG_SIZE.to_string())); +} + #[tokio::test] async fn get_config_returns_active_provider_model() -> Result<()> { let root = std::env::temp_dir().join(format!( diff --git a/crates/tui/src/runtime_threads.rs b/crates/tui/src/runtime_threads.rs index 0653175f30..97cc77487a 100644 --- a/crates/tui/src/runtime_threads.rs +++ b/crates/tui/src/runtime_threads.rs @@ -37,10 +37,12 @@ use crate::core::engine::{ use crate::core::events::{Event as EngineEvent, TurnOutcomeStatus}; use crate::core::ops::Op; use crate::cost_status::{ - EffectiveRouteEnvelope, EffectiveRouteUsage, RouteBillingMode, RuntimeUsageRecord, + EffectiveRouteEnvelope, EffectiveRouteUsage, RouteBillingMode, RuntimeUsageDropRecord, + RuntimeUsageRecord, }; use crate::models::Role; use crate::models::{ContentBlock, Message, SystemPrompt, Usage}; +use crate::provider_catalog_live::ProviderLivePricingQuote; use crate::route_budget::{ auto_compact_default_for_route, compaction_threshold_for_route_at_percent, known_route_limits, route_context_window_tokens, @@ -532,12 +534,32 @@ where { values .iter() - .map(|value| { - if value.len() == 64 && value.bytes().all(|byte| byte.is_ascii_hexdigit()) { - value.to_ascii_lowercase() - } else { - codewhale_config::catalog::base_url_fingerprint(value) - } + .map(|value| routed_usage_source_fingerprint(value)) + .collect::>() + .serialize(serializer) +} + +fn routed_usage_source_fingerprint(source_id: &str) -> String { + let source_id = source_id.trim(); + if source_id.len() == 64 && source_id.bytes().all(|byte| byte.is_ascii_hexdigit()) { + source_id.to_ascii_lowercase() + } else { + crate::cost_status::usage_source_fingerprint(source_id) + } +} + +fn serialize_routed_usage_drop_records( + values: &[RuntimeUsageDropRecord], + serializer: S, +) -> std::result::Result +where + S: serde::Serializer, +{ + values + .iter() + .map(|record| RuntimeUsageDropRecord { + source_id: routed_usage_source_fingerprint(&record.source_id), + route: record.route.sanitized_for_persistence(), }) .collect::>() .serialize(serializer) @@ -722,6 +744,11 @@ pub struct TurnRecord { pub duration_ms: Option, #[serde(skip_serializing_if = "Option::is_none")] pub usage: Option, + /// Portion of `usage` served by this turn's persisted effective route. + /// New records use this for parent cost/token accumulation, then add each + /// routed child independently. Legacy absence falls back to `usage`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub effective_route_usage: Option, /// Canonical posture that governed this turn. New records always carry /// this receipt; old records deserialize with no fabricated value. #[serde(default, skip_serializing_if = "Option::is_none")] @@ -756,6 +783,15 @@ pub struct TurnRecord { serialize_with = "serialize_endpoint_fingerprint_option" )] pub effective_endpoint_fingerprint: Option, + /// Immutable provider-live rate receipt frozen at CodeWhale's pre-permit + /// application-dispatch boundary. Missing on legacy records; Baseten then + /// stays unpriced while OpenRouter may use only its immutable bundled row. + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "crate::provider_catalog_live::deserialize_optional_provider_live_pricing" + )] + pub effective_provider_live_pricing: Option, /// Immutable billing classification captured before dispatch. #[serde(default, skip_serializing_if = "Option::is_none")] pub effective_billing_mode: Option, @@ -771,10 +807,20 @@ pub struct TurnRecord { )] pub effective_model: Option, /// Model calls made beneath this parent turn, each paired with its own - /// immutable route. These are exclusive of `usage`, which is only the - /// parent engine turn. + /// immutable route. These are exclusive of `effective_route_usage`; + /// `usage` remains the authoritative all-token turn total and may include + /// inline RLM/guardian calls. #[serde(default, skip_serializing_if = "Vec::is_empty")] pub routed_usage: Vec, + /// Exact missing-usage provider responses paired with their frozen route. + /// Sources are persisted only as fingerprints, routes are sanitized, and + /// the shared source ledger bounds usage plus missing-usage records. + #[serde( + default, + skip_serializing_if = "Vec::is_empty", + serialize_with = "serialize_routed_usage_drop_records" + )] + pub routed_usage_drop_records: Vec, /// Fingerprints of provider-call identities already appended to this turn. /// This durable ledger makes mailbox delivery, direct sinks, fallback /// recovery, and process restart idempotent without persisting raw ids. @@ -819,6 +865,7 @@ impl TurnRecord { self.effective_provider_id = Some(route.provider_identity); self.effective_billing_surface = route.billing_surface; self.effective_endpoint_fingerprint = route.endpoint_fingerprint; + self.effective_provider_live_pricing = route.provider_live_pricing; self.effective_billing_mode = Some(route.billing_mode); self.effective_dispatched_at = Some(route.dispatched_at); self.effective_model = Some(route.model); @@ -849,6 +896,7 @@ impl TurnRecord { model, billing_surface: self.effective_billing_surface.clone(), endpoint_fingerprint: self.effective_endpoint_fingerprint.clone(), + provider_live_pricing: self.effective_provider_live_pricing.clone(), billing_mode: self .effective_billing_mode .unwrap_or(RouteBillingMode::Unknown), @@ -868,7 +916,7 @@ fn append_routed_usage_record( source_id: &str, usage: EffectiveRouteUsage, ) -> bool { - let source_fingerprint = crate::cost_status::usage_source_fingerprint(source_id); + let source_fingerprint = routed_usage_source_fingerprint(source_id); if turn .routed_usage_source_ids .iter() @@ -876,11 +924,18 @@ fn append_routed_usage_record( { return false; } - turn.routed_usage_source_ids.push(source_fingerprint); - if turn.routed_usage.len() == MAX_ROUTED_USAGE_RECORDS_PER_TURN { - turn.routed_usage.remove(0); - turn.routed_usage_dropped_records = turn.routed_usage_dropped_records.saturating_add(1); + if turn.routed_usage_source_ids.len() == MAX_ROUTED_USAGE_RECORDS_PER_TURN { + // Stop admitting new records once both the record and exact-dedupe + // ledgers reach their shared bound. Retaining the first accepted set + // makes replay idempotent; one explicit incompleteness marker is safer + // than evicting fingerprints and later counting a replay twice. + if turn.routed_usage_dropped_records == 0 { + turn.routed_usage_dropped_records = 1; + return true; + } + return false; } + turn.routed_usage_source_ids.push(source_fingerprint); turn.routed_usage.push(EffectiveRouteUsage { route: usage.route.sanitized_for_persistence(), usage: usage.usage, @@ -888,6 +943,228 @@ fn append_routed_usage_record( true } +fn append_routed_usage_drop_record(turn: &mut TurnRecord, record: RuntimeUsageDropRecord) -> bool { + let source_fingerprint = routed_usage_source_fingerprint(&record.source_id); + if turn + .routed_usage_source_ids + .iter() + .any(|persisted| persisted == &source_fingerprint) + { + return false; + } + if turn.routed_usage_source_ids.len() == MAX_ROUTED_USAGE_RECORDS_PER_TURN { + if turn.routed_usage_dropped_records == 0 { + turn.routed_usage_dropped_records = 1; + return true; + } + return false; + } + turn.routed_usage_source_ids + .push(source_fingerprint.clone()); + turn.routed_usage_drop_records.push(RuntimeUsageDropRecord { + source_id: source_fingerprint, + route: record.route.sanitized_for_persistence(), + }); + true +} + +/// Bind pre-parent auxiliary calls to a reserved Runtime turn before the +/// engine accepts the parent operation. Exact records are persisted now so a +/// crash cannot erase them; the engine's terminal event remains the sole +/// owner of `dropped_records`, preventing the same incompleteness count from +/// being added both before and after the turn. +fn append_initial_routed_usage_to_turn( + turn: &mut TurnRecord, + batch: &crate::cost_status::RuntimeUsageBatch, +) { + for record in &batch.records { + append_routed_usage_record(turn, &record.source_id, record.usage.clone()); + } + for record in &batch.drop_records { + append_routed_usage_drop_record(turn, record.clone()); + } +} + +/// Displayed title of a settlement record. Deliberately fixed text: no turn +/// was accepted, so no prompt content belongs in a record that exists only to +/// keep an already-incurred auxiliary provider call. +const UNACCEPTED_TURN_SUMMARY: &str = "Routing failed before the turn started"; +const UNACCEPTED_TURN_REASON: &str = + "Turn was not accepted; this record keeps the completed pre-turn provider call"; + +fn routed_usage_batch_is_empty(batch: &crate::cost_status::RuntimeUsageBatch) -> bool { + batch.records.is_empty() && batch.drop_records.is_empty() && batch.dropped_records == 0 +} + +/// Durable identity of the settlement record for one completed pre-turn +/// provider call. +/// +/// Derived from the batch's own exact response fingerprints, so settling the +/// same completed call again — a raced operation replay, a retried start after +/// a restart, a reload of the same store — lands on the one record instead of +/// minting a second charge. A batch with no identifiable receipt has no +/// identity to be idempotent on, so it takes a fresh id rather than silently +/// collapsing two distinct coverage gaps into one. +fn unaccepted_routed_usage_turn_id( + thread_id: &str, + batch: &crate::cost_status::RuntimeUsageBatch, +) -> String { + let mut identities = batch + .records + .iter() + .map(|record| routed_usage_source_fingerprint(&record.source_id)) + .chain( + batch + .drop_records + .iter() + .map(|record| routed_usage_source_fingerprint(&record.source_id)), + ) + .collect::>(); + if identities.is_empty() { + return format!("turn_unaccepted_{}", &Uuid::new_v4().to_string()[..8]); + } + identities.sort_unstable(); + identities.dedup(); + let digest = routed_usage_source_fingerprint(&format!("{thread_id}:{}", identities.join(":"))); + format!("turn_unaccepted_{}", &digest[..32]) +} + +/// Settle a completed pre-turn provider call through the same durable receipt +/// authority an accepted turn uses: a Runtime turn record whose bounded +/// `routed_usage` / `routed_usage_drop_records` / `routed_usage_source_ids` +/// ledger is exact-once by source fingerprint. +/// +/// The record is terminal on creation and deliberately carries no parent +/// `usage`, no parent route, and no thread-pointer update: the parent turn was +/// never accepted, so only the auxiliary call's own frozen route may be +/// charged. Because no engine `TurnComplete` will ever arrive for it, this +/// record — unlike an accepted turn — also owns the coverage gap its bounded +/// ledger could not represent. Every persisted field is re-derived from the +/// batch, so replaying the same settlement rewrites the same values. +fn settle_unaccepted_routed_usage( + store: &RuntimeThreadStore, + thread_id: &str, + batch: &crate::cost_status::RuntimeUsageBatch, +) -> Result { + let turn_id = unaccepted_routed_usage_turn_id(thread_id, batch); + let _turn_mutation = store.turn_mutation.lock(); + let mut turn = if store.turn_path(&turn_id)?.exists() { + let existing = store.load_turn(&turn_id)?; + if existing.thread_id != thread_id { + bail!("settlement turn {turn_id} already belongs to another thread"); + } + existing + } else { + let now = Utc::now(); + TurnRecord { + schema_version: CURRENT_RUNTIME_SCHEMA_VERSION, + id: turn_id.clone(), + thread_id: thread_id.to_string(), + status: RuntimeTurnStatus::Failed, + input_summary: UNACCEPTED_TURN_SUMMARY.to_string(), + created_at: now, + started_at: Some(now), + ended_at: Some(now), + duration_ms: Some(0), + usage: None, + effective_route_usage: None, + permission_posture: None, + effective_provider: None, + effective_provider_id: None, + effective_billing_surface: None, + effective_endpoint_fingerprint: None, + effective_provider_live_pricing: None, + effective_billing_mode: None, + effective_dispatched_at: None, + effective_model: None, + routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), + routed_usage_source_ids: Vec::new(), + routed_usage_dropped_records: 0, + error: Some(UNACCEPTED_TURN_REASON.to_string()), + item_ids: Vec::new(), + steer_count: 0, + agent_mail_message_id: None, + } + }; + append_initial_routed_usage_to_turn(&mut turn, batch); + let reported = batch.records.len().saturating_add(batch.drop_records.len()); + let retained = turn + .routed_usage + .len() + .saturating_add(turn.routed_usage_drop_records.len()); + let unretained = u64::try_from(reported.saturating_sub(retained)).unwrap_or(u64::MAX); + let residual = batch + .dropped_records + .saturating_sub(u64::try_from(batch.drop_records.len()).unwrap_or(u64::MAX)); + turn.routed_usage_dropped_records = residual.saturating_add(unretained); + store.save_turn(&turn)?; + Ok(turn_id) +} + +struct InitialRoutedUsageSettlementGuard { + /// The Runtime store that owns the thread whose `start_turn` incurred the + /// call. It is the durability authority for runtime accounting, exactly as + /// it is for an accepted turn's routed usage. Ownerless in-process + /// accounting is only the last resort below: a headless Runtime has no + /// foreground session to drain it, and a closed cost scope (`/new`, + /// session load) rejects a stale report outright. + store: RuntimeThreadStore, + thread_id: String, + cost_scope: crate::cost_status::CostScopeToken, + batch: crate::cost_status::RuntimeUsageBatch, + armed: bool, +} + +impl InitialRoutedUsageSettlementGuard { + fn new( + store: RuntimeThreadStore, + thread_id: &str, + cost_scope: crate::cost_status::CostScopeToken, + batch: &crate::cost_status::RuntimeUsageBatch, + ) -> Self { + Self { + store, + thread_id: thread_id.to_string(), + cost_scope, + batch: batch.clone(), + armed: true, + } + } + + fn disarm(&mut self) { + self.armed = false; + } +} + +impl Drop for InitialRoutedUsageSettlementGuard { + fn drop(&mut self) { + if !self.armed || routed_usage_batch_is_empty(&self.batch) { + return; + } + match settle_unaccepted_routed_usage(&self.store, &self.thread_id, &self.batch) { + Ok(turn_id) => tracing::debug!( + thread_id = %self.thread_id, + turn_id = %turn_id, + "Settled a completed pre-turn provider call into the Runtime store" + ), + Err(error) => { + // The durable authority is unreachable. In-process accounting + // keeps the receipt for whatever is still draining this scope + // rather than discarding it, but it does not survive the + // process, so say so once at warn level. + tracing::warn!( + thread_id = %self.thread_id, + error = %error, + "Failed to persist a completed pre-turn provider call; \ + falling back to in-process accounting" + ); + crate::cost_status::report_runtime_usage_batch(self.cost_scope, None, &self.batch); + } + } + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct TurnItemRecord { #[serde(default = "default_runtime_schema_version")] @@ -2847,6 +3124,61 @@ fn accumulate_truncated_runtime_usage( .insert("runtime_usage_journal_truncated".to_string()); } +fn accumulate_exact_runtime_usage_drop( + totals: &mut UsageTotals, + buckets: &mut std::collections::BTreeMap, + group_by: UsageGroupBy, + record: &RuntimeUsageDropRecord, + turn: &TurnRecord, + thread: &ThreadRecord, +) { + let nonmetered = matches!( + record.route.billing_mode, + RouteBillingMode::Subscription | RouteBillingMode::Local + ); + totals.dropped_usage_records = totals.dropped_usage_records.saturating_add(1); + totals.turns = totals.turns.saturating_add(1); + if nonmetered { + totals.nonmetered_turns = totals.nonmetered_turns.saturating_add(1); + } else { + totals.unpriced_turns = totals.unpriced_turns.saturating_add(1); + totals.cny_unpriced_turns = totals.cny_unpriced_turns.saturating_add(1); + totals + .unpriced_reasons + .insert("provider_success_missing_usage".to_string()); + totals + .cny_unpriced_reasons + .insert("provider_success_missing_usage".to_string()); + } + let audit = record.route.audit(&Usage::default()); + totals + .route_receipts + .insert(format!("{} usage=missing", record.route.receipt(&audit))); + + let key = runtime_usage_bucket_key(group_by, Some(&record.route), turn, thread); + let bucket = buckets.entry(key.clone()).or_insert_with(|| UsageBucket { + key, + ..UsageBucket::default() + }); + bucket.dropped_usage_records = bucket.dropped_usage_records.saturating_add(1); + bucket.turns = bucket.turns.saturating_add(1); + if nonmetered { + bucket.nonmetered_turns = bucket.nonmetered_turns.saturating_add(1); + } else { + bucket.unpriced_turns = bucket.unpriced_turns.saturating_add(1); + bucket.cny_unpriced_turns = bucket.cny_unpriced_turns.saturating_add(1); + bucket + .unpriced_reasons + .insert("provider_success_missing_usage".to_string()); + bucket + .cny_unpriced_reasons + .insert("provider_success_missing_usage".to_string()); + } + bucket + .route_receipts + .insert(format!("{} usage=missing", record.route.receipt(&audit))); +} + fn resolve_runtime_thread_route( config: &Config, provider: ApiProvider, @@ -5220,7 +5552,7 @@ impl RuntimeThreadManager { let parent_dispatched_at = parent_route .as_ref() .map_or(turn.created_at, |route| route.dispatched_at); - if let Some(usage) = turn.usage.as_ref() + if let Some(usage) = turn.effective_route_usage.as_ref().or(turn.usage.as_ref()) && usage_timestamp_in_range(parent_dispatched_at, since, until) { accumulate_runtime_usage_record( @@ -5246,6 +5578,18 @@ impl RuntimeThreadManager { ); } } + for drop_record in &turn.routed_usage_drop_records { + if usage_timestamp_in_range(drop_record.route.dispatched_at, since, until) { + accumulate_exact_runtime_usage_drop( + &mut totals, + &mut buckets, + group_by, + drop_record, + &turn, + &thread, + ); + } + } // Dropped fallback receipts no longer carry a trustworthy // dispatch timestamp. Use the owning turn timestamp only to // decide whether the explicit incompleteness marker belongs @@ -5316,7 +5660,7 @@ impl RuntimeThreadManager { std::collections::BTreeMap::new(); for turn in &turns { let parent_route = turn.effective_route_envelope(); - if let Some(usage) = turn.usage.as_ref() { + if let Some(usage) = turn.effective_route_usage.as_ref().or(turn.usage.as_ref()) { accumulate_runtime_usage_record( &mut split.parent, &mut buckets, @@ -5338,6 +5682,16 @@ impl RuntimeThreadManager { &thread, ); } + for drop_record in &turn.routed_usage_drop_records { + accumulate_exact_runtime_usage_drop( + &mut split.routed_children, + &mut buckets, + UsageGroupBy::Thread, + drop_record, + turn, + &thread, + ); + } // Dropped fallback receipts are routed-child records, so the // incompleteness marker lands on the child side. accumulate_truncated_runtime_usage( @@ -6215,15 +6569,18 @@ impl RuntimeThreadManager { ended_at: Some(turn_at), duration_ms: Some(0), usage: None, + effective_route_usage: None, permission_posture: None, effective_provider: None, effective_provider_id: None, effective_billing_surface: None, effective_endpoint_fingerprint: None, + effective_provider_live_pricing: None, effective_billing_mode: None, effective_dispatched_at: None, effective_model: None, routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), routed_usage_source_ids: Vec::new(), routed_usage_dropped_records: 0, error: None, @@ -6447,9 +6804,15 @@ impl RuntimeThreadManager { for record in background_usage.records.iter().cloned() { append_routed_usage_record(&mut turn, &record.source_id, record.usage); } + for record in background_usage.drop_records.iter().cloned() { + append_routed_usage_drop_record(&mut turn, record); + } + let background_residual = background_usage.dropped_records.saturating_sub( + u64::try_from(background_usage.drop_records.len()).unwrap_or(u64::MAX), + ); turn.routed_usage_dropped_records = turn .routed_usage_dropped_records - .saturating_add(background_usage.dropped_records); + .saturating_add(background_residual); if turn.status == RuntimeTurnStatus::InProgress { turn.status = RuntimeTurnStatus::Failed; turn.ended_at = Some(now); @@ -6873,8 +7236,15 @@ impl RuntimeThreadManager { let mut thread_config = cfg_snapshot.clone(); thread_config.scope_to_provider_identity(&identity); let verbosity = thread_config.verbosity.clone(); - let (route, reasoning_effort, auto_controls_reasoning) = if auto_model { - let selection = crate::model_routing::resolve_auto_route_with_inventory( + let ( + route, + reasoning_effort, + auto_controls_reasoning, + initial_routed_usage, + mut initial_routed_usage_settlement, + ) = if auto_model { + let classifier_cost_scope = crate::cost_status::scope_token(); + let mut selection = crate::model_routing::resolve_auto_route_with_inventory( &thread_config, &prompt, "", @@ -6882,6 +7252,27 @@ impl RuntimeThreadManager { "auto", ) .await?; + // The classifier call has already completed. Move its immutable + // receipts out before resolving the selected parent route so a + // malformed/stale parent catalog entry cannot erase real + // auxiliary spend on the error path. + let initial_routed_usage = crate::cost_status::RuntimeUsageBatch { + records: std::mem::take(&mut selection.routed_usage), + drop_records: std::mem::take(&mut selection.routed_usage_drop_records), + dropped_records: std::mem::take( + &mut selection.routed_usage_dropped_records, + ), + }; + // Until the new turn is durably persisted and handed to the + // engine, every early return must settle this completed + // classifier call against its captured origin. The bounded clone + // is intentionally kept out of the original/replayed turn. + let settlement = InitialRoutedUsageSettlementGuard::new( + self.store.clone(), + thread_id, + classifier_cost_scope, + &initial_routed_usage, + ); let route = resolve_runtime_thread_route( &thread_config, selection.provider, @@ -6902,7 +7293,13 @@ impl RuntimeThreadManager { .as_setting() .to_string() }); - (route, reasoning_effort, auto_controls_reasoning) + ( + route, + reasoning_effort, + auto_controls_reasoning, + initial_routed_usage, + Some(settlement), + ) } else { let route = resolve_runtime_thread_route_for_identity( &cfg_snapshot, @@ -6930,7 +7327,13 @@ impl RuntimeThreadManager { .as_setting() .to_string() }); - (route, reasoning_effort, auto_controls_reasoning) + ( + route, + reasoning_effort, + auto_controls_reasoning, + crate::cost_status::RuntimeUsageBatch::default(), + None, + ) }; let route = if client_preflight_required { route @@ -6974,6 +7377,7 @@ impl RuntimeThreadManager { ended_at: None, duration_ms: None, usage: None, + effective_route_usage: None, permission_posture: Some(policy.permission_wire().to_string()), effective_provider: Some(provider.as_str().to_string()), effective_provider_id: provider_identity @@ -6982,10 +7386,12 @@ impl RuntimeThreadManager { .map(crate::cost_status::sanitize_persisted_route_label), effective_billing_surface: None, effective_endpoint_fingerprint: None, + effective_provider_live_pricing: None, effective_billing_mode: None, effective_dispatched_at: None, effective_model: Some(crate::cost_status::sanitize_persisted_route_label(&model)), routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), routed_usage_source_ids: Vec::new(), routed_usage_dropped_records: 0, error: None, @@ -6993,6 +7399,10 @@ impl RuntimeThreadManager { steer_count: 0, agent_mail_message_id: input_source.mail_message_id().map(str::to_string), }; + append_initial_routed_usage_to_turn(&mut turn, &initial_routed_usage); + // The engine's TurnComplete owns synchronous dropped coverage, + // including this classifier batch. Pre-persisting the count here and + // adding TurnComplete at settlement would count the same gap twice. let user_item_id = format!("item_{}", &Uuid::new_v4().to_string()[..8]); let user_item = TurnItemRecord { @@ -7015,6 +7425,7 @@ impl RuntimeThreadManager { mode, route: Box::new(route), compaction: Box::new(compaction), + initial_routed_usage: Box::new(initial_routed_usage), goal_objective: None, goal_token_budget: None, goal_status: crate::tools::goal::GoalStatus::Active, @@ -7144,6 +7555,9 @@ impl RuntimeThreadManager { configured_sandbox_mode, ); let _sender = permit.send(op); + if let Some(settlement) = initial_routed_usage_settlement.as_mut() { + settlement.disarm(); + } touch_lru(&mut active.lru, thread_id); self.spawn_claimed_turn_monitor( turn.clone(), @@ -7364,6 +7778,7 @@ impl RuntimeThreadManager { ended_at: None, duration_ms: None, usage: None, + effective_route_usage: None, permission_posture: Some( RuntimePolicyProjection::from_persisted( &thread.mode, @@ -7380,12 +7795,14 @@ impl RuntimeThreadManager { .map(crate::cost_status::sanitize_persisted_route_label), effective_billing_surface: None, effective_endpoint_fingerprint: None, + effective_provider_live_pricing: None, effective_billing_mode: None, effective_dispatched_at: None, effective_model: Some(crate::cost_status::sanitize_persisted_route_label( &route_model, )), routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), routed_usage_source_ids: Vec::new(), routed_usage_dropped_records: 0, error: None, @@ -8033,7 +8450,9 @@ impl RuntimeThreadManager { fn register_runtime_usage_sink(&self, turn_id: &str) { let store = self.store.clone(); let sink_turn_id = turn_id.to_string(); - crate::cost_status::register_runtime_usage_sink( + let drop_store = self.store.clone(); + let drop_turn_id = turn_id.to_string(); + crate::cost_status::register_runtime_usage_sink_with_drop( turn_id, Arc::new(move |record: RuntimeUsageRecord| { let _turn_mutation = store.turn_mutation.lock(); @@ -8045,6 +8464,16 @@ impl RuntimeThreadManager { } store.save_turn(&turn).is_ok() }), + Some(Arc::new(move |record: RuntimeUsageDropRecord| { + let _turn_mutation = drop_store.turn_mutation.lock(); + let Ok(mut turn) = drop_store.load_turn(&drop_turn_id) else { + return false; + }; + if !append_routed_usage_drop_record(&mut turn, record) { + return true; + } + drop_store.save_turn(&turn).is_ok() + })), ); } @@ -8059,6 +8488,8 @@ impl RuntimeThreadManager { let mut tool_items: HashMap = HashMap::new(); let mut compaction_items: HashMap = HashMap::new(); let mut turn_usage: Option = None; + let mut turn_effective_route_usage: Option = None; + let mut turn_routed_usage_dropped_records = 0_u64; let mut turn_status: Option = None; let mut turn_error: Option = None; let mut saw_engine_activity = false; @@ -8361,16 +8792,33 @@ impl RuntimeThreadManager { EngineEvent::ToolCallComplete { id, name, result } => { if let Ok(output) = &result && let Some(metadata) = output.metadata.as_ref() - && let Some(route) = - crate::cost_status::child_route_envelope_from_metadata(metadata) - && let Some(usage) = crate::cost_status::child_usage_from_metadata(metadata) { - let source = format!("tool:{id}"); - self.append_routed_usage_to_turn( - &turn_id, - &source, - EffectiveRouteUsage { route, usage }, - )?; + if let Some(batch) = + crate::cost_status::child_usage_records_from_metadata(metadata) + { + for record in batch.records { + self.append_routed_usage_to_turn( + &turn_id, + &record.source_id, + record.usage, + )?; + } + // Synchronous tool-batch coverage is already folded + // into TurnContext by the engine and arrives once in + // TurnComplete. Persisting it here as well would add + // the same dropped receipt twice at settlement. + } else if let Some(route) = + crate::cost_status::child_route_envelope_from_metadata(metadata) + && let Some(usage) = + crate::cost_status::child_usage_from_metadata(metadata) + { + let source = format!("tool:{id}"); + self.append_routed_usage_to_turn( + &turn_id, + &source, + EffectiveRouteUsage { route, usage }, + )?; + } } if let Some(item_id) = tool_items.remove(&id) { let mut item = self.store.load_item(&item_id)?; @@ -8447,7 +8895,10 @@ impl RuntimeThreadManager { self.append_routed_usage_to_turn( &turn_id, &source_id, - EffectiveRouteUsage { route, usage }, + EffectiveRouteUsage { + route: *route, + usage, + }, )?; } } @@ -9083,11 +9534,15 @@ impl RuntimeThreadManager { } EngineEvent::TurnComplete { usage, + parent_route_usage, + routed_usage_dropped_records, status, error, .. } => { turn_usage = Some(usage); + turn_effective_route_usage = Some(parent_route_usage); + turn_routed_usage_dropped_records = routed_usage_dropped_records; let reported_status = match status { TurnOutcomeStatus::Completed => RuntimeTurnStatus::Completed, TurnOutcomeStatus::Interrupted => RuntimeTurnStatus::Interrupted, @@ -9220,12 +9675,21 @@ impl RuntimeThreadManager { turn.ended_at = Some(ended_at); turn.duration_ms = turn.started_at.map(|start| duration_ms(start, ended_at)); turn.usage = turn_usage; + turn.effective_route_usage = turn_effective_route_usage; for record in background_usage.records { append_routed_usage_record(&mut turn, &record.source_id, record.usage); } + let background_exact_drop_count = background_usage.drop_records.len(); + for record in background_usage.drop_records { + append_routed_usage_drop_record(&mut turn, record); + } + let background_residual = background_usage + .dropped_records + .saturating_sub(u64::try_from(background_exact_drop_count).unwrap_or(u64::MAX)); turn.routed_usage_dropped_records = turn .routed_usage_dropped_records - .saturating_add(background_usage.dropped_records); + .saturating_add(background_residual) + .saturating_add(turn_routed_usage_dropped_records); turn.error = turn_error; turn }; diff --git a/crates/tui/src/runtime_threads/tests.rs b/crates/tui/src/runtime_threads/tests.rs index 04deca9eed..72b1eccf45 100644 --- a/crates/tui/src/runtime_threads/tests.rs +++ b/crates/tui/src/runtime_threads/tests.rs @@ -426,15 +426,18 @@ fn sample_turn(thread_id: &str, turn_id: &str, status: RuntimeTurnStatus) -> Tur ended_at: None, duration_ms: None, usage: None, + effective_route_usage: None, permission_posture: None, effective_provider: None, effective_provider_id: None, effective_billing_surface: None, effective_endpoint_fingerprint: None, + effective_provider_live_pricing: None, effective_billing_mode: None, effective_dispatched_at: None, effective_model: None, routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), routed_usage_source_ids: Vec::new(), routed_usage_dropped_records: 0, error: None, @@ -612,6 +615,7 @@ fn set_test_turn_route( model: model.to_string(), billing_surface: billing_surface.map(str::to_string), endpoint_fingerprint: None, + provider_live_pricing: None, billing_mode, dispatched_at: turn.created_at, }); @@ -1441,6 +1445,8 @@ async fn caller_cancellation_after_engine_acceptance_keeps_owned_turn_lifecycle( .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -1553,6 +1559,8 @@ async fn operation_key_replays_torn_response_survives_restart_and_rejects_mismat .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -1764,6 +1772,8 @@ async fn thread_updates_while_start_waits_for_capacity_survive_latest_turn_write .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -1926,6 +1936,8 @@ async fn compact_lifecycle_outlives_caller_and_preserves_concurrent_thread_updat .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -2734,6 +2746,7 @@ fn turn_record_persists_billing_surface_without_raw_endpoint() { model: "step-3.7-flash".to_string(), billing_surface: Some(crate::pricing::STEPFUN_PAYG_BILLING_SURFACE.to_string()), endpoint_fingerprint: Some(fingerprint.clone()), + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Metered, dispatched_at: turn.created_at, }); @@ -2750,6 +2763,159 @@ fn turn_record_persists_billing_surface_without_raw_endpoint() { assert!(value.get("effective_base_url").is_none()); } +#[test] +fn turn_record_round_trips_frozen_provider_live_pricing_and_drops_hostile_quotes() { + struct ProviderCatalogReset; + impl Drop for ProviderCatalogReset { + fn drop(&mut self) { + crate::provider_catalog_live::reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + } + } + + let _env = crate::test_support::lock_test_env(); + let _live = crate::provider_lake::lock_live_snapshot(); + let home = tempfile::tempdir().expect("test home"); + let _home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", home.path()); + let _reset = ProviderCatalogReset; + crate::provider_catalog_live::reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + + let dispatched_at = Utc::now(); + let fetched_at = u64::try_from(dispatched_at.timestamp()).expect("timestamp"); + let model = "synthetic-baseten-turn-record"; + let fingerprint = + codewhale_config::catalog::base_url_fingerprint(codewhale_config::BASETEN_BASE_URL); + let priced_delta = |input: f64, output: f64| codewhale_config::catalog::ProviderCatalogDelta { + provider: codewhale_config::BASETEN_TEMPLATE_ID.to_string(), + base_url_fingerprint: fingerprint.clone(), + fetched_at, + offerings: vec![codewhale_config::catalog::CatalogOffering { + provider: codewhale_config::BASETEN_TEMPLATE_ID.to_string(), + wire_model_id: model.to_string(), + endpoint_key: "chat".to_string(), + cost: Some(codewhale_config::models_dev::ModelsDevCost { + input: Some(input), + output: Some(output), + cache_read: Some(0.25), + cache_write: None, + }), + ..Default::default() + }], + }; + crate::provider_catalog_live::record_success(priced_delta(1.25, 5.0)); + let route = crate::cost_status::EffectiveRouteEnvelope::capture( + None, + ApiProvider::Custom, + codewhale_config::BASETEN_TEMPLATE_ID, + model, + Some(codewhale_config::BASETEN_BASE_URL), + dispatched_at, + ); + let valid_quote = route + .provider_live_pricing + .clone() + .expect("fresh Baseten scope freezes quote"); + + let mut turn = sample_turn("thr_quote", "turn_quote", RuntimeTurnStatus::Completed); + turn.persist_effective_route(&route); + let serialized = serde_json::to_string(&turn).expect("serialize quoted turn"); + for raw_secret in [codewhale_config::BASETEN_BASE_URL, "api_key", "Bearer "] { + // The assertion message must not itself log the credential fragment it + // checks for — name the check, not the secret. + assert!( + !serialized.contains(raw_secret), + "persisted turn serialization leaked a credential fragment" + ); + } + let restored: TurnRecord = serde_json::from_str(&serialized).expect("restore quoted turn"); + let restored_route = restored + .effective_route_envelope() + .expect("complete persisted route"); + assert_eq!( + restored_route.provider_live_pricing, + Some(valid_quote.clone()) + ); + + crate::provider_catalog_live::record_success(priced_delta(19.0, 29.0)); + let usage = Usage { + input_tokens: 1_000_000, + ..Usage::default() + }; + assert_eq!( + restored_route + .audit(&usage) + .estimate + .expect("frozen price survives refresh") + .usd, + 1.25 + ); + + // Model the strongest persisted attack: an adversary recomputes the + // unkeyed integrity revision after placing credential text in the quote. + // Structural validation, not the digest alone, must reject it. + let secret = "Authorization: Bearer persisted-secret"; + let mut hostile_quote = valid_quote; + hostile_quote.wire_model = secret.to_string(); + let revision_payload = serde_json::to_vec(&( + "codewhale-provider-live-pricing-quote-v1", + hostile_quote.provider, + &hostile_quote.provider_identity, + &hostile_quote.wire_model, + &hostile_quote.endpoint_fingerprint, + hostile_quote.catalog_fetched_at, + &hostile_quote.currency, + &hostile_quote.provenance, + &hostile_quote.input_per_million, + &hostile_quote.output_per_million, + &hostile_quote.cache_read_per_million, + &hostile_quote.cache_write_per_million, + )) + .expect("revision payload"); + hostile_quote.catalog_revision = + format!("sha256:{}", crate::hashing::sha256_hex(revision_payload)); + + let mut hostile_quote_value = serde_json::to_value( + restored_route + .provider_live_pricing + .as_ref() + .expect("valid persisted quote"), + ) + .expect("quote JSON"); + hostile_quote_value["wire_model"] = serde_json::json!(secret); + hostile_quote_value["catalog_revision"] = serde_json::json!(hostile_quote.catalog_revision); + + let mut hostile_turn_value = serde_json::to_value(&restored).expect("turn JSON"); + hostile_turn_value["effective_provider_live_pricing"] = hostile_quote_value.clone(); + let mut hostile_child_route = serde_json::to_value(&restored_route).expect("route JSON"); + hostile_child_route["provider_live_pricing"] = hostile_quote_value; + hostile_turn_value["routed_usage"] = serde_json::json!([{ + "route": hostile_child_route, + "usage": Usage::default(), + }]); + + let sanitized: TurnRecord = + serde_json::from_value(hostile_turn_value).expect("hostile optional quotes fail closed"); + assert!(sanitized.effective_provider_live_pricing.is_none()); + assert_eq!(sanitized.routed_usage.len(), 1); + assert!( + sanitized.routed_usage[0] + .route + .provider_live_pricing + .is_none() + ); + let sanitized_json = serde_json::to_string(&sanitized).expect("reserialize sanitized turn"); + assert!(!sanitized_json.contains("persisted-secret")); + assert_eq!( + sanitized + .effective_route_envelope() + .expect("route remains readable") + .audit(&usage) + .unpriced_reason, + Some(crate::pricing::UnpricedReason::UnverifiedLivePricing) + ); +} + #[test] fn serialized_turn_record_redacts_all_route_and_source_fields() { let mut value = serde_json::to_value(sample_turn( @@ -2907,6 +3073,7 @@ async fn aggregate_usage_for_thread_scopes_both_currencies_to_one_thread() -> Re model: "deepseek-v4-flash".to_string(), billing_surface: Some(crate::pricing::FIRST_PARTY_PAYG_BILLING_SURFACE.to_string()), endpoint_fingerprint: None, + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Metered, dispatched_at: turn.created_at, }, @@ -3279,6 +3446,7 @@ async fn aggregate_usage_includes_exclusive_child_calls_and_zero_usage_receipts( // the route audit will not price a metered route without one. billing_surface: Some(crate::pricing::FIRST_PARTY_PAYG_BILLING_SURFACE.to_string()), endpoint_fingerprint: None, + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Metered, dispatched_at: turn.created_at, }, @@ -3301,6 +3469,7 @@ async fn aggregate_usage_includes_exclusive_child_calls_and_zero_usage_receipts( crate::pricing::OAUTH_SUBSCRIPTION_BILLING_SURFACE.to_string(), ), endpoint_fingerprint: None, + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Subscription, dispatched_at: turn.created_at, }, @@ -3357,6 +3526,7 @@ async fn aggregate_usage_filters_each_call_by_its_dispatch_timestamp() -> Result model: "deepseek-v4-flash".to_string(), billing_surface: None, endpoint_fingerprint: None, + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Metered, dispatched_at: window, }, @@ -3528,7 +3698,7 @@ fn routed_usage_append_is_bounded_and_idempotent_for_every_delivery_path() { 1 => "mailbox", _ => "fallback", }; - assert!(append_routed_usage_record( + let changed = append_routed_usage_record( &mut turn, &format!("{path}:response:{index}"), crate::cost_status::EffectiveRouteUsage { @@ -3538,12 +3708,19 @@ fn routed_usage_append_is_bounded_and_idempotent_for_every_delivery_path() { ..Usage::default() }, }, - )); + ); + assert_eq!( + changed, + index <= u32::try_from(MAX_ROUTED_USAGE_RECORDS_PER_TURN).unwrap() + ); } assert_eq!(turn.routed_usage.len(), MAX_ROUTED_USAGE_RECORDS_PER_TURN); - assert_eq!(turn.routed_usage_dropped_records, 26); - assert_eq!(turn.routed_usage_source_ids.len(), 90); + assert_eq!(turn.routed_usage_dropped_records, 1); + assert_eq!( + turn.routed_usage_source_ids.len(), + MAX_ROUTED_USAGE_RECORDS_PER_TURN + ); let before = turn.clone(); assert!(!append_routed_usage_record( &mut turn, @@ -5218,6 +5395,12 @@ async fn thread_lifecycle_persists_across_restart() -> Result<()> { output_tokens: 12, ..Usage::default() }, + parent_route_usage: Usage { + input_tokens: 10, + output_tokens: 12, + ..Usage::default() + }, + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -5261,6 +5444,466 @@ async fn thread_lifecycle_persists_across_restart() -> Result<()> { Ok(()) } +#[tokio::test] +async fn initial_classifier_usage_is_persisted_before_terminal_and_merged_exactly_once() +-> Result<()> { + let manager = test_manager(test_runtime_dir())?; + let thread = manager + .create_thread(CreateThreadRequest::default()) + .await?; + let mut harness = install_mock_engine(&manager, &thread.id).await; + let turn = manager + .start_turn( + &thread.id, + StartTurnRequest { + prompt: "persist classifier receipt".to_string(), + ..StartTurnRequest::default() + }, + ) + .await?; + assert!(matches!( + harness.rx_op.recv().await, + Some(Op::SendMessage { .. }) + )); + + let classifier_usage = Usage { + input_tokens: 7, + output_tokens: 5, + ..Usage::default() + }; + let classifier_route = crate::cost_status::EffectiveRouteEnvelope::capture( + None, + ApiProvider::Openai, + "openai", + "classifier-model", + Some(ApiProvider::Openai.default_base_url()), + Utc::now(), + ); + let classifier_batch = crate::cost_status::RuntimeUsageBatch { + records: vec![crate::cost_status::RuntimeUsageRecord { + source_id: "auto-router:runtime-fixture".to_string(), + usage: crate::cost_status::EffectiveRouteUsage { + route: classifier_route.clone(), + usage: classifier_usage.clone(), + }, + }], + drop_records: vec![crate::cost_status::RuntimeUsageDropRecord { + source_id: "auto-router:runtime-missing-usage".to_string(), + route: classifier_route, + }], + // One exact drop plus two residual/unidentifiable gaps. + dropped_records: 3, + }; + { + let _turn_mutation = manager.store.turn_mutation.lock(); + let mut reserved = manager.store.load_turn(&turn.id)?; + append_initial_routed_usage_to_turn(&mut reserved, &classifier_batch); + append_initial_routed_usage_to_turn(&mut reserved, &classifier_batch); + manager.store.save_turn(&reserved)?; + } + let preterminal = manager.store.load_turn(&turn.id)?; + assert_eq!(preterminal.routed_usage.len(), 1); + assert_eq!(preterminal.routed_usage_drop_records.len(), 1); + assert_eq!(preterminal.routed_usage_drop_records[0].source_id.len(), 64); + assert!( + !preterminal.routed_usage_drop_records[0] + .source_id + .contains("runtime-missing-usage") + ); + assert_eq!(preterminal.routed_usage[0].usage, classifier_usage); + assert_eq!(preterminal.routed_usage_source_ids.len(), 2); + assert!( + preterminal + .routed_usage_source_ids + .contains(&preterminal.routed_usage_drop_records[0].source_id), + "the persisted exact-drop receipt must retain the same durable source fingerprint as the dedupe ledger" + ); + assert_eq!( + preterminal.routed_usage_dropped_records, 0, + "TurnComplete, not pre-persistence, owns classifier drop coverage" + ); + { + let _turn_mutation = manager.store.turn_mutation.lock(); + let mut replayed = manager.store.load_turn(&turn.id)?; + append_initial_routed_usage_to_turn(&mut replayed, &classifier_batch); + manager.store.save_turn(&replayed)?; + } + let replayed = manager.store.load_turn(&turn.id)?; + assert_eq!(replayed.routed_usage.len(), 1); + assert_eq!(replayed.routed_usage_drop_records.len(), 1); + assert_eq!(replayed.routed_usage_source_ids.len(), 2); + + harness + .tx_event + .send(EngineEvent::TurnStarted { + turn_id: "engine_classifier_receipt".to_string(), + created_at: Utc::now(), + route: None, + }) + .await?; + harness + .tx_event + .send(EngineEvent::TurnComplete { + usage: Usage { + input_tokens: 18, + output_tokens: 8, + ..Usage::default() + }, + parent_route_usage: Usage { + input_tokens: 11, + output_tokens: 3, + ..Usage::default() + }, + routed_usage_dropped_records: classifier_batch.dropped_records.saturating_sub( + u64::try_from(classifier_batch.drop_records.len()).unwrap_or(u64::MAX), + ), + status: TurnOutcomeStatus::Completed, + error: None, + tool_catalog: None, + base_url: None, + }) + .await?; + + let completed = wait_for_terminal_turn(&manager, &turn.id, Duration::from_secs(2)).await?; + assert_eq!(completed.routed_usage.len(), 1); + assert_eq!(completed.routed_usage_drop_records.len(), 1); + assert_eq!(completed.routed_usage_source_ids.len(), 2); + assert_eq!(completed.routed_usage[0].usage, classifier_usage); + assert_eq!(completed.routed_usage_dropped_records, 2); + assert_eq!( + completed.usage, + Some(Usage { + input_tokens: 18, + output_tokens: 8, + ..Usage::default() + }) + ); + assert_eq!( + completed.effective_route_usage, + Some(Usage { + input_tokens: 11, + output_tokens: 3, + ..Usage::default() + }) + ); + let aggregate = manager + .aggregate_usage(None, None, UsageGroupBy::Thread) + .await?; + assert_eq!(aggregate.totals.dropped_usage_records, 3); + assert!( + aggregate.totals.route_receipts.iter().any( + |receipt| receipt.contains("classifier-model") && receipt.contains("usage=missing") + ) + ); + Ok(()) +} + +/// One completed auxiliary provider call, on its own frozen route, as the +/// classifier hands it to `start_turn` before the parent route is resolved. +fn classifier_settlement_batch( + identity: &str, + model: &str, + source_prefix: &str, + dropped_records: u64, +) -> crate::cost_status::RuntimeUsageBatch { + let mut route = crate::cost_status::EffectiveRouteEnvelope::capture( + None, + ApiProvider::Openrouter, + identity, + model, + Some(ApiProvider::Openrouter.default_base_url()), + Utc::now(), + ); + route.billing_mode = crate::cost_status::RouteBillingMode::Subscription; + crate::cost_status::RuntimeUsageBatch { + records: vec![crate::cost_status::RuntimeUsageRecord { + source_id: format!("auto-router:{source_prefix}-usage"), + usage: crate::cost_status::EffectiveRouteUsage { + route: route.clone(), + usage: Usage { + input_tokens: 9, + output_tokens: 4, + ..Usage::default() + }, + }, + }], + drop_records: vec![crate::cost_status::RuntimeUsageDropRecord { + source_id: format!("auto-router:{source_prefix}-drop"), + route, + }], + dropped_records, + } +} + +/// A classifier call that completed before the parent route failed must land +/// in the Runtime store, not in ownerless in-process accounting: headless +/// Runtime/API execution has no foreground session draining that pool, so a +/// process loss erased real provider spend. Settling the same completed call +/// again must not charge a second time, and the failed parent route must never +/// become the billing route. +#[tokio::test] +async fn failed_runtime_parent_route_resolve_and_preflight_settle_batch_once_without_repricing() +-> Result<()> { + let _cost_scope = crate::cost_status::test_scope(); + let scope = crate::cost_status::scope_token(); + let runtime_dir = test_runtime_dir(); + let manager = test_manager(runtime_dir.clone())?; + let thread = manager + .create_thread(CreateThreadRequest::default()) + .await?; + // One exact missing-usage route plus two residual coverage gaps. + let batch = classifier_settlement_batch( + "classifier-openrouter", + "classifier/frozen-route", + "runtime-parent-failure", + 3, + ); + let frozen_batch = batch.clone(); + + for error_message in [ + "selected parent provider/model is invalid", + "selected parent route failed preflight", + ] { + let error = { + let _settlement = InitialRoutedUsageSettlementGuard::new( + manager.store.clone(), + &thread.id, + scope, + &batch, + ); + anyhow::anyhow!(error_message) + }; + assert_eq!( + error.to_string(), + error_message, + "settlement must preserve the original route failure" + ); + } + + assert_eq!( + batch, frozen_batch, + "settlement must not coerce classifier receipts onto the failed parent provider/model" + ); + + let turns = manager.store.list_turns_for_thread(&thread.id)?; + assert_eq!( + turns.len(), + 1, + "settling the same completed call twice must reuse one durable record: {turns:?}" + ); + let settled = &turns[0]; + assert_eq!(settled.status, RuntimeTurnStatus::Failed); + assert_eq!(settled.usage, None); + assert_eq!(settled.effective_route_usage, None); + assert_eq!(settled.effective_provider, None); + assert_eq!(settled.effective_model, None); + assert!( + settled.effective_route_envelope().is_none(), + "the parent route never dispatched and must never be charged" + ); + assert_eq!(settled.routed_usage.len(), 1); + assert_eq!( + settled.routed_usage[0].route.model, + "classifier/frozen-route" + ); + assert_eq!( + settled.routed_usage[0].usage, + Usage { + input_tokens: 9, + output_tokens: 4, + ..Usage::default() + } + ); + assert_eq!(settled.routed_usage_drop_records.len(), 1); + assert_eq!(settled.routed_usage_drop_records[0].source_id.len(), 64); + assert_eq!(settled.routed_usage_source_ids.len(), 2); + assert_eq!( + settled.routed_usage_dropped_records, 2, + "no engine TurnComplete will ever arrive, so this record owns its residual gap" + ); + assert!( + crate::cost_status::drain().is_empty(), + "runtime spend must not be charged to whichever session happens to be live" + ); + + // Process loss: a fresh manager over the same store still prices the call + // exactly once, under the classifier's own frozen route. + drop(manager); + let reopened = test_manager(runtime_dir)?; + let aggregate = reopened + .aggregate_usage(None, None, UsageGroupBy::Provider) + .await?; + assert_eq!(aggregate.totals.input_tokens, 9); + assert_eq!(aggregate.totals.output_tokens, 4); + assert_eq!(aggregate.totals.dropped_usage_records, 3); + assert_eq!( + aggregate + .buckets + .iter() + .map(|bucket| bucket.key.as_str()) + .collect::>(), + vec!["classifier-openrouter", "unknown-truncated"], + "exact receipts bill the classifier identity; only the residual gap is unattributed" + ); + assert!(aggregate.totals.route_receipts.iter().any(|receipt| { + receipt.contains("identity=classifier-openrouter") + && receipt.contains("model=classifier/frozen-route") + })); + assert!( + aggregate + .totals + .route_receipts + .iter() + .all(|receipt| !receipt.contains("selected_parent")), + "the failed parent route must never become the classifier billing route" + ); + + // Restart recovery treats the record as the terminal turn it is: it emits + // the one missing receipt and re-prices nothing. + let recovered = reopened.get_thread(&thread.id).await?; + assert_eq!( + recovered.latest_turn_id.as_deref(), + Some(settled.id.as_str()) + ); + assert_eq!( + serde_json::to_value( + reopened + .aggregate_usage(None, None, UsageGroupBy::Provider) + .await? + .totals + )?, + serde_json::to_value(aggregate.totals)?, + "flushing the recovery receipt must not count the call again" + ); + Ok(()) +} + +/// `/new` and session load close the cost scope captured at dispatch, and +/// ownerless in-process settlement rejects a stale scope outright — which is +/// how a real classifier call vanished. The durable record must not depend on +/// the scope, and must not charge the replacement scope either. +#[tokio::test] +async fn unaccepted_classifier_settlement_survives_a_closed_cost_scope() -> Result<()> { + let _cost_scope = crate::cost_status::test_scope(); + let dispatch_scope = crate::cost_status::scope_token(); + let manager = test_manager(test_runtime_dir())?; + let thread = manager + .create_thread(CreateThreadRequest::default()) + .await?; + let batch = classifier_settlement_batch( + "stale-scope-openrouter", + "classifier/stale-scope", + "stale-scope", + 1, + ); + + let _closed = crate::cost_status::close_current_scope(); + assert_ne!(dispatch_scope, crate::cost_status::scope_token()); + drop(InitialRoutedUsageSettlementGuard::new( + manager.store.clone(), + &thread.id, + dispatch_scope, + &batch, + )); + + let turns = manager.store.list_turns_for_thread(&thread.id)?; + assert_eq!(turns.len(), 1); + assert_eq!(turns[0].routed_usage.len(), 1); + assert_eq!( + turns[0].routed_usage[0].route.model, + "classifier/stale-scope" + ); + assert_eq!(turns[0].routed_usage_drop_records.len(), 1); + assert_eq!(turns[0].routed_usage_dropped_records, 0); + assert!( + crate::cost_status::drain().is_empty(), + "the replacement scope must not inherit the closed scope's spend" + ); + Ok(()) +} + +/// An exact operation-key retry replays the original turn, but its own +/// classifier call really happened. It settles into its own record, and the +/// replayed turn stays exactly as its owner left it. +#[tokio::test] +async fn raced_runtime_operation_replay_settles_second_classifier_without_mutating_original_turn() +-> Result<()> { + let _cost_scope = crate::cost_status::test_scope(); + let scope = crate::cost_status::scope_token(); + let manager = test_manager(test_runtime_dir())?; + let thread = manager + .create_thread(CreateThreadRequest::default()) + .await?; + let original_turn = sample_turn( + &thread.id, + "turn_original_operation", + RuntimeTurnStatus::Completed, + ); + manager.store.save_turn(&original_turn)?; + let second_classifier_batch = classifier_settlement_batch( + "second-classifier-openrouter", + "classifier/second-actual-call", + "raced-replay-second-call", + 1, + ); + + let replayed_turn = { + let _settlement = InitialRoutedUsageSettlementGuard::new( + manager.store.clone(), + &thread.id, + scope, + &second_classifier_batch, + ); + manager.store.load_turn(&original_turn.id)? + }; + assert_eq!(replayed_turn.id, original_turn.id); + assert_eq!(replayed_turn.status, original_turn.status); + assert!(replayed_turn.routed_usage.is_empty()); + assert!(replayed_turn.routed_usage_drop_records.is_empty()); + + let persisted_original = manager.store.load_turn(&original_turn.id)?; + assert!(persisted_original.routed_usage.is_empty()); + assert!(persisted_original.routed_usage_drop_records.is_empty()); + assert_eq!(persisted_original.routed_usage_dropped_records, 0); + + let settled = manager + .store + .list_turns_for_thread(&thread.id)? + .into_iter() + .filter(|turn| turn.id != original_turn.id) + .collect::>(); + assert_eq!(settled.len(), 1); + assert_eq!(settled[0].routed_usage.len(), 1); + assert_eq!( + settled[0].routed_usage[0].route.model, + "classifier/second-actual-call" + ); + assert_eq!(settled[0].routed_usage_drop_records.len(), 1); + Ok(()) +} + +/// Auto routing that reported nothing has nothing to settle. Minting a record +/// anyway would put a phantom turn on every thread whose parent route failed. +#[tokio::test] +async fn pre_turn_settlement_with_nothing_to_keep_writes_no_record() -> Result<()> { + let _cost_scope = crate::cost_status::test_scope(); + let manager = test_manager(test_runtime_dir())?; + let thread = manager + .create_thread(CreateThreadRequest::default()) + .await?; + drop(InitialRoutedUsageSettlementGuard::new( + manager.store.clone(), + &thread.id, + crate::cost_status::scope_token(), + &crate::cost_status::RuntimeUsageBatch::default(), + )); + assert!( + manager.store.list_turns_for_thread(&thread.id)?.is_empty(), + "a classifier that reported nothing must not mint a phantom turn" + ); + Ok(()) +} + #[tokio::test] async fn monitor_separates_lifecycle_start_from_billing_dispatch_and_child_usage() -> Result<()> { let manager = test_manager(test_runtime_dir())?; @@ -5306,6 +5949,7 @@ async fn monitor_separates_lifecycle_start_from_billing_dispatch_and_child_usage billing: Some(crate::core::events::RouteBillingEnvelope { billing_surface: Some(crate::pricing::STEPFUN_PAYG_BILLING_SURFACE.to_string()), endpoint_fingerprint: Some(endpoint_fingerprint.clone()), + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Metered, dispatched_at, }), @@ -5323,7 +5967,7 @@ async fn monitor_separates_lifecycle_start_from_billing_dispatch_and_child_usage message: crate::tools::subagent::MailboxMessage::TokenUsage { agent_id: "agent_child".to_string(), source_id: "response-child".to_string(), - route: crate::cost_status::EffectiveRouteEnvelope { + route: Box::new(crate::cost_status::EffectiveRouteEnvelope { provider: ApiProvider::OpenaiCodex, provider_identity: "codex-child".to_string(), model: "gpt-5.5".to_string(), @@ -5331,9 +5975,10 @@ async fn monitor_separates_lifecycle_start_from_billing_dispatch_and_child_usage crate::pricing::OAUTH_SUBSCRIPTION_BILLING_SURFACE.to_string(), ), endpoint_fingerprint: None, + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Subscription, dispatched_at, - }, + }), usage: Usage { input_tokens: 3, output_tokens: 2, @@ -5343,6 +5988,27 @@ async fn monitor_separates_lifecycle_start_from_billing_dispatch_and_child_usage }, }) .await?; + harness + .tx_event + .send(EngineEvent::ToolCallStarted { + id: "tool-routed-coverage".to_string(), + name: "rlm".to_string(), + input: json!({"action": "eval"}), + }) + .await?; + harness + .tx_event + .send(EngineEvent::ToolCallComplete { + id: "tool-routed-coverage".to_string(), + name: "rlm".to_string(), + result: Ok( + crate::tools::spec::ToolResult::success("covered").with_metadata(json!({ + "child_usage_records": [], + "child_usage_dropped_records": 2, + })), + ), + }) + .await?; harness .tx_event .send(EngineEvent::TurnComplete { @@ -5356,6 +6022,20 @@ async fn monitor_separates_lifecycle_start_from_billing_dispatch_and_child_usage }), ..Usage::default() }, + parent_route_usage: Usage { + input_tokens: 10, + output_tokens: 4, + reasoning_replay_tokens: Some(6), + server_tool_use: Some(crate::models::ServerToolUsage { + code_execution_requests: Some(2), + tool_search_requests: Some(3), + }), + ..Usage::default() + }, + // The engine already folded the synchronous tool metadata into + // this authoritative turn count. Runtime must not persist the + // ToolCallComplete copy and then add it again here. + routed_usage_dropped_records: 2, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -5385,6 +6065,7 @@ async fn monitor_separates_lifecycle_start_from_billing_dispatch_and_child_usage assert_eq!(completed.effective_dispatched_at, Some(dispatched_at)); assert_eq!(completed.started_at, Some(started_at)); assert_eq!(completed.routed_usage.len(), 1); + assert_eq!(completed.routed_usage_dropped_records, 2); assert_eq!(completed.routed_usage[0].usage.reasoning_tokens, Some(2)); let persisted_usage = completed.usage.expect("parent usage"); assert_eq!(persisted_usage.reasoning_replay_tokens, Some(6)); @@ -5430,7 +6111,7 @@ async fn monitor_separates_lifecycle_start_from_billing_dispatch_and_child_usage message: crate::tools::subagent::MailboxMessage::TokenUsage { agent_id: "agent-child-second".to_string(), source_id: "response-child-second".to_string(), - route: crate::cost_status::EffectiveRouteEnvelope { + route: Box::new(crate::cost_status::EffectiveRouteEnvelope { provider: ApiProvider::OpenaiCodex, provider_identity: "codex-child".to_string(), model: "gpt-5.5".to_string(), @@ -5438,9 +6119,10 @@ async fn monitor_separates_lifecycle_start_from_billing_dispatch_and_child_usage crate::pricing::OAUTH_SUBSCRIPTION_BILLING_SURFACE.to_string(), ), endpoint_fingerprint: None, + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Subscription, dispatched_at: Utc::now(), - }, + }), usage: Usage { input_tokens: 5, output_tokens: 1, @@ -5453,6 +6135,8 @@ async fn monitor_separates_lifecycle_start_from_billing_dispatch_and_child_usage .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -5502,6 +6186,12 @@ async fn completed_turn_without_engine_output_fails() -> Result<()> { output_tokens: 0, ..Usage::default() }, + parent_route_usage: Usage { + input_tokens: 10, + output_tokens: 0, + ..Usage::default() + }, + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -5585,6 +6275,8 @@ async fn preturn_control_status_does_not_make_empty_turn_succeed() -> Result<()> let _ = tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -5645,6 +6337,8 @@ async fn engine_error_remains_failed_after_nominal_turn_complete() -> Result<()> let _ = tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -6184,6 +6878,8 @@ async fn compact_interrupt_persists_canceled_item_for_the_exact_request() -> Res .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Interrupted, error: None, tool_catalog: None, @@ -6421,6 +7117,12 @@ async fn multi_turn_continuity_same_thread() -> Result<()> { output_tokens: 5, ..Usage::default() }, + parent_route_usage: Usage { + input_tokens: 5, + output_tokens: 5, + ..Usage::default() + }, + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -6747,6 +7449,8 @@ async fn approval_required_with_stale_active_turn_is_denied() -> Result<()> { output_tokens: 0, ..Usage::default() }, + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -6862,6 +7566,8 @@ async fn approval_required_awaits_external_decision_allow() -> Result<()> { .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -6974,6 +7680,8 @@ async fn user_input_snapshot_survives_reload_and_clears_after_submission() -> Re .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -7321,6 +8029,8 @@ async fn thread_detail_cursor_precedes_projection_reads_at_terminal_boundary() - .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -7571,6 +8281,8 @@ async fn thread_detail_materializes_stream_prefixes_before_their_delta_cursor() .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Interrupted, error: None, tool_catalog: None, @@ -7709,6 +8421,8 @@ async fn thread_detail_delta_boundary_is_replay_idempotent() -> Result<()> { .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Interrupted, error: None, tool_catalog: None, @@ -7778,6 +8492,8 @@ async fn terminal_turn_cancels_pending_user_input_and_clears_snapshot() -> Resul .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -7966,6 +8682,8 @@ async fn dynamic_tool_result_settles_snapshot_and_emits_one_safe_resolution() -> .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -8150,6 +8868,8 @@ async fn dynamic_tool_result_receipt_outlives_canceled_delivery_future() -> Resu .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -8978,6 +9698,8 @@ async fn dynamic_tool_timeout_clears_snapshot_and_emits_once() -> Result<()> { .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -9050,6 +9772,8 @@ async fn terminal_turn_cancels_pending_dynamic_tool_exactly_once() -> Result<()> .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Interrupted, error: None, tool_catalog: None, @@ -9164,6 +9888,8 @@ async fn approval_required_external_deny_is_denied() -> Result<()> { .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -9234,6 +9960,8 @@ async fn auto_review_force_prompt_is_denied_without_opening_a_modal() -> Result< .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -9332,6 +10060,8 @@ async fn approval_timeout_denies_clears_ui_and_next_turn_can_start() -> Result<( .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -9422,6 +10152,8 @@ async fn thinking_delta_emits_agent_reasoning_item() -> Result<()> { .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -9559,6 +10291,8 @@ async fn approval_required_remember_flips_thread_auto_approve() -> Result<()> { .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -9643,6 +10377,8 @@ async fn elevation_required_with_stale_active_turn_is_denied() -> Result<()> { output_tokens: 0, ..Usage::default() }, + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -9709,6 +10445,12 @@ async fn steer_turn_on_active_turn_records_item_and_event() -> Result<()> { output_tokens: 9, ..Usage::default() }, + parent_route_usage: Usage { + input_tokens: 8, + output_tokens: 9, + ..Usage::default() + }, + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -9866,6 +10608,8 @@ async fn steer_receipts_outlive_caller_cancellation_after_engine_acceptance() -> tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -9953,6 +10697,8 @@ async fn steer_rejects_a_terminal_durable_turn_without_dispatch_or_item() -> Res tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -10226,6 +10972,12 @@ async fn compaction_lifecycle_emits_item_events_with_compaction_counts() -> Resu output_tokens: 3, ..Usage::default() }, + parent_route_usage: Usage { + input_tokens: 3, + output_tokens: 3, + ..Usage::default() + }, + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -10263,6 +11015,12 @@ async fn compaction_lifecycle_emits_item_events_with_compaction_counts() -> Resu output_tokens: 1, ..Usage::default() }, + parent_route_usage: Usage { + input_tokens: 1, + output_tokens: 1, + ..Usage::default() + }, + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, @@ -10473,15 +11231,18 @@ fn opening_manager_recovers_stale_queued_and_in_progress_work() -> Result<()> { ended_at: None, duration_ms: None, usage: None, + effective_route_usage: None, permission_posture: None, effective_provider: None, effective_provider_id: None, effective_billing_surface: None, effective_endpoint_fingerprint: None, + effective_provider_live_pricing: None, effective_billing_mode: None, effective_dispatched_at: None, effective_model: None, routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), routed_usage_source_ids: Vec::new(), routed_usage_dropped_records: 0, error: None, @@ -10500,15 +11261,18 @@ fn opening_manager_recovers_stale_queued_and_in_progress_work() -> Result<()> { ended_at: None, duration_ms: None, usage: None, + effective_route_usage: None, permission_posture: None, effective_provider: None, effective_provider_id: None, effective_billing_surface: None, effective_endpoint_fingerprint: None, + effective_provider_live_pricing: None, effective_billing_mode: None, effective_dispatched_at: None, effective_model: None, routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), routed_usage_source_ids: Vec::new(), routed_usage_dropped_records: 0, error: None, @@ -10773,15 +11537,18 @@ fn seed_turns_with_user_messages( ended_at: Some(created_at), duration_ms: Some(0), usage: None, + effective_route_usage: None, permission_posture: None, effective_provider: None, effective_provider_id: None, effective_billing_surface: None, effective_endpoint_fingerprint: None, + effective_provider_live_pricing: None, effective_billing_mode: None, effective_dispatched_at: None, effective_model: None, routed_usage: Vec::new(), + routed_usage_drop_records: Vec::new(), routed_usage_source_ids: Vec::new(), routed_usage_dropped_records: 0, error: None, @@ -11129,6 +11896,8 @@ async fn agent_mail_release_acceptance_two_task_matrix() -> Result<()> { .tx_event .send(EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: TurnOutcomeStatus::Completed, error: None, tool_catalog: None, diff --git a/crates/tui/src/runtime_web/app.mjs b/crates/tui/src/runtime_web/app.mjs index 2ea0576aec..1902c8e3eb 100644 --- a/crates/tui/src/runtime_web/app.mjs +++ b/crates/tui/src/runtime_web/app.mjs @@ -632,6 +632,56 @@ function appendItemDelta(state, itemId, payload) { }); } +const PROVIDER_MODELS_PAGE_SIZE = 250; +const MAX_PROVIDER_MODELS = 10_000; +const MAX_PROVIDER_MODEL_PAGES = Math.ceil( + MAX_PROVIDER_MODELS / PROVIDER_MODELS_PAGE_SIZE, +); + +/** + * Load every bounded page of one provider catalog. + * + * `fetchPage` is injected so the browser client can retain its authenticated + * Runtime API boundary and tests can prove catalogs larger than one page are + * not silently truncated. Cursors are opaque and may never repeat. + */ +export async function collectProviderModelPages(providerId, fetchPage) { + const provider = String(providerId || "").trim(); + if (!provider || typeof fetchPage !== "function") { + throw new Error("A provider and page loader are required."); + } + + const entries = []; + const seenCursors = new Set(); + let cursor = ""; + for (let page = 0; page < MAX_PROVIDER_MODEL_PAGES; page += 1) { + const query = new URLSearchParams({ limit: String(PROVIDER_MODELS_PAGE_SIZE) }); + if (cursor) query.set("cursor", cursor); + const response = await fetchPage( + `/v1/providers/${encodeURIComponent(provider)}/models?${query.toString()}`, + ); + if (String(response?.provider || "") !== provider) { + throw new Error("The Runtime returned a model page for a different provider."); + } + const pageEntries = Array.isArray(response?.models) ? response.models : []; + if (entries.length + pageEntries.length > MAX_PROVIDER_MODELS) { + throw new Error(`The provider catalog exceeds ${MAX_PROVIDER_MODELS} models.`); + } + entries.push(...pageEntries); + + const nextCursor = typeof response?.nextCursor === "string" + ? response.nextCursor.trim() + : ""; + if (!nextCursor) return entries; + if (pageEntries.length === 0 || seenCursors.has(nextCursor)) { + throw new Error("The Runtime returned a non-progressing model cursor."); + } + seenCursors.add(nextCursor); + cursor = nextCursor; + } + throw new Error(`The provider catalog exceeds ${MAX_PROVIDER_MODELS} models.`); +} + function startBrowserClient() { const dom = { shell: document.querySelector("#app-shell"), @@ -1856,11 +1906,17 @@ function startBrowserClient() { setNewThreadStatus("Loading models…"); syncNewThreadControls(); try { - const response = await api(`/v1/providers/${encodeURIComponent(provider.id)}/models`); + const modelEntries = await collectProviderModelPages(provider.id, async (path) => { + const page = await api(path); + if (generation !== app.newThreadGeneration || !dom.newThreadDialog.open) { + throw new Error("The model request was superseded."); + } + return page; + }); if (generation !== app.newThreadGeneration || !dom.newThreadDialog.open) return; const seen = new Set(); const models = []; - for (const entry of Array.isArray(response?.models) ? response.models : []) { + for (const entry of modelEntries) { const id = String(entry?.id || "").trim(); const key = id.toLowerCase(); if (!id || seen.has(key)) continue; diff --git a/crates/tui/src/runtime_web/app.test.mjs b/crates/tui/src/runtime_web/app.test.mjs index dfcae9e40d..07f9ac1914 100644 --- a/crates/tui/src/runtime_web/app.test.mjs +++ b/crates/tui/src/runtime_web/app.test.mjs @@ -16,6 +16,7 @@ import { describe, expect, it } from "vitest"; import { NO_TARGET, canReply, + collectProviderModelPages, refusalMessage, receiptPresentation, resolveApprovalTarget, @@ -26,6 +27,49 @@ import { workflowReceiptPresentation, } from "./app.mjs"; +describe("collectProviderModelPages", () => { + it("loads a 600-model catalog through every opaque page", async () => { + const all = Array.from({ length: 600 }, (_, index) => ({ + id: `openrouter/model-${String(index).padStart(3, "0")}`, + image_input: "unknown", + })); + const cursors = new Map([ + ["", { start: 0, nextCursor: "page-2" }], + ["page-2", { start: 250, nextCursor: "page-3" }], + ["page-3", { start: 500, nextCursor: "" }], + ]); + const paths = []; + const models = await collectProviderModelPages("openrouter", async (path) => { + paths.push(path); + const url = new URL(path, "http://runtime.local"); + const cursor = url.searchParams.get("cursor") || ""; + const page = cursors.get(cursor); + expect(url.searchParams.get("limit")).toBe("250"); + expect(page).toBeDefined(); + return { + provider: "openrouter", + models: all.slice(page.start, page.start + 250), + total: all.length, + ...(page.nextCursor ? { nextCursor: page.nextCursor } : {}), + }; + }); + + expect(models).toEqual(all); + expect(paths).toHaveLength(3); + }); + + it("rejects a repeated or non-progressing cursor", async () => { + await expect( + collectProviderModelPages("openrouter", async () => ({ + provider: "openrouter", + models: [{ id: "model-a", image_input: "unknown" }], + total: 2, + nextCursor: "same-page", + })), + ).rejects.toThrow("non-progressing model cursor"); + }); +}); + describe("receiptPresentation", () => { it("keeps a failed MCP transport compact while preserving the raw receipt", () => { const raw = "Failed to connect MCP server 'github': Stdio transport closed MCP server stderr (last 1 line): Docker is not running"; diff --git a/crates/tui/src/session_manager.rs b/crates/tui/src/session_manager.rs index 4be9274bdf..327b67c43b 100644 --- a/crates/tui/src/session_manager.rs +++ b/crates/tui/src/session_manager.rs @@ -22,7 +22,7 @@ use crate::work_graph::ReasoningEffortTier; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use std::collections::{BTreeMap, BTreeSet}; -use std::fs; +use std::fs::{self, OpenOptions}; use std::io; use std::path::{Component, Path, PathBuf}; use uuid::Uuid; @@ -39,6 +39,45 @@ const MAX_SESSION_GOAL_OBJECTIVE_CHARS: usize = 8_192; const MAX_SESSION_GOAL_FILE_BYTES: u64 = 64 * 1_024; const CURRENT_SESSION_SCHEMA_VERSION: u32 = 1; const CURRENT_QUEUE_SCHEMA_VERSION: u32 = 1; +const LATE_USAGE_DIR: &str = ".late-usage"; +const CURRENT_LATE_USAGE_SCHEMA_VERSION: u32 = 1; +const MAX_LATE_USAGE_RECORDS_PER_SESSION: usize = 64; +const MAX_LATE_USAGE_LEDGER_BYTES: u64 = 1024 * 1024; + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct LateUsageRecord { + source_fingerprint: String, + turn_fingerprint: String, + route: crate::cost_status::EffectiveRouteEnvelope, + #[serde(default, skip_serializing_if = "Option::is_none")] + usage: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct LateUsageLedger { + schema_version: u32, + #[serde(default)] + records: Vec, + #[serde(default)] + overflowed: bool, +} + +impl Default for LateUsageLedger { + fn default() -> Self { + Self { + schema_version: CURRENT_LATE_USAGE_SCHEMA_VERSION, + records: Vec::new(), + overflowed: false, + } + } +} + +fn is_sha256_fingerprint(value: &str) -> bool { + value.len() == 64 + && value + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) +} const fn default_session_schema_version() -> u32 { CURRENT_SESSION_SCHEMA_VERSION @@ -73,6 +112,119 @@ fn normalize_managed_dir(path: PathBuf) -> std::io::Result { std::env::current_dir().map(|cwd| cwd.join(path)) } +fn open_private_lock_file(path: &Path) -> io::Result { + let mut options = OpenOptions::new(); + options.create(true).read(true).write(true); + #[cfg(unix)] + { + use std::os::unix::fs::{OpenOptionsExt, PermissionsExt}; + options + .mode(0o600) + .custom_flags(libc::O_NOFOLLOW | libc::O_CLOEXEC | libc::O_NONBLOCK); + let file = options.open(path)?; + validate_private_regular_file(&file, path)?; + file.set_permissions(fs::Permissions::from_mode(0o600))?; + Ok(file) + } + #[cfg(windows)] + { + use std::os::windows::fs::OpenOptionsExt as _; + use windows_sys::Win32::Storage::FileSystem::FILE_FLAG_OPEN_REPARSE_POINT; + options.custom_flags(FILE_FLAG_OPEN_REPARSE_POINT); + let file = options.open(path)?; + validate_private_regular_file(&file, path)?; + Ok(file) + } + #[cfg(all(not(unix), not(windows)))] + { + let file = options.open(path)?; + validate_private_regular_file(&file, path)?; + Ok(file) + } +} + +fn open_private_read_file(path: &Path) -> io::Result { + let mut options = OpenOptions::new(); + options.read(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt as _; + options.custom_flags(libc::O_NOFOLLOW | libc::O_CLOEXEC | libc::O_NONBLOCK); + } + #[cfg(windows)] + { + use std::os::windows::fs::OpenOptionsExt as _; + use windows_sys::Win32::Storage::FileSystem::FILE_FLAG_OPEN_REPARSE_POINT; + options.custom_flags(FILE_FLAG_OPEN_REPARSE_POINT); + } + let file = options.open(path)?; + validate_private_regular_file(&file, path)?; + Ok(file) +} + +#[cfg(unix)] +fn validate_private_regular_file(file: &fs::File, path: &Path) -> io::Result<()> { + use std::os::unix::fs::MetadataExt as _; + + let metadata = file.metadata()?; + if !metadata.is_file() || metadata.nlink() != 1 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "private sidecar file {} must be one regular filesystem link", + path.display() + ), + )); + } + Ok(()) +} + +#[cfg(windows)] +fn validate_private_regular_file(file: &fs::File, path: &Path) -> io::Result<()> { + use std::os::windows::fs::MetadataExt as _; + use std::os::windows::io::AsRawHandle as _; + use windows_sys::Win32::Storage::FileSystem::{ + BY_HANDLE_FILE_INFORMATION, FILE_ATTRIBUTE_REPARSE_POINT, GetFileInformationByHandle, + }; + + let metadata = file.metadata()?; + if !metadata.is_file() || metadata.file_attributes() & FILE_ATTRIBUTE_REPARSE_POINT != 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "private sidecar file {} must be a non-reparse regular file", + path.display() + ), + )); + } + let mut info = BY_HANDLE_FILE_INFORMATION::default(); + // SAFETY: `file` keeps the handle valid and `info` is writable for the call. + if unsafe { GetFileInformationByHandle(file.as_raw_handle(), &mut info) } == 0 { + return Err(io::Error::last_os_error()); + } + if info.nNumberOfLinks != 1 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "private sidecar file {} must have exactly one filesystem link", + path.display() + ), + )); + } + Ok(()) +} + +#[cfg(all(not(unix), not(windows)))] +fn validate_private_regular_file(file: &fs::File, path: &Path) -> io::Result<()> { + if !file.metadata()?.is_file() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("private sidecar file {} must be regular", path.display()), + )); + } + Ok(()) +} + /// Persisted queued message for offline/degraded mode. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct QueuedSessionMessage { @@ -436,6 +588,45 @@ pub struct SessionCostSnapshot { } impl SessionCostSnapshot { + fn absorb_late_background_cost(&mut self, pool: &crate::cost_status::PendingBackgroundCost) { + let estimate = crate::pricing::CostEstimate { + usd: self.subagent_cost_usd, + cny: self.subagent_cost_cny, + } + .saturating_add(pool.estimate); + self.subagent_cost_usd = estimate.usd; + self.subagent_cost_cny = estimate.cny; + self.priced_turns = self.priced_turns.saturating_add(pool.priced_turns); + self.unpriced_turns = self.unpriced_turns.saturating_add(pool.unpriced_turns); + self.cny_priced_turns = self.cny_priced_turns.saturating_add(pool.cny_priced_turns); + self.cny_unpriced_turns = self + .cny_unpriced_turns + .saturating_add(pool.cny_unpriced_turns); + self.unpriced_reasons + .extend(pool.unpriced_reasons.iter().map(ToString::to_string)); + self.cny_unpriced_reasons + .extend(pool.cny_unpriced_reasons.iter().map(ToString::to_string)); + self.unpriced_classes + .extend(pool.unpriced_classes.iter().map(ToString::to_string)); + self.pricing_provenances + .extend(pool.pricing_provenances.iter().map(ToString::to_string)); + self.live_pricing_defects + .extend(pool.live_pricing_defects.iter().map(ToString::to_string)); + self.live_pricing_unusable_defects.extend( + pool.live_pricing_unusable_defects + .iter() + .map(ToString::to_string), + ); + self.route_receipts + .extend(pool.route_receipts.iter().cloned()); + self.usage_source_fingerprints + .extend(pool.usage_source_fingerprints.iter().cloned()); + self.coverage_recorded = true; + let total = self.total_estimate(); + self.displayed_cost_high_water_usd = self.displayed_cost_high_water_usd.max(total.usd); + self.displayed_cost_high_water_cny = self.displayed_cost_high_water_cny.max(total.cny); + } + /// Session + subagent spend as **one** dual-currency accumulator. /// /// The persisted USD and CNY columns are projections of per-turn @@ -1039,6 +1230,224 @@ impl SessionManager { &self.sessions_dir } + fn late_usage_paths(&self, session_id: &str) -> io::Result<(PathBuf, PathBuf)> { + let session_id = self.validated_session_id(session_id)?; + let dir = self.sessions_dir.join(LATE_USAGE_DIR); + fs::create_dir_all(&dir)?; + let metadata = fs::symlink_metadata(&dir)?; + if metadata.file_type().is_symlink() || !metadata.is_dir() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "late usage store {} must be a real directory", + dir.display() + ), + )); + } + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + fs::set_permissions(&dir, fs::Permissions::from_mode(0o700))?; + } + Ok(( + dir.join(format!("{session_id}.json")), + dir.join(format!("{session_id}.lock")), + )) + } + + fn load_late_usage_unlocked(path: &Path) -> io::Result { + let file = match open_private_read_file(path) { + Ok(file) => file, + Err(error) if error.kind() == io::ErrorKind::NotFound => { + return Ok(LateUsageLedger::default()); + } + Err(error) => return Err(error), + }; + let metadata = file.metadata()?; + if metadata.len() > MAX_LATE_USAGE_LEDGER_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "late usage ledger {} exceeds its size bound", + path.display() + ), + )); + } + use std::io::Read as _; + let mut raw = Vec::with_capacity( + usize::try_from(metadata.len().min(MAX_LATE_USAGE_LEDGER_BYTES)).unwrap_or(0), + ); + file.take(MAX_LATE_USAGE_LEDGER_BYTES.saturating_add(1)) + .read_to_end(&mut raw)?; + if u64::try_from(raw.len()).unwrap_or(u64::MAX) > MAX_LATE_USAGE_LEDGER_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "late usage ledger {} exceeds its size bound", + path.display() + ), + )); + } + let ledger: LateUsageLedger = serde_json::from_slice(&raw) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + if ledger.schema_version != CURRENT_LATE_USAGE_SCHEMA_VERSION + || ledger.records.len() > MAX_LATE_USAGE_RECORDS_PER_SESSION + || ledger.records.iter().any(|record| { + !is_sha256_fingerprint(&record.source_fingerprint) + || !is_sha256_fingerprint(&record.turn_fingerprint) + }) + { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "late usage ledger has an unsupported or unbounded shape", + )); + } + Ok(ledger) + } + + fn persist_late_usage_record( + &self, + session_id: &str, + turn_id: &str, + source_id: &str, + route: &crate::cost_status::EffectiveRouteEnvelope, + usage: Option<&crate::models::Usage>, + ) -> io::Result { + let (path, lock_path) = self.late_usage_paths(session_id)?; + let lock_file = open_private_lock_file(&lock_path)?; + let mut lock = fd_lock::RwLock::new(lock_file); + let _guard = lock.write()?; + let mut ledger = Self::load_late_usage_unlocked(&path)?; + let source_fingerprint = crate::cost_status::usage_source_fingerprint(source_id); + if ledger + .records + .iter() + .any(|record| record.source_fingerprint == source_fingerprint) + { + return Ok(true); + } + if ledger.records.len() == MAX_LATE_USAGE_RECORDS_PER_SESSION { + if !ledger.overflowed { + ledger.overflowed = true; + let bytes = serde_json::to_vec(&ledger) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + write_atomic(&path, &bytes)?; + } + return Ok(true); + } + ledger.records.push(LateUsageRecord { + source_fingerprint, + turn_fingerprint: crate::cost_status::usage_source_fingerprint(turn_id), + route: route.sanitized_for_persistence(), + usage: usage.cloned(), + }); + let bytes = serde_json::to_vec(&ledger) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + write_atomic(&path, &bytes)?; + Ok(true) + } + + pub(crate) fn persist_late_runtime_usage( + &self, + session_id: &str, + turn_id: &str, + record: &crate::cost_status::RuntimeUsageRecord, + ) -> io::Result { + self.persist_late_usage_record( + session_id, + turn_id, + &record.source_id, + &record.usage.route, + Some(&record.usage.usage), + ) + } + + pub(crate) fn persist_late_runtime_drop( + &self, + session_id: &str, + turn_id: &str, + record: &crate::cost_status::RuntimeUsageDropRecord, + ) -> io::Result { + self.persist_late_usage_record(session_id, turn_id, &record.source_id, &record.route, None) + } + + fn load_late_usage(&self, session_id: &str) -> io::Result { + let (path, lock_path) = self.late_usage_paths(session_id)?; + let lock_file = open_private_lock_file(&lock_path)?; + let lock = fd_lock::RwLock::new(lock_file); + let _guard = lock.read()?; + Self::load_late_usage_unlocked(&path) + } + + fn apply_late_usage_to_metadata(&self, metadata: &mut SessionMetadata) -> io::Result<()> { + let ledger = self.load_late_usage(&metadata.id)?; + for record in ledger.records { + let source_fingerprint = record.source_fingerprint.clone(); + let source_id = format!("late:{}", record.source_fingerprint); + let mut pending = if let Some(usage) = record.usage.as_ref() { + crate::cost_status::background_cost_for_runtime_usage( + &crate::cost_status::RuntimeUsageRecord { + source_id, + usage: crate::cost_status::EffectiveRouteUsage { + route: record.route, + usage: usage.clone(), + }, + }, + ) + } else { + crate::cost_status::background_cost_for_runtime_drop( + &crate::cost_status::RuntimeUsageDropRecord { + source_id, + route: record.route, + }, + ) + }; + // The sidecar already stores the canonical SHA-256 identity. Do + // not hash it again while projecting the receipt into the saved + // session, or a concurrent main-snapshot writer that already + // contains the response would not dedupe against this overlay. + pending.usage_source_fingerprints.clear(); + pending + .usage_source_fingerprints + .insert(source_fingerprint.clone()); + if metadata + .cost + .usage_source_fingerprints + .contains(&source_fingerprint) + { + continue; + } + if let Some(usage) = record.usage { + metadata.total_tokens = metadata + .total_tokens + .saturating_add(u64::from(usage.input_tokens)) + .saturating_add(u64::from(usage.output_tokens)); + } + metadata.cost.absorb_late_background_cost(&pending); + } + if ledger.overflowed { + let fingerprint = crate::cost_status::usage_source_fingerprint(&format!( + "late-usage-overflow:{}", + crate::cost_status::usage_source_fingerprint(&metadata.id) + )); + if metadata.cost.usage_source_fingerprints.insert(fingerprint) { + metadata.cost.unpriced_turns = metadata.cost.unpriced_turns.saturating_add(1); + metadata.cost.cny_unpriced_turns = + metadata.cost.cny_unpriced_turns.saturating_add(1); + metadata + .cost + .unpriced_reasons + .insert("late_usage_ledger_overflow".to_string()); + metadata + .cost + .cny_unpriced_reasons + .insert("late_usage_ledger_overflow".to_string()); + metadata.cost.coverage_recorded = true; + } + } + Ok(()) + } + /// Persist the bounded goal control state for one saved session. /// `None` is the canonical clear operation and is idempotent. pub fn save_session_goal( @@ -1448,6 +1857,7 @@ impl SessionManager { session.system_prompt = strip_legacy_truncation_note(session.system_prompt); session.ensure_journal(); self.hydrate_approval_receipts(&mut session)?; + self.apply_late_usage_to_metadata(&mut session.metadata)?; Ok(session) } @@ -1531,8 +1941,9 @@ impl SessionManager { let path = entry.path(); if path.extension().is_some_and(|ext| ext == "json") - && let Ok(session) = Self::load_session_metadata(&path) + && let Ok(mut session) = Self::load_session_metadata(&path) { + self.apply_late_usage_to_metadata(&mut session)?; sessions.push(session); } } @@ -2502,6 +2913,284 @@ mod tests { } } + #[test] + fn late_usage_sidecar_survives_stale_session_save_and_replays_once() { + let tmp = tempdir().expect("tempdir"); + let manager = SessionManager::new(tmp.path().join("sessions")).expect("manager"); + let old_id = "old-session"; + let new_id = "new-session"; + let old = create_saved_session_with_id_and_mode( + old_id.to_string(), + &[make_test_message("user", "old session")], + "deepseek-v4-flash", + tmp.path(), + 0, + None, + Some("agent"), + ); + let new = create_saved_session_with_id_and_mode( + new_id.to_string(), + &[make_test_message("user", "new session")], + "deepseek-v4-flash", + tmp.path(), + 0, + None, + Some("agent"), + ); + manager.save_session(&old).expect("save old"); + manager.save_session(&new).expect("save new"); + + let priced_route = crate::cost_status::EffectiveRouteEnvelope::capture( + None, + ApiProvider::Deepseek, + "deepseek", + "deepseek-v4-flash", + Some(crate::config::DEFAULT_DEEPSEEK_BASE_URL), + Utc::now(), + ); + let usage = crate::models::Usage { + input_tokens: 17, + output_tokens: 5, + ..crate::models::Usage::default() + }; + let usage_record = crate::cost_status::RuntimeUsageRecord { + source_id: "translation:old-turn:assistant:1".to_string(), + usage: crate::cost_status::EffectiveRouteUsage { + route: priced_route.clone(), + usage: usage.clone(), + }, + }; + let missing_record = crate::cost_status::RuntimeUsageDropRecord { + source_id: "advisor:old-turn:provider-response:0".to_string(), + route: priced_route, + }; + let mut subscription_route = missing_record.route.clone(); + subscription_route.billing_mode = crate::cost_status::RouteBillingMode::Subscription; + let subscription_missing = crate::cost_status::RuntimeUsageDropRecord { + source_id: "translation:old-turn:thinking:2".to_string(), + route: subscription_route, + }; + + for _ in 0..2 { + assert!( + manager + .persist_late_runtime_usage(old_id, "old-turn", &usage_record) + .expect("persist late usage") + ); + assert!( + manager + .persist_late_runtime_drop(old_id, "old-turn", &missing_record) + .expect("persist missing usage") + ); + assert!( + manager + .persist_late_runtime_drop(old_id, "old-turn", &subscription_missing) + .expect("persist subscription missing usage") + ); + } + + // A concurrent stale whole-session writer cannot erase the independent + // origin ledger. Loading overlays it once by stable response identity. + manager.save_session(&old).expect("stale old-session save"); + let first = manager.load_session_snapshot(old_id).expect("load old"); + let second = manager.load_session_snapshot(old_id).expect("replay old"); + for loaded in [&first, &second] { + assert_eq!(loaded.metadata.total_tokens, 22); + assert_eq!(loaded.metadata.cost.unpriced_turns, 1); + assert_eq!(loaded.metadata.cost.cny_unpriced_turns, 1); + assert_eq!(loaded.metadata.cost.usage_source_fingerprints.len(), 3); + assert!( + loaded + .metadata + .cost + .unpriced_reasons + .contains("provider_success_missing_usage") + ); + } + assert_eq!(first.metadata.cost.priced_turns, 1); + + let clean = manager.load_session_snapshot(new_id).expect("load new"); + assert_eq!(clean.metadata.total_tokens, 0); + assert_eq!(clean.metadata.cost.priced_turns, 0); + assert_eq!(clean.metadata.cost.unpriced_turns, 0); + assert!(clean.metadata.cost.usage_source_fingerprints.is_empty()); + + let ledger = fs::read_to_string( + manager + .sessions_dir() + .join(LATE_USAGE_DIR) + .join(format!("{old_id}.json")), + ) + .expect("late ledger"); + assert!(!ledger.contains("translation:old-turn")); + assert!(!ledger.contains(crate::config::DEFAULT_DEEPSEEK_BASE_URL)); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let ledger_dir = manager.sessions_dir().join(LATE_USAGE_DIR); + assert_eq!( + fs::metadata(&ledger_dir) + .expect("private sidecar directory") + .permissions() + .mode() + & 0o777, + 0o700 + ); + for path in [ + ledger_dir.join(format!("{old_id}.json")), + ledger_dir.join(format!("{old_id}.lock")), + ] { + assert_eq!( + fs::metadata(path) + .expect("private sidecar metadata") + .permissions() + .mode() + & 0o777, + 0o600 + ); + } + } + } + + #[test] + fn late_usage_sidecar_has_a_bounded_fail_closed_overflow() { + let tmp = tempdir().expect("tempdir"); + let manager = SessionManager::new(tmp.path().join("sessions")).expect("manager"); + let session_id = "bounded-session"; + let session = create_saved_session_with_id_and_mode( + session_id.to_string(), + &[make_test_message("user", "bounded session")], + "local-model", + tmp.path(), + 0, + None, + Some("agent"), + ); + manager.save_session(&session).expect("save bounded"); + let mut route = crate::cost_status::EffectiveRouteEnvelope::capture( + None, + ApiProvider::Custom, + "local-provider", + "local-model", + Some("http://127.0.0.1:11434/v1"), + Utc::now(), + ); + route.billing_mode = crate::cost_status::RouteBillingMode::Local; + for index in 0..=MAX_LATE_USAGE_RECORDS_PER_SESSION { + manager + .persist_late_runtime_usage( + session_id, + "bounded-turn", + &crate::cost_status::RuntimeUsageRecord { + source_id: format!("late-bounded:{index}"), + usage: crate::cost_status::EffectiveRouteUsage { + route: route.clone(), + usage: crate::models::Usage { + input_tokens: 1, + ..crate::models::Usage::default() + }, + }, + }, + ) + .expect("bounded append"); + } + + let loaded = manager + .load_session_snapshot(session_id) + .expect("load bounded"); + assert_eq!( + loaded.metadata.total_tokens, + u64::try_from(MAX_LATE_USAGE_RECORDS_PER_SESSION).unwrap_or(u64::MAX) + ); + assert_eq!(loaded.metadata.cost.unpriced_turns, 1); + assert!( + loaded + .metadata + .cost + .unpriced_reasons + .contains("late_usage_ledger_overflow") + ); + let ledger = manager.load_late_usage(session_id).expect("bounded ledger"); + assert_eq!(ledger.records.len(), MAX_LATE_USAGE_RECORDS_PER_SESSION); + assert!(ledger.overflowed); + } + + #[cfg(unix)] + #[test] + fn late_usage_sidecar_rejects_linked_lock_and_ledger_leaves() { + use std::os::unix::fs::symlink; + + let tmp = tempdir().expect("tempdir"); + let manager = SessionManager::new(tmp.path().join("sessions")).expect("manager"); + let session_id = "linked-sidecar-session"; + let (ledger_path, lock_path) = manager.late_usage_paths(session_id).expect("paths"); + let route = crate::cost_status::EffectiveRouteEnvelope::capture( + None, + ApiProvider::Deepseek, + "deepseek", + "deepseek-v4-flash", + Some(crate::config::DEFAULT_DEEPSEEK_BASE_URL), + Utc::now(), + ); + let record = crate::cost_status::RuntimeUsageRecord { + source_id: "linked-sidecar-response".to_string(), + usage: crate::cost_status::EffectiveRouteUsage { + route, + usage: crate::models::Usage { + input_tokens: 1, + ..crate::models::Usage::default() + }, + }, + }; + + let outside_lock = tmp.path().join("outside.lock"); + fs::write(&outside_lock, b"outside-lock").expect("outside lock"); + symlink(&outside_lock, &lock_path).expect("symlink lock"); + assert!( + manager + .persist_late_runtime_usage(session_id, "turn", &record) + .is_err(), + "a symlink lock leaf must fail closed" + ); + assert_eq!( + fs::read(&outside_lock).expect("outside lock unchanged"), + b"outside-lock" + ); + fs::remove_file(&lock_path).expect("remove lock symlink"); + + fs::hard_link(&outside_lock, &lock_path).expect("hard-linked lock"); + assert!( + manager + .persist_late_runtime_usage(session_id, "turn", &record) + .is_err(), + "a multiply linked lock leaf must fail closed" + ); + fs::remove_file(&lock_path).expect("remove hard-linked lock"); + + let outside_ledger = tmp.path().join("outside.json"); + fs::write( + &outside_ledger, + br#"{"schema_version":1,"records":[],"overflowed":false}"#, + ) + .expect("outside ledger"); + symlink(&outside_ledger, &ledger_path).expect("symlink ledger"); + assert!( + manager.load_late_usage(session_id).is_err(), + "a symlink ledger leaf must fail closed" + ); + fs::remove_file(&ledger_path).expect("remove ledger symlink"); + + fs::hard_link(&outside_ledger, &ledger_path).expect("hard-linked ledger"); + assert!( + manager.load_late_usage(session_id).is_err(), + "a multiply linked ledger leaf must fail closed" + ); + assert_eq!( + fs::read(&outside_ledger).expect("outside ledger unchanged"), + br#"{"schema_version":1,"records":[],"overflowed":false}"# + ); + } + #[test] fn session_goal_sidecar_round_trips_control_state_without_model_output() { let tmp = tempdir().expect("tempdir"); diff --git a/crates/tui/src/tools/rlm.rs b/crates/tui/src/tools/rlm.rs index 86e6b89f3a..2a1b03b3a2 100644 --- a/crates/tui/src/tools/rlm.rs +++ b/crates/tui/src/tools/rlm.rs @@ -45,8 +45,7 @@ const ALL_ACTIONS: &[&str] = &["session_objects", "open", "eval", "configure", " fn rlm_kernel_error_result( error: &str, elapsed: Duration, - route: &crate::cost_status::EffectiveRouteEnvelope, - usage: &crate::models::Usage, + usage_batch: &crate::cost_status::RuntimeUsageBatch, ) -> ToolResult { let mut metadata = json!({ // The registered tool is `rlm`; `eval` is its action. Naming a @@ -57,7 +56,7 @@ fn rlm_kernel_error_result( "duration_ms": elapsed.as_millis() as u64, "kernel_error": true, }); - crate::cost_status::attach_child_usage_metadata(&mut metadata, route, usage); + crate::cost_status::attach_child_usage_batch_metadata(&mut metadata, usage_batch); ToolResult::error(format!("rlm action='eval': {error}")).with_metadata(metadata) } @@ -463,16 +462,14 @@ impl RlmTool { }; let started = Instant::now(); - let (round, child_usage, child_route) = if let Some(client) = self.client.clone() { - let route = client.effective_route_envelope(&self.root_model, chrono::Utc::now()); + let (round, child_usage_batch) = if let Some(client) = self.client.clone() { let bridge = RlmBridge::new( Arc::new(client), self.root_model.clone(), config.sub_rlm_max_depth.min(HARD_SUB_RLM_DEPTH_CAP), ); - let usage_handle = bridge.usage_handle(); let round_result = kernel.run(code, Some(&bridge)).await; - let usage = usage_handle.lock().await.clone(); + let usage = bridge.usage_snapshot().await; let round = match round_result { Ok(round) => round, Err(error) => { @@ -485,18 +482,28 @@ impl RlmTool { return Ok(rlm_kernel_error_result( &error.to_string(), started.elapsed(), - &route, - &usage, + &crate::cost_status::RuntimeUsageBatch { + records: usage.records, + drop_records: usage.drop_records, + dropped_records: usage.dropped_records, + }, )); } }; - (round, usage, Some(route)) + ( + round, + crate::cost_status::RuntimeUsageBatch { + records: usage.records, + drop_records: usage.drop_records, + dropped_records: usage.dropped_records, + }, + ) } else { let round = kernel .run(code, None::<&RlmBridge>) .await .map_err(|e| ToolError::execution_failed(format!("rlm_eval: {e}")))?; - (round, Default::default(), None) + (round, crate::cost_status::RuntimeUsageBatch::default()) }; session.rpc_count = session.rpc_count.saturating_add(round.rpc_count); @@ -595,11 +602,10 @@ impl RlmTool { "tool": "rlm_eval", "duration_ms": started.elapsed().as_millis() as u64, }); - // RLM fans out dozens of child rounds, so an undercounted class here - // scales; report every billable class from the shared producer (#4318). - if let Some(route) = child_route.as_ref() { - crate::cost_status::attach_child_usage_metadata(&mut metadata, route, &child_usage); - } + // Every RLM provider call keeps its own dispatch timestamp and frozen + // quote. The preferred batch format prevents a fan-out from being + // retroactively priced as one aggregate call on the first route. + crate::cost_status::attach_child_usage_batch_metadata(&mut metadata, &child_usage_batch); Ok(ToolResult::json(&output) .map_err(|e| ToolError::execution_failed(e.to_string()))? @@ -964,25 +970,39 @@ mod tests { reasoning_replay_tokens: Some(7), ..Default::default() }; + let record = crate::cost_status::RuntimeUsageRecord { + source_id: "rlm:test:request:0".to_string(), + usage: crate::cost_status::EffectiveRouteUsage { + route: route.clone(), + usage: usage.clone(), + }, + }; + let drop_record = crate::cost_status::RuntimeUsageDropRecord { + source_id: "rlm:test:request:1".to_string(), + route: route.clone(), + }; let result = rlm_kernel_error_result( "kernel stdout closed", Duration::from_millis(11), - &route, - &usage, + &crate::cost_status::RuntimeUsageBatch { + records: vec![record], + drop_records: vec![drop_record], + dropped_records: 1, + }, ); assert!(!result.success); let metadata = result .metadata .expect("usage metadata on failed tool result"); - assert_eq!( - crate::cost_status::child_route_envelope_from_metadata(&metadata), - Some(route) - ); - assert_eq!( - crate::cost_status::child_usage_from_metadata(&metadata), - Some(usage) - ); + let batch = crate::cost_status::child_usage_records_from_metadata(&metadata) + .expect("preferred routed batch"); + assert_eq!(batch.dropped_records, 1); + assert_eq!(batch.records.len(), 1); + assert_eq!(batch.drop_records.len(), 1); + assert_eq!(batch.records[0].usage.route, route); + assert_eq!(batch.records[0].usage.usage, usage); + assert_eq!(batch.drop_records[0].route, route); } #[test] diff --git a/crates/tui/src/tools/subagent/advisor.rs b/crates/tui/src/tools/subagent/advisor.rs index 0132097b84..39153708e7 100644 --- a/crates/tui/src/tools/subagent/advisor.rs +++ b/crates/tui/src/tools/subagent/advisor.rs @@ -2,7 +2,7 @@ //! //! When enabled, the advisor wakes on turn boundaries, reads a bounded slice //! of recent tool calls from the session transcript, makes a concise LLM -//! advisory call (reusing the same `DeepSeekClient` as the parent turn), and +//! advisory call on an exactly resolved provider/model client, and //! emits an [`Event::AdvisoryNote`] fire-and-forget. //! //! Key design properties: @@ -26,6 +26,7 @@ use tokio::sync::mpsc; use tracing::debug; use crate::client::DeepSeekClient; +use crate::config::Config; use crate::core::events::Event; use crate::llm_client::LlmClient; use crate::models::Role; @@ -109,6 +110,60 @@ pub struct EmissionGuard { last_note_hash_at: Option, } +/// Accounting ownership captured while the originating turn is still live. +/// +/// Runtime turns retain their synchronous durable sink through the lease; +/// ordinary interactive turns fall back to the exact session cost generation +/// captured here. Neither path can spill into a later session. +#[derive(Debug)] +pub(crate) struct AdvisorUsageContext { + cost_scope: crate::cost_status::CostScopeToken, + runtime_usage_lease: Option, +} + +impl AdvisorUsageContext { + #[must_use] + pub(crate) fn capture(runtime_owner: Option<&str>) -> Self { + Self { + cost_scope: crate::cost_status::scope_token(), + runtime_usage_lease: runtime_owner + .and_then(crate::cost_status::acquire_runtime_usage_lease), + } + } + + fn report( + &self, + source_id: &str, + route: &crate::cost_status::EffectiveRouteEnvelope, + usage: &crate::models::Usage, + ) { + crate::cost_status::report_effective_route_for_runtime( + self.cost_scope, + self.runtime_usage_lease + .as_ref() + .map(crate::cost_status::RuntimeUsageLease::owner), + source_id, + route, + usage, + ); + } + + fn report_unreceipted( + &self, + source_id: &str, + route: &crate::cost_status::EffectiveRouteEnvelope, + ) { + crate::cost_status::report_unreceipted_provider_success( + self.cost_scope, + self.runtime_usage_lease + .as_ref() + .map(crate::cost_status::RuntimeUsageLease::owner), + source_id, + route, + ); + } +} + impl EmissionGuard { /// Create a fresh guard with no emission history. #[must_use] @@ -259,7 +314,9 @@ pub async fn run_advisor_for_turn( messages: Vec, config: AdvisorConfig, client: DeepSeekClient, + route_config: Config, session_model: String, + usage_context: AdvisorUsageContext, guard: std::sync::Arc>, tx_event: mpsc::Sender, ) { @@ -292,6 +349,15 @@ pub async fn run_advisor_for_turn( .clone() .unwrap_or_else(|| session_model.clone()); + let (client, model) = match exact_advisor_client(&route_config, client, &session_model, &model) + { + Ok(route) => route, + Err(error) => { + tracing::warn!(target: "advisor", "advisor route resolution failed for turn {turn_id}: {error}"); + return; + } + }; + let route = client.effective_route_envelope(&model, chrono::Utc::now()); let request = MessageRequest { model: model.clone(), messages: vec![Message { @@ -323,6 +389,19 @@ pub async fn run_advisor_for_turn( } }; + // A decoded provider response is billable even when its partial/empty + // content is rejected below or the emission guard suppresses a duplicate. + let usage_source_id = format!("advisor:{turn_id}:provider-response:0"); + if response.usage == crate::models::Usage::default() { + usage_context.report_unreceipted(&usage_source_id, &route); + tracing::warn!( + target: "advisor", + "advisor provider response omitted usage for turn {turn_id}; cost coverage is unknown" + ); + } else { + usage_context.report(&usage_source_id, &route, &response.usage); + } + if crate::models::is_incomplete_stop_reason(response.stop_reason.as_deref()) { tracing::warn!( target: "advisor", @@ -374,6 +453,62 @@ pub async fn run_advisor_for_turn( debug!(target: "advisor", "advisory note emitted for turn {turn_id} ({tool_call_count} tool calls reviewed)"); } +fn exact_advisor_client( + config: &Config, + parent_client: DeepSeekClient, + session_model: &str, + requested_model: &str, +) -> anyhow::Result<(DeepSeekClient, String)> { + if requested_model + .trim() + .eq_ignore_ascii_case(session_model.trim()) + { + return Ok((parent_client, session_model.trim().to_string())); + } + + if config.providers.as_ref().is_some_and(|providers| { + providers.custom.values().any(|provider| { + provider + .model + .as_deref() + .is_some_and(|model| model.trim().eq_ignore_ascii_case(requested_model.trim())) + }) + }) { + anyhow::bail!( + "advisor model `{}` belongs to a custom provider but no exact provider identity is carried", + requested_model.trim() + ); + } + + let selection = + crate::model_routing::resolve_explicit_route_with_inventory(config, requested_model); + let (provider, model) = if let Some(selection) = selection { + if selection.provider == crate::config::ApiProvider::Custom { + anyhow::bail!( + "advisor model `{}` resolved only to a custom provider kind without an exact provider identity", + requested_model.trim() + ); + } + (selection.provider, selection.model) + } else { + let candidates = + crate::model_routing::explicit_route_candidate_providers(config, requested_model); + if !candidates.is_empty() && !candidates.contains(&config.api_provider()) { + anyhow::bail!( + "advisor model `{}` is not owned by the originating provider and has no unique exact route", + requested_model.trim() + ); + } + (config.api_provider(), requested_model.trim().to_string()) + }; + let client = crate::route_runtime::resolve_runtime_route(config, provider, Some(&model)) + .map_err(anyhow::Error::msg)? + .validate() + .map(|route| route.client) + .map_err(anyhow::Error::msg)?; + Ok((client, model)) +} + fn hash_str(s: &str) -> u64 { let mut h = DefaultHasher::new(); s.hash(&mut h); @@ -385,7 +520,10 @@ fn hash_str(s: &str) -> u64 { #[cfg(test)] mod tests { use super::*; + use crate::config::{ProviderConfig, ProvidersConfig}; use std::time::Duration; + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; fn test_config() -> AdvisorConfig { AdvisorConfig { @@ -603,4 +741,228 @@ mod tests { "prompt must include the tool input" ); } + + #[test] + fn advisor_model_override_builds_the_owning_provider_client() { + let config = Config { + provider: Some("deepseek".to_string()), + providers: Some(ProvidersConfig { + deepseek: ProviderConfig { + api_key: Some("sk-deepseek-advisor-test".to_string()), + model: Some("deepseek-chat".to_string()), + ..ProviderConfig::default() + }, + zai: ProviderConfig { + api_key: Some("zai-advisor-test-key".to_string()), + model: Some(crate::config::DEFAULT_ZAI_MODEL.to_string()), + ..ProviderConfig::default() + }, + ..ProvidersConfig::default() + }), + ..Config::default() + }; + let parent = DeepSeekClient::new(&config).expect("parent client"); + let (advisor, resolved_model) = exact_advisor_client( + &config, + parent, + "deepseek-chat", + crate::config::DEFAULT_ZAI_MODEL, + ) + .expect("cross-provider advisor route"); + let route = advisor.effective_route_envelope(&resolved_model, chrono::Utc::now()); + + assert_eq!(route.provider, crate::config::ApiProvider::Zai); + assert_eq!(route.provider_identity, "zai"); + assert_eq!(route.model, crate::config::DEFAULT_ZAI_MODEL); + } + + #[test] + fn advisor_foreign_custom_override_fails_closed_without_exact_identity() { + let config = Config { + provider: Some("deepseek".to_string()), + providers: Some(ProvidersConfig { + deepseek: ProviderConfig { + api_key: Some("sk-deepseek-advisor-test".to_string()), + model: Some("deepseek-chat".to_string()), + ..ProviderConfig::default() + }, + custom: [( + "private-route".to_string(), + ProviderConfig { + api_key: Some("custom-advisor-test-key".to_string()), + base_url: Some("https://custom.invalid/v1".to_string()), + model: Some("private-advisor-model".to_string()), + ..ProviderConfig::default() + }, + )] + .into_iter() + .collect(), + ..ProvidersConfig::default() + }), + ..Config::default() + }; + let parent = DeepSeekClient::new(&config).expect("parent client"); + let error = + match exact_advisor_client(&config, parent, "deepseek-chat", "private-advisor-model") { + Ok(_) => panic!("generic custom kind cannot identify the exact foreign route"), + Err(error) => error, + }; + assert!( + error.to_string().contains("exact provider identity"), + "{error}" + ); + } + + #[test] + fn advisor_named_custom_a_cannot_route_model_owned_by_custom_b() { + let config = Config { + provider: Some("custom-a".to_string()), + providers: Some(ProvidersConfig { + custom: [ + ( + "custom-a".to_string(), + ProviderConfig { + api_key: Some("custom-a-advisor-test-key".to_string()), + base_url: Some("https://custom-a.invalid/v1".to_string()), + model: Some("custom-a-model".to_string()), + kind: Some("openai-compatible".to_string()), + ..ProviderConfig::default() + }, + ), + ( + "custom-b".to_string(), + ProviderConfig { + api_key: Some("custom-b-advisor-test-key".to_string()), + base_url: Some("https://custom-b.invalid/v1".to_string()), + model: Some("custom-b-model".to_string()), + kind: Some("openai-compatible".to_string()), + ..ProviderConfig::default() + }, + ), + ] + .into_iter() + .collect(), + ..ProvidersConfig::default() + }), + ..Config::default() + }; + let parent = DeepSeekClient::new(&config).expect("active custom-a client"); + let error = match exact_advisor_client(&config, parent, "custom-a-model", "custom-b-model") + { + Ok(_) => panic!("custom-b must not reuse custom-a's endpoint or credential"), + Err(error) => error, + }; + assert!( + error.to_string().contains("exact provider identity"), + "{error}" + ); + } + + async fn run_billed_advisor_fixture( + note: &str, + stop_reason: &str, + suppress_as_duplicate: bool, + include_usage: bool, + ) -> (crate::cost_status::PendingBackgroundCost, Option) { + let _scope = crate::cost_status::test_scope(); + let server = MockServer::start().await; + let mut provider_response = serde_json::json!({ + "id": "advisor-provider-response", + "model": "deepseek-chat", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": note}, + "finish_reason": stop_reason + }] + }); + if include_usage { + provider_response["usage"] = serde_json::json!({ + "prompt_tokens": 9, + "completion_tokens": 3, + "total_tokens": 12 + }); + } + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(provider_response)) + .expect(1) + .mount(&server) + .await; + let route_config = Config { + provider: Some("deepseek".to_string()), + providers: Some(ProvidersConfig { + deepseek: ProviderConfig { + api_key: Some("sk-deepseek-advisor-test".to_string()), + model: Some("deepseek-chat".to_string()), + ..ProviderConfig::default() + }, + ..ProvidersConfig::default() + }), + ..Config::default() + }; + let mut client = DeepSeekClient::new(&route_config).expect("advisor client"); + client.set_test_chat_transport_base_url(server.uri()); + let mut emission_guard = EmissionGuard::new(); + if suppress_as_duplicate { + emission_guard.record_emission(note); + } + let guard = std::sync::Arc::new(tokio::sync::Mutex::new(emission_guard)); + let (tx, mut rx) = mpsc::channel(1); + run_advisor_for_turn( + "advisor-turn".to_string(), + make_messages_with_n_tool_calls(1), + AdvisorConfig { + enabled: true, + max_tool_calls: 5, + rate_limit: Duration::ZERO, + dedup_window: Duration::from_secs(60), + model: None, + }, + client, + route_config, + "deepseek-chat".to_string(), + AdvisorUsageContext::capture(None), + guard, + tx, + ) + .await; + (crate::cost_status::drain(), rx.try_recv().ok()) + } + + #[tokio::test] + async fn advisor_incomplete_and_dedup_suppressed_responses_are_each_billed_once() { + let (incomplete, incomplete_event) = + run_billed_advisor_fixture("partial note", "max_tokens", false, true).await; + assert!(incomplete_event.is_none()); + assert_eq!( + incomplete + .priced_turns + .saturating_add(incomplete.unpriced_turns), + 1 + ); + + let (dedup, dedup_event) = + run_billed_advisor_fixture("same advisory", "stop", true, true).await; + assert!(dedup_event.is_none()); + assert_eq!(dedup.priced_turns.saturating_add(dedup.unpriced_turns), 1); + } + + #[tokio::test] + async fn advisor_provider_success_without_usage_marks_unknown_once() { + let (pending, event) = + run_billed_advisor_fixture("use a smaller focused slice", "stop", false, false).await; + + assert!( + event.is_some(), + "the semantic advisor response remains usable" + ); + assert_eq!(pending.priced_turns, 0); + assert_eq!(pending.unpriced_turns, 1); + assert_eq!(pending.cny_unpriced_turns, 1); + assert!( + pending + .unpriced_reasons + .contains("provider_success_missing_usage") + ); + } } diff --git a/crates/tui/src/tools/subagent/mailbox.rs b/crates/tui/src/tools/subagent/mailbox.rs index 0ad37788a8..793feaaa60 100644 --- a/crates/tui/src/tools/subagent/mailbox.rs +++ b/crates/tui/src/tools/subagent/mailbox.rs @@ -84,8 +84,9 @@ pub enum MailboxMessage { /// this across direct durability, mailbox replay, and restart dedupe. source_id: String, /// Immutable provider/model/billing evidence captured before the - /// child request was sent. - route: crate::cost_status::EffectiveRouteEnvelope, + /// child request was sent. Boxed: the envelope dwarfs every other + /// variant, and mailboxes queue many messages. + route: Box, /// Provider usage payload, including cache-hit/cache-miss fields. usage: Usage, }, @@ -141,7 +142,7 @@ impl MailboxMessage { Self::TokenUsage { agent_id: agent_id.into(), source_id: source_id.into(), - route, + route: Box::new(route), usage, } } @@ -618,7 +619,7 @@ mod tests { MailboxMessage::TokenUsage { agent_id: "a9".into(), source_id: "response-a9".into(), - route: test_route(ApiProvider::Deepseek, "deepseek-v4-flash"), + route: Box::new(test_route(ApiProvider::Deepseek, "deepseek-v4-flash")), usage: Usage { input_tokens: 100, output_tokens: 50, @@ -641,6 +642,7 @@ mod tests { model: "k3".to_string(), billing_surface: Some(crate::pricing::MOONSHOT_KIMI_CODE_BILLING_SURFACE.to_string()), endpoint_fingerprint: Some("a".repeat(64)), + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Subscription, dispatched_at: chrono::DateTime::::from_timestamp(1_234, 0) .expect("timestamp"), diff --git a/crates/tui/src/tools/subagent/mod.rs b/crates/tui/src/tools/subagent/mod.rs index 2d8b1fca16..ebe68e38ed 100644 --- a/crates/tui/src/tools/subagent/mod.rs +++ b/crates/tui/src/tools/subagent/mod.rs @@ -1071,6 +1071,17 @@ fn usage_total_tokens(usage: &Usage) -> u64 { u64::from(usage.input_tokens).saturating_add(u64::from(usage.output_tokens)) } +fn usage_has_reported_data(usage: &Usage) -> bool { + usage.input_tokens > 0 + || usage.output_tokens > 0 + || usage.prompt_cache_hit_tokens.is_some() + || usage.prompt_cache_miss_tokens.is_some() + || usage.prompt_cache_write_tokens.is_some() + || usage.reasoning_tokens.is_some() + || usage.reasoning_replay_tokens.is_some() + || usage.server_tool_use.is_some() +} + /// Convert an authoritative USD audit into the workflow IR's integer /// microdollar receipt. Route coverage stays on the cost-status path; this /// narrow projection deliberately preserves only a priced subtotal. @@ -1086,9 +1097,85 @@ fn priced_usd_microusd(audit: &crate::pricing::TurnCostAudit) -> Option { Some(microusd.round() as u64) } +/// Immutable accounting ownership of one sub-agent dispatch. +/// +/// Captured once, when the runtime is built (`SubAgentRuntime::new`: the +/// engine's per-turn runtime, its off-turn continuation runtime, the direct +/// Workflow runtime), and inherited unchanged by every `child_runtime` / +/// `background_runtime` descendant and by the child tool registry's guardian +/// runtime. It replaces the completion-time `cost_status::scope_token()` read +/// `record_provider_response_usage` used to perform: a provider response that +/// lands after `/new` or a session load settles against the session it was +/// dispatched from, never against whichever session is live when it arrives. +/// +/// Runtime-owned dispatches (a `RuntimeUsageLease`) already carry their origin +/// inside the owner sink. This origin governs the ownerless paths: off-turn +/// continuations, direct Workflow runtimes, and their missing-usage receipts. +#[derive(Debug, Clone)] +pub(crate) struct SubAgentAccountingOrigin { + /// Cost-scope generation at dispatch. The live pool rejects a stale token; + /// it is never re-derived at completion. + cost_scope: crate::cost_status::CostScopeToken, + /// Root session of the dispatch: the tool context's state namespace, the + /// same identity the terminal fan-in and coordination paths use. A receipt + /// whose scope has been retired is appended to this session's durable + /// late-usage sidecar. + session_id: String, +} + +impl SubAgentAccountingOrigin { + fn capture(context: &ToolContext) -> Self { + Self { + cost_scope: crate::cost_status::scope_token(), + session_id: context.state_namespace.clone(), + } + } + + /// Settle one reported provider response that has no runtime owner. A + /// live origin scope charges the interactive pool exactly once per source + /// id; a retired one appends the exact frozen receipt to the origin + /// session's sidecar, labelled by the child run that dispatched it. Either + /// way the replacement session never sees it. + fn report_ownerless_usage( + &self, + agent_id: &str, + source_id: &str, + route: &crate::cost_status::EffectiveRouteEnvelope, + usage: &Usage, + ) { + crate::cost_status::report_effective_route_for_interactive_origin( + self.cost_scope, + &self.session_id, + agent_id, + source_id, + route, + usage, + ); + } + + /// Same settlement for a provider-success response without a usage + /// payload: one route-aware missing-coverage receipt for the origin. + fn report_ownerless_missing_usage( + &self, + agent_id: &str, + source_id: &str, + route: &crate::cost_status::EffectiveRouteEnvelope, + ) { + crate::cost_status::report_unreceipted_for_interactive_origin( + self.cost_scope, + &self.session_id, + agent_id, + source_id, + route, + ); + } +} + /// Publish one child provider response into every projection that owns it. /// The stable source id is the shared exactly-once key: runtime/session cost /// and the durable worker record hash it with the same canonical function. +/// Accounting ownership is the runtime's dispatch-time +/// [`SubAgentAccountingOrigin`]; nothing here consults the live cost scope. async fn record_provider_response_usage( runtime: &SubAgentRuntime, agent_id: &str, @@ -1096,23 +1183,53 @@ async fn record_provider_response_usage( route: crate::cost_status::EffectiveRouteEnvelope, usage: &Usage, ) { - let priced_cost_microusd = priced_usd_microusd(&route.audit(usage)); - if let Some(lease) = runtime.runtime_usage_lease.as_ref() { - crate::cost_status::report_effective_route_for_runtime( - crate::cost_status::scope_token(), - Some(lease.owner()), - source_id, - &route, - usage, - ); - } - if let Some(mailbox) = runtime.mailbox.as_ref() { - let _ = mailbox.send(MailboxMessage::token_usage( - agent_id, - source_id, - route, - usage.clone(), - )); + let has_reported_usage = usage_has_reported_data(usage); + let priced_cost_microusd = has_reported_usage + .then(|| priced_usd_microusd(&route.audit(usage))) + .flatten(); + let origin = &runtime.accounting_origin; + let runtime_owner = runtime + .runtime_usage_lease + .as_ref() + .map(crate::cost_status::RuntimeUsageLease::owner); + if has_reported_usage { + if let Some(owner) = runtime_owner { + crate::cost_status::report_effective_route_for_runtime( + origin.cost_scope, + Some(owner), + source_id, + &route, + usage, + ); + } else if runtime.mailbox.is_none() { + // Off-turn continuations and direct Workflow runtimes have no + // turn mailbox or durable turn owner. Their provider responses + // belong to the session that dispatched them; do not silently + // retain tokens only in the worker projection. + origin.report_ownerless_usage(agent_id, source_id, &route, usage); + } + if let Some(mailbox) = runtime.mailbox.as_ref() { + let _ = mailbox.send(MailboxMessage::token_usage( + agent_id, + source_id, + route, + usage.clone(), + )); + } + } else { + // A decoded provider-success response with the default Usage shape is + // indistinguishable from an omitted payload. Preserve its exact frozen + // route as one missing-coverage receipt; do not also publish a legacy + // priced-zero mailbox message under the same logical response. + match runtime_owner { + Some(owner) => crate::cost_status::report_unreceipted_provider_success( + origin.cost_scope, + Some(owner), + source_id, + &route, + ), + None => origin.report_ownerless_missing_usage(agent_id, source_id, &route), + } } runtime.manager.write().await.record_worker_usage( agent_id, @@ -1122,6 +1239,27 @@ async fn record_provider_response_usage( ); } +/// One logical held child-tool call gets one guardian usage identity even if a +/// mailbox/monitor replays its receipt. Raw agent/tool ids can be model-owned, +/// so only this fixed-length digest crosses telemetry or persistence seams. +fn child_guardian_usage_source_id(agent_id: &str, tool_id: &str) -> String { + use sha2::{Digest as _, Sha256}; + + let mut digest = Sha256::new(); + for part in [ + b"codewhale:subagent-auto-review-guardian:v1".as_slice(), + agent_id.as_bytes(), + tool_id.as_bytes(), + ] { + digest.update((part.len() as u64).to_le_bytes()); + digest.update(part); + } + format!( + "subagent-guardian:{}", + crate::hashing::hex_bytes(digest.finalize()) + ) +} + fn refresh_usage_note(usage: &mut AgentRunUsage) { let worker_total = usage.total_tokens.unwrap_or(0); if let Some(limit) = usage.token_budget { @@ -2528,6 +2666,10 @@ pub struct SubAgentRuntime { /// runtimes clone this guard, keeping the sink alive after the parent UI /// mailbox closes until the final child response has been persisted. pub(crate) runtime_usage_lease: Option, + /// Dispatch-time accounting ownership for every provider response this + /// runtime or its descendants settle. Captured in [`Self::new`], cloned + /// by [`Self::child_runtime`], never refreshed from the live cost scope. + pub(crate) accounting_origin: SubAgentAccountingOrigin, /// Wakeup channel for this runtime's immediate parent (issue #756). For /// the engine's direct children this points at the engine turn loop. While /// a sub-agent is running, its tool registry swaps this for a local inbox @@ -2603,6 +2745,9 @@ impl SubAgentRuntime { event_tx: Option>, manager: SharedSubAgentManager, ) -> Self { + // Accounting ownership is fixed here, at dispatch, before any child + // provider request exists; see `SubAgentAccountingOrigin`. + let accounting_origin = SubAgentAccountingOrigin::capture(&context); Self { client, api_config: None, @@ -2630,6 +2775,7 @@ impl SubAgentRuntime { foreground_children: None, mailbox: None, runtime_usage_lease: None, + accounting_origin, parent_completion_tx: None, fork_context: None, mcp_pool: None, @@ -2967,6 +3113,7 @@ impl SubAgentRuntime { foreground_children: self.foreground_children.clone(), mailbox: self.mailbox.clone(), runtime_usage_lease: self.runtime_usage_lease.clone(), + accounting_origin: self.accounting_origin.clone(), parent_completion_tx: self.parent_completion_tx.clone(), fork_context: self.fork_context.clone(), mcp_pool: self.mcp_pool.clone(), @@ -14681,12 +14828,34 @@ impl SubAgentToolRegistry { let context_text = crate::tui::auto_review::build_reviewer_context(review_context, held_reason, input); + // Capture the child guardian's own immutable quote immediately before + // its provider future starts. It is a routed auxiliary call, never + // billable under the parent turn's model/quote. + let review_route = self + .gate_runtime + .client + .effective_route_envelope(self.gate_runtime.client.model(), chrono::Utc::now()); let review = consult_reviewer( &self.gate_runtime.client, &context_text, &self.gate_runtime.cancel_token, ) .await; + // A provider-success reply carries usage even when it is incomplete or + // semantically invalid. Record before interpreting the verdict so the + // fail-closed path cannot erase spend. Pre-dispatch cancellation and + // transport failure expose no usage and therefore mint no receipt. + if let Some(usage) = review.usage.as_ref() { + let source_id = child_guardian_usage_source_id(agent_id, tool_id); + record_provider_response_usage( + &self.gate_runtime, + agent_id, + &source_id, + review_route, + usage, + ) + .await; + } let risk = review.outcome.audit_risk(); let (verdict, reason) = match &review.outcome { ReviewerOutcome::Allow { reason, .. } => (ToolGateVerdict::Allowed, reason.clone()), diff --git a/crates/tui/src/tools/subagent/tests.rs b/crates/tui/src/tools/subagent/tests.rs index cc67ef1acd..56b28cc76d 100644 --- a/crates/tui/src/tools/subagent/tests.rs +++ b/crates/tui/src/tools/subagent/tests.rs @@ -2199,6 +2199,381 @@ async fn detached_interactive_usage_after_mailbox_seal_reaches_session_accountin } } +#[tokio::test] +async fn child_guardian_usage_source_is_sanitized_and_replay_idempotent() { + let _cost_scope = crate::cost_status::test_scope(); + let tmp = tempdir().expect("tempdir"); + let manager = Arc::new(RwLock::new(SubAgentManager::new( + tmp.path().to_path_buf(), + 1, + ))); + manager.write().await.register_worker_for_session( + make_worker_spec("agent_guardian", tmp.path().to_path_buf()), + "guardian-usage-session", + ); + + let runtime_owner = "interactive:guardian-usage-session:turn-parent"; + crate::cost_status::register_interactive_runtime_usage_sink( + runtime_owner, + crate::cost_status::scope_token(), + ); + let mut runtime = stub_runtime(); + runtime.manager = Arc::clone(&manager); + runtime.runtime_usage_lease = crate::cost_status::acquire_runtime_usage_lease(runtime_owner); + + let source_id = child_guardian_usage_source_id("agent_guardian", "tool-fixed"); + assert_eq!( + source_id, + child_guardian_usage_source_id("agent_guardian", "tool-fixed"), + "replaying one logical held call must preserve its accounting identity" + ); + assert_ne!( + source_id, + child_guardian_usage_source_id("agent_guardian", "tool-other") + ); + assert_eq!(source_id.len(), "subagent-guardian:".len() + 64); + assert!(!source_id.contains("agent_guardian")); + assert!(!source_id.contains("tool-fixed")); + + let route = crate::cost_status::EffectiveRouteEnvelope::capture( + None, + ApiProvider::Deepseek, + "deepseek-direct", + "deepseek-v4-flash", + Some(ApiProvider::Deepseek.default_base_url()), + chrono::Utc::now(), + ); + let usage = Usage { + input_tokens: 7, + output_tokens: 5, + ..Usage::default() + }; + record_provider_response_usage( + &runtime, + "agent_guardian", + &source_id, + route.clone(), + &usage, + ) + .await; + // A mailbox or durable-record replay must not charge the routed guardian + // call a second time in either the session or worker projection. + record_provider_response_usage(&runtime, "agent_guardian", &source_id, route, &usage).await; + + crate::cost_status::finish_runtime_usage_owner(runtime_owner); + let session_usage = crate::cost_status::drain(); + let fingerprint = crate::cost_status::usage_source_fingerprint(&source_id); + assert_eq!( + session_usage.usage_source_fingerprints, + [fingerprint.clone()].into() + ); + assert_eq!( + session_usage + .priced_turns + .saturating_add(session_usage.unpriced_turns), + 1, + "the routed guardian response contributes exactly one cost receipt" + ); + let worker = manager + .read() + .await + .get_worker_record("agent_guardian") + .expect("guardian worker record") + .clone(); + assert_eq!(worker.usage.input_tokens, Some(7)); + assert_eq!(worker.usage.output_tokens, Some(5)); + assert_eq!(worker.usage_source_fingerprints, [fingerprint].into()); +} + +#[tokio::test] +async fn ownerless_no_mailbox_provider_usage_reaches_accounting_once() { + let _cost_scope = crate::cost_status::test_scope(); + let tmp = tempdir().expect("tempdir"); + let manager = Arc::new(RwLock::new(SubAgentManager::new( + tmp.path().to_path_buf(), + 1, + ))); + manager.write().await.register_worker_for_session( + make_worker_spec("agent_direct", tmp.path().to_path_buf()), + "direct-usage-session", + ); + let mut runtime = stub_runtime(); + runtime.manager = Arc::clone(&manager); + assert!(runtime.runtime_usage_lease.is_none()); + assert!(runtime.mailbox.is_none()); + + let source_id = "subagent:agent_direct:step:1:response:direct"; + let route = crate::cost_status::EffectiveRouteEnvelope::capture( + None, + ApiProvider::Deepseek, + "deepseek-direct", + "deepseek-v4-flash", + Some(ApiProvider::Deepseek.default_base_url()), + chrono::Utc::now(), + ); + let usage = Usage { + input_tokens: 19, + output_tokens: 7, + ..Usage::default() + }; + for _ in 0..2 { + record_provider_response_usage(&runtime, "agent_direct", source_id, route.clone(), &usage) + .await; + } + + let session_usage = crate::cost_status::drain(); + let fingerprint = crate::cost_status::usage_source_fingerprint(source_id); + assert_eq!( + session_usage.usage_source_fingerprints, + [fingerprint.clone()].into() + ); + assert_eq!( + session_usage + .priced_turns + .saturating_add(session_usage.unpriced_turns), + 1, + "a replayed ownerless provider response must have one cost receipt" + ); + assert!(session_usage.estimate.is_positive()); + + let worker = manager + .read() + .await + .get_worker_record("agent_direct") + .expect("direct worker record") + .clone(); + assert_eq!(worker.usage.input_tokens, Some(19)); + assert_eq!(worker.usage.output_tokens, Some(7)); + assert_eq!(worker.usage_source_fingerprints, [fingerprint].into()); +} + +/// Dispatch-time accounting ownership must survive `/new`. An off-turn +/// continuation runtime (no turn owner, no mailbox) is dispatched in the +/// origin session; a later provider response, its monitor replay, and a +/// guardian reply without a usage payload all land after the origin scope has +/// been retired. Each must settle exactly once against the origin session's +/// durable sidecar and never against the replacement scope. +#[tokio::test] +async fn ownerless_child_usage_crossing_new_settles_to_its_dispatch_origin_once() { + let _env = crate::test_support::lock_test_env(); + let _cost_scope = crate::cost_status::test_scope(); + let tmp = tempdir().expect("tempdir"); + let home = tmp.path().join("home"); + let _codewhale_home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", &home); + let sessions = + crate::session_manager::SessionManager::default_location().expect("session store"); + assert!( + sessions.sessions_dir().starts_with(&home), + "the sidecar must resolve into the guarded temporary home, not the real store: {}", + sessions.sessions_dir().display() + ); + let origin_session_id = "origin-session"; + let replacement_session_id = "replacement-session"; + for session_id in [origin_session_id, replacement_session_id] { + let session = crate::session_manager::create_saved_session_with_id_and_mode( + session_id.to_string(), + &[], + "deepseek-v4-flash", + tmp.path(), + 0, + None, + Some("agent"), + ); + sessions.save_session(&session).expect("save session"); + } + + let agent_id = "agent_off_turn"; + let manager = Arc::new(RwLock::new(SubAgentManager::new( + tmp.path().to_path_buf(), + 1, + ))); + manager.write().await.register_worker_for_session( + make_worker_spec(agent_id, tmp.path().to_path_buf()), + origin_session_id, + ); + + // Dispatch in the origin session: the engine's off-turn continuation + // runtime carries neither a runtime owner lease nor a turn mailbox. + let mut root = stub_runtime(); + root.manager = Arc::clone(&manager); + root.context = ToolContext::new(tmp.path()).with_state_namespace(origin_session_id); + root.accounting_origin = SubAgentAccountingOrigin::capture(&root.context); + let child = root.background_runtime(); + assert!(child.runtime_usage_lease.is_none()); + assert!(child.mailbox.is_none()); + + let route = crate::cost_status::EffectiveRouteEnvelope::capture( + None, + ApiProvider::Deepseek, + "deepseek-direct", + "deepseek-v4-flash", + Some(ApiProvider::Deepseek.default_base_url()), + chrono::Utc::now(), + ); + let usage = Usage { + input_tokens: 19, + output_tokens: 7, + ..Usage::default() + }; + let first_source = "subagent:agent_off_turn:step:1:response:before-new"; + let late_source = "subagent:agent_off_turn:step:2:response:after-new"; + let missing_source = child_guardian_usage_source_id(agent_id, "tool-after-new"); + let first_fingerprint = crate::cost_status::usage_source_fingerprint(first_source); + let late_fingerprint = crate::cost_status::usage_source_fingerprint(late_source); + let missing_fingerprint = crate::cost_status::usage_source_fingerprint(&missing_source); + + // Step 1 settles while the origin scope is live. + record_provider_response_usage(&child, agent_id, first_source, route.clone(), &usage).await; + // `/new` retires the origin scope exactly as the TUI does, while the + // child keeps running. + let settled_origin = crate::cost_status::close_current_scope(); + assert_eq!( + settled_origin.usage_source_fingerprints, + [first_fingerprint.clone()].into() + ); + assert_eq!(settled_origin.priced_turns, 1); + assert!(settled_origin.estimate.is_positive()); + + // The next response, its monitor replay, and a guardian reply without a + // usage payload all arrive after the scope generation moved on. + for _ in 0..2 { + record_provider_response_usage(&child, agent_id, late_source, route.clone(), &usage).await; + record_provider_response_usage( + &child, + agent_id, + &missing_source, + route.clone(), + &Usage::default(), + ) + .await; + } + + let replacement_live = crate::cost_status::drain(); + assert!( + replacement_live.is_empty(), + "late origin receipts charged the replacement scope: {replacement_live:?}" + ); + + let origin = sessions + .load_session_snapshot(origin_session_id) + .expect("origin session"); + assert_eq!(origin.metadata.total_tokens, 26); + assert_eq!(origin.metadata.cost.priced_turns, 1); + assert_eq!(origin.metadata.cost.unpriced_turns, 1); + assert!( + origin + .metadata + .cost + .unpriced_reasons + .contains("provider_success_missing_usage") + ); + assert_eq!( + origin.metadata.cost.usage_source_fingerprints, + [late_fingerprint.clone(), missing_fingerprint.clone()].into() + ); + + let replacement = sessions + .load_session_snapshot(replacement_session_id) + .expect("replacement session"); + assert_eq!(replacement.metadata.total_tokens, 0); + assert_eq!(replacement.metadata.cost.priced_turns, 0); + assert_eq!(replacement.metadata.cost.unpriced_turns, 0); + assert!( + replacement + .metadata + .cost + .usage_source_fingerprints + .is_empty() + ); + + let worker = manager + .read() + .await + .get_worker_record(agent_id) + .expect("off-turn worker record") + .clone(); + assert_eq!(worker.usage.input_tokens, Some(38)); + assert_eq!(worker.usage.output_tokens, Some(14)); + assert_eq!( + worker.usage_source_fingerprints, + [first_fingerprint, late_fingerprint, missing_fingerprint].into() + ); +} + +#[tokio::test] +async fn provider_success_without_usage_records_one_route_aware_gap_and_no_zero_mail() { + let _cost_scope = crate::cost_status::test_scope(); + let tmp = tempdir().expect("tempdir"); + let manager = Arc::new(RwLock::new(SubAgentManager::new( + tmp.path().to_path_buf(), + 1, + ))); + manager.write().await.register_worker_for_session( + make_worker_spec("agent_missing_usage", tmp.path().to_path_buf()), + "missing-usage-session", + ); + + let runtime_owner = "interactive:missing-usage-session:turn-parent"; + crate::cost_status::register_interactive_runtime_usage_sink( + runtime_owner, + crate::cost_status::scope_token(), + ); + let (mailbox, mut mailbox_rx) = Mailbox::new(CancellationToken::new()); + let mut runtime = stub_runtime(); + runtime.manager = Arc::clone(&manager); + runtime.mailbox = Some(mailbox); + runtime.runtime_usage_lease = crate::cost_status::acquire_runtime_usage_lease(runtime_owner); + + let source_id = child_guardian_usage_source_id("agent_missing_usage", "tool-fixed"); + let route = crate::cost_status::EffectiveRouteEnvelope::capture( + None, + ApiProvider::Deepseek, + "deepseek-direct", + "deepseek-v4-flash", + Some(ApiProvider::Deepseek.default_base_url()), + chrono::Utc::now(), + ); + for _ in 0..2 { + record_provider_response_usage( + &runtime, + "agent_missing_usage", + &source_id, + route.clone(), + &Usage::default(), + ) + .await; + } + + assert!( + !mailbox_rx.has_pending(), + "missing usage must not also publish a priced-zero TokenUsage message" + ); + crate::cost_status::finish_runtime_usage_owner(runtime_owner); + let session_usage = crate::cost_status::drain(); + let fingerprint = crate::cost_status::usage_source_fingerprint(&source_id); + assert_eq!( + session_usage.usage_source_fingerprints, + [fingerprint.clone()].into() + ); + assert_eq!(session_usage.priced_turns, 0); + assert_eq!(session_usage.unpriced_turns, 1); + assert!( + session_usage + .unpriced_reasons + .contains("provider_success_missing_usage") + ); + + let worker = manager + .read() + .await + .get_worker_record("agent_missing_usage") + .expect("missing-usage worker record") + .clone(); + assert_eq!(worker.usage.total_tokens, Some(0)); + assert_eq!(worker.usage.cost_microusd, None); + assert_eq!(worker.usage_source_fingerprints, [fingerprint].into()); +} + /// Like [`delayed_chat_client`] but delays *every* attempt, so the per-step /// API timeout fires on the first call and on every retry — the shape needed /// to drive the timeout-retry budget to exhaustion. @@ -12947,6 +13322,7 @@ pub(crate) fn stub_runtime() -> SubAgentRuntime { api_key: Some("test-key".to_string()), ..crate::config::Config::default() }; + let accounting_origin = SubAgentAccountingOrigin::capture(&context); SubAgentRuntime { client: stub_client(), api_config: Some(std::sync::Arc::new(stub_config)), @@ -12971,6 +13347,7 @@ pub(crate) fn stub_runtime() -> SubAgentRuntime { foreground_children: None, mailbox: None, runtime_usage_lease: None, + accounting_origin, parent_agent_id: None, parent_completion_tx: None, fork_context: None, @@ -20518,7 +20895,9 @@ mod child_permission_gate { if let Some(client) = client { runtime.client = client; } - runtime.context = ToolContext::new(tmp.path().to_path_buf()); + let workspace = tmp.path().to_path_buf(); + runtime.context = + ToolContext::new(workspace.clone()).with_state_namespace("guardian-test-session"); runtime.context.auto_approve = auto_approve; runtime.allow_shell = true; runtime.event_tx = Some(tx); @@ -20528,6 +20907,15 @@ mod child_permission_gate { parent_can_prompt, ); let manager = Arc::clone(&runtime.manager); + { + let mut manager = manager + .try_write() + .expect("guardian test worker registry is uncontended"); + manager.register_worker_for_session( + make_worker_spec("agent_gate", workspace), + "guardian-test-session", + ); + } // Keep the tempdir alive for the registry's lifetime by leaking it // into the workspace path (tests are short-lived). std::mem::forget(tmp); @@ -20562,8 +20950,12 @@ mod child_permission_gate { out } - /// A chat-completions mock that answers every request with `content`. - async fn guardian_mock(content: &str) -> (wiremock::MockServer, DeepSeekClient) { + /// A chat-completions mock that answers every request with `content` and + /// the requested semantic completion state. + async fn guardian_mock_with_stop( + content: &str, + finish_reason: &str, + ) -> (wiremock::MockServer, DeepSeekClient) { use wiremock::matchers::method; use wiremock::{Mock, MockServer, ResponseTemplate}; let _ = rustls::crypto::ring::default_provider().install_default(); @@ -20576,7 +20968,7 @@ mod child_permission_gate { "choices": [{ "index": 0, "message": {"role": "assistant", "content": content}, - "finish_reason": "stop" + "finish_reason": finish_reason }], "usage": {"prompt_tokens": 9, "completion_tokens": 3, "total_tokens": 12} }))) @@ -20591,6 +20983,37 @@ mod child_permission_gate { (server, client) } + async fn guardian_mock(content: &str) -> (wiremock::MockServer, DeepSeekClient) { + guardian_mock_with_stop(content, "stop").await + } + + async fn guardian_mock_without_usage(content: &str) -> (wiremock::MockServer, DeepSeekClient) { + use wiremock::matchers::method; + use wiremock::{Mock, MockServer, ResponseTemplate}; + let _ = rustls::crypto::ring::default_provider().install_default(); + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "chatcmpl-guardian-missing-usage", + "object": "chat.completion", + "model": "deepseek-v4-pro", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": content}, + "finish_reason": "stop" + }] + }))) + .mount(&server) + .await; + let config = crate::config::Config { + api_key: Some("test-key".to_string()), + base_url: Some(server.uri()), + ..crate::config::Config::default() + }; + let client = DeepSeekClient::new(&config).expect("mock-backed client"); + (server, client) + } + fn unreachable_client() -> DeepSeekClient { let _ = rustls::crypto::ring::default_provider().install_default(); let config = crate::config::Config { @@ -20730,6 +21153,19 @@ mod child_permission_gate { #[cfg(not(windows))] const GUARDIAN_PIPELINE: &str = "echo built | cat"; + /// [`GUARDIAN_PIPELINE`] with a caller-chosen marker, for tests that + /// need to recognize their own output. + fn guardian_marker_pipeline(marker: &str) -> String { + #[cfg(windows)] + { + format!("echo {marker} | Out-String") + } + #[cfg(not(windows))] + { + format!("echo {marker} | cat") + } + } + #[tokio::test] async fn auto_review_consults_the_guardian_and_runs_an_allowed_call_with_a_receipt() { let (_server, client) = guardian_mock( @@ -20774,6 +21210,145 @@ mod child_permission_gate { assert_eq!(receipts[0].2.as_deref(), Some("high")); } + #[tokio::test] + async fn auto_review_guardian_semantic_failures_still_record_provider_usage() { + for (content, finish_reason, expected_reason) in [ + ("not a guardian verdict", "stop", "answer was unparseable"), + ( + r#"{"risk_level":"low","decision":"allow","reason":"truncated"}"#, + "length", + "answer was incomplete", + ), + ] { + let (_server, client) = guardian_mock_with_stop(content, finish_reason).await; + let (registry, mut rx, manager) = + worker_registry(ApprovalMode::Auto, false, true, Some(client)); + // Keep each provider-success case distinct from the process-wide + // deterministic response cache: this test is proving two actual + // semantic outcomes, not replaying the first cached answer. + let command = format!("{GUARDIAN_PIPELINE} # {finish_reason}"); + let err = registry + .execute("agent_gate", "bash", json!({"command": command})) + .await + .expect_err("a semantically unusable guardian response fails closed"); + assert!(err.to_string().contains(expected_reason), "{err}"); + + let worker = manager + .read() + .await + .get_worker_record("agent_gate") + .expect("guardian usage reaches the durable worker record") + .clone(); + assert_eq!(worker.usage.input_tokens, Some(9)); + assert_eq!(worker.usage.output_tokens, Some(3)); + assert_eq!( + worker.usage_source_fingerprints.len(), + 1, + "one provider-success guardian call records exactly once" + ); + + let receipts = drain_gate_receipts(&mut rx); + assert_eq!(receipts.len(), 1, "{receipts:?}"); + assert_eq!(receipts[0].0, ToolGate::AutoReviewGuardian); + assert_eq!(receipts[0].1, ToolGateVerdict::Unavailable); + } + } + + #[tokio::test] + async fn auto_review_guardian_success_without_usage_records_one_missing_coverage_receipt() { + let _cost_scope = crate::cost_status::test_scope(); + let (_server, client) = guardian_mock_without_usage( + r#"{"risk_level":"low","decision":"allow","reason":"missing usage regression"}"#, + ) + .await; + let (registry, mut rx, manager) = + worker_registry(ApprovalMode::Auto, false, true, Some(client)); + + let output = registry + .execute( + "agent_gate", + "bash", + json!({"command": guardian_marker_pipeline("guardian-missing-usage")}), + ) + .await + .expect("semantic guardian success remains usable"); + assert!(output.contains("guardian-missing-usage"), "{output}"); + + let pending = crate::cost_status::drain(); + assert_eq!(pending.priced_turns, 0); + assert_eq!(pending.unpriced_turns, 1); + assert!( + pending + .unpriced_reasons + .contains("provider_success_missing_usage") + ); + assert_eq!(pending.usage_source_fingerprints.len(), 1); + + let worker = manager + .read() + .await + .get_worker_record("agent_gate") + .expect("guardian missing usage reaches the worker ledger") + .clone(); + assert_eq!(worker.usage.total_tokens, Some(0)); + assert_eq!(worker.usage.cost_microusd, None); + assert_eq!(worker.usage_source_fingerprints.len(), 1); + + let receipts = drain_gate_receipts(&mut rx); + assert_eq!(receipts.len(), 1, "{receipts:?}"); + assert_eq!(receipts[0].0, ToolGate::AutoReviewGuardian); + assert_eq!(receipts[0].1, ToolGateVerdict::Allowed); + } + + #[tokio::test] + async fn auto_review_guardian_cache_hit_does_not_mint_a_second_usage_receipt() { + let (server, client) = guardian_mock( + r#"{"risk_level":"low","decision":"allow","reason":"bounded cache regression"}"#, + ) + .await; + let (registry, mut rx, manager) = + worker_registry(ApprovalMode::Auto, false, true, Some(client)); + let input = json!({"command": guardian_marker_pipeline("guardian-cache-regression")}); + + for _ in 0..2 { + let output = registry + .execute("agent_gate", "bash", input.clone()) + .await + .expect("both the provider verdict and its cache replay allow the call"); + assert!(output.contains("guardian-cache-regression"), "{output}"); + } + + let requests = server + .received_requests() + .await + .expect("guardian provider requests are recorded"); + assert_eq!( + requests.len(), + 1, + "the second verdict came from response cache" + ); + let worker = manager + .read() + .await + .get_worker_record("agent_gate") + .expect("guardian usage reaches the durable worker record") + .clone(); + assert_eq!(worker.usage.input_tokens, Some(9)); + assert_eq!(worker.usage.output_tokens, Some(3)); + assert_eq!( + worker.usage_source_fingerprints.len(), + 1, + "an all-zero cache replay is not a second provider call" + ); + let receipts = drain_gate_receipts(&mut rx); + assert_eq!(receipts.len(), 2, "both gate decisions remain auditable"); + assert!( + receipts + .iter() + .all(|receipt| receipt.1 == ToolGateVerdict::Allowed) + ); + } + #[tokio::test] async fn auto_review_with_an_unreachable_guardian_fails_closed_with_a_receipt() { let (registry, mut rx, _) = diff --git a/crates/tui/src/tui/app/tests.rs b/crates/tui/src/tui/app/tests.rs index 3d041ad652..5624ec8eb7 100644 --- a/crates/tui/src/tui/app/tests.rs +++ b/crates/tui/src/tui/app/tests.rs @@ -1030,6 +1030,7 @@ fn active_turn_zai_receipt_overrides_all_mutable_parallel_route_metadata() { billing: Some(crate::core::events::RouteBillingEnvelope { billing_surface: None, endpoint_fingerprint: None, + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Unknown, dispatched_at: chrono::Utc::now(), }), diff --git a/crates/tui/src/tui/model_picker.rs b/crates/tui/src/tui/model_picker.rs index 1860a30f25..8d03049dfb 100644 --- a/crates/tui/src/tui/model_picker.rs +++ b/crates/tui/src/tui/model_picker.rs @@ -20,7 +20,7 @@ use ratatui::{ widgets::{Block, Paragraph, Widget}, }; -use codewhale_config::catalog::CatalogSource; +use codewhale_config::catalog::{CatalogRefreshError, CatalogSource, CatalogStatus}; use codewhale_config::model_reference::ModelReferenceCard; use codewhale_config::pricing::OfferingPricing; @@ -37,7 +37,8 @@ use crate::model_registry; use crate::models_dev_live::{self, ModelsDevFreshness}; use crate::palette; use crate::provider_lake::{ - all_catalog_models_for_provider, catalog_offering_for_model, configured_providers, + all_catalog_models_for_provider_identity, catalog_offering_for_model, + catalog_offering_for_model_identity, configured_providers, }; use crate::settings::PinnedModel; use crate::tui::app::{App, ReasoningEffort}; @@ -1300,6 +1301,7 @@ fn picker_model_rows_for_app(app: &App, config: &Config) -> Vec push_provider_model_rows( &mut rows, app.api_provider, + (app.api_provider == ApiProvider::Custom).then(|| app.provider_identity_for_persistence()), active_model_ids, app.api_provider, config, @@ -1335,6 +1337,7 @@ fn picker_model_rows_for_app(app: &App, config: &Config) -> Vec push_provider_model_rows( &mut rows, provider, + None, model_ids, app.api_provider, config, @@ -1343,21 +1346,10 @@ fn picker_model_rows_for_app(app: &App, config: &Config) -> Vec ); } - // `ApiProvider::Custom` is shared by every named custom route. Preserve - // the concrete active route key on rows so exact pins cannot collide. - let active_custom_identity = (app.api_provider == ApiProvider::Custom) - .then(|| app.provider_identity_for_persistence().to_string()); - for row in &mut rows { - if row.provider == Some(ApiProvider::Custom) { - row.provider_identity = active_custom_identity.clone(); - } - } - for row in &mut rows { row.enabled = model_row_enabled_for_app(app, config, row); if let Some(pin) = app.pinned_models.iter().find(|pin| { - row_provider_identity(row) - .is_some_and(|provider| provider.eq_ignore_ascii_case(&pin.provider)) + row_provider_identity(row).is_some_and(|provider| provider == pin.provider) && row.id.eq_ignore_ascii_case(&pin.model) }) { let label = pin.label.as_deref().unwrap_or("pinned"); @@ -1371,13 +1363,17 @@ fn picker_model_rows_for_app(app: &App, config: &Config) -> Vec for pin in &app.pinned_models { let provider = ApiProvider::parse(&pin.provider).unwrap_or(ApiProvider::Custom); if rows.iter().any(|row| { - row_provider_identity(row) - .is_some_and(|identity| identity.eq_ignore_ascii_case(&pin.provider)) + row_provider_identity(row).is_some_and(|identity| identity == pin.provider) && row.id.eq_ignore_ascii_case(&pin.model) }) { continue; } - let metadata = effective_picker_metadata(config, Some(provider), &pin.model); + let metadata = effective_picker_metadata_for_identity( + config, + Some(provider), + Some(&pin.provider), + &pin.model, + ); // Bypass the ordinary `(enum provider, model)` de-duplication here: // two named Custom routes may intentionally expose the same model id. rows.push(ModelPickerRow { @@ -1409,11 +1405,13 @@ fn model_row_enabled_for_app(app: &App, config: &Config, row: &ModelPickerRow) - return true; } } - let provider_identity = if provider == app.api_provider { - app.provider_identity_for_persistence() - } else { - provider.as_str() - }; + let provider_identity = row_provider_identity(row).unwrap_or_else(|| { + if provider == app.api_provider { + app.provider_identity_for_persistence() + } else { + provider.as_str() + } + }); if app.provider_model_is_enabled(provider_identity, &row.id) || app .provider_models @@ -1453,6 +1451,7 @@ fn model_row_enabled_for_app(app: &App, config: &Config, row: &ModelPickerRow) - fn push_provider_model_rows( rows: &mut Vec, provider: ApiProvider, + provider_identity: Option<&str>, model_ids: Vec, active_provider: ApiProvider, config: &Config, @@ -1478,9 +1477,26 @@ fn push_provider_model_rows( None }; let codex_freshness = roster_entry.map(|_| codex_roster.freshness); - let metadata = - effective_picker_metadata_with_codex(config, Some(provider), &id, codex_metadata); - let mut hint = render_picker_model_hint(&id, Some(provider), &metadata, codex_freshness); + let metadata = effective_picker_metadata_with_codex( + config, + Some(provider), + provider_identity, + &id, + codex_metadata, + ); + let provider_catalog_receipt = provider_catalog_receipt_for_route( + provider, + provider_identity, + config, + metadata.source.as_ref(), + ); + let mut hint = render_picker_model_hint( + &id, + Some(provider), + &metadata, + codex_freshness, + provider_catalog_receipt.as_ref(), + ); hint = format!("{readiness_label} · {hint}"); if provider != active_provider { hint = format!("switch route · {hint}"); @@ -1490,6 +1506,7 @@ fn push_provider_model_rows( rows, id.clone(), Some(provider), + provider_identity.map(str::to_string), hint, metadata, selectable, @@ -1498,6 +1515,40 @@ fn push_provider_model_rows( } } +fn provider_catalog_receipt_for_route( + provider: ApiProvider, + provider_identity: Option<&str>, + config: &Config, + source: Option<&CatalogSource>, +) -> Option<(CatalogStatus, bool)> { + let identity = provider_identity.unwrap_or_else(|| provider.as_str()); + let owns_provider_catalog = matches!( + provider, + ApiProvider::Openrouter | ApiProvider::Telecomjs | ApiProvider::Edenai + ) || (provider == ApiProvider::Custom + && codewhale_config::provider_setup_template(identity) + .is_some_and(|template| template.is_compatible())); + if !owns_provider_catalog { + return None; + } + + let base_url = config.base_url_for_route_identity(provider, identity); + let endpoint_matches = match source { + Some(CatalogSource::Live { + base_url_fingerprint, + .. + }) => *base_url_fingerprint == codewhale_config::catalog::base_url_fingerprint(&base_url), + // A bundled/template fallback has no endpoint claim to compare. Its + // exact-scope status still matters: a first refresh failure must be + // visible even though no live row exists yet. + _ => true, + }; + Some(( + crate::provider_catalog_live::status_for_scope(identity, &base_url), + endpoint_matches, + )) +} + fn push_auto_model_row(rows: &mut Vec, app: &App, config: &Config, hint: &str) { let readiness = crate::provider_readiness::resolve_for_model( config, @@ -1512,6 +1563,7 @@ fn push_auto_model_row(rows: &mut Vec, app: &App, config: &Confi rows, "auto".to_string(), None, + None, format!("{} · {hint}", readiness.label()), metadata, selectable, @@ -1578,8 +1630,15 @@ fn push_configured_provider_model( } fn provider_catalog_model_ids(provider: ApiProvider) -> Vec { + provider_catalog_model_ids_for_identity(provider, None) +} + +fn provider_catalog_model_ids_for_identity( + provider: ApiProvider, + provider_identity: Option<&str>, +) -> Vec { let mut models = Vec::new(); - for id in all_catalog_models_for_provider(provider) { + for id in all_catalog_models_for_provider_identity(provider, provider_identity) { // The catalog describes the built-in provider route. A custom route's // endpoint-owned current/configured model is appended separately. push_model_id(&mut models, picker_visible_model_id(provider, &id, false)); @@ -1592,7 +1651,9 @@ fn provider_scoped_model_ids_for_app(app: &App, include_current_model: bool) -> // separate custom/current-model row. let mut models = Vec::new(); push_model_id(&mut models, "auto"); - for id in provider_catalog_model_ids(app.api_provider) { + let provider_identity = + (app.api_provider == ApiProvider::Custom).then(|| app.provider_identity_for_persistence()); + for id in provider_catalog_model_ids_for_identity(app.api_provider, provider_identity) { push_model_id(&mut models, &id); } @@ -1674,21 +1735,30 @@ fn push_model_row( rows: &mut Vec, id: String, provider: Option, + provider_identity: Option, hint: String, metadata: EffectivePickerMetadata, selectable: bool, blocked_reason: Option, ) { - if rows - .iter() - .any(|row| row.id == id && row.provider == provider) - { + if rows.iter().any(|row| { + row.id == id + && row.provider == provider + && match ( + row.provider_identity.as_deref(), + provider_identity.as_deref(), + ) { + (Some(left), Some(right)) => left == right, + (None, None) => true, + _ => false, + } + }) { return; } rows.push(ModelPickerRow { id, provider, - provider_identity: None, + provider_identity, hint, metadata, selectable, @@ -1770,9 +1840,13 @@ fn normalize_picker_search_text(text: &str) -> String { /// offer, the provider's own id — the `[providers.]` key the user would /// edit — supplies the discriminator, with the leading run it already shares /// with the display name removed so the suffix is the part that differs. -fn route_labels_for_rows(rows: &[&ModelPickerRow]) -> BTreeMap<&'static str, String> { +fn route_labels_for_rows(rows: &[&ModelPickerRow]) -> BTreeMap { let mut by_display: BTreeMap<&'static str, Vec> = BTreeMap::new(); - for provider in rows.iter().filter_map(|row| row.provider) { + for provider in rows + .iter() + .filter_map(|row| row.provider) + .filter(|provider| *provider != ApiProvider::Custom) + { let bucket = by_display.entry(provider.display_name()).or_default(); if !bucket.contains(&provider) { bucket.push(provider); @@ -1790,9 +1864,21 @@ fn route_labels_for_rows(rows: &[&ModelPickerRow]) -> BTreeMap<&'static str, Str // distinct. Some(None) | None => display.to_string(), }; - labels.insert(provider.as_str(), label); + labels.insert(provider.as_str().to_string(), label); } } + for row in rows + .iter() + .filter(|row| row.provider == Some(ApiProvider::Custom)) + { + let Some(identity) = row_provider_identity(row) else { + continue; + }; + let label = codewhale_config::provider_setup_template(identity) + .map(|template| template.display_name.to_string()) + .unwrap_or_else(|| identity.to_string()); + labels.entry(identity.to_string()).or_insert(label); + } labels } @@ -1845,11 +1931,15 @@ fn route_discriminator(display: &str, provider_id: &str) -> Option { /// DeepSeek has published both `deepseek` and `deepseek-thinking` as family /// values for its current V4 models, so keep its picker heading stable and /// provider-facing rather than exposing either implementation detail. -fn catalog_family_for(provider: ApiProvider, model_id: &str) -> Option { +fn catalog_family_for_identity( + provider: ApiProvider, + provider_identity: Option<&str>, + model_id: &str, +) -> Option { if provider == ApiProvider::Deepseek { return Some(provider.display_name().to_string()); } - crate::provider_lake::catalog_offering_for_model(provider, model_id) + catalog_offering_for_model_identity(provider, provider_identity, model_id) .and_then(|offering| offering.family) } @@ -2014,7 +2104,7 @@ fn row_provider_identity(row: &ModelPickerRow) -> Option<&str> { fn offering_for_row(row: &ModelPickerRow) -> Option { let provider = row.provider?; - catalog_offering_for_model(provider, &row.id) + catalog_offering_for_model_identity(provider, row.provider_identity.as_deref(), &row.id) } fn offering_fetched_at(row: &ModelPickerRow) -> u64 { @@ -2068,16 +2158,27 @@ fn effective_picker_metadata( provider: Option, id: &str, ) -> EffectivePickerMetadata { - effective_picker_metadata_with_codex(config, provider, id, None) + effective_picker_metadata_for_identity(config, provider, None, id) +} + +fn effective_picker_metadata_for_identity( + config: &Config, + provider: Option, + provider_identity: Option<&str>, + id: &str, +) -> EffectivePickerMetadata { + effective_picker_metadata_with_codex(config, provider, provider_identity, id, None) } fn effective_picker_metadata_with_codex( config: &Config, provider: Option, + provider_identity: Option<&str>, id: &str, codex_metadata: Option<&CodexModelMetadata>, ) -> EffectivePickerMetadata { - let offering = provider.and_then(|provider| catalog_offering_for_model(provider, id)); + let offering = provider + .and_then(|provider| catalog_offering_for_model_identity(provider, provider_identity, id)); let card = offering.as_ref().map(ModelReferenceCard::from_offering); let registry = model_registry::lookup(id); @@ -2228,6 +2329,7 @@ fn render_picker_model_hint( provider: Option, metadata: &EffectivePickerMetadata, codex_freshness: Option, + provider_catalog_receipt: Option<&(CatalogStatus, bool)>, ) -> String { debug_assert_ne!(id, "auto", "Auto rows use the context-aware picker hint"); @@ -2315,16 +2417,26 @@ fn render_picker_model_hint( PickerPricing::Known(label) => parts.push(label.clone()), PickerPricing::Unknown => parts.push("price unknown".to_string()), } + let provider_live_source = matches!(metadata.source.as_ref(), Some(CatalogSource::Live { .. })); match metadata.source.as_ref() { - Some(CatalogSource::Live { .. } | CatalogSource::ModelsDevLive { .. }) => { - parts.push("live".to_string()) + Some(CatalogSource::Live { .. }) => { + parts.push(provider_catalog_source_label(provider_catalog_receipt)) } + Some(CatalogSource::ModelsDevLive { .. }) => parts.push("live".to_string()), Some(CatalogSource::Bundled) => parts.push("bundled".to_string()), Some(CatalogSource::ConfigOverride | CatalogSource::UserOverride) => { parts.push("override".to_string()) } None => {} } + if !provider_live_source + && let Some((CatalogStatus::Failed { reason }, _)) = provider_catalog_receipt + { + parts.push(format!( + "refresh failed ({})", + catalog_refresh_error_label(*reason) + )); + } if provider == Some(ApiProvider::OpenaiCodex) { parts.push(match codex_freshness { Some(freshness) => freshness.picker_label().to_string(), @@ -2339,6 +2451,38 @@ fn render_picker_model_hint( } } +fn provider_catalog_source_label(receipt: Option<&(CatalogStatus, bool)>) -> String { + let Some((status, endpoint_matches)) = receipt else { + return "catalog freshness unknown".to_string(); + }; + if !endpoint_matches { + return "catalog from different endpoint".to_string(); + } + match status { + CatalogStatus::Fresh => "live".to_string(), + CatalogStatus::Stale { age_secs } => { + let age_hours = age_secs.saturating_add(3_599) / 3_600; + format!("stale catalog ({age_hours}h)") + } + CatalogStatus::Failed { reason } => { + format!("refresh failed ({})", catalog_refresh_error_label(*reason)) + } + CatalogStatus::Unknown => "catalog freshness unknown".to_string(), + } +} + +fn catalog_refresh_error_label(error: CatalogRefreshError) -> &'static str { + match error { + CatalogRefreshError::Unauthorized => "unauthorized", + CatalogRefreshError::Forbidden => "forbidden", + CatalogRefreshError::NotFound => "not found", + CatalogRefreshError::RateLimited => "rate limited", + CatalogRefreshError::InvalidResponse => "invalid response", + CatalogRefreshError::EmptyList => "empty list", + CatalogRefreshError::Network => "network error", + } +} + pub(crate) fn format_picker_context_window(tokens: u64) -> String { if tokens >= 1_000_000 { if tokens.is_multiple_of(1_000_000) { @@ -2736,11 +2880,15 @@ impl ModelPickerView { Some(provider) => PaneRow { primary: row.id.clone(), route: route_labels - .get(provider.as_str()) + .get(row_provider_identity(row).unwrap_or(provider.as_str())) .cloned() .unwrap_or_else(|| provider.display_name().to_string()), meta: model_row_meta_chips(row), - family: catalog_family_for(provider, &row.id), + family: catalog_family_for_identity( + provider, + row.provider_identity.as_deref(), + &row.id, + ), active, }, } @@ -2995,7 +3143,7 @@ mod tests { #[test] fn deepseek_picker_heading_hides_legacy_family_metadata() { assert_eq!( - catalog_family_for(ApiProvider::Deepseek, "deepseek-v4-pro").as_deref(), + catalog_family_for_identity(ApiProvider::Deepseek, None, "deepseek-v4-pro").as_deref(), Some("DeepSeek") ); } @@ -3007,4 +3155,125 @@ mod tests { assert!(model_row_visible_by_default(&row, ApiProvider::Deepseek)); assert!(!model_row_visible_by_default(&row, ApiProvider::Openai)); } + + #[test] + fn baseten_picker_models_use_exact_identity_and_direct_provider_label() { + let _live = crate::provider_lake::lock_live_snapshot(); + crate::provider_lake::clear_live_snapshot(); + + let models = provider_catalog_model_ids_for_identity( + ApiProvider::Custom, + Some(codewhale_config::BASETEN_TEMPLATE_ID), + ); + assert_eq!(models.len(), codewhale_config::BASETEN_MODELS.len()); + assert!(models.contains(&codewhale_config::BASETEN_DEFAULT_MODEL.to_string())); + + let row = ModelPickerRow { + id: codewhale_config::BASETEN_DEFAULT_MODEL.to_string(), + provider: Some(ApiProvider::Custom), + provider_identity: Some(codewhale_config::BASETEN_TEMPLATE_ID.to_string()), + hint: String::new(), + metadata: EffectivePickerMetadata::default(), + selectable: true, + blocked_reason: None, + enabled: true, + }; + let labels = route_labels_for_rows(&[&row]); + assert_eq!(labels.get("baseten").map(String::as_str), Some("Baseten")); + } + + #[test] + fn provider_catalog_hint_never_calls_failed_or_mismatched_rows_live() { + assert_eq!( + provider_catalog_source_label(Some(&(CatalogStatus::Fresh, true))), + "live" + ); + assert_eq!( + provider_catalog_source_label(Some(&( + CatalogStatus::Failed { + reason: CatalogRefreshError::Unauthorized, + }, + true, + ))), + "refresh failed (unauthorized)" + ); + assert_eq!( + provider_catalog_source_label(Some(&(CatalogStatus::Fresh, false))), + "catalog from different endpoint" + ); + assert_eq!( + provider_catalog_source_label(None), + "catalog freshness unknown" + ); + } + + #[test] + fn first_provider_catalog_failure_is_visible_on_bundled_fallback_rows() { + let _env = crate::test_support::lock_test_env(); + let _live = crate::provider_lake::lock_live_snapshot(); + let home = tempfile::tempdir().expect("test home"); + let _home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", home.path()); + crate::provider_catalog_live::reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + + let config = Config { + provider: Some("openrouter".to_string()), + ..Config::default() + }; + let base_url = config.base_url_for_route_identity(ApiProvider::Openrouter, "openrouter"); + let fingerprint = codewhale_config::catalog::base_url_fingerprint(&base_url); + crate::provider_catalog_live::record_failure( + "openrouter", + &fingerprint, + CatalogRefreshError::Unauthorized, + ); + + let model = provider_catalog_model_ids(ApiProvider::Openrouter) + .into_iter() + .next() + .expect("bundled OpenRouter fallback"); + let mut rows = Vec::new(); + let codex_roster = CodexModelRoster { + models: Vec::new(), + freshness: CodexModelCacheFreshness::Missing, + fetched_at: None, + }; + push_provider_model_rows( + &mut rows, + ApiProvider::Openrouter, + None, + vec![model], + ApiProvider::Openrouter, + &config, + &codex_roster, + &crate::provider_readiness::ProviderReadinessSnapshot::default(), + ); + assert_eq!(rows.len(), 1); + assert!( + rows[0].hint.contains("refresh failed (unauthorized)"), + "{}", + rows[0].hint + ); + + crate::provider_catalog_live::reset_cache_for_test(); + crate::provider_lake::clear_live_snapshot(); + } + + #[test] + fn model_rows_keep_case_distinct_custom_identities() { + let mut rows = Vec::new(); + for identity in ["CustomA", "customa"] { + push_model_row( + &mut rows, + "shared-model".to_string(), + Some(ApiProvider::Custom), + Some(identity.to_string()), + String::new(), + EffectivePickerMetadata::default(), + true, + None, + ); + } + assert_eq!(rows.len(), 2); + } } diff --git a/crates/tui/src/tui/prompt_suggestion.rs b/crates/tui/src/tui/prompt_suggestion.rs index 4204e27f28..e5de1fcd8a 100644 --- a/crates/tui/src/tui/prompt_suggestion.rs +++ b/crates/tui/src/tui/prompt_suggestion.rs @@ -1235,6 +1235,7 @@ mod tests { billing: Some(crate::core::events::RouteBillingEnvelope { billing_surface: None, endpoint_fingerprint: None, + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Unknown, dispatched_at: chrono::Utc::now(), }), @@ -1583,6 +1584,7 @@ mod tests { billing: Some(crate::core::events::RouteBillingEnvelope { billing_surface: None, endpoint_fingerprint: None, + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Unknown, dispatched_at: chrono::Utc::now(), }), diff --git a/crates/tui/src/tui/subagent_routing.rs b/crates/tui/src/tui/subagent_routing.rs index 3cb1d10700..2cbcde42c2 100644 --- a/crates/tui/src/tui/subagent_routing.rs +++ b/crates/tui/src/tui/subagent_routing.rs @@ -1331,7 +1331,10 @@ mod tests { &MailboxMessage::TokenUsage { agent_id: "agent_route".to_string(), source_id: "response-route".to_string(), - route: test_route(crate::config::ApiProvider::Openrouter, "vendor/model-real"), + route: Box::new(test_route( + crate::config::ApiProvider::Openrouter, + "vendor/model-real", + )), usage: crate::models::Usage::default(), }, ); @@ -1353,7 +1356,7 @@ mod tests { &MailboxMessage::TokenUsage { agent_id: "agent_spend".to_string(), source_id: "response-1".to_string(), - route: route.clone(), + route: Box::new(route.clone()), usage: crate::models::Usage { input_tokens: 1_000, output_tokens: 40, @@ -1367,7 +1370,7 @@ mod tests { &MailboxMessage::TokenUsage { agent_id: "agent_spend".to_string(), source_id: "response-2".to_string(), - route, + route: Box::new(route), usage: crate::models::Usage { input_tokens: 2_000, output_tokens: 60, diff --git a/crates/tui/src/tui/tool_routing.rs b/crates/tui/src/tui/tool_routing.rs index 92ed7b4b53..f3cf445b72 100644 --- a/crates/tui/src/tui/tool_routing.rs +++ b/crates/tui/src/tui/tool_routing.rs @@ -454,6 +454,32 @@ fn accrue_child_token_cost_if_any(app: &mut App, result: &Result 0 { + let dropped = u32::try_from(residual_dropped_records).unwrap_or(u32::MAX); + app.session.cost_unpriced_turns = + app.session.cost_unpriced_turns.saturating_add(dropped); + app.session.cost_cny_unpriced_turns = + app.session.cost_cny_unpriced_turns.saturating_add(dropped); + app.session + .cost_unpriced_reasons + .insert("routed_usage_receipt_missing".to_string()); + app.session + .cost_cny_unpriced_reasons + .insert("routed_usage_receipt_missing".to_string()); + } + return; + } let Some(route) = crate::cost_status::child_route_envelope_from_metadata(metadata) else { return; }; @@ -464,6 +490,13 @@ fn accrue_child_token_cost_if_any(app: &mut App, result: &Result bool { ) } -/// Translate text to the requested target language using a dedicated -/// translation agent. -/// -/// This is a lightweight, focused API call — no streaming, no tool calls, -/// no conversation history. The agent's only role is translation. -/// -/// # Errors -/// -/// Returns an error if the API call fails or the response is malformed. -pub async fn translate_text( - text: &str, - client: &DeepSeekClient, - model: &str, - target_language: &str, -) -> Result { - client.translate(text, model, target_language).await -} - /// Status of a translation operation for a single message. #[derive(Debug, Clone, PartialEq, Eq)] #[allow(dead_code)] diff --git a/crates/tui/src/tui/ui.rs b/crates/tui/src/tui/ui.rs index ab77f78597..17f118b092 100644 --- a/crates/tui/src/tui/ui.rs +++ b/crates/tui/src/tui/ui.rs @@ -270,15 +270,21 @@ type PendingToolUses = Vec<(String, String, serde_json::Value)>; #[derive(Debug)] enum TranslationEvent { AssistantMessage { + origin_session_fingerprint: Option, + origin_turn_fingerprint: Option, history_index: Option, original_text: String, translated: anyhow::Result, + usage: Option, thinking: Option, tool_uses: PendingToolUses, }, Thinking { + origin_session_fingerprint: Option, + origin_turn_fingerprint: Option, placeholder: String, translated: anyhow::Result, + usage: Option, }, } diff --git a/crates/tui/src/tui/ui/dispatch.rs b/crates/tui/src/tui/ui/dispatch.rs index b05befc5e2..9f7c848072 100644 --- a/crates/tui/src/tui/ui/dispatch.rs +++ b/crates/tui/src/tui/ui/dispatch.rs @@ -680,6 +680,10 @@ pub(crate) async fn spawned_dispatch_execute( completion_permit.send(apply); } +pub(super) fn settle_failed_dispatch_routed_usage(batch: &crate::cost_status::RuntimeUsageBatch) { + crate::cost_status::report_runtime_usage_batch(crate::cost_status::scope_token(), None, batch); +} + pub(crate) async fn spawned_dispatch_inner( prepare: UserDispatchPrepare, recovery: DispatchRecovery, @@ -722,6 +726,7 @@ pub(crate) async fn spawned_dispatch_inner( effective_reasoning_effort, auto_controls_reasoning, auto_selection, + initial_routed_usage, routing_source: _, } = planned; let effective_reasoning_tier = selected_reasoning_effort @@ -738,12 +743,18 @@ pub(crate) async fn spawned_dispatch_inner( &turn_route.model, ); + // Retain one fallback copy until mailbox acceptance. The classifier has + // already run; if the engine channel closes now, its exact spend still + // belongs to this interactive session rather than disappearing with the + // unstarted parent turn. + let send_failure_routed_usage = initial_routed_usage.clone(); if let Err(err) = engine_handle .send(Op::SendMessage { content: prepare.content.clone(), mode: prepare.mode, route: Box::new(turn_route), compaction: Box::new(turn_compaction.clone()), + initial_routed_usage: Box::new(initial_routed_usage), goal_objective: prepare.goal_objective.clone(), goal_token_budget: prepare.goal_token_budget, goal_status: prepare.goal_status, @@ -763,6 +774,7 @@ pub(crate) async fn spawned_dispatch_inner( }) .await { + settle_failed_dispatch_routed_usage(&send_failure_routed_usage); return build_dispatch_error_closure(prepare, recovery, err.to_string()); } diff --git a/crates/tui/src/tui/ui/event_loop.rs b/crates/tui/src/tui/ui/event_loop.rs index c02bf14a2e..a23817c5e8 100644 --- a/crates/tui/src/tui/ui/event_loop.rs +++ b/crates/tui/src/tui/ui/event_loop.rs @@ -20,6 +20,198 @@ pub(super) fn event_owner_is_active( !owner_session_id.is_empty() && current_session_id == Some(owner_session_id) } +#[derive(Debug)] +struct TranslationAccountingContext { + cost_scope: crate::cost_status::CostScopeToken, + origin_session_id: Option, + origin_turn_id: Option, + source_id: String, +} + +struct SettledTranslation { + translated: anyhow::Result, + usage: Option, +} + +impl TranslationAccountingContext { + fn capture(app: &App, kind: &str, sequence: u64) -> Self { + let raw_source = format!( + "translation:{}:{}:{kind}:{sequence}", + app.current_session_id.as_deref().unwrap_or("no-session"), + app.runtime_turn_id.as_deref().unwrap_or("no-turn") + ); + Self { + cost_scope: crate::cost_status::scope_token(), + origin_session_id: app.current_session_id.clone(), + origin_turn_id: app.runtime_turn_id.clone(), + source_id: format!( + "translation:{}", + crate::cost_status::usage_source_fingerprint(&raw_source) + ), + } + } + + fn settle( + self, + response: anyhow::Result, + ) -> SettledTranslation { + let response = match response { + Ok(response) => response, + Err(error) => { + return SettledTranslation { + translated: Err(error), + usage: None, + }; + } + }; + if let Some(usage) = response.usage.as_ref() { + if let (Some(session_id), Some(turn_id)) = ( + self.origin_session_id.as_deref(), + self.origin_turn_id.as_deref(), + ) { + crate::cost_status::report_effective_route_for_interactive_origin( + self.cost_scope, + session_id, + turn_id, + &self.source_id, + &response.route, + usage, + ); + } else { + crate::cost_status::report_effective_route_for_runtime( + self.cost_scope, + None, + &self.source_id, + &response.route, + usage, + ); + } + } else { + if let (Some(session_id), Some(turn_id)) = ( + self.origin_session_id.as_deref(), + self.origin_turn_id.as_deref(), + ) { + crate::cost_status::report_unreceipted_for_interactive_origin( + self.cost_scope, + session_id, + turn_id, + &self.source_id, + &response.route, + ); + } else { + crate::cost_status::report_unreceipted_provider_success( + self.cost_scope, + None, + &self.source_id, + &response.route, + ); + } + } + SettledTranslation { + translated: response.translated, + usage: response.usage, + } + } +} + +fn accrue_translation_usage(app: &mut App, usage: &crate::models::Usage) { + let turn_tokens = usage.input_tokens.saturating_add(usage.output_tokens); + app.session.total_tokens = app.session.total_tokens.saturating_add(turn_tokens); + app.session.total_conversation_tokens = app + .session + .total_conversation_tokens + .saturating_add(turn_tokens); + app.session.total_input_tokens = app + .session + .total_input_tokens + .saturating_add(usage.input_tokens); + app.session.total_output_tokens = app + .session + .total_output_tokens + .saturating_add(usage.output_tokens); + if usage.prompt_cache_hit_tokens.is_some() + || usage.prompt_cache_miss_tokens.is_some() + || usage.prompt_cache_write_tokens.is_some() + { + let classes = crate::pricing::token_usage_for_pricing(usage); + app.session.total_cache_hit_tokens = app + .session + .total_cache_hit_tokens + .saturating_add(u32::try_from(classes.cache_read).unwrap_or(u32::MAX)); + app.session.total_cache_miss_tokens = app + .session + .total_cache_miss_tokens + .saturating_add(u32::try_from(classes.input).unwrap_or(u32::MAX)); + app.session.total_cache_write_tokens = app + .session + .total_cache_write_tokens + .saturating_add(u32::try_from(classes.cache_write).unwrap_or(u32::MAX)); + } +} + +fn translation_origin(app: &App) -> (Option, Option) { + // Fixed-size one-way identities avoid retaining raw imported ids in a + // detached completion envelope without introducing truncation aliases. + let fingerprint = |value: Option<&str>| value.map(crate::cost_status::usage_source_fingerprint); + ( + fingerprint(app.current_session_id.as_deref()), + fingerprint(app.runtime_turn_id.as_deref()), + ) +} + +fn translation_origin_is_current( + app: &App, + origin_session_fingerprint: Option<&str>, + origin_turn_fingerprint: Option<&str>, +) -> bool { + let current = translation_origin(app); + current.0.as_deref() == origin_session_fingerprint + && current.1.as_deref() == origin_turn_fingerprint +} + +fn translation_session_is_current(app: &App, origin_session_fingerprint: Option<&str>) -> bool { + translation_origin(app).0.as_deref() == origin_session_fingerprint +} + +fn exact_translation_client( + config: &Config, + route: &crate::core::events::TurnRoute, +) -> anyhow::Result> { + let identity = config + .resolve_persisted_provider_identity( + Some(route.provider.as_str()), + Some(&route.provider_identity), + ) + .map_err(anyhow::Error::msg)?; + let validated = crate::route_runtime::resolve_runtime_route_for_identity( + config, + &identity, + Some(&route.model), + ) + .map_err(anyhow::Error::msg)? + .validate() + .map_err(anyhow::Error::msg)?; + if validated.identity.key != route.provider_identity + || validated.model != route.model + || validated.candidate.endpoint().base_url != route.base_url + { + anyhow::bail!( + "translation route changed after turn dispatch; refusing to reuse a different provider client" + ); + } + if let Some(receipt) = route.receipt.as_ref() + && &validated + .client + .turn_route_receipt(&route.provider_identity) + != receipt + { + anyhow::bail!( + "translation credential or endpoint changed after turn dispatch; refusing stale completion ownership" + ); + } + Ok(Arc::new(validated.client)) +} + /// Bind the Runtime thread store to a session before the process-owner lock /// is taken, so a second Codewhale on the same machine does not collide on /// the default root (#5630). Resume reuses the loaded id; a fresh session @@ -884,6 +1076,10 @@ pub(crate) async fn run_event_loop( let mut stream_display_clock = StreamDisplayClock::default(); let (translation_tx, mut translation_rx) = tokio::sync::mpsc::unbounded_channel::(); + let fallback_translation_client = translation_client; + let mut active_translation_client = fallback_translation_client.clone(); + let mut active_translation_route: Option = None; + let mut translation_sequence = 0_u64; let mut pending_translations = 0usize; let mut pending_thinking_translations = 0usize; let mut last_queue_state = (app.queued_messages.clone(), app.queued_draft.clone()); @@ -1059,14 +1255,31 @@ pub(crate) async fn run_event_loop( while let Ok(event) = translation_rx.try_recv() { match event { TranslationEvent::AssistantMessage { + origin_session_fingerprint, + origin_turn_fingerprint, history_index, original_text, translated, + usage, thinking, tool_uses, } => { pending_translations = pending_translations.saturating_sub(1); - pending_thinking_translations = pending_thinking_translations.saturating_sub(1); + if translation_session_is_current(app, origin_session_fingerprint.as_deref()) + && let Some(usage) = usage.as_ref() + { + accrue_translation_usage(app, usage); + } + if !translation_origin_is_current( + app, + origin_session_fingerprint.as_deref(), + origin_turn_fingerprint.as_deref(), + ) { + tracing::debug!( + "discarded assistant translation completed for a stale session/turn" + ); + continue; + } let text = match translated { Ok(text) => { app.status_message = Some( @@ -1111,10 +1324,29 @@ pub(crate) async fn run_event_loop( app.needs_redraw = true; } TranslationEvent::Thinking { + origin_session_fingerprint, + origin_turn_fingerprint, placeholder, translated, + usage, } => { pending_translations = pending_translations.saturating_sub(1); + pending_thinking_translations = pending_thinking_translations.saturating_sub(1); + if translation_session_is_current(app, origin_session_fingerprint.as_deref()) + && let Some(usage) = usage.as_ref() + { + accrue_translation_usage(app, usage); + } + if !translation_origin_is_current( + app, + origin_session_fingerprint.as_deref(), + origin_turn_fingerprint.as_deref(), + ) { + tracing::debug!( + "discarded thinking translation completed for a stale session/turn" + ); + continue; + } let text = match translated { Ok(text) => { app.status_message = Some( @@ -1374,7 +1606,7 @@ pub(crate) async fn run_event_loop( if app.translation_enabled && !current_streaming_text.is_empty() && crate::tui::translation::needs_translation(¤t_streaming_text) - && let Some(translation_client) = translation_client.as_ref() + && let Some(translation_client) = active_translation_client.as_ref() { app.status_message = Some( crate::localization::tr( @@ -1388,24 +1620,38 @@ pub(crate) async fn run_event_loop( let tx = translation_tx.clone(); let client = translation_client.clone(); let original_text = current_streaming_text.clone(); - let translation_model = app - .last_effective_model - .clone() + let translation_model = active_translation_route + .as_ref() + .map(|route| route.model.clone()) + .or_else(|| app.last_effective_model.clone()) .unwrap_or_else(|| app.model.clone()); + translation_sequence = translation_sequence.saturating_add(1); + let accounting = TranslationAccountingContext::capture( + app, + "assistant", + translation_sequence, + ); + let (origin_session_fingerprint, origin_turn_fingerprint) = + translation_origin(app); let target_language = app.ui_locale.translation_target_name().to_string(); tokio::spawn(async move { - let translated = crate::tui::translation::translate_text( - &original_text, - &client, - &translation_model, - &target_language, - ) - .await; + let settled = accounting.settle( + client + .translate_with_usage( + &original_text, + &translation_model, + &target_language, + ) + .await, + ); let _ = tx.send(TranslationEvent::AssistantMessage { + origin_session_fingerprint, + origin_turn_fingerprint, history_index, original_text, - translated, + translated: settled.translated, + usage: settled.usage, thinking, tool_uses, }); @@ -1462,7 +1708,7 @@ pub(crate) async fn run_event_loop( } if !original_thinking.is_empty() && crate::tui::translation::needs_translation(&original_thinking) - && let Some(translation_client) = translation_client.as_ref() + && let Some(translation_client) = active_translation_client.as_ref() { app.status_message = Some( crate::localization::thinking_translation_in_progress( @@ -1476,10 +1722,19 @@ pub(crate) async fn run_event_loop( pending_thinking_translations.saturating_add(1); let tx = translation_tx.clone(); let client = translation_client.clone(); - let translation_model = app - .last_effective_model - .clone() + let translation_model = active_translation_route + .as_ref() + .map(|route| route.model.clone()) + .or_else(|| app.last_effective_model.clone()) .unwrap_or_else(|| app.model.clone()); + translation_sequence = translation_sequence.saturating_add(1); + let accounting = TranslationAccountingContext::capture( + app, + "thinking", + translation_sequence, + ); + let (origin_session_fingerprint, origin_turn_fingerprint) = + translation_origin(app); let placeholder = crate::localization::thinking_translation_placeholder( app.ui_locale, @@ -1488,16 +1743,21 @@ pub(crate) async fn run_event_loop( let target_language = app.ui_locale.translation_target_name().to_string(); tokio::spawn(async move { - let translated = crate::tui::translation::translate_text( - &original_thinking, - &client, - &translation_model, - &target_language, - ) - .await; + let settled = accounting.settle( + client + .translate_with_usage( + &original_thinking, + &translation_model, + &target_language, + ) + .await, + ); let _ = tx.send(TranslationEvent::Thinking { + origin_session_fingerprint, + origin_turn_fingerprint, placeholder, - translated, + translated: settled.translated, + usage: settled.usage, }); }); } else { @@ -1655,7 +1915,7 @@ pub(crate) async fn run_event_loop( subagent_list_refresh_requested = true; } } - EngineEvent::TurnStarted { turn_id, .. } => { + EngineEvent::TurnStarted { turn_id, route, .. } => { // A prior turn that died without its `TurnComplete` // must not leak its provisional estimate into this one. app.clear_pending_turn_cost(); @@ -1696,6 +1956,19 @@ pub(crate) async fn run_event_loop( if app.status_message.is_none() { app.status_message = Some("Press Esc or Ctrl+C to cancel".to_string()); } + active_translation_client = match route.as_ref() { + Some(route) => match exact_translation_client(config, route) { + Ok(client) => Some(client), + Err(error) => { + tracing::warn!( + "translation client rejected the frozen turn route: {error}" + ); + None + } + }, + None => fallback_translation_client.clone(), + }; + active_translation_route = route; app.runtime_turn_id = Some(turn_id); app.runtime_turn_status = Some("in_progress".to_string()); app.turn_counter = app.turn_counter.saturating_add(1); @@ -1725,9 +1998,26 @@ pub(crate) async fn run_event_loop( EngineEvent::ToolRequestSnapshot { snapshot } => { app.session.last_tool_request_snapshot = Some(snapshot); } - EngineEvent::RouteDispatched { .. } => {} + EngineEvent::RouteDispatched { turn_id, route } => { + if app.runtime_turn_id.as_deref() == Some(turn_id.as_str()) { + active_translation_client = match exact_translation_client( + config, &route, + ) { + Ok(client) => Some(client), + Err(error) => { + tracing::warn!( + "translation client rejected the dispatched turn route: {error}" + ); + None + } + }; + active_translation_route = Some(route); + } + } EngineEvent::TurnComplete { usage, + parent_route_usage, + routed_usage_dropped_records, status, error, tool_catalog, @@ -1944,19 +2234,19 @@ pub(crate) async fn run_event_loop( .as_ref() .and_then(|turn| turn.route.as_ref()) .and_then(crate::core::events::TurnRoute::cost_envelope) - .map(|route| route.audit(&usage)); + .map(|route| route.audit(&parent_route_usage)); app.push_turn_cache_record(crate::tui::app::TurnCacheRecord { provider, provider_identity, model, auto_model, - input_tokens: usage.input_tokens, - output_tokens: usage.output_tokens, - cache_hit_tokens: usage.prompt_cache_hit_tokens, - cache_miss_tokens: usage.prompt_cache_miss_tokens, - reasoning_replay_tokens: usage.reasoning_replay_tokens, - cache_write_tokens: usage.prompt_cache_write_tokens, - reasoning_tokens: usage.reasoning_tokens, + input_tokens: parent_route_usage.input_tokens, + output_tokens: parent_route_usage.output_tokens, + cache_hit_tokens: parent_route_usage.prompt_cache_hit_tokens, + cache_miss_tokens: parent_route_usage.prompt_cache_miss_tokens, + reasoning_replay_tokens: parent_route_usage.reasoning_replay_tokens, + cache_write_tokens: parent_route_usage.prompt_cache_write_tokens, + reasoning_tokens: parent_route_usage.reasoning_tokens, cost_audit: cost_audit.clone(), recorded_at: Instant::now(), }); @@ -1975,8 +2265,9 @@ pub(crate) async fn run_event_loop( // *not* cover so `/cost` can stay honest about it. // // `cost_audit` above came from `cost_envelope()`, i.e. - // the billing envelope stamped at the wire boundary - // and classified from this turn's frozen receipt. It + // the billing envelope frozen at CodeWhale's + // pre-permit application-dispatch boundary and + // classified from this turn's frozen receipt. It // is `None` for a route that was never dispatched, and // a route whose receipt named no product classified as // Unknown — either way nothing accrues. A `/provider` @@ -1999,6 +2290,20 @@ pub(crate) async fn run_event_loop( if let Some(cost) = turn_cost { app.accrue_session_cost_estimate(cost); } + if routed_usage_dropped_records > 0 { + let dropped = + u32::try_from(routed_usage_dropped_records).unwrap_or(u32::MAX); + app.session.cost_unpriced_turns = + app.session.cost_unpriced_turns.saturating_add(dropped); + app.session.cost_cny_unpriced_turns = + app.session.cost_cny_unpriced_turns.saturating_add(dropped); + app.session + .cost_unpriced_reasons + .insert("routed_usage_receipt_missing".to_string()); + app.session + .cost_cny_unpriced_reasons + .insert("routed_usage_receipt_missing".to_string()); + } // Emit OSC 9 / BEL desktop notification for long turns, and // always stop the title animation that began on TurnStarted. @@ -3283,6 +3588,24 @@ pub(crate) async fn run_event_loop( } app.session.accrue_pending_turn_usage(&usage); } + EngineEvent::RoutedTurnUsage { + usage, + duration_ms, + first_token_ms, + request_ms, + } => { + // Routed calls own separate immutable cost receipts. + // Preserve model-call telemetry without pricing them + // provisionally under the active parent route or + // incrementally adding tokens that TurnComplete will + // reconcile authoritatively. + app.session_metrics.record_model_call( + usage.output_tokens, + duration_ms, + first_token_ms, + request_ms, + ); + } EngineEvent::AdvisoryNote { note, .. } => { // Advisor background watcher note. Display as a // concise system message in the transcript so the @@ -6281,4 +6604,175 @@ mod session_boot_event_tests { Some("fresh") ); } + + fn translation_test_route() -> crate::cost_status::EffectiveRouteEnvelope { + crate::cost_status::EffectiveRouteEnvelope { + provider: crate::config::ApiProvider::Deepseek, + provider_identity: "deepseek".to_string(), + model: "deepseek-chat".to_string(), + billing_surface: crate::pricing::billing_surface_for_route( + crate::config::ApiProvider::Deepseek, + Some("https://api.deepseek.com/v1"), + ) + .map(str::to_string), + endpoint_fingerprint: crate::cost_status::endpoint_fingerprint( + "https://api.deepseek.com/v1", + ), + provider_live_pricing: None, + billing_mode: crate::cost_status::RouteBillingMode::Metered, + dispatched_at: chrono::Utc::now(), + } + } + + #[test] + fn assistant_and_thinking_translation_usage_each_accrue_once() { + let _scope = crate::cost_status::test_scope(); + let mut app = test_app(); + app.current_session_id = Some("session-translation".to_string()); + app.runtime_turn_id = Some("turn-translation".to_string()); + let usage_a = crate::models::Usage { + input_tokens: 5, + output_tokens: 2, + ..crate::models::Usage::default() + }; + let usage_b = crate::models::Usage { + input_tokens: 3, + output_tokens: 1, + ..crate::models::Usage::default() + }; + + let assistant = TranslationAccountingContext::capture(&app, "assistant", 1).settle(Ok( + crate::client::TranslationProviderResponse { + translated: Ok("助理".to_string()), + route: translation_test_route(), + usage: Some(usage_a.clone()), + }, + )); + let thinking = TranslationAccountingContext::capture(&app, "thinking", 2).settle(Ok( + crate::client::TranslationProviderResponse { + translated: Err(anyhow::anyhow!("incomplete: max_tokens")), + route: translation_test_route(), + usage: Some(usage_b.clone()), + }, + )); + + assert_eq!(assistant.usage.as_ref(), Some(&usage_a)); + assert_eq!(thinking.usage.as_ref(), Some(&usage_b)); + assert!( + thinking.translated.is_err(), + "semantic rejection is preserved" + ); + accrue_translation_usage(&mut app, assistant.usage.as_ref().expect("assistant usage")); + accrue_translation_usage(&mut app, thinking.usage.as_ref().expect("thinking usage")); + assert_eq!(app.session.total_input_tokens, 8); + assert_eq!(app.session.total_output_tokens, 3); + assert_eq!(app.session.total_tokens, 11); + + let pending = crate::cost_status::drain(); + assert_eq!( + pending.priced_turns.saturating_add(pending.unpriced_turns), + 2, + "each decoded provider response is audited exactly once" + ); + } + + #[test] + fn translation_unreceipted_success_is_marked_once_but_transport_failure_is_not() { + let _scope = crate::cost_status::test_scope(); + let mut app = test_app(); + app.current_session_id = Some("session-translation-missing-usage".to_string()); + app.runtime_turn_id = Some("turn-translation-missing-usage".to_string()); + + for _ in 0..2 { + let settled = TranslationAccountingContext::capture(&app, "assistant", 7).settle(Ok( + crate::client::TranslationProviderResponse { + translated: Ok("translation remains usable".to_string()), + route: translation_test_route(), + usage: None, + }, + )); + assert_eq!( + settled.translated.expect("semantic output remains usable"), + "translation remains usable" + ); + assert_eq!(settled.usage, None); + } + let transport = TranslationAccountingContext::capture(&app, "assistant", 8) + .settle(Err(anyhow::anyhow!("HTTP 429"))); + assert!(transport.translated.is_err()); + assert_eq!(transport.usage, None); + + let pending = crate::cost_status::drain(); + assert_eq!(pending.priced_turns, 0); + assert_eq!( + pending.unpriced_turns, 1, + "stable response id dedupes replay" + ); + assert_eq!(pending.cny_unpriced_turns, 1); + assert!( + pending + .unpriced_reasons + .contains("provider_success_missing_usage") + ); + } + + #[test] + fn late_translation_delivery_isolated_from_new_session_or_turn() { + let mut app = test_app(); + app.current_session_id = Some("session-a".to_string()); + app.runtime_turn_id = Some("turn-a".to_string()); + let (session, turn) = translation_origin(&app); + assert!(translation_origin_is_current( + &app, + session.as_deref(), + turn.as_deref() + )); + + app.current_session_id = Some("session-b".to_string()); + assert!(!translation_session_is_current(&app, session.as_deref())); + assert!(!translation_origin_is_current( + &app, + session.as_deref(), + turn.as_deref() + )); + app.current_session_id = Some("session-a".to_string()); + app.runtime_turn_id = Some("turn-b".to_string()); + assert!( + translation_session_is_current(&app, session.as_deref()), + "same-session late usage still belongs in session totals" + ); + assert!(!translation_origin_is_current( + &app, + session.as_deref(), + turn.as_deref() + )); + + let usage = crate::models::Usage { + input_tokens: 4, + output_tokens: 2, + ..crate::models::Usage::default() + }; + if translation_session_is_current(&app, session.as_deref()) { + accrue_translation_usage(&mut app, &usage); + } + assert_eq!(app.session.total_tokens, 6); + app.current_session_id = Some("session-b".to_string()); + if translation_session_is_current(&app, session.as_deref()) { + accrue_translation_usage(&mut app, &usage); + } + assert_eq!( + app.session.total_tokens, 6, + "cross-session late usage must not pollute the new session" + ); + + let shared_prefix = "x".repeat(300); + app.current_session_id = Some(format!("{shared_prefix}:old")); + app.runtime_turn_id = Some("turn-long".to_string()); + let (long_session, long_turn) = translation_origin(&app); + app.current_session_id = Some(format!("{shared_prefix}:new")); + assert!( + !translation_origin_is_current(&app, long_session.as_deref(), long_turn.as_deref()), + "fixed fingerprints must distinguish ids with the same long prefix" + ); + } } diff --git a/crates/tui/src/tui/ui/tests.rs b/crates/tui/src/tui/ui/tests.rs index 8cdee9865d..bfb6b6d701 100644 --- a/crates/tui/src/tui/ui/tests.rs +++ b/crates/tui/src/tui/ui/tests.rs @@ -52,6 +52,42 @@ use crate::models::Role; use crate::tui::selection::{SelectionAutoscroll, TranscriptSelectionPoint}; use tempfile::TempDir; +#[test] +fn failed_engine_channel_settles_classifier_batch_once() { + let _cost_scope = crate::cost_status::test_scope(); + let route = crate::cost_status::EffectiveRouteEnvelope::capture( + None, + ApiProvider::Deepseek, + "deepseek", + "classifier-model", + Some(ApiProvider::Deepseek.default_base_url()), + chrono::Utc::now(), + ); + let batch = crate::cost_status::RuntimeUsageBatch { + records: vec![crate::cost_status::RuntimeUsageRecord { + source_id: "auto-router:dispatch-usage".to_string(), + usage: crate::cost_status::EffectiveRouteUsage { + route: route.clone(), + usage: crate::models::Usage { + input_tokens: 6, + output_tokens: 2, + ..Default::default() + }, + }, + }], + drop_records: vec![crate::cost_status::RuntimeUsageDropRecord { + source_id: "auto-router:dispatch-drop".to_string(), + route, + }], + dropped_records: 1, + }; + + super::dispatch::settle_failed_dispatch_routed_usage(&batch); + super::dispatch::settle_failed_dispatch_routed_usage(&batch); + let pending = crate::cost_status::drain(); + assert_eq!(pending.usage_source_fingerprints.len(), 2); +} + #[test] fn session_shell_area_fills_the_host_terminal_at_every_width() { // #5322: transcript / composer share the full host width — no wide-terminal @@ -520,6 +556,7 @@ fn completed_turn_cost_receipt_uses_the_captured_effective_route() { // receipt fails closed on anything that is not one, so a // hand-written placeholder here would have tested nothing. endpoint_fingerprint: served_endpoint_fingerprint(), + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Metered, dispatched_at: chrono::Utc::now(), }), @@ -10939,6 +10976,7 @@ fn turn_liveness_recovers_stalled_in_progress_turn() { billing: Some(crate::core::events::RouteBillingEnvelope { billing_surface: Some(crate::pricing::FIRST_PARTY_PAYG_BILLING_SURFACE.to_string()), endpoint_fingerprint: Some("openai-endpoint".to_string()), + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Metered, dispatched_at: chrono::Utc::now(), }), @@ -10990,6 +11028,7 @@ fn engine_event_disconnect_recovers_live_turn_immediately() { billing: Some(crate::core::events::RouteBillingEnvelope { billing_surface: Some(crate::pricing::FIRST_PARTY_PAYG_BILLING_SURFACE.to_string()), endpoint_fingerprint: Some("openai-endpoint".to_string()), + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Metered, dispatched_at: chrono::Utc::now(), }), @@ -11065,6 +11104,7 @@ fn engine_event_disconnect_cleans_cancelled_turn_metadata() { billing: Some(crate::core::events::RouteBillingEnvelope { billing_surface: Some(crate::pricing::FIRST_PARTY_PAYG_BILLING_SURFACE.to_string()), endpoint_fingerprint: Some("openai-endpoint".to_string()), + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Metered, dispatched_at: chrono::Utc::now(), }), @@ -11955,7 +11995,10 @@ fn subagent_token_usage_updates_live_cost_counter_without_card_change() { &crate::tools::subagent::MailboxMessage::TokenUsage { agent_id: "agent-a".to_string(), source_id: "response-a".to_string(), - route: test_mailbox_route(ApiProvider::Deepseek, "deepseek-v4-flash"), + route: Box::new(test_mailbox_route( + ApiProvider::Deepseek, + "deepseek-v4-flash", + )), usage: crate::models::Usage { input_tokens: 10_000, output_tokens: 1_000, @@ -11982,7 +12025,7 @@ fn subagent_token_usage_prices_the_child_route_not_the_parent_route() { &crate::tools::subagent::MailboxMessage::TokenUsage { agent_id: "agent-codex".to_string(), source_id: "response-codex".to_string(), - route: test_mailbox_route(ApiProvider::OpenaiCodex, "gpt-5.5"), + route: Box::new(test_mailbox_route(ApiProvider::OpenaiCodex, "gpt-5.5")), usage: crate::models::Usage { input_tokens: 10_000, output_tokens: 1_000, @@ -12003,7 +12046,10 @@ fn subagent_token_usage_is_deduped_by_response_source() { let usage = crate::tools::subagent::MailboxMessage::TokenUsage { agent_id: "agent-a".to_string(), source_id: "response-a".to_string(), - route: test_mailbox_route(ApiProvider::Deepseek, "deepseek-v4-flash"), + route: Box::new(test_mailbox_route( + ApiProvider::Deepseek, + "deepseek-v4-flash", + )), usage: crate::models::Usage { input_tokens: 10_000, output_tokens: 1_000, @@ -12030,7 +12076,10 @@ fn subagent_token_usage_source_is_stable_across_engine_turns() { let usage = crate::tools::subagent::MailboxMessage::TokenUsage { agent_id: "agent-a".to_string(), source_id: "response-a".to_string(), - route: test_mailbox_route(ApiProvider::Deepseek, "deepseek-v4-flash"), + route: Box::new(test_mailbox_route( + ApiProvider::Deepseek, + "deepseek-v4-flash", + )), usage: crate::models::Usage { input_tokens: 10_000, output_tokens: 1_000, @@ -13196,6 +13245,8 @@ fn local_cancel_marks_late_stream_events_for_suppression() { assert!(!suppress_engine_event_after_local_cancel( &EngineEvent::TurnComplete { usage: Usage::default(), + parent_route_usage: Usage::default(), + routed_usage_dropped_records: 0, status: crate::core::events::TurnOutcomeStatus::Interrupted, error: None, tool_catalog: None, @@ -13240,6 +13291,7 @@ fn turn_started_route_is_captured_before_cancel_suppression() { billing: Some(crate::core::events::RouteBillingEnvelope { billing_surface: Some(crate::pricing::FIRST_PARTY_PAYG_BILLING_SURFACE.to_string()), endpoint_fingerprint: Some("openai-endpoint".to_string()), + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Metered, dispatched_at: created_at, }), @@ -13304,6 +13356,7 @@ fn turn_started_suggestion_authority_comes_from_the_route_receipt_not_config() { billing: Some(crate::core::events::RouteBillingEnvelope { billing_surface: None, endpoint_fingerprint: None, + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Unknown, dispatched_at: chrono::Utc::now(), }), @@ -13346,6 +13399,7 @@ fn turn_started_without_a_route_receipt_captures_no_suggestion_authority() { billing: Some(crate::core::events::RouteBillingEnvelope { billing_surface: None, endpoint_fingerprint: None, + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Unknown, dispatched_at: chrono::Utc::now(), }), @@ -13382,6 +13436,7 @@ fn engine_error_health_accounting_uses_active_turn_route() { billing: Some(crate::core::events::RouteBillingEnvelope { billing_surface: Some(crate::pricing::FIRST_PARTY_PAYG_BILLING_SURFACE.to_string()), endpoint_fingerprint: Some("openai-endpoint".to_string()), + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Metered, dispatched_at: chrono::Utc::now(), }), @@ -16366,6 +16421,7 @@ fn legacy_child_usage_metadata_fails_closed_without_parent_route_fallback() { endpoint_fingerprint: crate::cost_status::endpoint_fingerprint( crate::config::DEFAULT_DEEPSEEK_BASE_URL, ), + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Metered, dispatched_at: chrono::Utc::now(), }), @@ -16429,6 +16485,7 @@ fn child_usage_metadata_carries_cache_write_and_reasoning_end_to_end() { endpoint_fingerprint: crate::cost_status::endpoint_fingerprint( "https://api.anthropic.com/v1", ), + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Metered, dispatched_at: chrono::DateTime::::from_timestamp(0, 0).expect("epoch"), }; @@ -16482,6 +16539,7 @@ fn child_usage_metadata_carries_cache_write_and_reasoning_end_to_end() { model: "kimi-k2.7-code".to_string(), billing_surface: Some(crate::pricing::FIRST_PARTY_PAYG_BILLING_SURFACE.to_string()), endpoint_fingerprint: Some("test-moonshot-endpoint".to_string()), + provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Metered, dispatched_at: chrono::DateTime::::from_timestamp(0, 0).expect("epoch"), }; @@ -16558,6 +16616,53 @@ fn zero_usage_model_child_still_records_priced_receipt() { assert_eq!(app.session.cost_route_receipts.len(), 1); } +#[test] +fn routed_missing_usage_batch_prices_exact_routes_and_only_residual_as_generic() { + let mut app = create_test_app(); + let mut metered_route = test_mailbox_route(ApiProvider::Deepseek, "deepseek-v4-flash"); + metered_route.billing_mode = crate::cost_status::RouteBillingMode::Metered; + let mut local_route = metered_route.clone(); + local_route.provider_identity = "local-computer".to_string(); + local_route.billing_mode = crate::cost_status::RouteBillingMode::Local; + let batch = crate::cost_status::RuntimeUsageBatch { + records: Vec::new(), + drop_records: vec![ + crate::cost_status::RuntimeUsageDropRecord { + source_id: "rlm:missing:metered".to_string(), + route: metered_route, + }, + crate::cost_status::RuntimeUsageDropRecord { + source_id: "rlm:missing:local".to_string(), + route: local_route, + }, + ], + // Two exact routes plus one overflowed/route-less residual. + dropped_records: 3, + }; + let mut metadata = serde_json::json!({}); + crate::cost_status::attach_child_usage_batch_metadata(&mut metadata, &batch); + let result = Ok(crate::tools::spec::ToolResult::success("ok").with_metadata(metadata)); + + handle_tool_call_complete(&mut app, "rlm-missing", "rlm", &result); + + assert_eq!(app.session.cost_priced_turns, 0); + assert_eq!( + app.session.cost_unpriced_turns, 2, + "metered exact drop and one residual gap count; local exact drop does not" + ); + assert!( + app.session + .cost_unpriced_reasons + .contains("provider_success_missing_usage") + ); + assert!( + app.session + .cost_unpriced_reasons + .contains("routed_usage_receipt_missing") + ); + assert_eq!(app.session.subagent_usage_sources.len(), 2); +} + #[test] fn picker_renamed_active_title_survives_automatic_snapshot() { let mut app = create_test_app(); @@ -21177,7 +21282,10 @@ fn duplicate_mailbox_token_usage_does_not_regress_displayed_cost() { let usage = crate::tools::subagent::MailboxMessage::TokenUsage { agent_id: "agent-x".to_string(), source_id: "response-x".to_string(), - route: test_mailbox_route(ApiProvider::Deepseek, "deepseek-v4-flash"), + route: Box::new(test_mailbox_route( + ApiProvider::Deepseek, + "deepseek-v4-flash", + )), usage: crate::models::Usage { input_tokens: 10_000, output_tokens: 1_000, diff --git a/crates/tui/src/tui/work_surface/mod.rs b/crates/tui/src/tui/work_surface/mod.rs index 896630e470..ebcda35b6c 100644 --- a/crates/tui/src/tui/work_surface/mod.rs +++ b/crates/tui/src/tui/work_surface/mod.rs @@ -1326,7 +1326,7 @@ mod tests { let usage = |source_id: &str, input_tokens, output_tokens| MailboxMessage::TokenUsage { agent_id: "agent_stream".to_string(), source_id: source_id.to_string(), - route: route.clone(), + route: Box::new(route.clone()), usage: crate::models::Usage { input_tokens, output_tokens, diff --git a/crates/tui/src/turn_route_plan.rs b/crates/tui/src/turn_route_plan.rs index f7eb514810..8222a9782f 100644 --- a/crates/tui/src/turn_route_plan.rs +++ b/crates/tui/src/turn_route_plan.rs @@ -69,6 +69,11 @@ pub(crate) struct PlannedTurnRoute { pub(crate) effective_reasoning_effort: Option, pub(crate) auto_controls_reasoning: bool, pub(crate) auto_selection: Option, + /// Bounded auxiliary classifier usage that must enter the accepted turn + /// under its own frozen routes. It is moved out of `auto_selection` so a + /// UI-only receipt consumer cannot accidentally become the accounting + /// owner or price it under the parent route. + pub(crate) initial_routed_usage: crate::cost_status::RuntimeUsageBatch, /// Why this concrete route was selected. This is captured by the planner, /// not inferred later from the resulting provider/model pair. pub(crate) routing_source: TurnRoutingSource, @@ -108,6 +113,18 @@ fn reasoning_effort_for_route_selection( } } +fn settle_failed_parent_route( + error: String, + initial_routed_usage: &crate::cost_status::RuntimeUsageBatch, +) -> String { + crate::cost_status::report_runtime_usage_batch( + crate::cost_status::scope_token(), + None, + initial_routed_usage, + ); + error +} + /// Resolve the route for one turn. /// /// This is *the* route planner (#1004). `spawned_dispatch_inner` calls it to @@ -125,7 +142,7 @@ fn reasoning_effort_for_route_selection( pub(crate) async fn plan_turn_route( request: TurnRoutePlanRequest<'_>, ) -> Result { - let auto_selection = if request.should_auto_resolve { + let mut auto_selection = if request.should_auto_resolve { Some( crate::model_routing::resolve_auto_route_with_inventory_for_session_and_cache_policy( request.route_config, @@ -163,6 +180,18 @@ pub(crate) async fn plan_turn_route( request.app_model.to_string() }; + // Move classifier accounting out immediately. Every later parent-route + // failure must settle this already-incurred auxiliary call instead of + // returning an error that silently drops its exact quote/usage. + let initial_routed_usage = auto_selection + .as_mut() + .map(|selection| crate::cost_status::RuntimeUsageBatch { + records: std::mem::take(&mut selection.routed_usage), + drop_records: std::mem::take(&mut selection.routed_usage_drop_records), + dropped_records: std::mem::take(&mut selection.routed_usage_dropped_records), + }) + .unwrap_or_default(); + let turn_route = if effective_provider == request.app_route_identity.provider { resolve_runtime_route_for_identity( request.route_config, @@ -177,9 +206,22 @@ pub(crate) async fn plan_turn_route( ) }; - let turn_route = turn_route.map_err(|err| err.to_string())?; + let turn_route = match turn_route { + Ok(route) => route, + Err(err) => { + return Err(settle_failed_parent_route( + err.to_string(), + &initial_routed_usage, + )); + } + }; let turn_route = if request.preflight_required { - turn_route.preflight()? + match turn_route.preflight() { + Ok(route) => route, + Err(err) => { + return Err(settle_failed_parent_route(err, &initial_routed_usage)); + } + } } else { turn_route }; @@ -262,6 +304,7 @@ pub(crate) async fn plan_turn_route( effective_reasoning_effort, auto_controls_reasoning, auto_selection, + initial_routed_usage, routing_source, }) } @@ -280,6 +323,49 @@ mod tests { } } + #[test] + fn failed_parent_route_settles_classifier_batch_once() { + let _cost_scope = crate::cost_status::test_scope(); + let route = crate::cost_status::EffectiveRouteEnvelope::capture( + None, + ApiProvider::Deepseek, + "deepseek", + "classifier-model", + Some(ApiProvider::Deepseek.default_base_url()), + chrono::Utc::now(), + ); + let batch = crate::cost_status::RuntimeUsageBatch { + records: vec![crate::cost_status::RuntimeUsageRecord { + source_id: "auto-router:plan-usage".to_string(), + usage: crate::cost_status::EffectiveRouteUsage { + route: route.clone(), + usage: crate::models::Usage { + input_tokens: 4, + output_tokens: 2, + ..Default::default() + }, + }, + }], + drop_records: vec![crate::cost_status::RuntimeUsageDropRecord { + source_id: "auto-router:plan-drop".to_string(), + route, + }], + dropped_records: 1, + }; + + assert_eq!( + settle_failed_parent_route("route failed".to_string(), &batch), + "route failed" + ); + settle_failed_parent_route("route failed".to_string(), &batch); + let pending = crate::cost_status::drain(); + assert_eq!( + pending.usage_source_fingerprints.len(), + 2, + "both exact classifier outcomes persist, and replay is idempotent" + ); + } + #[test] fn auto_model_route_selection_keeps_raw_reasoning_preference() { assert_eq!( diff --git a/crates/tui/tests/runtime_web_client.test.mjs b/crates/tui/tests/runtime_web_client.test.mjs index 189b42c801..cb2eb05349 100644 --- a/crates/tui/tests/runtime_web_client.test.mjs +++ b/crates/tui/tests/runtime_web_client.test.mjs @@ -407,7 +407,10 @@ test("new thread dialog labels exact vision capability without exposing attachme assert.match(html, /does not change your Runtime defaults/); assert.doesNotMatch(html, /type="file"/); assert.match(source, /api\("\/v1\/providers"\)/); - assert.match(source, /\/v1\/providers\/\$\{encodeURIComponent\(provider\.id\)\}\/models/); + // The dialog loads the catalog through the bounded, paginated collector + // keyed by provider.id; the wire endpoint stays /v1/providers//models. + assert.match(source, /collectProviderModelPages\(provider\.id/); + assert.match(source, /\/v1\/providers\/\$\{encodeURIComponent\(provider\)\}\/models\?\$\{query\.toString\(\)\}/); assert.match(source, /body: JSON\.stringify\(request\)/); assert.match(source, /function trapFocusWithin\(event, container\)/); assert.match(source, /dom\.newThreadCancel\.focus\(\{ preventScroll: true \}\)/);