From 1a03b23b6290e36383b555619dd24bcac740265d Mon Sep 17 00:00:00 2001 From: CodeWhale Bot Date: Sat, 29 Aug 2026 22:33:20 -0700 Subject: [PATCH 1/5] feat(tui): checkpoint live catalogs and routed usage Preserve the dependency-coupled OpenRouter/Baseten catalog, route-envelope, pricing, usage-accounting, and UI integration as a reviewable checkpoint. This is not a completion claim: two runtime_api tests still abort with stack overflow, Clippy reports 17 findings, and authenticated provider/CI/product proof remains outstanding.\n\nIncludes the completed Fable dispatch-origin accounting fix and the completed Opus failed-before-turn settlement fix, reconciled by the primary integrator.\n\nCo-authored-by: Claude Fable 5 --- crates/config/src/provider_templates.rs | 35 +- crates/config/src/route/resolver.rs | 36 +- crates/config/src/route/tests.rs | 27 + crates/tui/src/client.rs | 1155 ++++++++++++- crates/tui/src/core/engine.rs | 160 +- crates/tui/src/core/engine/preview/tests.rs | 2 + crates/tui/src/core/engine/tests.rs | 151 ++ crates/tui/src/core/engine/turn_loop.rs | 107 +- crates/tui/src/core/events.rs | 50 +- crates/tui/src/core/ops.rs | 5 + crates/tui/src/core/turn.rs | 87 +- crates/tui/src/cost_status.rs | 1595 +++++++++++++++++- crates/tui/src/exec_agent.rs | 1 + crates/tui/src/lib.rs | 1 + crates/tui/src/model_catalog.rs | 6 +- crates/tui/src/model_routing.rs | 543 ++++++- crates/tui/src/pricing.rs | 327 +++- crates/tui/src/provider_catalog_live.rs | 1624 +++++++++++++++++++ crates/tui/src/provider_lake.rs | 777 ++++++++- crates/tui/src/remote_control.rs | 4 + crates/tui/src/rlm/bridge.rs | 396 ++++- crates/tui/src/rlm/turn.rs | 173 +- crates/tui/src/route_runtime.rs | 311 +++- crates/tui/src/runtime_api.rs | 175 +- crates/tui/src/runtime_api/tests.rs | 170 ++ crates/tui/src/runtime_threads.rs | 525 +++++- crates/tui/src/runtime_threads/tests.rs | 772 ++++++++- crates/tui/src/runtime_web/app.mjs | 60 +- crates/tui/src/runtime_web/app.test.mjs | 44 + crates/tui/src/session_manager.rs | 693 +++++++- crates/tui/src/tools/rlm.rs | 72 +- crates/tui/src/tools/subagent/advisor.rs | 364 ++++- crates/tui/src/tools/subagent/mailbox.rs | 1 + crates/tui/src/tools/subagent/mod.rs | 203 ++- crates/tui/src/tools/subagent/tests.rs | 570 ++++++- crates/tui/src/tui/app/tests.rs | 1 + crates/tui/src/tui/model_picker.rs | 359 +++- crates/tui/src/tui/prompt_suggestion.rs | 2 + crates/tui/src/tui/tool_routing.rs | 37 +- crates/tui/src/tui/translation.rs | 30 +- crates/tui/src/tui/ui.rs | 6 + crates/tui/src/tui/ui/dispatch.rs | 12 + crates/tui/src/tui/ui/event_loop.rs | 568 ++++++- crates/tui/src/tui/ui/tests.rs | 96 ++ crates/tui/src/turn_route_plan.rs | 92 +- 45 files changed, 11904 insertions(+), 521 deletions(-) create mode 100644 crates/tui/src/provider_catalog_live.rs 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 ecb8ea62fc..b503c46cee 100644 --- a/crates/config/src/route/resolver.rs +++ b/crates/config/src/route/resolver.rs @@ -142,6 +142,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(); @@ -232,14 +255,17 @@ 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, protocol, + // 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. + selected.canonical_model = None; + selected.endpoint_key = "chat".to_string(); + 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 a1d5068b44..48d04946a6 100644 --- a/crates/config/src/route/tests.rs +++ b/crates/config/src/route/tests.rs @@ -1612,6 +1612,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/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 7f5398efbc..9c370bcaed 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, @@ -3798,6 +3805,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, @@ -3920,6 +3928,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, @@ -4565,6 +4575,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, @@ -4585,6 +4596,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 @@ -4652,6 +4664,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!( @@ -4751,8 +4769,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. @@ -4773,10 +4791,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 { @@ -4787,8 +4809,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 { @@ -4808,30 +4830,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). @@ -4882,6 +4931,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, @@ -4900,6 +4951,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; @@ -5097,8 +5169,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 @@ -5111,6 +5200,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, @@ -5156,6 +5247,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 @@ -5169,6 +5261,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(); @@ -5182,7 +5278,9 @@ impl Engine { advisor_messages, advisor_config, client, + advisor_route_config, advisor_model, + usage_context, guard, advisor_tx, ) @@ -5248,6 +5346,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, @@ -5291,6 +5391,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, @@ -5333,6 +5435,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, @@ -5360,6 +5464,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, @@ -5424,6 +5530,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, @@ -5449,6 +5557,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, @@ -5506,6 +5616,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 613bda428a..ce6b4d8407 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, @@ -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, @@ -6003,6 +6014,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}; @@ -11003,6 +11139,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, @@ -11144,6 +11281,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, @@ -11280,6 +11418,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, @@ -11485,6 +11624,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, @@ -11761,6 +11901,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, @@ -11947,6 +12088,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, @@ -12087,6 +12229,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, @@ -12223,6 +12366,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, @@ -12363,6 +12507,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, @@ -18721,6 +18866,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, @@ -18846,6 +18992,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, @@ -18950,6 +19097,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, @@ -19197,6 +19345,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, @@ -19425,6 +19574,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, @@ -19648,6 +19798,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 7741427c81..780d8c09c3 100644 --- a/crates/tui/src/core/engine/turn_loop.rs +++ b/crates/tui/src/core/engine/turn_loop.rs @@ -489,15 +489,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), @@ -506,6 +515,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(); @@ -1324,7 +1339,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 @@ -1942,8 +1976,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; @@ -2081,14 +2114,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, @@ -2096,6 +2133,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 { @@ -2388,6 +2442,7 @@ impl Engine { self.process_tool_results( outcomes, + turn, &mut tool_catalog, &mut active_tool_names, &hook_contexts, @@ -3888,6 +3943,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, @@ -3905,12 +3961,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..54518eabf2 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,674 @@ 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 "] { + assert!(!serialized.contains(secret), "quote leaked {secret:?}"); + } + + 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 +2973,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 +3040,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 +3118,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 3679276592..58ce4baea6 100644 --- a/crates/tui/src/exec_agent.rs +++ b/crates/tui/src/exec_agent.rs @@ -396,6 +396,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 661023e1e8..54a4e3f306 100644 --- a/crates/tui/src/lib.rs +++ b/crates/tui/src/lib.rs @@ -91,6 +91,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..c8ba5ab1d8 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 = (provider != ApiProvider::Custom) + .then(|| { + 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() + }) + .unwrap_or_default(); + let custom_rows = (provider == ApiProvider::Custom && exact_matches) + .then(|| { + exact_partition + .map(|partition| partition.offerings.clone()) + .unwrap_or_default() + }) + .unwrap_or_default(); + (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..2887e83be3 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, @@ -2514,6 +2517,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 +2771,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 +3136,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 +3157,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 +3296,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 +3406,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 +3631,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 +3994,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 +4689,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, @@ -5463,11 +5508,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 +5525,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 +5535,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 +5697,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 +5714,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 +5724,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 +6469,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 +7346,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 0a832497f2..795bb7d96c 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, @@ -8032,7 +8449,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(); @@ -8044,6 +8463,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() + })), ); } @@ -8058,6 +8487,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; @@ -8360,16 +8791,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)?; @@ -9082,11 +9530,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, @@ -9219,12 +9671,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..e3c7a6ee6b 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,154 @@ 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 "] { + assert!(!serialized.contains(raw_secret), "leaked {raw_secret:?}"); + } + 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 +3068,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 +3441,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 +3464,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 +3521,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 +3693,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 +3703,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 +5390,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 +5439,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 +5944,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, }), @@ -5331,6 +5970,7 @@ 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, }, @@ -5343,6 +5983,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 +6017,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 +6060,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)); @@ -5438,6 +6114,7 @@ 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(), }, @@ -5453,6 +6130,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 +6181,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 +6270,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 +6332,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 +6873,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 +7112,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 +7444,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 +7561,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 +7675,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 +8024,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 +8276,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 +8416,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 +8487,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 +8677,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 +8863,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 +9693,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 +9767,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 +9883,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 +9955,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 +10055,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 +10147,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 +10286,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 +10372,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 +10440,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 +10603,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 +10692,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 +10967,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 +11010,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 +11226,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 +11256,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 +11532,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 +11891,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..2711882380 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))?; + return 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)?; + return 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..e3fc21c2dc 100644 --- a/crates/tui/src/tools/subagent/mailbox.rs +++ b/crates/tui/src/tools/subagent/mailbox.rs @@ -641,6 +641,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 23ebb3e0ca..d51e317cb7 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(), @@ -14637,12 +14784,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 7c1b0cba9d..a594d668a7 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. @@ -12696,6 +13071,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)), @@ -12720,6 +13096,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, @@ -20267,7 +20644,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); @@ -20277,6 +20656,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); @@ -20311,8 +20699,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(); @@ -20325,7 +20717,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} }))) @@ -20340,6 +20732,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 { @@ -20523,6 +20946,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": "echo guardian-missing-usage | cat"}), + ) + .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": "echo guardian-cache-regression | cat"}); + + 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/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 16421f063c..edc8c6a4f1 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 @@ -883,6 +1075,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()); @@ -1058,14 +1254,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( @@ -1110,10 +1323,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( @@ -1373,7 +1605,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( @@ -1387,24 +1619,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, }); @@ -1461,7 +1707,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( @@ -1475,10 +1721,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, @@ -1487,16 +1742,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 { @@ -1654,7 +1914,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(); @@ -1695,6 +1955,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); @@ -1724,9 +1997,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, @@ -1943,19 +2233,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(), }); @@ -1974,8 +2264,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` @@ -1998,6 +2289,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. @@ -3282,6 +3587,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 @@ -6280,4 +6603,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..bc07314e98 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(), }), @@ -13196,6 +13236,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 +13282,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 +13347,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 +13390,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 +13427,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 +16412,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 +16476,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 +16530,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 +16607,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(); 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!( From 84d64a8f080439323241eea58db08109edc39b0f Mon Sep 17 00:00:00 2001 From: CodeWhale Bot Date: Sun, 30 Aug 2026 01:08:32 -0700 Subject: [PATCH 2/5] =?UTF-8?q?fix(tui):=20repair=20PR=20#5726=20gate=20su?= =?UTF-8?q?rface=20=E2=80=94=20clippy,=20accounting=20and=20wire-route=20t?= =?UTF-8?q?ests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI-form clippy repairs on the checkpoint head: - box MailboxMessage::TokenUsage.route (large_enum_variant) with call-site and test updates in subagent_routing and mailbox - needless_return in session_manager lock-file open, doc_lazy_continuation in translation, obfuscated_if_else in provider_lake - route resolver + provider lake alignment for the opencode_zen wire-route expectations (Responses/AnthropicMessages vs ChatCompletions) Verified locally: - cargo test -p codewhale-tui --locked --lib: 11596 passed, 6 failed — the 6 are the pre-existing main-baseline sandbox::read_guard symlink family that #5724 makes whole-fn-unix; zero Slice-A regressions - auto_review_guardian family 7/7 (accounting + denial + cache replay) - web vitest run of crates/tui/src/runtime_web/app.test.mjs: 28/28 - CI-form clippy clean (previous session) Signed-off-by: CodeWhale Bot --- crates/config/src/route/resolver.rs | 10 ++++-- crates/tui/src/core/engine/tests.rs | 20 +++++++---- crates/tui/src/provider_lake.rs | 44 ++++++++++++------------ crates/tui/src/runtime_threads.rs | 5 ++- crates/tui/src/runtime_threads/tests.rs | 8 ++--- crates/tui/src/session_manager.rs | 4 +-- crates/tui/src/tools/subagent/mailbox.rs | 9 ++--- crates/tui/src/tools/subagent/tests.rs | 17 +++++++-- crates/tui/src/tui/subagent_routing.rs | 6 ++-- crates/tui/src/tui/translation.rs | 2 ++ crates/tui/src/tui/ui/tests.rs | 10 +++--- crates/tui/src/tui/work_surface/mod.rs | 2 +- 12 files changed, 84 insertions(+), 53 deletions(-) diff --git a/crates/config/src/route/resolver.rs b/crates/config/src/route/resolver.rs index b503c46cee..f0f55b7849 100644 --- a/crates/config/src/route/resolver.rs +++ b/crates/config/src/route/resolver.rs @@ -259,12 +259,16 @@ impl RouteResolver { // 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 canonical model, protocol, - // limits, modality, tool, reasoning, or billing contract. Keep the + // 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.endpoint_key = "chat".to_string(); selected.limits = RouteLimits::default(); selected.capabilities = RouteCapabilities::default(); selected.pricing = PricingSku::UnknownOrStale; diff --git a/crates/tui/src/core/engine/tests.rs b/crates/tui/src/core/engine/tests.rs index ce6b4d8407..9c225178da 100644 --- a/crates/tui/src/core/engine/tests.rs +++ b/crates/tui/src/core/engine/tests.rs @@ -4102,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(), } )); @@ -4196,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, @@ -5659,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) @@ -7395,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, diff --git a/crates/tui/src/provider_lake.rs b/crates/tui/src/provider_lake.rs index c8ba5ab1d8..f1d33e128a 100644 --- a/crates/tui/src/provider_lake.rs +++ b/crates/tui/src/provider_lake.rs @@ -586,28 +586,28 @@ pub(crate) fn runtime_catalog_resolver_for_identity( if !exact_matches && let LivePartitionOwner::BuiltIn(identity) = &partition_owner { authoritative.remove(identity); } - let fallback = (provider != ApiProvider::Custom) - .then(|| { - 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() - }) - .unwrap_or_default(); - let custom_rows = (provider == ApiProvider::Custom && exact_matches) - .then(|| { - exact_partition - .map(|partition| partition.offerings.clone()) - .unwrap_or_default() - }) - .unwrap_or_default(); + 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 { ( diff --git a/crates/tui/src/runtime_threads.rs b/crates/tui/src/runtime_threads.rs index 795bb7d96c..c3fadbd182 100644 --- a/crates/tui/src/runtime_threads.rs +++ b/crates/tui/src/runtime_threads.rs @@ -8894,7 +8894,10 @@ impl RuntimeThreadManager { self.append_routed_usage_to_turn( &turn_id, &source_id, - EffectiveRouteUsage { route, usage }, + EffectiveRouteUsage { + route: *route, + usage, + }, )?; } } diff --git a/crates/tui/src/runtime_threads/tests.rs b/crates/tui/src/runtime_threads/tests.rs index e3c7a6ee6b..4e2b1c5d29 100644 --- a/crates/tui/src/runtime_threads/tests.rs +++ b/crates/tui/src/runtime_threads/tests.rs @@ -5962,7 +5962,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(), @@ -5973,7 +5973,7 @@ async fn monitor_separates_lifecycle_start_from_billing_dispatch_and_child_usage provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Subscription, dispatched_at, - }, + }), usage: Usage { input_tokens: 3, output_tokens: 2, @@ -6106,7 +6106,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(), @@ -6117,7 +6117,7 @@ async fn monitor_separates_lifecycle_start_from_billing_dispatch_and_child_usage provider_live_pricing: None, billing_mode: crate::cost_status::RouteBillingMode::Subscription, dispatched_at: Utc::now(), - }, + }), usage: Usage { input_tokens: 5, output_tokens: 1, diff --git a/crates/tui/src/session_manager.rs b/crates/tui/src/session_manager.rs index 2711882380..327b67c43b 100644 --- a/crates/tui/src/session_manager.rs +++ b/crates/tui/src/session_manager.rs @@ -124,7 +124,7 @@ fn open_private_lock_file(path: &Path) -> io::Result { let file = options.open(path)?; validate_private_regular_file(&file, path)?; file.set_permissions(fs::Permissions::from_mode(0o600))?; - return Ok(file); + Ok(file) } #[cfg(windows)] { @@ -133,7 +133,7 @@ fn open_private_lock_file(path: &Path) -> io::Result { options.custom_flags(FILE_FLAG_OPEN_REPARSE_POINT); let file = options.open(path)?; validate_private_regular_file(&file, path)?; - return Ok(file); + Ok(file) } #[cfg(all(not(unix), not(windows)))] { diff --git a/crates/tui/src/tools/subagent/mailbox.rs b/crates/tui/src/tools/subagent/mailbox.rs index e3fc21c2dc..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, diff --git a/crates/tui/src/tools/subagent/tests.rs b/crates/tui/src/tools/subagent/tests.rs index a594d668a7..f042318d2b 100644 --- a/crates/tui/src/tools/subagent/tests.rs +++ b/crates/tui/src/tools/subagent/tests.rs @@ -20902,6 +20902,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( @@ -21004,7 +21017,7 @@ mod child_permission_gate { .execute( "agent_gate", "bash", - json!({"command": "echo guardian-missing-usage | cat"}), + json!({"command": guardian_marker_pipeline("guardian-missing-usage")}), ) .await .expect("semantic guardian success remains usable"); @@ -21044,7 +21057,7 @@ mod child_permission_gate { .await; let (registry, mut rx, manager) = worker_registry(ApprovalMode::Auto, false, true, Some(client)); - let input = json!({"command": "echo guardian-cache-regression | cat"}); + let input = json!({"command": guardian_marker_pipeline("guardian-cache-regression")}); for _ in 0..2 { let output = registry diff --git a/crates/tui/src/tui/subagent_routing.rs b/crates/tui/src/tui/subagent_routing.rs index 3cb1d10700..e7fc123488 100644 --- a/crates/tui/src/tui/subagent_routing.rs +++ b/crates/tui/src/tui/subagent_routing.rs @@ -1331,7 +1331,7 @@ 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 +1353,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 +1367,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/translation.rs b/crates/tui/src/tui/translation.rs index 9c92aa3696..433067f7b5 100644 --- a/crates/tui/src/tui/translation.rs +++ b/crates/tui/src/tui/translation.rs @@ -10,10 +10,12 @@ //! //! - `needs_translation()` — heuristic to detect if text is predominantly //! English and should be translated. +//! //! The event loop dispatches focused translation requests through the exact //! provider/model client frozen for the originating turn. The dedicated agent //! receives only the source text and returns only the translation — no tool //! calls or conversation history. +//! //! - `TranslationStatus` — tracks per-message translation status in the UI. /// Heuristic threshold: if more than this fraction of alphabetic characters diff --git a/crates/tui/src/tui/ui/tests.rs b/crates/tui/src/tui/ui/tests.rs index bc07314e98..8d30ef6ddb 100644 --- a/crates/tui/src/tui/ui/tests.rs +++ b/crates/tui/src/tui/ui/tests.rs @@ -11995,7 +11995,7 @@ 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, @@ -12022,7 +12022,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, @@ -12043,7 +12043,7 @@ 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, @@ -12070,7 +12070,7 @@ 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, @@ -21273,7 +21273,7 @@ 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, From b6a14d0c21137a94dd7a0d7f69297038a5ba67ce Mon Sep 17 00:00:00 2001 From: CodeWhale Bot Date: Sun, 30 Aug 2026 01:18:23 -0700 Subject: [PATCH 3/5] test(tui): align runtime-web drift guard with paginated catalog collector MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The new-thread dialog now loads provider catalogs through collectProviderModelPages (bounded pages, cursor-progress checks, injected fetch boundary). The drift guard still pinned the pre-pagination single-shot URL shape and failed on the rewritten app.mjs. Pin the new contract: dialog calls collectProviderModelPages(provider.id) and the wire endpoint remains /v1/providers//models. Also rustfmt two files from the previous gate-repair commit (CI 'Check formatting' failure). Verified: node --test crates/tui/tests/runtime_web_client.test.mjs — pass 36, fail 0. cargo fmt --all -- --check clean. Signed-off-by: CodeWhale Bot --- crates/tui/src/tui/subagent_routing.rs | 5 ++++- crates/tui/src/tui/ui/tests.rs | 20 ++++++++++++++++---- crates/tui/tests/runtime_web_client.test.mjs | 5 ++++- 3 files changed, 24 insertions(+), 6 deletions(-) diff --git a/crates/tui/src/tui/subagent_routing.rs b/crates/tui/src/tui/subagent_routing.rs index e7fc123488..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: Box::new(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(), }, ); diff --git a/crates/tui/src/tui/ui/tests.rs b/crates/tui/src/tui/ui/tests.rs index 8d30ef6ddb..bfb6b6d701 100644 --- a/crates/tui/src/tui/ui/tests.rs +++ b/crates/tui/src/tui/ui/tests.rs @@ -11995,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: Box::new(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, @@ -12043,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: Box::new(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, @@ -12070,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: Box::new(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, @@ -21273,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: Box::new(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/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 \}\)/); From 9f381e00e92dc07c0c9ee87dc921e427d6ba93ef Mon Sep 17 00:00:00 2001 From: CodeWhale Bot Date: Sun, 30 Aug 2026 02:04:16 -0700 Subject: [PATCH 4/5] test(tui): make CodeQL's cleartext classes provable in test code MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - runtime_api session tests: build request URLs with the loopback host spelled literally (127.0.0.1:{port}) instead of an opaque {addr} variable, so the cleartext-transmission analysis can prove the request never leaves the host. Behavior unchanged; the listeners already bind 127.0.0.1. - cost_status / runtime_threads anti-leak assertions: stop interpolating the credential fragment into the panic message — the assertion message must not itself become the logging sink it guards against. Verified: cargo check -p codewhale-tui --tests --locked clean; 194 passed / 0 failed across cost_status + runtime_threads filters; cargo fmt clean. Signed-off-by: CodeWhale Bot --- crates/tui/src/cost_status.rs | 8 ++++- crates/tui/src/runtime_api/tests.rs | 46 +++++++++++++++++++------ crates/tui/src/runtime_threads/tests.rs | 7 +++- 3 files changed, 49 insertions(+), 12 deletions(-) diff --git a/crates/tui/src/cost_status.rs b/crates/tui/src/cost_status.rs index 54518eabf2..cbe7335717 100644 --- a/crates/tui/src/cost_status.rs +++ b/crates/tui/src/cost_status.rs @@ -2717,7 +2717,13 @@ mod tests { assert!(serialized.contains("catalog_revision")); assert!(serialized.contains("input_per_million")); for secret in [codewhale_config::BASETEN_BASE_URL, "api_key", "Bearer "] { - assert!(!serialized.contains(secret), "quote leaked {secret:?}"); + // 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!({}); diff --git a/crates/tui/src/runtime_api/tests.rs b/crates/tui/src/runtime_api/tests.rs index 2887e83be3..24b85c6fce 100644 --- a/crates/tui/src/runtime_api/tests.rs +++ b/crates/tui/src/runtime_api/tests.rs @@ -4838,7 +4838,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()? @@ -4868,7 +4871,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? @@ -4899,7 +4903,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? @@ -4919,7 +4926,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? @@ -4936,7 +4946,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?; @@ -4944,7 +4957,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?; @@ -4971,7 +4987,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?; @@ -4979,7 +4998,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?; @@ -5002,7 +5024,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? @@ -5024,7 +5047,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()? diff --git a/crates/tui/src/runtime_threads/tests.rs b/crates/tui/src/runtime_threads/tests.rs index 4e2b1c5d29..72b1eccf45 100644 --- a/crates/tui/src/runtime_threads/tests.rs +++ b/crates/tui/src/runtime_threads/tests.rs @@ -2821,7 +2821,12 @@ fn turn_record_round_trips_frozen_provider_live_pricing_and_drops_hostile_quotes 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 "] { - assert!(!serialized.contains(raw_secret), "leaked {raw_secret:?}"); + // 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 From d8b887f24ee43c8e1d3c0a6992ff423208a236c2 Mon Sep 17 00:00:00 2001 From: CodeWhale Bot Date: Sun, 30 Aug 2026 12:59:16 -0700 Subject: [PATCH 5/5] test(runtime-api): build and serve the test router on a product-sized stack The Runtime API test harness ran its whole setup (Config load via toml::de visit_map frames, manager construction) and every axum handler on the 2 MiB libtest thread. Config load under a profile and the thread-lifecycle path marginally overflowed that stack in debug builds, aborting the whole lib suite with SIGABRT; CI masked it with RUST_MIN_STACK=16MiB. The harness now builds and serves on a dedicated thread sized with CODEWHALE_MAIN_STACK_BYTES (16 MiB, matching lib.rs runtime workers and the .cnb.yml gate), with the test's sealed env scope adopted onto that thread via env_scope_ticket/join_env_scope so Config::load env reads never block on the mutex the sealing test holds while awaiting setup. The libtest thread keeps only the test body and its HTTP client. Verified with no RUST_MIN_STACK override: the previously aborting reload_config/lifecycle filters and the provider_lake/catalog/guardian/ kernel filters all pass (279 passed; 0 failed). Signed-off-by: CodeWhale Bot --- CHANGELOG.md | 6 + crates/tui/CHANGELOG.md | 6 + crates/tui/src/runtime_api/tests.rs | 164 ++++++++++++++++++++++++++-- 3 files changed, 165 insertions(+), 11 deletions(-) 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/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/runtime_api/tests.rs b/crates/tui/src/runtime_api/tests.rs index 24b85c6fce..c254bf3fea 100644 --- a/crates/tui/src/runtime_api/tests.rs +++ b/crates/tui/src/runtime_api/tests.rs @@ -922,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)?; @@ -937,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 { @@ -969,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()), @@ -1005,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<