diff --git a/Dockerfile b/Dockerfile index ad0ca7079..9cfc30c6b 100644 --- a/Dockerfile +++ b/Dockerfile @@ -47,10 +47,13 @@ FROM docker.io/envoyproxy/envoy:${ENVOY_VERSION} AS envoy FROM python:3.14-slim AS arch +# Install runtime deps first, then upgrade — so security patches also land on +# packages pulled in transitively here (e.g. curl -> libssh2-1t64), not just those +# already present in the base image. RUN set -eux; \ apt-get update; \ - apt-get upgrade -y; \ apt-get install -y --no-install-recommends gettext-base curl procps; \ + apt-get upgrade -y; \ apt-get clean; rm -rf /var/lib/apt/lists/* RUN pip install --no-cache-dir supervisor diff --git a/config/envoy.template.yaml b/config/envoy.template.yaml index b2b9fb1f6..d3c778305 100644 --- a/config/envoy.template.yaml +++ b/config/envoy.template.yaml @@ -269,6 +269,13 @@ static_resources: auto_host_rewrite: true cluster: {{ cluster_name }} timeout: 300s + {% if cluster.prefix_affinity %} + # Feeds the cluster's RING_HASH lb_policy: requests with the + # same prompt-prefix hash stick to the same replica. + hash_policy: + - header: + header_name: x-plano-prefix-hash + {% endif %} {% endfor %} http_filters: - name: envoy.filters.http.router @@ -977,9 +984,18 @@ static_resources: {% else -%} connect_timeout: {{ upstream_connect_timeout | default('5s') }} {% endif -%} + {% if cluster.prefix_affinity -%} + # KV-aware replica stickiness: resolve every replica address and consistent-hash + # requests (by x-plano-prefix-hash, see the route-level hash_policy) so the same + # prompt prefix lands on the replica holding its warm KV cache. + type: STRICT_DNS + dns_lookup_family: V4_ONLY + lb_policy: RING_HASH + {% else -%} type: LOGICAL_DNS dns_lookup_family: V4_ONLY lb_policy: ROUND_ROBIN + {% endif -%} load_assignment: cluster_name: {{ cluster_name }} endpoints: diff --git a/config/plano_config_schema.yaml b/config/plano_config_schema.yaml index f356b5735..d4549d57a 100644 --- a/config/plano_config_schema.yaml +++ b/config/plano_config_schema.yaml @@ -155,6 +155,13 @@ properties: - https http_host: type: string + prefix_affinity: + type: boolean + description: > + For self-hosted multi-replica backends (e.g. vLLM): consistent-hash requests + to replicas by the x-plano-prefix-hash header so the same prompt prefix lands + on the same replica and reuses its KV cache. Uses ring-hash load balancing + across all resolved endpoint addresses. Default false. additionalProperties: false required: - endpoint @@ -316,6 +323,33 @@ properties: orchestrator_model_context_length: type: integer description: "Maximum token length for the orchestrator/routing model context window. Default is 8192." + prompt_caching: + type: object + description: > + Automatic provider prompt caching, configured once for the whole Plano instance. + Disabled by default; set enabled: true to opt in. Prompt caching never changes + which model routing selects — it only keeps a conversation on the same + model/provider so the upstream prompt cache stays warm across turns, and + auto-injects provider cache-control markers where supported. + properties: + enabled: + type: boolean + description: "Master switch. Default false (opt-in). Applies across the entire instance." + session_affinity: + type: boolean + description: "Auto-derive a session key from the prompt prefix (system + tools + first user message) when X-Model-Affinity is absent, so follow-up turns reuse the same warm cache. Default true when enabled." + inject_cache_control: + type: boolean + description: "Auto-insert ephemeral cache_control breakpoints for providers that need explicit markers (Anthropic). Default true when enabled." + min_prefix_tokens: + type: integer + minimum: 1 + description: "Skip breakpoint injection when the estimated stable prefix is below this many tokens. Default 1024." + session_ttl_seconds: + type: integer + minimum: 1 + description: "Pin lifetime for implicit/explicit sessions; align with the provider cache window (e.g. 3600 for Anthropic 1h caching). Defaults to routing.session_ttl_seconds." + additionalProperties: false system_prompt: type: string prompt_targets: @@ -510,6 +544,51 @@ properties: Optional HTTP header name whose value is used as a tenant prefix in the cache key. When set, keys are scoped as plano:affinity:{tenant_id}:{session_id}. additionalProperties: false + routing_budget: + type: object + description: > + Per-session cost gate on model switching. Independent of prompt caching — it + applies whenever configured (presence of this block turns it on). The default + posture is to stick to the model a session is warm on. When routing proposes a + different model while that session's provider cache is plausibly still warm + (inferred from the idle gap vs. the provider's cache window), the actual + input-token cost of abandoning the cache — + context_tokens x (candidate_uncached_input_rate - anchor_cached_input_rate), + output cost deliberately excluded — is drawn from the session's budget: a paid + switch is allowed only while budget remains, and an outright-cheaper switch is + free and can credit the budget back. Requires a cost source in + model_metrics_sources. + properties: + seed_usd: + type: number + minimum: 0 + description: > + Initial cumulative budget (USD) granted to a session. 0 means never pay to + switch (only outright-cheaper switches are allowed); larger values buy more + quality-driven switches over the session. + replenish_on_rebind: + type: boolean + description: "Re-seed the budget to seed_usd when a cold session re-binds. Default true." + credit_negative: + type: boolean + description: "Credit the budget back when a switch is outright cheaper (negative cost). Default true." + cache_read_discount: + type: number + minimum: 0 + maximum: 1 + description: > + Fallback used to estimate a model's cached input rate when the pricing feed + doesn't publish one (cached_rate = input_rate x discount). A pricing detail, + not a cost policy. Default 0.1. + record_counterfactual: + type: boolean + description: > + When true, a vetoed switch records the route the gate would have taken had + the switch been allowed, as the plano.switch.counterfactual_route span + attribute. Telemetry only — the counterfactual model is never dispatched. + Useful for evals/benchmarks. Default false. + required: ["seed_usd"] + additionalProperties: false additionalProperties: false state_storage: type: object diff --git a/crates/brightstaff/src/affinity.rs b/crates/brightstaff/src/affinity.rs new file mode 100644 index 000000000..bf05d9565 --- /dev/null +++ b/crates/brightstaff/src/affinity.rs @@ -0,0 +1,181 @@ +//! Implicit session-affinity key derivation. +//! +//! When a request carries no explicit `X-Model-Affinity` header, Plano derives a +//! stable session key from the parts of the prompt that repeat verbatim at the head +//! of every turn — the same bytes the provider's prompt cache is keyed on: +//! +//! ```text +//! session_key = hash(system + tools + first_user_message) +//! prefix_hash = hash(system + tools) +//! ``` +//! +//! The session key is constant for the life of a conversation (history grows at the +//! tail, not the head), so turns 2+ reuse the same pin without any client changes. +//! The prefix hash covers only the fully-stable segment and is stored with the pin +//! for drift detection: if it changes, the provider cache is already lost and +//! re-routing fresh is safe. +//! +//! Only salted hashes are ever stored — never prompt content. + +use hermesllm::apis::openai::{Message, Role}; + +/// Salt folded into every hash so stored keys can't be trivially correlated with +/// prompt content across systems. Deterministic across processes/replicas so a +/// shared Redis session cache keys consistently. +const HASH_SALT: &str = "plano-affinity-v1"; + +/// Prefix distinguishing derived keys from client-supplied `X-Model-Affinity` ids. +const IMPLICIT_KEY_PREFIX: &str = "implicit:"; + +/// Derived affinity identifiers for one request. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ImplicitAffinity { + /// Session-cache key: `implicit:{hex}` over system + tools + first user message. + pub session_key: String, + /// Hash of the stable prefix only (system + tools), for drift detection. + pub prefix_hash: u64, +} + +/// FNV-1a 64-bit — stable across processes and Rust versions (unlike `DefaultHasher`), +/// dependency-free, and plenty for cache keying (collisions merely over-pin, which is +/// cache-friendly). +fn fnv1a64(chunks: &[&str]) -> u64 { + const OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325; + const PRIME: u64 = 0x0000_0100_0000_01b3; + let mut hash = OFFSET_BASIS; + let mut feed = |bytes: &[u8]| { + for &b in bytes { + hash ^= b as u64; + hash = hash.wrapping_mul(PRIME); + } + // Field separator so ("ab","c") and ("a","bc") hash differently. + hash ^= 0x1f; + hash = hash.wrapping_mul(PRIME); + }; + feed(HASH_SALT.as_bytes()); + for chunk in chunks { + feed(chunk.as_bytes()); + } + hash +} + +/// Derive the implicit affinity key from parsed request messages and tool names. +/// +/// Returns `None` when there is no user message to anchor on (nothing distinguishes +/// the conversation, so pinning would be meaningless). +pub fn derive_implicit_affinity( + messages: &[Message], + tool_names: Option<&[String]>, + tenant_id: Option<&str>, +) -> Option { + let system_text: String = messages + .iter() + .filter(|m| matches!(m.role, Role::System | Role::Developer)) + .filter_map(|m| m.content.as_ref().map(|c| c.to_string())) + .collect::>() + .join("\n"); + + let first_user = messages + .iter() + .find(|m| matches!(m.role, Role::User)) + .and_then(|m| m.content.as_ref().map(|c| c.to_string()))?; + + let tools_text = tool_names.map(|names| names.join(",")).unwrap_or_default(); + let tenant = tenant_id.unwrap_or_default(); + + let prefix_hash = fnv1a64(&[tenant, &system_text, &tools_text]); + let session_hash = fnv1a64(&[tenant, &system_text, &tools_text, &first_user]); + + Some(ImplicitAffinity { + session_key: format!("{IMPLICIT_KEY_PREFIX}{session_hash:016x}"), + prefix_hash, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use hermesllm::apis::openai::MessageContent; + + fn msg(role: Role, text: &str) -> Message { + Message { + role, + content: Some(MessageContent::Text(text.to_string())), + name: None, + tool_call_id: None, + tool_calls: None, + } + } + + #[test] + fn key_is_stable_as_history_grows() { + let turn1 = vec![ + msg(Role::System, "You are a coding agent."), + msg(Role::User, "Fix the bug in main.rs"), + ]; + let turn5 = vec![ + msg(Role::System, "You are a coding agent."), + msg(Role::User, "Fix the bug in main.rs"), + msg(Role::Assistant, "Done. Anything else?"), + msg(Role::User, "Now add tests"), + msg(Role::Assistant, "Added."), + ]; + + let a1 = derive_implicit_affinity(&turn1, None, None).unwrap(); + let a5 = derive_implicit_affinity(&turn5, None, None).unwrap(); + assert_eq!(a1.session_key, a5.session_key); + assert_eq!(a1.prefix_hash, a5.prefix_hash); + assert!(a1.session_key.starts_with("implicit:")); + } + + #[test] + fn different_first_user_message_yields_different_session() { + let base = msg(Role::System, "You are a coding agent."); + let a = derive_implicit_affinity( + &[base.clone(), msg(Role::User, "conversation A")], + None, + None, + ) + .unwrap(); + let b = derive_implicit_affinity(&[base, msg(Role::User, "conversation B")], None, None) + .unwrap(); + // Same stable prefix, different conversations. + assert_eq!(a.prefix_hash, b.prefix_hash); + assert_ne!(a.session_key, b.session_key); + } + + #[test] + fn changed_system_prompt_changes_prefix_hash() { + let a = derive_implicit_affinity( + &[msg(Role::System, "v1 prompt"), msg(Role::User, "hi")], + None, + None, + ) + .unwrap(); + let b = derive_implicit_affinity( + &[msg(Role::System, "v2 prompt"), msg(Role::User, "hi")], + None, + None, + ) + .unwrap(); + assert_ne!(a.prefix_hash, b.prefix_hash); + assert_ne!(a.session_key, b.session_key); + } + + #[test] + fn tools_and_tenant_are_part_of_the_key() { + let messages = [msg(Role::System, "s"), msg(Role::User, "u")]; + let plain = derive_implicit_affinity(&messages, None, None).unwrap(); + let with_tools = + derive_implicit_affinity(&messages, Some(&["get_weather".to_string()]), None).unwrap(); + let with_tenant = derive_implicit_affinity(&messages, None, Some("acme")).unwrap(); + assert_ne!(plain.session_key, with_tools.session_key); + assert_ne!(plain.session_key, with_tenant.session_key); + } + + #[test] + fn no_user_message_yields_none() { + assert!(derive_implicit_affinity(&[msg(Role::System, "s")], None, None).is_none()); + assert!(derive_implicit_affinity(&[], None, None).is_none()); + } +} diff --git a/crates/brightstaff/src/app_state.rs b/crates/brightstaff/src/app_state.rs index de4393f3e..da8d15b95 100644 --- a/crates/brightstaff/src/app_state.rs +++ b/crates/brightstaff/src/app_state.rs @@ -1,7 +1,10 @@ use std::collections::HashMap; use std::sync::Arc; -use common::configuration::{Agent, FilterPipeline, Listener, ModelAlias, SpanAttributes}; +use common::configuration::{ + Agent, EffectivePromptCaching, EffectiveRoutingBudget, FilterPipeline, Listener, ModelAlias, + SpanAttributes, +}; use common::llm_providers::LlmProviders; use tokio::sync::RwLock; @@ -31,4 +34,10 @@ pub struct AppState { /// When false, agentic signal analysis is skipped on LLM responses to save CPU. /// Controlled by `overrides.disable_signals` in plano config. pub signals_enabled: bool, + /// Instance-wide automatic prompt-caching settings, resolved once from the + /// top-level `prompt_caching` config. Disabled by default (opt-in). + pub prompt_caching: EffectivePromptCaching, + /// Per-session model-switch cost gate, resolved from `routing.routing_budget`. + /// Independent of prompt caching; `None` when not configured (off by default). + pub routing_budget: Option, } diff --git a/crates/brightstaff/src/handlers/function_calling.rs b/crates/brightstaff/src/handlers/function_calling.rs index 3e2543bc6..24ca5c524 100644 --- a/crates/brightstaff/src/handlers/function_calling.rs +++ b/crates/brightstaff/src/handlers/function_calling.rs @@ -1239,6 +1239,7 @@ impl ArchFunctionHandler { total_tokens: 0, prompt_tokens_details: None, completion_tokens_details: None, + ..Default::default() }, system_fingerprint: None, service_tier: None, diff --git a/crates/brightstaff/src/handlers/llm/mod.rs b/crates/brightstaff/src/handlers/llm/mod.rs index a1d3ce973..eee96d5d1 100644 --- a/crates/brightstaff/src/handlers/llm/mod.rs +++ b/crates/brightstaff/src/handlers/llm/mod.rs @@ -1,6 +1,9 @@ use bytes::Bytes; use common::configuration::{FilterPipeline, ModelAlias}; -use common::consts::{ARCH_IS_STREAMING_HEADER, ARCH_PROVIDER_HINT_HEADER, MODEL_AFFINITY_HEADER}; +use common::consts::{ + ARCH_IS_STREAMING_HEADER, ARCH_PROVIDER_HINT_HEADER, MODEL_AFFINITY_HEADER, PLANO_CACHE_HEADER, + PLANO_PREFIX_HASH_HEADER, +}; use common::llm_providers::LlmProviders; use hermesllm::apis::openai::Message; use hermesllm::apis::openai_responses::InputParam; @@ -19,6 +22,8 @@ use tokio::sync::RwLock; use tracing::{debug, info, info_span, warn, Instrument}; pub(crate) mod model_selection; +pub(crate) mod prompt_caching; +pub(crate) mod session_router; use crate::app_state::AppState; use crate::handlers::agents::pipeline::PipelineProcessor; @@ -31,7 +36,7 @@ use crate::state::{ }; use crate::streaming::{ create_streaming_response, create_streaming_response_with_output_filter, truncate_message, - LlmMetricsCtx, ObservableStreamProcessor, StreamProcessor, + LlmMetricsCtx, ObservableStreamProcessor, SessionUpdateCtx, StreamProcessor, }; use crate::tracing::{ collect_custom_trace_attributes, llm as tracing_llm, operation_component, @@ -112,8 +117,7 @@ async fn llm_chat_inner( } } - // Session pinning: extract session ID and check cache before routing - let session_id: Option = request_headers + let explicit_session_id: Option = request_headers .get(MODEL_AFFINITY_HEADER) .and_then(|h| h.to_str().ok()) .map(|s| s.to_string()); @@ -123,40 +127,17 @@ async fn llm_chat_inner( .and_then(|hdr| request_headers.get(hdr)) .and_then(|v| v.to_str().ok()) .map(|s| s.to_string()); - let cached_route = if let Some(ref sid) = session_id { - state - .orchestrator_service - .get_cached_route(sid, tenant_id.as_deref()) - .await - } else { - None - }; - let (pinned_model, pinned_route_name): (Option, Option) = match cached_route { - Some(c) => (Some(c.model_name), c.route_name), - None => (None, None), - }; - - // Record session id on the LLM span for the observability console. - if let Some(ref sid) = session_id { - get_active_span(|span| { - span.set_attribute(opentelemetry::KeyValue::new( - tracing_plano::SESSION_ID, - sid.clone(), - )); - }); - } - if let Some(ref route_name) = pinned_route_name { - get_active_span(|span| { - span.set_attribute(opentelemetry::KeyValue::new( - tracing_plano::ROUTE_NAME, - route_name.clone(), - )); - }); - } + // `X-Plano-Cache: off` disables implicit pinning + marker injection for one call. + let cache_off_for_request = request_headers + .get(PLANO_CACHE_HEADER) + .and_then(|h| h.to_str().ok()) + .is_some_and(|v| v.eq_ignore_ascii_case("off")); let full_qualified_llm_provider_url = format!("{}{}", state.llm_provider_url, request_path); // --- Phase 1: Parse and validate the incoming request --- + // (Parsing happens before session-key derivation: the implicit affinity key is + // computed from the request body's stable prompt prefix.) let parsed = match parse_and_validate_request( request, &request_path, @@ -187,6 +168,14 @@ async fn llm_chat_inner( provider_id, } = parsed; + // Prompt caching is configured once for the whole instance and never influences + // which model routing selects — it only keeps a conversation on the same + // model/provider so the upstream prompt cache stays warm across turns. Note: + // session affinity and pin lookup are derived later (Phase 2a), *after* input + // filters and Responses-API state merge, so the prefix hash reflects the exact + // stable prefix sent upstream rather than the pre-filter request body. + let prompt_caching = state.prompt_caching; + // Record LLM-specific span attributes let span = tracing::Span::current(); if let Some(temp) = temperature { @@ -283,6 +272,18 @@ async fn llm_chat_inner( *bad_request.status_mut() = StatusCode::BAD_REQUEST; return Ok(bad_request); } + + // Auto-inject prompt-cache markers (grouped in `prompt_caching`). No-op unless + // caching is enabled for this request and the model needs explicit markers. + prompt_caching::inject_cache_markers( + &mut client_request, + provider_id, + &model_name_only, + &upstream_api, + &prompt_caching, + cache_off_for_request, + &alias_resolved_model, + ); } // --- Phase 2: Resolve conversation state (v1/responses API) --- @@ -301,6 +302,40 @@ async fn llm_chat_inner( Err(response) => return Ok(response), }; + // --- Phase 2a: Session affinity (see `session_router`) --- + // Derived here, after input filters and Responses-API state merge, so the + // implicit session key and prefix hash are computed over the exact stable + // prefix (system + tools + first user message) that is actually sent upstream — + // the same bytes the provider's prompt cache is keyed on. The prefix hash is + // derived even when caching is off, so the `x-plano-prefix-hash` RING_HASH + // replica-stickiness header still works. + let routing_budget = state.routing_budget; + // Derive the implicit session key when either prompt-caching affinity or the + // routing budget is active — the budget needs a session anchor even with caching off. + let implicit_affinity_enabled = prompt_caching.session_affinity || routing_budget.is_some(); + let request_messages = client_request.get_messages(); + let session_router::SessionResolution { + request_prefix_hash, + session_id, + } = session_router::resolve_session( + explicit_session_id, + &request_messages, + tool_names.as_deref(), + tenant_id.as_deref(), + implicit_affinity_enabled, + cache_off_for_request, + ); + + // Record session id on the LLM span for the observability console. + if let Some(ref sid) = session_id { + get_active_span(|span| { + span.set_attribute(opentelemetry::KeyValue::new( + tracing_plano::SESSION_ID, + sid.clone(), + )); + }); + } + // Serialize request for upstream BEFORE router consumes it let client_request_bytes_for_upstream: Bytes = match ProviderRequestType::to_bytes(&client_request) { @@ -313,76 +348,85 @@ async fn llm_chat_inner( } }; - // --- Phase 3: Route the request (or use pinned model from session cache) --- - let resolved_model = if let Some(cached_model) = pinned_model { - info!( - session_id = %session_id.as_deref().unwrap_or(""), - model = %cached_model, - "using pinned routing decision from cache" - ); - cached_model + // Context size for the switch-cost estimate, computed before the router consumes + // the request. Only needed when stickiness could act on a session. + let est_context_tokens: u64 = if session_id.is_some() && routing_budget.is_some() { + session_router::estimate_context_tokens(&request_messages, &model_name_only) } else { - let routing_span = info_span!( - "routing", - component = "routing", - http.method = "POST", - http.target = %request_path, - model.requested = %model_from_request, - model.alias_resolved = %alias_resolved_model, - route.selected_model = tracing::field::Empty, - routing.determination_ms = tracing::field::Empty, - ); - let routing_result = match async { - set_service_name(operation_component::ROUTING); - router_chat_get_upstream_model( - Arc::clone(&state.orchestrator_service), - client_request, - &request_path, - &request_id, - inline_routing_preferences, - ) - .await - } - .instrument(routing_span) + 0 + }; + + // --- Phase 3: Route the request (quality), then apply the session-cache decision --- + // Routing stays cache-blind: the quality router always picks a candidate. The + // session router then honors it or sticks to the warm anchor (see `session_router`). + let routing_span = info_span!( + "routing", + component = "routing", + http.method = "POST", + http.target = %request_path, + model.requested = %model_from_request, + model.alias_resolved = %alias_resolved_model, + route.selected_model = tracing::field::Empty, + routing.determination_ms = tracing::field::Empty, + ); + let routing_result = match async { + set_service_name(operation_component::ROUTING); + router_chat_get_upstream_model( + Arc::clone(&state.orchestrator_service), + client_request, + &request_path, + &request_id, + inline_routing_preferences, + ) .await - { - Ok(result) => result, - Err(err) => { - let mut internal_error = Response::new(full(err.message)); - *internal_error.status_mut() = err.status_code; - return Ok(internal_error); - } - }; + } + .instrument(routing_span) + .await + { + Ok(result) => result, + Err(err) => { + let mut internal_error = Response::new(full(err.message)); + *internal_error.status_mut() = err.status_code; + return Ok(internal_error); + } + }; - let (router_selected_model, route_name) = - (routing_result.model_name, routing_result.route_name); - let model = if router_selected_model != "none" { - router_selected_model - } else { - alias_resolved_model.clone() - }; + let candidate_model = if routing_result.model_name != "none" { + routing_result.model_name + } else { + alias_resolved_model.clone() + }; + let candidate_route = routing_result.route_name; + + let decision = session_router::route( + &state.orchestrator_service, + routing_budget.as_ref(), + session_router::RouteFacts { + session_id: session_id.as_deref(), + tenant_id: tenant_id.as_deref(), + prefix_hash: request_prefix_hash, + est_context_tokens, + candidate_model: &candidate_model, + candidate_route: candidate_route.as_deref(), + }, + ) + .await; - // Record route name on the LLM span (only when the orchestrator produced one). - if let Some(ref rn) = route_name { - if !rn.is_empty() && rn != "none" { - get_active_span(|span| { - span.set_attribute(opentelemetry::KeyValue::new( - tracing_plano::ROUTE_NAME, - rn.clone(), - )); - }); - } - } + let resolved_model = decision.model; + let resolved_route_name = decision.route_name; - if let Some(ref sid) = session_id { - state - .orchestrator_service - .cache_route(sid.clone(), tenant_id.as_deref(), model.clone(), route_name) - .await; + // Record route name on the LLM span (only when a real route was produced). + if let Some(ref rn) = resolved_route_name { + if !rn.is_empty() && rn != "none" { + get_active_span(|span| { + span.set_attribute(opentelemetry::KeyValue::new( + tracing_plano::ROUTE_NAME, + rn.clone(), + )); + }); } + } - model - }; tracing::Span::current().record(tracing_llm::MODEL_NAME, resolved_model.as_str()); // Record the provider (derived from the `provider/model` prefix) so @@ -398,6 +442,34 @@ async fn llm_chat_inner( }); } + // Response-side refresh: update `last_used` + the context-size estimate from the + // real response. The routing decision itself was already persisted by `route()`. + let session_update_ctx: Option = + session_id.as_ref().map(|sid| SessionUpdateCtx { + orchestrator: Arc::clone(&state.orchestrator_service), + session_id: sid.clone(), + tenant_id: tenant_id.clone(), + anchor_model: resolved_model.clone(), + route_name: resolved_route_name.clone(), + prefix_hash: request_prefix_hash, + switch_budget_usd: decision.switch_budget_usd, + switches: decision.switches, + est_context_tokens: decision.cached_tokens, + gc_ttl: decision.gc_ttl, + }); + + // Forward the prefix hash so self-hosted multi-replica backends can do + // KV-aware replica stickiness (consistent-hash on this header at the + // cluster layer). + if let Some(hash) = request_prefix_hash { + if let Ok(val) = header::HeaderValue::from_str(&format!("{hash:016x}")) { + request_headers.insert( + header::HeaderName::from_static(PLANO_PREFIX_HASH_HEADER), + val, + ); + } + } + // --- Phase 4: Forward to upstream and stream back --- send_upstream( &state.http_client, @@ -415,6 +487,7 @@ async fn llm_chat_inner( state.state_storage.clone(), request_id, &state.filter_pipeline, + session_update_ctx, ) .await } @@ -692,6 +765,7 @@ async fn send_upstream( state_storage: Option>, request_id: String, filter_pipeline: &Arc, + session_update_ctx: Option, ) -> Result>, hyper::Error> { let span_name = if model_from_request == resolved_model { format!("POST {} {}", request_path, resolved_model) @@ -807,7 +881,7 @@ async fn send_upstream( let byte_stream = llm_response.bytes_stream(); // Create base processor for metrics and tracing - let base_processor = ObservableStreamProcessor::new( + let mut base_processor = ObservableStreamProcessor::new( operation_component::LLM, span_name, request_start_time, @@ -818,6 +892,13 @@ async fn send_upstream( model: metric_model.clone(), upstream_status: upstream_status.as_u16(), }); + // Only refresh the session binding for successful responses — errors carry no + // usage block and shouldn't reset the warmth clock. + if upstream_status.is_success() { + if let Some(update_ctx) = session_update_ctx { + base_processor = base_processor.with_session_update(update_ctx); + } + } let output_filter_request_headers = if filter_pipeline.has_output_filters() { Some(request_headers.clone()) diff --git a/crates/brightstaff/src/handlers/llm/prompt_caching.rs b/crates/brightstaff/src/handlers/llm/prompt_caching.rs new file mode 100644 index 000000000..2d0b5a08c --- /dev/null +++ b/crates/brightstaff/src/handlers/llm/prompt_caching.rs @@ -0,0 +1,70 @@ +//! Prompt-caching request handling for the LLM path. +//! +//! This module owns the "make the upstream provider's prompt cache work" concern: +//! resolving the correct cache-marking strategy for the `(gateway × model family × +//! upstream API)` combination and injecting the markers into the outbound request. +//! It never influences routing — see [`super::session_router`] for the session-cache +//! lookup and the routing-budget switch-cost concern. + +use common::configuration::EffectivePromptCaching; +use hermesllm::clients::SupportedUpstreamAPIs; +use hermesllm::{cache_marker_strategy, CacheMarkerStrategy, ProviderId, ProviderRequestType}; +use tracing::debug; + +/// Auto-inject prompt-cache markers into `client_request`. +/// +/// The strategy is resolved from `(gateway × model family × upstream API)` so that +/// Anthropic-family models cache whether they arrive over the native Messages API or +/// an OpenAI-compatible gateway (DigitalOcean, OpenRouter), while OpenAI-family models +/// (which cache automatically) and unimplemented backends are left untouched. +/// +/// A no-op when caching is disabled, `inject_cache_control` is off, the request opted +/// out (`X-Plano-Cache: off`), or the provider caches automatically. Injection is +/// idempotent (client-supplied markers are respected) and threshold-guarded against +/// the provider's minimum cacheable prefix. +pub fn inject_cache_markers( + client_request: &mut ProviderRequestType, + provider_id: ProviderId, + model_name_only: &str, + upstream_api: &SupportedUpstreamAPIs, + prompt_caching: &EffectivePromptCaching, + cache_off_for_request: bool, + alias_resolved_model: &str, +) { + if !(prompt_caching.enabled && prompt_caching.inject_cache_control && !cache_off_for_request) { + return; + } + + match cache_marker_strategy(provider_id, model_name_only, upstream_api) { + CacheMarkerStrategy::AnthropicMessagesBreakpoints { + min_prefix_tokens, .. + } => { + if let ProviderRequestType::MessagesRequest(req) = client_request { + let threshold = prompt_caching.min_prefix_tokens.max(min_prefix_tokens); + if req.inject_cache_breakpoints(threshold) { + debug!( + model = %alias_resolved_model, + min_prefix_tokens = threshold, + "injected anthropic ephemeral cache breakpoints" + ); + } + } + } + CacheMarkerStrategy::OpenAiContentPartCacheControl { + min_prefix_tokens, + ttl, + } => { + if let ProviderRequestType::ChatCompletionsRequest(req) = client_request { + let threshold = prompt_caching.min_prefix_tokens.max(min_prefix_tokens); + if req.inject_cache_control(ttl, threshold) { + debug!( + model = %alias_resolved_model, + min_prefix_tokens = threshold, + "injected openai content-part cache_control" + ); + } + } + } + CacheMarkerStrategy::Automatic | CacheMarkerStrategy::None => {} + } +} diff --git a/crates/brightstaff/src/handlers/llm/session_router.rs b/crates/brightstaff/src/handlers/llm/session_router.rs new file mode 100644 index 000000000..ecbb78a09 --- /dev/null +++ b/crates/brightstaff/src/handlers/llm/session_router.rs @@ -0,0 +1,640 @@ +//! Session-cache-aware routing. +//! +//! Routing itself stays cache-blind: the `llm_router` (quality) still picks a +//! candidate model for every request. This module then decides whether to *honor* +//! that candidate or stick to the session's warm anchor, based on: +//! +//! * **Cache warmth** — inferred structurally from how long ago the session was last +//! used vs. the provider's cache window ([`hermesllm::provider_cache_capability`]), +//! so it works on the decision path with no provider response in hand. +//! * **A cumulative per-session switch budget** — a paid switch (the candidate must +//! re-ingest the context at its uncached rate) is allowed only while budget remains; +//! an outright-cheaper switch is free and can credit the budget back. +//! +//! The default posture is to stick. Quality and cost stay separate: the router decides +//! whether a switch *improves quality*; the budget decides whether it is *affordable*. +//! +//! Prompt-cache *marker injection* is a separate concern — see [`super::prompt_caching`]. + +use std::time::{Duration, SystemTime}; + +use common::configuration::EffectiveRoutingBudget; +use hermesllm::apis::openai::Message; +use hermesllm::{provider_cache_capability, ProviderCacheCapability, ProviderId}; +use opentelemetry::trace::get_active_span; +use opentelemetry::KeyValue; +use tracing::{debug, info}; + +use crate::affinity::derive_implicit_affinity; +use crate::metrics as bs_metrics; +use crate::metrics::labels as metric_labels; +use crate::router::orchestrator::OrchestratorService; +use crate::session_cache::SessionBinding; +use crate::tracing::plano as tracing_plano; + +/// Resolved session identity for one request. +pub struct SessionResolution { + /// Stable prefix hash (system + tools + first user message), independent of + /// `prompt_caching.enabled` so it can still drive the `x-plano-prefix-hash` + /// RING_HASH replica-stickiness header. `None` when the request opted out or has + /// no anchorable prompt. + pub request_prefix_hash: Option, + /// Session key: the explicit `X-Model-Affinity` value, or the implicit prefix-hash + /// key when implicit affinity is active. `None` when there is nothing to anchor to. + pub session_id: Option, +} + +/// Resolve the session key and prefix hash from the (already filtered / state-merged) +/// request. An explicit affinity header always anchors; the implicit key is derived +/// when `implicit_affinity_enabled` is set — true when either prompt caching's +/// `session_affinity` or the routing budget is active, so stickiness works whether or +/// not prompt caching is enabled. The prefix hash is derived regardless (only +/// `X-Plano-Cache: off` or an unanchorable prompt suppresses it) so the +/// `x-plano-prefix-hash` RING_HASH replica-stickiness header still works. +pub fn resolve_session( + explicit_session_id: Option, + messages: &[Message], + tool_names: Option<&[String]>, + tenant_id: Option<&str>, + implicit_affinity_enabled: bool, + cache_off_for_request: bool, +) -> SessionResolution { + let implicit_affinity = if cache_off_for_request { + None + } else { + derive_implicit_affinity(messages, tool_names, tenant_id) + }; + let request_prefix_hash = implicit_affinity.as_ref().map(|a| a.prefix_hash); + + let session_id = match explicit_session_id { + Some(sid) => Some(sid), + None if implicit_affinity_enabled && !cache_off_for_request => { + implicit_affinity.as_ref().map(|a| a.session_key.clone()) + } + None => None, + }; + + SessionResolution { + request_prefix_hash, + session_id, + } +} + +/// Extra memory retention beyond the warmth window, so a still-warm binding is never +/// GC'd out from under the router before it could plausibly go cold. +const GC_SLACK: Duration = Duration::from_secs(60); + +/// Stable request facts the router reasons about. Independent of the transport (full +/// proxy vs. decision endpoint) so both paths route identically. +pub struct RouteFacts<'a> { + /// Session key (explicit `X-Model-Affinity` or the implicit prefix key). `None` + /// disables stickiness for this request (nothing to anchor to). + pub session_id: Option<&'a str>, + pub tenant_id: Option<&'a str>, + /// Stable prompt-prefix hash; a mismatch vs. the stored binding means the provider + /// cache is already lost, so a switch is free. + pub prefix_hash: Option, + /// Estimated context size in tokens (the tokens a switch would re-ingest). + pub est_context_tokens: u64, + /// The model the quality router picked for this request. + pub candidate_model: &'a str, + pub candidate_route: Option<&'a str>, +} + +/// The routing decision plus the session state to carry into the response side. +pub struct RouteDecision { + /// The model to actually dispatch to (the anchor when a switch was vetoed). + pub model: String, + pub route_name: Option, + /// Whether the session's cache was inferred warm at decision time. + pub warm: bool, + /// Remaining switch budget after this decision. + pub switch_budget_usd: f64, + /// Cumulative switches taken this session (after this decision). + pub switches: u32, + /// Context-token estimate persisted with the binding (refined later from usage). + pub cached_tokens: u64, + /// GC bound the binding was stored with (reused when the response side refreshes). + pub gc_ttl: Duration, +} + +/// Estimate the request's context size in tokens. Uses the tiktoken-based counter when +/// available, falling back to the chars/4 heuristic. Precision is not critical — it only +/// scales the switch-cost estimate, and both sides of the comparison scale with it. +pub fn estimate_context_tokens(messages: &[Message], model: &str) -> u64 { + let text: String = messages + .iter() + .filter_map(|m| m.content.as_ref().map(|c| c.to_string())) + .collect::>() + .join("\n"); + match common::tokenizer::token_count(model, &text) { + Ok(count) => count as u64, + Err(_) => (text.len() / 4) as u64, + } +} + +/// Resolve a provider-qualified model id (e.g. `openai/gpt-4o`) to its cache window. +/// Unknown providers fall back to the conservative default. +fn capability_for_model(model: &str) -> ProviderCacheCapability { + let provider_part = model.split_once('/').map(|(p, _)| p).unwrap_or(model); + ProviderId::try_from(provider_part) + .map(provider_cache_capability) + .unwrap_or_default() +} + +/// How long a binding on this model can sit idle before its cache is certainly cold. +fn warmth_window(cap: &ProviderCacheCapability) -> Duration { + if cap.extended_retention { + cap.extended_ttl + } else { + cap.idle_ttl.min(cap.hard_ttl) + } +} + +/// Whether the session's provider cache is plausibly still warm given how long ago it +/// was last used. Returns the warmth verdict and the measured idle gap. +fn warmth( + binding: &SessionBinding, + cap: &ProviderCacheCapability, + now: SystemTime, +) -> (bool, Duration) { + let idle = now + .duration_since(binding.last_used) + .unwrap_or(Duration::ZERO); + let warm = if cap.extended_retention { + idle <= cap.extended_ttl + } else { + idle <= cap.idle_ttl && idle <= cap.hard_ttl + }; + (warm, idle) +} + +/// Decide the final model for this request and persist the updated session binding. +/// +/// Never overrides the router on a *cold* session — it only protects a warm cache. The +/// returned [`RouteDecision`] carries the model to dispatch plus the session state the +/// response side reuses when it refreshes the binding from real usage. +pub async fn route( + orchestrator: &OrchestratorService, + routing_budget: Option<&EffectiveRoutingBudget>, + facts: RouteFacts<'_>, +) -> RouteDecision { + let now = SystemTime::now(); + let candidate_gc_ttl = warmth_window(&capability_for_model(facts.candidate_model)) + GC_SLACK; + + // No session to anchor to: honor the candidate, persist nothing. + let Some(session_id) = facts.session_id else { + return RouteDecision { + model: facts.candidate_model.to_string(), + route_name: facts.candidate_route.map(str::to_string), + warm: false, + switch_budget_usd: 0.0, + switches: 0, + cached_tokens: facts.est_context_tokens, + gc_ttl: candidate_gc_ttl, + }; + }; + + let existing = orchestrator.get_binding(session_id, facts.tenant_id).await; + + // Warmth + prefix drift. A drifted prefix means the cache is already cold. + let (warm, idle) = match &existing { + Some(b) => warmth(b, &capability_for_model(&b.anchor_model), now), + None => (false, Duration::ZERO), + }; + let drifted = match ( + existing.as_ref().and_then(|b| b.prefix_hash), + facts.prefix_hash, + ) { + (Some(stored), Some(current)) => stored != current, + _ => false, + }; + let effective_warm = warm && !drifted; + + let seed = routing_budget.map(|s| s.seed_usd).unwrap_or(0.0); + + // Resolve the final model, budget, switch count, and decision telemetry. + let mut model = facts.candidate_model.to_string(); + let mut route_name = facts.candidate_route.map(str::to_string); + let budget_before; + let mut budget; + let mut switches; + let mut cost_opt: Option = None; + let mut counterfactual: Option = None; + let decision_label: &'static str; + let reason: &'static str; + + match existing.as_ref() { + Some(b) if effective_warm => { + budget_before = b.switch_budget_usd; + budget = b.switch_budget_usd; + switches = b.switches; + if facts.candidate_model == b.anchor_model { + // Router agrees with the anchor — stick, no cost. + decision_label = metric_labels::SWITCH_DECISION_ALLOWED; + reason = metric_labels::SWITCH_REASON_SAME_ANCHOR; + } else if let Some(cfg) = routing_budget { + let context_tokens = if b.cached_tokens > 0 { + b.cached_tokens + } else { + facts.est_context_tokens + }; + match orchestrator + .estimate_switch_cost_in_usd( + context_tokens, + &b.anchor_model, + facts.candidate_model, + cfg.cache_read_discount, + ) + .await + { + // No pricing for one side — fail open (switch freely) rather than + // veto the router on guesswork. + None => { + switches += 1; + decision_label = metric_labels::SWITCH_DECISION_ALLOWED; + reason = metric_labels::SWITCH_REASON_NO_PRICING; + debug!( + anchor = %b.anchor_model, + candidate = %facts.candidate_model, + "switch allowed — missing pricing data, cannot gate" + ); + } + Some(cost) => { + cost_opt = Some(cost); + if cost <= 0.0 { + // Outright cheaper: free switch, credit the budget back. + if cfg.credit_negative { + budget -= cost; // cost is negative → budget grows + } + switches += 1; + decision_label = metric_labels::SWITCH_DECISION_ALLOWED; + reason = metric_labels::SWITCH_REASON_FREE; + info!( + anchor = %b.anchor_model, + candidate = %facts.candidate_model, + switch_cost_in_usd = cost, + "switch allowed — candidate undercuts the cached rate" + ); + } else if cost <= budget { + budget -= cost; + switches += 1; + decision_label = metric_labels::SWITCH_DECISION_ALLOWED; + reason = metric_labels::SWITCH_REASON_WITHIN_BUDGET; + info!( + anchor = %b.anchor_model, + candidate = %facts.candidate_model, + switch_cost_in_usd = cost, + budget_remaining_in_usd = budget, + "switch allowed — within session switch budget" + ); + } else { + // Unaffordable: retain the warm anchor. + if cfg.record_counterfactual { + counterfactual = Some(match route_name.as_deref() { + Some(rn) if !rn.is_empty() && rn != "none" => { + format!("{} ({rn})", facts.candidate_model) + } + _ => facts.candidate_model.to_string(), + }); + } + model = b.anchor_model.clone(); + route_name = b.route_name.clone(); + decision_label = metric_labels::SWITCH_DECISION_RETAINED; + reason = metric_labels::SWITCH_REASON_OVER_BUDGET; + info!( + anchor = %b.anchor_model, + candidate = %facts.candidate_model, + switch_cost_in_usd = cost, + budget_remaining_in_usd = budget, + "switch vetoed — cost exceeds remaining budget, retaining anchor" + ); + } + } + } + } else { + // Warm but no budget configured — follow the router freely. + switches += 1; + decision_label = metric_labels::SWITCH_DECISION_ALLOWED; + reason = metric_labels::SWITCH_REASON_FREE; + } + bs_metrics::record_session_switch_decision(decision_label, reason); + } + _ => { + // Cold (or no binding, or drifted): honor the candidate and (re)seed a + // fresh warm episode. Switches reset — this is a new cache lifetime. + budget_before = seed; + budget = match (routing_budget, existing.as_ref()) { + (Some(cfg), Some(b)) if !cfg.replenish_on_rebind => b.switch_budget_usd, + _ => seed, + }; + switches = 0; + } + } + + // Context estimate persisted with the binding (refined later from real usage). + let cached_tokens = if facts.est_context_tokens > 0 { + facts.est_context_tokens + } else { + existing.as_ref().map(|b| b.cached_tokens).unwrap_or(0) + }; + let gc_ttl = warmth_window(&capability_for_model(&model)) + GC_SLACK; + + // Observability: cache warmth + budget/switch state on the current span. + get_active_span(|span| { + span.set_attribute(KeyValue::new(tracing_plano::CACHE_WARM, effective_warm)); + span.set_attribute(KeyValue::new( + tracing_plano::CACHE_IDLE_MS, + idle.as_millis() as i64, + )); + if routing_budget.is_some() { + span.set_attribute(KeyValue::new( + tracing_plano::SESSION_BUDGET_REMAINING_IN_USD, + budget, + )); + span.set_attribute(KeyValue::new( + tracing_plano::SESSION_SWITCHES, + switches as i64, + )); + } + if let Some(cost) = cost_opt { + span.set_attribute(KeyValue::new(tracing_plano::SWITCH_COST_IN_USD, cost)); + span.set_attribute(KeyValue::new( + tracing_plano::SWITCH_THRESHOLD_IN_USD, + budget_before, + )); + span.set_attribute(KeyValue::new( + tracing_plano::SWITCH_DECISION, + if model == facts.candidate_model { + metric_labels::SWITCH_DECISION_ALLOWED + } else { + metric_labels::SWITCH_DECISION_RETAINED + }, + )); + } + if let Some(ref cf) = counterfactual { + span.set_attribute(KeyValue::new( + tracing_plano::SWITCH_COUNTERFACTUAL_ROUTE, + cf.clone(), + )); + } + }); + + orchestrator + .store_binding( + session_id, + facts.tenant_id, + SessionBinding { + anchor_model: model.clone(), + route_name: route_name.clone(), + prefix_hash: facts.prefix_hash, + last_used: now, + cached_tokens, + switch_budget_usd: budget, + switches, + }, + Some(gc_ttl), + ) + .await; + + RouteDecision { + model, + route_name, + warm: effective_warm, + switch_budget_usd: budget, + switches, + cached_tokens, + gc_ttl, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn cap_5m_1h() -> ProviderCacheCapability { + ProviderCacheCapability { + idle_ttl: Duration::from_secs(300), + hard_ttl: Duration::from_secs(3600), + extended_retention: false, + extended_ttl: Duration::from_secs(3600), + } + } + + fn binding_used_ago(secs: u64) -> SessionBinding { + SessionBinding { + anchor_model: "anthropic/claude-sonnet-4-5".to_string(), + route_name: None, + prefix_hash: Some(1), + last_used: SystemTime::now() - Duration::from_secs(secs), + cached_tokens: 100_000, + switch_budget_usd: 0.10, + switches: 0, + } + } + + #[test] + fn warm_within_idle_window() { + let (warm, _) = warmth(&binding_used_ago(60), &cap_5m_1h(), SystemTime::now()); + assert!(warm); + } + + #[test] + fn cold_past_idle_window() { + let (warm, _) = warmth(&binding_used_ago(600), &cap_5m_1h(), SystemTime::now()); + assert!(!warm); + } + + #[test] + fn extended_retention_keeps_warm_past_idle() { + let cap = ProviderCacheCapability { + extended_retention: true, + ..cap_5m_1h() + }; + // 10 minutes idle: cold under 5m, warm under the 1h extended window. + let (warm, _) = warmth(&binding_used_ago(600), &cap, SystemTime::now()); + assert!(warm); + } + + #[test] + fn capability_resolves_from_model_prefix() { + // Known provider prefix resolves; unknown falls back to the default. + let anthropic = capability_for_model("anthropic/claude-sonnet-4-5"); + assert_eq!(anthropic, ProviderCacheCapability::default()); + let unknown = capability_for_model("madeup/model-x"); + assert_eq!(unknown, ProviderCacheCapability::default()); + } + + // ---- route() budget behavior ---- + + use crate::router::model_metrics::{ModelMetricsService, ModelRates}; + use crate::session_cache::memory::MemorySessionCache; + use std::collections::HashMap; + use std::sync::Arc; + + // Anchor cached rate 0.3, candidate `pricey` input 5.0, candidate `cheap` input 0.1. + // With a 100k-token context the paid switch costs 0.1 * (5.0 - 0.3) = $0.47 and the + // cheap switch is 0.1 * (0.1 - 0.3) = -$0.02 (a credit). + fn orch_with_rates() -> OrchestratorService { + let mut rates = HashMap::new(); + rates.insert( + "anthropic/expensive".to_string(), + ModelRates { + input_per_million: 3.0, + output_per_million: 15.0, + cache_read_per_million: Some(0.3), + }, + ); + rates.insert( + "openai/pricey".to_string(), + ModelRates { + input_per_million: 5.0, + output_per_million: 15.0, + cache_read_per_million: Some(0.5), + }, + ); + rates.insert( + "google/cheap".to_string(), + ModelRates { + input_per_million: 0.1, + output_per_million: 0.4, + cache_read_per_million: Some(0.01), + }, + ); + let metrics = Arc::new(ModelMetricsService::from_rates_for_test(rates)); + let cache = Arc::new(MemorySessionCache::new(100)); + OrchestratorService::with_routing( + "http://localhost/v1/chat/completions".to_string(), + "m".to_string(), + "p".to_string(), + None, + Some(metrics), + Some(600), + cache, + None, + 8192, + ) + } + + fn routing_budget(seed: f64) -> EffectiveRoutingBudget { + EffectiveRoutingBudget { + seed_usd: seed, + replenish_on_rebind: true, + credit_negative: true, + cache_read_discount: 0.1, + record_counterfactual: false, + } + } + + async fn seed_warm_binding(orch: &OrchestratorService, budget: f64, idle_secs: u64) { + orch.store_binding( + "s1", + None, + SessionBinding { + anchor_model: "anthropic/expensive".to_string(), + route_name: None, + prefix_hash: Some(1), + last_used: SystemTime::now() - Duration::from_secs(idle_secs), + cached_tokens: 100_000, + switch_budget_usd: budget, + switches: 0, + }, + Some(Duration::from_secs(3600)), + ) + .await; + } + + fn facts_for<'a>(candidate: &'a str) -> RouteFacts<'a> { + RouteFacts { + session_id: Some("s1"), + tenant_id: None, + prefix_hash: Some(1), + est_context_tokens: 0, + candidate_model: candidate, + candidate_route: None, + } + } + + #[tokio::test] + async fn paid_switch_within_budget_is_allowed_and_debits() { + let orch = orch_with_rates(); + seed_warm_binding(&orch, 1.0, 30).await; + let st = routing_budget(1.0); + let d = route(&orch, Some(&st), facts_for("openai/pricey")).await; + + assert_eq!(d.model, "openai/pricey"); + assert!(d.warm); + assert_eq!(d.switches, 1); + assert!( + (d.switch_budget_usd - 0.53).abs() < 1e-6, + "budget {} != 0.53", + d.switch_budget_usd + ); + } + + #[tokio::test] + async fn paid_switch_over_budget_retains_anchor() { + let orch = orch_with_rates(); + seed_warm_binding(&orch, 0.10, 30).await; + let st = routing_budget(0.10); + let d = route(&orch, Some(&st), facts_for("openai/pricey")).await; + + assert_eq!(d.model, "anthropic/expensive"); + assert!(d.warm); + assert_eq!(d.switches, 0); + assert!((d.switch_budget_usd - 0.10).abs() < 1e-6); + } + + #[tokio::test] + async fn cheaper_switch_is_free_and_credits_budget() { + let orch = orch_with_rates(); + seed_warm_binding(&orch, 0.10, 30).await; + let st = routing_budget(0.10); + let d = route(&orch, Some(&st), facts_for("google/cheap")).await; + + assert_eq!(d.model, "google/cheap"); + assert!(d.warm); + assert_eq!(d.switches, 1); + // -(-0.02) credited back: 0.10 + 0.02 = 0.12. + assert!( + (d.switch_budget_usd - 0.12).abs() < 1e-6, + "budget {} != 0.12", + d.switch_budget_usd + ); + } + + #[tokio::test] + async fn cold_session_reseeds_budget_and_follows_router() { + let orch = orch_with_rates(); + // 10 minutes idle: past Anthropic's 5m idle window -> cold. Budget was drained. + seed_warm_binding(&orch, 0.0, 600).await; + let st = routing_budget(1.0); + let d = route(&orch, Some(&st), facts_for("openai/pricey")).await; + + assert_eq!(d.model, "openai/pricey"); + assert!(!d.warm); + assert_eq!(d.switches, 0); + assert!( + (d.switch_budget_usd - 1.0).abs() < 1e-6, + "budget {} != seed", + d.switch_budget_usd + ); + } + + #[tokio::test] + async fn no_session_honors_candidate() { + let orch = orch_with_rates(); + let st = routing_budget(1.0); + let facts = RouteFacts { + session_id: None, + tenant_id: None, + prefix_hash: Some(1), + est_context_tokens: 0, + candidate_model: "openai/pricey", + candidate_route: None, + }; + let d = route(&orch, Some(&st), facts).await; + assert_eq!(d.model, "openai/pricey"); + assert!(!d.warm); + } +} diff --git a/crates/brightstaff/src/handlers/routing_service.rs b/crates/brightstaff/src/handlers/routing_service.rs index b93b14222..45f76b8e8 100644 --- a/crates/brightstaff/src/handlers/routing_service.rs +++ b/crates/brightstaff/src/handlers/routing_service.rs @@ -1,9 +1,11 @@ use bytes::Bytes; -use common::configuration::{SpanAttributes, TopLevelRoutingPreference}; +use common::configuration::{ + EffectivePromptCaching, EffectiveRoutingBudget, SpanAttributes, TopLevelRoutingPreference, +}; use common::consts::{MODEL_AFFINITY_HEADER, REQUEST_ID_HEADER}; use common::errors::BrightStaffError; use hermesllm::clients::SupportedAPIsFromClient; -use hermesllm::ProviderRequestType; +use hermesllm::{ProviderRequest, ProviderRequestType}; use http_body_util::combinators::BoxBody; use http_body_util::{BodyExt, Full}; use hyper::{Request, Response, StatusCode}; @@ -12,6 +14,7 @@ use tracing::{debug, info, info_span, warn, Instrument}; use super::extract_or_generate_traceparent; use crate::handlers::llm::model_selection::router_chat_get_upstream_model; +use crate::handlers::llm::session_router; use crate::metrics as bs_metrics; use crate::metrics::labels as metric_labels; use crate::router::orchestrator::OrchestratorService; @@ -60,11 +63,14 @@ struct RoutingDecisionResponse { pinned: bool, } +#[allow(clippy::too_many_arguments)] pub async fn routing_decision( request: Request, orchestrator_service: Arc, request_path: String, span_attributes: &Option, + prompt_caching: EffectivePromptCaching, + routing_budget: Option, ) -> Result>, hyper::Error> { let request_headers = request.headers().clone(); let request_id: String = request_headers @@ -73,7 +79,7 @@ pub async fn routing_decision( .map(|s| s.to_string()) .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - let session_id: Option = request_headers + let explicit_session_id: Option = request_headers .get(MODEL_AFFINITY_HEADER) .and_then(|h| h.to_str().ok()) .map(|s| s.to_string()); @@ -101,8 +107,10 @@ pub async fn routing_decision( request_path, request_headers, custom_attrs, - session_id, + explicit_session_id, tenant_id, + prompt_caching, + routing_budget, ) .instrument(request_span) .await @@ -116,8 +124,10 @@ async fn routing_decision_inner( request_path: String, request_headers: hyper::HeaderMap, custom_attrs: std::collections::HashMap, - session_id: Option, + explicit_session_id: Option, tenant_id: Option, + prompt_caching: EffectivePromptCaching, + routing_budget: Option, ) -> Result>, hyper::Error> { set_service_name(operation_component::ROUTING); opentelemetry::trace::get_active_span(|span| { @@ -135,37 +145,9 @@ async fn routing_decision_inner( .unwrap_or("unknown") .to_string(); - if let Some(ref sid) = session_id { - if let Some(cached) = orchestrator_service - .get_cached_route(sid, tenant_id.as_deref()) - .await - { - info!( - session_id = %sid, - model = %cached.model_name, - route = ?cached.route_name, - "returning pinned routing decision from cache" - ); - let response = RoutingDecisionResponse { - models: vec![cached.model_name], - route: cached.route_name, - trace_id, - session_id: Some(sid.clone()), - pinned: true, - }; - let json = serde_json::to_string(&response).unwrap(); - let body = Full::new(Bytes::from(json)) - .map_err(|never| match never {}) - .boxed(); - return Ok(Response::builder() - .status(StatusCode::OK) - .header("Content-Type", "application/json") - .body(body) - .unwrap()); - } - } - - // Parse request body + // Parse the request body up front so a pin can be validated against prefix drift + // and a fresh pin can be stored with its prefix hash. This endpoint shares the + // session cache with the LLM handler, so both must key drift the same way. let raw_bytes = request.collect().await?.to_bytes(); debug!( @@ -202,6 +184,39 @@ async fn routing_decision_inner( } }; + // `X-Plano-Cache: off` opts this request out of implicit affinity (same sentinel + // the LLM handler honors), so callers can bypass stickiness per request. + let cache_off_for_request = request_headers + .get(common::consts::PLANO_CACHE_HEADER) + .and_then(|h| h.to_str().ok()) + .is_some_and(|v| v.eq_ignore_ascii_case("off")); + + let request_messages = client_request.get_messages(); + let tool_names = client_request.get_tool_names(); + + // Session key + prefix hash resolved identically to the LLM handler so pins + // interoperate across the full-proxy and decision paths. + // Derive the implicit session key when either prompt-caching affinity or the routing + // budget is active, so the budget works the same way with caching off. + let implicit_affinity_enabled = prompt_caching.session_affinity || routing_budget.is_some(); + let session_router::SessionResolution { + request_prefix_hash, + session_id, + } = session_router::resolve_session( + explicit_session_id, + &request_messages, + tool_names.as_deref(), + tenant_id.as_deref(), + implicit_affinity_enabled, + cache_off_for_request, + ); + + let est_context_tokens: u64 = if session_id.is_some() && routing_budget.is_some() { + session_router::estimate_context_tokens(&request_messages, client_request.model()) + } else { + 0 + }; + let routing_result = router_chat_get_upstream_model( Arc::clone(&orchestrator_service), client_request, @@ -213,23 +228,37 @@ async fn routing_decision_inner( match routing_result { Ok(result) => { - if let Some(ref sid) = session_id { - orchestrator_service - .cache_route( - sid.clone(), - tenant_id.as_deref(), - result.model_name.clone(), - result.route_name.clone(), - ) - .await; + let candidate_model = result.model_name.clone(); + let decision = session_router::route( + &orchestrator_service, + routing_budget.as_ref(), + session_router::RouteFacts { + session_id: session_id.as_deref(), + tenant_id: tenant_id.as_deref(), + prefix_hash: request_prefix_hash, + est_context_tokens, + candidate_model: &candidate_model, + candidate_route: result.route_name.as_deref(), + }, + ) + .await; + + // Front the ranked fallback list with the decided model (the anchor, when a + // switch was vetoed), so 429/5xx fallbacks still work. + let mut models = result.models; + if models.first() != Some(&decision.model) { + models.retain(|m| m != &decision.model); + models.insert(0, decision.model.clone()); } let response = RoutingDecisionResponse { - models: result.models, - route: result.route_name, + models, + route: decision.route_name, trace_id, session_id, - pinned: false, + // `pinned` signals a warm, stuck session — safe for callers to treat as + // "keep this provider's cache warm". + pinned: decision.warm, }; // Distinguish "decision served" (a concrete model picked) from @@ -247,6 +276,7 @@ async fn routing_decision_inner( primary_model = %response.models.first().map(|s| s.as_str()).unwrap_or("none"), total_models = response.models.len(), route = ?response.route, + pinned = response.pinned, "routing decision completed" ); diff --git a/crates/brightstaff/src/lib.rs b/crates/brightstaff/src/lib.rs index 66c6eadf7..f0b6b13dc 100644 --- a/crates/brightstaff/src/lib.rs +++ b/crates/brightstaff/src/lib.rs @@ -1,3 +1,4 @@ +pub mod affinity; pub mod app_state; pub mod handlers; pub mod metrics; diff --git a/crates/brightstaff/src/main.rs b/crates/brightstaff/src/main.rs index c9e8b9bc7..b69284ce9 100644 --- a/crates/brightstaff/src/main.rs +++ b/crates/brightstaff/src/main.rs @@ -343,6 +343,36 @@ async fn init_app_state( let signals_enabled = !overrides.disable_signals.unwrap_or(false); + let prompt_caching = + common::configuration::EffectivePromptCaching::from_config(config.prompt_caching.as_ref())?; + + // Routing budget lives under `routing` and is independent of prompt caching. + let routing_budget = common::configuration::EffectiveRoutingBudget::from_config( + config + .routing + .as_ref() + .and_then(|r| r.routing_budget.as_ref()), + )?; + + // The routing-budget cost gate needs per-model pricing to compute switch cost. + if routing_budget.is_some() { + use common::configuration::MetricsSource; + let has_cost_source = config + .model_metrics_sources + .as_deref() + .unwrap_or_default() + .iter() + .any(|s| matches!(s, MetricsSource::Cost(_))); + if !has_cost_source { + return Err( + "routing.routing_budget is configured but no cost metrics source is \ + configured — add a cost source (e.g. models.dev) to model_metrics_sources so \ + per-model input/cached rates are available for the switch-cost calculation" + .into(), + ); + } + } + Ok(AppState { orchestrator_service, model_aliases: config.model_aliases.clone(), @@ -356,6 +386,8 @@ async fn init_app_state( http_client: reqwest::Client::new(), filter_pipeline, signals_enabled, + prompt_caching, + routing_budget, }) } @@ -527,6 +559,8 @@ async fn dispatch( Arc::clone(&state.orchestrator_service), stripped, &state.span_attributes, + state.prompt_caching, + state.routing_budget, ) .with_context(parent_cx) .await; diff --git a/crates/brightstaff/src/metrics/labels.rs b/crates/brightstaff/src/metrics/labels.rs index 4eaf3e59a..9b9474ad0 100644 --- a/crates/brightstaff/src/metrics/labels.rs +++ b/crates/brightstaff/src/metrics/labels.rs @@ -18,6 +18,10 @@ pub const ROUTE_LLM: &str = "llm"; // Token kind for brightstaff_llm_tokens_total. pub const TOKEN_KIND_PROMPT: &str = "prompt"; pub const TOKEN_KIND_COMPLETION: &str = "completion"; +/// Input tokens served from the provider's prompt cache (billed at the cached rate). +pub const TOKEN_KIND_CACHE_READ: &str = "cache_read"; +/// Input tokens written into the provider's prompt cache (cache-creation surcharge). +pub const TOKEN_KIND_CACHE_WRITE: &str = "cache_write"; // LLM error_class values (match docstring in metrics/mod.rs). pub const LLM_ERR_NONE: &str = "none"; @@ -36,3 +40,38 @@ pub const ROUTING_SVC_POLICY_ERROR: &str = "policy_error"; pub const SESSION_CACHE_HIT: &str = "hit"; pub const SESSION_CACHE_MISS: &str = "miss"; pub const SESSION_CACHE_STORE: &str = "store"; + +// Prompt cache outcome values (brightstaff_prompt_cache_requests_total). +pub const PROMPT_CACHE_HIT: &str = "hit"; +pub const PROMPT_CACHE_MISS: &str = "miss"; + +// Session pin lifecycle events (brightstaff_session_pin_events_total). +/// Implicit session committed its pin after the first observed cache activity. +pub const PIN_EVENT_IMPLICIT_COMMIT: &str = "implicit_commit"; +/// An existing pin was refreshed (TTL extended, observed-hit state updated). +pub const PIN_EVENT_REFRESH: &str = "refresh"; +/// A pinned request's prefix hash no longer matched — cache already lost, re-routed. +pub const PIN_EVENT_PREFIX_DRIFT: &str = "prefix_drift"; +/// A logically-expired pin was used as a soft switch-penalty hint for routing. +pub const PIN_EVENT_STALE_HINT: &str = "stale_hint"; +/// A pinned session that previously produced cache hits stopped producing them. +pub const PIN_EVENT_VALIDATION_FAILED: &str = "validation_failed"; + +// Session-stickiness decisions (brightstaff_session_switch_decisions_total). +// `decision` label — the coarse outcome: +/// The proposed switch was honored (free, within budget, or unpriced fail-open). +pub const SWITCH_DECISION_ALLOWED: &str = "allowed"; +/// The switch cost exceeded the remaining budget — the warm anchor was retained. +pub const SWITCH_DECISION_RETAINED: &str = "retained"; + +// `reason` label — why the decision was made: +/// The router agreed with the warm anchor; no switch was needed. +pub const SWITCH_REASON_SAME_ANCHOR: &str = "same_anchor"; +/// The candidate was outright cheaper (negative cost) — a free switch. +pub const SWITCH_REASON_FREE: &str = "free"; +/// A paid switch that fit within the session's remaining switch budget. +pub const SWITCH_REASON_WITHIN_BUDGET: &str = "within_budget"; +/// A paid switch that exceeded the remaining budget — retained the anchor. +pub const SWITCH_REASON_OVER_BUDGET: &str = "over_budget"; +/// Pricing was missing for one side, so the switch was allowed without a cost gate. +pub const SWITCH_REASON_NO_PRICING: &str = "no_pricing"; diff --git a/crates/brightstaff/src/metrics/mod.rs b/crates/brightstaff/src/metrics/mod.rs index 34679ccac..7b00702b7 100644 --- a/crates/brightstaff/src/metrics/mod.rs +++ b/crates/brightstaff/src/metrics/mod.rs @@ -18,6 +18,9 @@ //! `brightstaff_router_decision_duration_seconds`, //! `brightstaff_routing_service_requests_total`, //! `brightstaff_session_cache_events_total`. +//! - Prompt caching: `brightstaff_prompt_cache_requests_total`, +//! `brightstaff_session_pin_events_total` (cache read/write tokens are emitted as +//! `kind=cache_read|cache_write` on `brightstaff_llm_tokens_total`). //! - Process: via `metrics-process`. //! - Build: `brightstaff_build_info`. @@ -172,6 +175,17 @@ fn describe_all() { "brightstaff_session_cache_events_total", "Session affinity cache lookups and stores, by outcome." ); + describe_counter!( + "brightstaff_prompt_cache_requests_total", + "LLM responses with usage data, by provider, model and prompt-cache outcome \ + (hit = cache read/creation tokens reported, miss = none). The miss rate is \ + the silent cache-miss baseline." + ); + describe_counter!( + "brightstaff_session_pin_events_total", + "Session pin lifecycle events: implicit_commit, refresh, prefix_drift, \ + stale_hint, validation_failed." + ); describe_gauge!( "brightstaff_build_info", @@ -375,3 +389,37 @@ pub fn record_session_cache_event(outcome: &'static str) { ) .increment(1); } + +/// Record whether a completed LLM response reported prompt-cache activity. +/// `hit` means the provider billed cache-read or cache-creation tokens. +pub fn record_prompt_cache_outcome(provider: &str, model: &str, outcome: &'static str) { + counter!( + "brightstaff_prompt_cache_requests_total", + "provider" => provider.to_string(), + "model" => model.to_string(), + "outcome" => outcome, + ) + .increment(1); +} + +/// Record a session pin lifecycle event (see `metrics::labels::PIN_EVENT_*`). +pub fn record_session_pin_event(event: &'static str) { + counter!( + "brightstaff_session_pin_events_total", + "event" => event, + ) + .increment(1); +} + +/// Record a session-stickiness decision on a proposed model switch. `decision` is the +/// coarse outcome (`allowed`/`retained`, see `metrics::labels::SWITCH_DECISION_*`) and +/// `reason` explains why (`same_anchor`/`free`/`within_budget`/`over_budget`/`no_pricing`, +/// see `metrics::labels::SWITCH_REASON_*`). +pub fn record_session_switch_decision(decision: &'static str, reason: &'static str) { + counter!( + "brightstaff_session_switch_decisions_total", + "decision" => decision, + "reason" => reason, + ) + .increment(1); +} diff --git a/crates/brightstaff/src/router/model_metrics.rs b/crates/brightstaff/src/router/model_metrics.rs index a4b3df48b..a75d88875 100644 --- a/crates/brightstaff/src/router/model_metrics.rs +++ b/crates/brightstaff/src/router/model_metrics.rs @@ -11,14 +11,55 @@ use tracing::{debug, info, warn}; const DO_PRICING_URL: &str = "https://api.digitalocean.com/v2/gen-ai/models/catalog"; const MODELS_DEV_URL: &str = "https://models.dev/api.json"; +/// DigitalOcean publishes prices per token; scale to the per-million convention +/// that `ModelRates` and the absolute-USD consumers (switch-cost gate) expect. +const TOKENS_PER_MILLION: f64 = 1_000_000.0; + +/// Structured per-million-token USD rates for one model, as published by the cost +/// feed. Kept alongside the blended ranking metric so cost-sensitive features (e.g. +/// the session-stickiness switch-cost gate) can reason about input vs cached-input +/// pricing without touching `rank_models`. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct ModelRates { + pub input_per_million: f64, + pub output_per_million: f64, + /// Cached (prompt-cache read) input rate. Present when the feed publishes it + /// (models.dev `cost.cache_read`); absent for feeds that don't (DO catalog). + pub cache_read_per_million: Option, +} + +impl ModelRates { + /// Blended input+output metric used for `prefer: cheapest` ranking. + fn blended(&self) -> f64 { + self.input_per_million + self.output_per_million + } + + /// Cached input rate, falling back to `input * cache_read_discount` when the + /// feed doesn't publish a cached rate. + pub fn cached_input_rate(&self, cache_read_discount: f64) -> f64 { + self.cache_read_per_million + .unwrap_or(self.input_per_million * cache_read_discount) + } +} + +/// Derive the blended ranking map from the structured rates map. +fn blended_costs(rates: &HashMap) -> HashMap { + rates + .iter() + .map(|(k, r)| (k.clone(), r.blended())) + .collect() +} + pub struct ModelMetricsService { cost: Arc>>, + rates: Arc>>, latency: Arc>>, } impl ModelMetricsService { pub async fn new(sources: &[MetricsSource], client: reqwest::Client) -> Self { let cost_data = Arc::new(RwLock::new(HashMap::new())); + let rates_data = Arc::new(RwLock::new(HashMap::new())); let latency_data = Arc::new(RwLock::new(HashMap::new())); for source in sources { @@ -34,10 +75,12 @@ impl ModelMetricsService { let data = fetch_cost_pricing(&provider, &url, &client, &aliases).await; info!(models = data.len(), provider = provider_name, url = %url, "fetched cost pricing"); - *cost_data.write().await = data; + *cost_data.write().await = blended_costs(&data); + *rates_data.write().await = data; if let Some(interval_secs) = cfg.refresh_interval { let cost_clone = Arc::clone(&cost_data); + let rates_clone = Arc::clone(&rates_data); let client_clone = client.clone(); let interval = Duration::from_secs(interval_secs); tokio::spawn(async move { @@ -47,7 +90,8 @@ impl ModelMetricsService { fetch_cost_pricing(&provider, &url, &client_clone, &aliases) .await; info!(models = data.len(), provider = provider_name, url = %url, "refreshed cost pricing"); - *cost_clone.write().await = data; + *cost_clone.write().await = blended_costs(&data); + *rates_clone.write().await = data; } }); } @@ -81,10 +125,34 @@ impl ModelMetricsService { ModelMetricsService { cost: cost_data, + rates: rates_data, latency: latency_data, } } + /// Build a service directly from a rates map (no network). Test-only: lets + /// handler/router tests exercise switch-cost math with deterministic pricing. + #[cfg(test)] + pub fn from_rates_for_test(rates: HashMap) -> Self { + ModelMetricsService { + cost: Arc::new(RwLock::new(blended_costs(&rates))), + rates: Arc::new(RwLock::new(rates)), + latency: Arc::new(RwLock::new(HashMap::new())), + } + } + + /// Structured per-million rates for a model, if the cost feed published them. + /// Falls back to the bare model id (without `provider/` prefix) like the + /// blended cost map does. + pub async fn model_rates(&self, model: &str) -> Option { + let rates = self.rates.read().await; + rates.get(model).copied().or_else(|| { + model + .split_once('/') + .and_then(|(_, bare)| rates.get(bare).copied()) + }) + } + /// Rank `models` by `policy`, returning them in preference order. /// Models with no metric data are appended at the end in their original order. pub async fn rank_models(&self, models: &[String], policy: &SelectionPolicy) -> Vec { @@ -167,10 +235,14 @@ struct DoModel { pricing: Option, } +/// DigitalOcean catalog pricing. Despite the `_per_million` field names, the DO +/// API returns these as USD **per token** (e.g. gpt-4o input `2.5e-6`), so they +/// are scaled to per-million at ingestion to match `ModelRates`' convention. #[derive(serde::Deserialize)] struct DoPricing { input_price_per_million: Option, output_price_per_million: Option, + cache_read_input_price_per_million: Option, } #[derive(serde::Deserialize)] @@ -188,6 +260,7 @@ struct ModelsDevModel { struct ModelsDevCost { input: Option, output: Option, + cache_read: Option, } fn default_cost_url(provider: &CostProvider) -> &'static str { @@ -209,7 +282,7 @@ async fn fetch_cost_pricing( url: &str, client: &reqwest::Client, aliases: &HashMap, -) -> HashMap { +) -> HashMap { match provider { CostProvider::Digitalocean => fetch_do_pricing(url, client, aliases).await, CostProvider::ModelsDev => fetch_models_dev_pricing(url, client, aliases).await, @@ -220,21 +293,10 @@ async fn fetch_do_pricing( url: &str, client: &reqwest::Client, aliases: &HashMap, -) -> HashMap { +) -> HashMap { match client.get(url).send().await { Ok(resp) => match resp.json::().await { - Ok(list) => list - .data - .into_iter() - .filter_map(|m| { - let pricing = m.pricing?; - let raw_key = m.model_id.clone(); - let key = aliases.get(&raw_key).cloned().unwrap_or(raw_key); - let cost = pricing.input_price_per_million.unwrap_or(0.0) - + pricing.output_price_per_million.unwrap_or(0.0); - Some((key, cost)) - }) - .collect(), + Ok(list) => parse_do_pricing(list, aliases), Err(err) => { warn!(error = %err, url = %url, "failed to parse digitalocean pricing response"); HashMap::new() @@ -247,16 +309,45 @@ async fn fetch_do_pricing( } } +/// Map the DO catalog into `ModelRates`, scaling DO's per-token prices into the +/// per-million convention the rest of the router uses. +fn parse_do_pricing( + list: DoModelList, + aliases: &HashMap, +) -> HashMap { + list.data + .into_iter() + .filter_map(|m| { + let pricing = m.pricing?; + let raw_key = m.model_id.clone(); + let key = aliases.get(&raw_key).cloned().unwrap_or(raw_key); + // DO reports prices per token; scale to per-million so the + // absolute-USD consumers (the switch-cost gate) are correct. + // Relative ranking via `blended()` is unaffected either way. + let rates = ModelRates { + input_per_million: pricing.input_price_per_million.unwrap_or(0.0) + * TOKENS_PER_MILLION, + output_per_million: pricing.output_price_per_million.unwrap_or(0.0) + * TOKENS_PER_MILLION, + cache_read_per_million: pricing + .cache_read_input_price_per_million + .map(|r| r * TOKENS_PER_MILLION), + }; + Some((key, rates)) + }) + .collect() +} + /// models.dev publishes a top-level object keyed by provider id; each provider /// carries a `models` map whose keys are `creator/model` ids and whose `cost` -/// block holds per-million USD rates. We sum input + output (mirroring the DO -/// ranking metric) and key the result by `creator/model_id` so it lines up with -/// Plano's `provider/model` routing names. +/// block holds per-million USD rates. We keep the structured rates (including +/// `cache_read` when published) and key the result by `creator/model_id` so it +/// lines up with Plano's `provider/model` routing names. async fn fetch_models_dev_pricing( url: &str, client: &reqwest::Client, aliases: &HashMap, -) -> HashMap { +) -> HashMap { match client.get(url).send().await { Ok(resp) => match resp.json::>().await { Ok(providers) => parse_models_dev_pricing(providers, aliases), @@ -275,7 +366,7 @@ async fn fetch_models_dev_pricing( fn parse_models_dev_pricing( providers: HashMap, aliases: &HashMap, -) -> HashMap { +) -> HashMap { let mut out = HashMap::new(); for (provider_id, provider) in providers { for (model_key, model) in provider.models { @@ -286,11 +377,15 @@ fn parse_models_dev_pricing( // First-party providers use bare model keys (`claude-opus-4-5`), // so compose `provider/model` to line up with Plano routing names. let raw_key = format!("{provider_id}/{model_key}"); - let total = input + output; + let rates = ModelRates { + input_per_million: input, + output_per_million: output, + cache_read_per_million: cost.cache_read, + }; let key = aliases.get(&raw_key).cloned().unwrap_or(raw_key); - out.insert(key, total); + out.insert(key, rates); // Also register the bare model id as a fallback lookup. - out.entry(model_key).or_insert(total); + out.entry(model_key).or_insert(rates); } } out @@ -356,6 +451,17 @@ mod tests { SelectionPolicy { prefer } } + fn service_with( + cost: HashMap, + latency: HashMap, + ) -> ModelMetricsService { + ModelMetricsService { + cost: Arc::new(RwLock::new(cost)), + rates: Arc::new(RwLock::new(HashMap::new())), + latency: Arc::new(RwLock::new(latency)), + } + } + #[test] fn test_rank_by_ascending_metric_picks_lowest_first() { let models = vec!["a".to_string(), "b".to_string(), "c".to_string()]; @@ -386,15 +492,15 @@ mod tests { #[tokio::test] async fn test_rank_models_cheapest() { - let service = ModelMetricsService { - cost: Arc::new(RwLock::new({ + let service = service_with( + { let mut m = HashMap::new(); m.insert("gpt-4o".to_string(), 0.005); m.insert("gpt-4o-mini".to_string(), 0.0001); m - })), - latency: Arc::new(RwLock::new(HashMap::new())), - }; + }, + HashMap::new(), + ); let models = vec!["gpt-4o".to_string(), "gpt-4o-mini".to_string()]; let result = service .rank_models(&models, &make_policy(SelectionPreference::Cheapest)) @@ -404,15 +510,12 @@ mod tests { #[tokio::test] async fn test_rank_models_fastest() { - let service = ModelMetricsService { - cost: Arc::new(RwLock::new(HashMap::new())), - latency: Arc::new(RwLock::new({ - let mut m = HashMap::new(); - m.insert("gpt-4o".to_string(), 200.0); - m.insert("claude-sonnet".to_string(), 120.0); - m - })), - }; + let service = service_with(HashMap::new(), { + let mut m = HashMap::new(); + m.insert("gpt-4o".to_string(), 200.0); + m.insert("claude-sonnet".to_string(), 120.0); + m + }); let models = vec!["gpt-4o".to_string(), "claude-sonnet".to_string()]; let result = service .rank_models(&models, &make_policy(SelectionPreference::Fastest)) @@ -422,10 +525,7 @@ mod tests { #[tokio::test] async fn test_rank_models_fallback_no_metrics() { - let service = ModelMetricsService { - cost: Arc::new(RwLock::new(HashMap::new())), - latency: Arc::new(RwLock::new(HashMap::new())), - }; + let service = service_with(HashMap::new(), HashMap::new()); let models = vec!["model-a".to_string(), "model-b".to_string()]; let result = service .rank_models(&models, &make_policy(SelectionPreference::Cheapest)) @@ -435,14 +535,14 @@ mod tests { #[tokio::test] async fn test_rank_models_partial_data_appended_last() { - let service = ModelMetricsService { - cost: Arc::new(RwLock::new({ + let service = service_with( + { let mut m = HashMap::new(); m.insert("gpt-4o".to_string(), 0.005); m - })), - latency: Arc::new(RwLock::new(HashMap::new())), - }; + }, + HashMap::new(), + ); let models = vec!["gpt-4o-mini".to_string(), "gpt-4o".to_string()]; let result = service .rank_models(&models, &make_policy(SelectionPreference::Cheapest)) @@ -452,15 +552,15 @@ mod tests { #[tokio::test] async fn test_rank_models_none_preserves_order() { - let service = ModelMetricsService { - cost: Arc::new(RwLock::new({ + let service = service_with( + { let mut m = HashMap::new(); m.insert("gpt-4o-mini".to_string(), 0.0001); m.insert("gpt-4o".to_string(), 0.005); m - })), - latency: Arc::new(RwLock::new(HashMap::new())), - }; + }, + HashMap::new(), + ); let models = vec!["gpt-4o".to_string(), "gpt-4o-mini".to_string()]; let result = service .rank_models(&models, &make_policy(SelectionPreference::None)) @@ -469,12 +569,61 @@ mod tests { assert_eq!(result, vec!["gpt-4o", "gpt-4o-mini"]); } + #[test] + fn test_parse_do_pricing_scales_per_token_to_per_million() { + // DO's catalog returns USD *per token* despite the `_per_million` names. + // These are the real values from the live catalog. + let json = r#"{ + "data": [ + { + "model_id": "digitalocean/anthropic-claude-4.6-sonnet", + "pricing": { + "input_price_per_million": 3e-6, + "output_price_per_million": 15e-6, + "cache_read_input_price_per_million": 3e-7 + } + }, + { + "model_id": "digitalocean/openai-gpt-4o", + "pricing": { + "input_price_per_million": 2.5e-6, + "output_price_per_million": 10e-6 + } + } + ] + }"#; + let list: DoModelList = serde_json::from_str(json).unwrap(); + let rates = parse_do_pricing(list, &HashMap::new()); + + // Scaled to per-million so the switch-cost gate compares real dollars. + let sonnet = rates + .get("digitalocean/anthropic-claude-4.6-sonnet") + .unwrap(); + assert!((sonnet.input_per_million - 3.0).abs() < 1e-9); + assert!((sonnet.output_per_million - 15.0).abs() < 1e-9); + // DO *does* publish a cached-read rate — it must be captured, not dropped. + assert_eq!(sonnet.cache_read_per_million, Some(0.3)); + + let gpt4o = rates.get("digitalocean/openai-gpt-4o").unwrap(); + assert!((gpt4o.input_per_million - 2.5).abs() < 1e-9); + assert_eq!(gpt4o.cache_read_per_million, None); + + // Regression: a ~5k-token switch off warm sonnet to cold gpt-4o must now + // cost real dollars, not ~1e-8. This is the bug the live gate test caught. + let switch_cost = + 5_000.0 / 1_000_000.0 * (gpt4o.input_per_million - sonnet.cached_input_rate(0.1)); + assert!( + switch_cost > 0.01, + "expected ~$0.011 of switch cost, got {switch_cost}" + ); + } + #[test] fn test_parse_models_dev_pricing_composes_provider_keys() { let json = r#"{ "anthropic": { "models": { - "claude-opus-4-5": {"cost": {"input": 5.0, "output": 25.0}} + "claude-opus-4-5": {"cost": {"input": 5.0, "output": 25.0, "cache_read": 0.5}} } }, "groq": { @@ -486,14 +635,21 @@ mod tests { }"#; let providers: HashMap = serde_json::from_str(json).unwrap(); let aliases = HashMap::new(); - let prices = parse_models_dev_pricing(providers, &aliases); - - assert_eq!(prices.get("anthropic/claude-opus-4-5"), Some(&30.0)); - assert_eq!(prices.get("groq/llama-3.3-70b-versatile"), Some(&1.38)); + let rates = parse_models_dev_pricing(providers, &aliases); + + let opus = rates.get("anthropic/claude-opus-4-5").unwrap(); + assert_eq!(opus.blended(), 30.0); + assert_eq!(opus.cache_read_per_million, Some(0.5)); + let llama = rates.get("groq/llama-3.3-70b-versatile").unwrap(); + assert_eq!(llama.blended(), 1.38); + assert_eq!(llama.cache_read_per_million, None); // bare fallback also registered - assert_eq!(prices.get("claude-opus-4-5"), Some(&30.0)); + assert_eq!( + rates.get("claude-opus-4-5").map(|r| r.blended()), + Some(30.0) + ); // models with no cost block are skipped - assert!(!prices.contains_key("groq/whisper-large-v3-turbo")); + assert!(!rates.contains_key("groq/whisper-large-v3-turbo")); } #[test] @@ -507,10 +663,53 @@ mod tests { "openai/gpt-oss-120b".to_string(), "openai/gpt-4o".to_string(), ); - let prices = parse_models_dev_pricing(providers, &aliases); + let rates = parse_models_dev_pricing(providers, &aliases); + + assert_eq!(rates.get("openai/gpt-4o").map(|r| r.blended()), Some(3.0)); + assert!(!rates.contains_key("openai/gpt-oss-120b")); + } + + #[test] + fn test_cached_input_rate_prefers_published_rate() { + let with_feed = ModelRates { + input_per_million: 3.0, + output_per_million: 15.0, + cache_read_per_million: Some(0.3), + }; + // Published cache_read wins; the discount is ignored. + assert_eq!(with_feed.cached_input_rate(0.5), 0.3); + + let without_feed = ModelRates { + input_per_million: 3.0, + output_per_million: 15.0, + cache_read_per_million: None, + }; + // Falls back to input * discount. + assert!((without_feed.cached_input_rate(0.1) - 0.3).abs() < 1e-9); + } - assert_eq!(prices.get("openai/gpt-4o"), Some(&3.0)); - assert!(!prices.contains_key("openai/gpt-oss-120b")); + #[tokio::test] + async fn test_model_rates_falls_back_to_bare_model_id() { + let service = ModelMetricsService { + cost: Arc::new(RwLock::new(HashMap::new())), + rates: Arc::new(RwLock::new({ + let mut m = HashMap::new(); + m.insert( + "claude-sonnet-4-5".to_string(), + ModelRates { + input_per_million: 3.0, + output_per_million: 15.0, + cache_read_per_million: Some(0.3), + }, + ); + m + })), + latency: Arc::new(RwLock::new(HashMap::new())), + }; + // provider-prefixed lookup falls back to the bare id + let rates = service.model_rates("vercel/claude-sonnet-4-5").await; + assert_eq!(rates.map(|r| r.input_per_million), Some(3.0)); + assert!(service.model_rates("unknown/model").await.is_none()); } #[test] diff --git a/crates/brightstaff/src/router/orchestrator.rs b/crates/brightstaff/src/router/orchestrator.rs index 2d7b25dee..78c3a1d51 100644 --- a/crates/brightstaff/src/router/orchestrator.rs +++ b/crates/brightstaff/src/router/orchestrator.rs @@ -20,9 +20,26 @@ use crate::metrics::labels as metric_labels; use crate::router::orchestrator_model_v1; use crate::session_cache::SessionCache; -pub use crate::session_cache::CachedRoute; +pub use crate::session_cache::SessionBinding; const DEFAULT_SESSION_TTL_SECONDS: u64 = 600; +const TOKENS_PER_MILLION: f64 = 1_000_000.0; + +/// Input-cost of abandoning a plausibly-warm cache to switch to `candidate`. +/// +/// Deliberately input-only: output-token savings are unknowable before the response +/// is generated (reasoning models can emit 5-10x the tokens for the same task), so they +/// are never credited. Negative when the candidate's uncached input rate undercuts the +/// anchor's cached rate — those switches are outright cheaper. Rates are USD per million +/// tokens; the caller draws a positive cost down from the session's switch budget. +pub fn switch_cost_in_usd( + est_context_tokens: u64, + anchor_cached_rate: f64, + candidate_uncached_rate: f64, +) -> f64 { + let context_millions = est_context_tokens as f64 / TOKENS_PER_MILLION; + context_millions * (candidate_uncached_rate - anchor_cached_rate) +} pub struct OrchestratorService { orchestrator_url: String, @@ -126,43 +143,80 @@ impl OrchestratorService { } } - pub async fn get_cached_route( + /// Look up a session binding. Warmth is the caller's concern (time since + /// `last_used`); this only reports whether a binding exists. + pub async fn get_binding( &self, session_id: &str, tenant_id: Option<&str>, - ) -> Option { + ) -> Option { let cache = self.session_cache.as_ref()?; let result = cache.get(&Self::session_key(tenant_id, session_id)).await; - bs_metrics::record_session_cache_event(if result.is_some() { - metric_labels::SESSION_CACHE_HIT - } else { - metric_labels::SESSION_CACHE_MISS + bs_metrics::record_session_cache_event(match result { + Some(_) => metric_labels::SESSION_CACHE_HIT, + None => metric_labels::SESSION_CACHE_MISS, }); result } - pub async fn cache_route( + /// The GC bound for a session binding: the per-scope override when provided, + /// otherwise the global `routing.session_ttl_seconds`. This only governs when an + /// idle binding is reclaimed from memory — not whether its cache is warm. + pub fn effective_session_ttl(&self, ttl_override_seconds: Option) -> Duration { + ttl_override_seconds + .map(Duration::from_secs) + .unwrap_or(self.session_ttl) + } + + /// Persist a session binding with a GC bound (defaults to `routing.session_ttl_seconds` + /// when `gc_ttl` is `None`). + pub async fn store_binding( &self, - session_id: String, + session_id: &str, tenant_id: Option<&str>, - model_name: String, - route_name: Option, + binding: SessionBinding, + gc_ttl: Option, ) { if let Some(ref cache) = self.session_cache { cache .put( - &Self::session_key(tenant_id, &session_id), - CachedRoute { - model_name, - route_name, - }, - self.session_ttl, + &Self::session_key(tenant_id, session_id), + binding, + gc_ttl.unwrap_or(self.session_ttl), ) .await; bs_metrics::record_session_cache_event(metric_labels::SESSION_CACHE_STORE); } } + /// Structured per-million pricing for a model, from the configured cost feed. + /// `None` when no cost source is configured or the model is unknown to the feed. + pub async fn model_rates(&self, model: &str) -> Option { + self.metrics_service.as_ref()?.model_rates(model).await + } + + /// Estimate the input-cost (USD) of switching a warm session from `anchor_model` + /// to `candidate_model`. Fetches per-model rates from the configured cost feed; + /// returns `None` when pricing is missing for either side so the caller can fail + /// open (switch freely) rather than veto the router on guesswork. + /// `cache_read_discount` estimates the anchor's cached-read rate when the feed + /// doesn't publish one. Negative when the switch is outright cheaper. + pub async fn estimate_switch_cost_in_usd( + &self, + est_context_tokens: u64, + anchor_model: &str, + candidate_model: &str, + cache_read_discount: f64, + ) -> Option { + let anchor = self.model_rates(anchor_model).await?; + let candidate = self.model_rates(candidate_model).await?; + Some(switch_cost_in_usd( + est_context_tokens, + anchor.cached_input_rate(cache_read_discount), + candidate.input_per_million, + )) + } + // ---- LLM routing ---- pub async fn determine_route( @@ -348,86 +402,163 @@ mod tests { ) } + fn binding(model: &str, route_name: Option<&str>) -> SessionBinding { + SessionBinding { + anchor_model: model.to_string(), + route_name: route_name.map(|r| r.to_string()), + prefix_hash: None, + last_used: std::time::SystemTime::now(), + cached_tokens: 0, + switch_budget_usd: 0.0, + switches: 0, + } + } + #[tokio::test] async fn test_cache_miss_returns_none() { let svc = make_orchestrator_service(600, 100); - assert!(svc - .get_cached_route("unknown-session", None) - .await - .is_none()); + assert!(svc.get_binding("unknown-session", None).await.is_none()); } #[tokio::test] - async fn test_cache_hit_returns_cached_route() { + async fn test_cache_hit_returns_binding() { let svc = make_orchestrator_service(600, 100); - svc.cache_route( - "s1".to_string(), - None, - "gpt-4o".to_string(), - Some("code".to_string()), - ) - .await; + svc.store_binding("s1", None, binding("gpt-4o", Some("code")), None) + .await; - let cached = svc.get_cached_route("s1", None).await.unwrap(); - assert_eq!(cached.model_name, "gpt-4o"); + let cached = svc.get_binding("s1", None).await.unwrap(); + assert_eq!(cached.anchor_model, "gpt-4o"); assert_eq!(cached.route_name, Some("code".to_string())); } #[tokio::test] async fn test_cache_expired_entry_returns_none() { let svc = make_orchestrator_service(0, 100); - svc.cache_route("s1".to_string(), None, "gpt-4o".to_string(), None) + svc.store_binding("s1", None, binding("gpt-4o", None), None) .await; - assert!(svc.get_cached_route("s1", None).await.is_none()); + assert!(svc.get_binding("s1", None).await.is_none()); } #[tokio::test] async fn test_expired_entries_not_returned() { let svc = make_orchestrator_service(0, 100); - svc.cache_route("s1".to_string(), None, "gpt-4o".to_string(), None) + svc.store_binding("s1", None, binding("gpt-4o", None), None) .await; - svc.cache_route("s2".to_string(), None, "claude".to_string(), None) + svc.store_binding("s2", None, binding("claude", None), None) .await; - assert!(svc.get_cached_route("s1", None).await.is_none()); - assert!(svc.get_cached_route("s2", None).await.is_none()); + assert!(svc.get_binding("s1", None).await.is_none()); + assert!(svc.get_binding("s2", None).await.is_none()); } #[tokio::test] async fn test_cache_evicts_oldest_when_full() { let svc = make_orchestrator_service(600, 2); - svc.cache_route("s1".to_string(), None, "model-a".to_string(), None) + svc.store_binding("s1", None, binding("model-a", None), None) .await; tokio::time::sleep(Duration::from_millis(10)).await; - svc.cache_route("s2".to_string(), None, "model-b".to_string(), None) + svc.store_binding("s2", None, binding("model-b", None), None) .await; - svc.cache_route("s3".to_string(), None, "model-c".to_string(), None) + svc.store_binding("s3", None, binding("model-c", None), None) .await; - assert!(svc.get_cached_route("s1", None).await.is_none()); - assert!(svc.get_cached_route("s2", None).await.is_some()); - assert!(svc.get_cached_route("s3", None).await.is_some()); + assert!(svc.get_binding("s1", None).await.is_none()); + assert!(svc.get_binding("s2", None).await.is_some()); + assert!(svc.get_binding("s3", None).await.is_some()); } #[tokio::test] async fn test_cache_update_existing_session_does_not_evict() { let svc = make_orchestrator_service(600, 2); - svc.cache_route("s1".to_string(), None, "model-a".to_string(), None) + svc.store_binding("s1", None, binding("model-a", None), None) + .await; + svc.store_binding("s2", None, binding("model-b", None), None) .await; - svc.cache_route("s2".to_string(), None, "model-b".to_string(), None) + + svc.store_binding("s1", None, binding("model-a-updated", Some("route")), None) .await; - svc.cache_route( - "s1".to_string(), + let s1 = svc.get_binding("s1", None).await.unwrap(); + assert_eq!(s1.anchor_model, "model-a-updated"); + assert!(svc.get_binding("s2", None).await.is_some()); + } + + #[tokio::test] + async fn test_gc_ttl_override_extends_binding_lifetime() { + // Global GC bound of 0 would reclaim immediately; the per-call override keeps it. + let svc = make_orchestrator_service(0, 100); + svc.store_binding( + "s1", None, - "model-a-updated".to_string(), - Some("route".to_string()), + binding("gpt-4o", None), + Some(Duration::from_secs(600)), ) .await; + let cached = svc.get_binding("s1", None).await.unwrap(); + assert_eq!(cached.anchor_model, "gpt-4o"); + } + + #[tokio::test] + async fn test_binding_fields_round_trip_through_cache() { + let svc = make_orchestrator_service(600, 100); + let mut b = binding("gpt-4o", None); + b.prefix_hash = Some(0xdead_beef); + b.cached_tokens = 12_345; + b.switch_budget_usd = 0.42; + b.switches = 3; + svc.store_binding("s1", None, b, None).await; + + let cached = svc.get_binding("s1", None).await.unwrap(); + assert_eq!(cached.prefix_hash, Some(0xdead_beef)); + assert_eq!(cached.cached_tokens, 12_345); + assert!((cached.switch_budget_usd - 0.42).abs() < 1e-9); + assert_eq!(cached.switches, 3); + } + + // ---- switch-cost math ---- + // + // Real models.dev rates (USD per million input tokens): + // claude-opus-4-1: input 15, cache_read 1.5 + // claude-sonnet-4-5: input 3, cache_read 0.3 + // claude-haiku-4-5: input 1, cache_read 0.1 + // gpt-4.1: input 2, cache_read 0.5 + + #[test] + fn negative_cost_when_candidate_undercuts_cached_rate() { + // Anchor opus (cached 1.5) -> haiku (uncached 1.0) over 100k context: + // cost = 0.1M x (1.0 - 1.5) = -$0.05 — cheaper even after re-reading. + let cost = switch_cost_in_usd(100_000, 1.5, 1.0); + assert!((cost - (-0.05)).abs() < 1e-9); + } + + #[test] + fn positive_cost_when_candidate_pricier_than_cached_rate() { + // Anchor opus (cached 1.5) -> gpt-4.1 (uncached 2.0) over 100k: + // cost = 0.1M x (2.0 - 1.5) = +$0.05. + let cost = switch_cost_in_usd(100_000, 1.5, 2.0); + assert!((cost - 0.05).abs() < 1e-9); + } + + #[test] + fn large_context_amplifies_cost() { + // Anchor sonnet (cached 0.3) -> gpt-5.5-class (uncached 5.0) over 150k: + // cost = 0.15M x (5.0 - 0.3) = +$0.705. + let cost = switch_cost_in_usd(150_000, 0.3, 5.0); + assert!((cost - 0.705).abs() < 1e-9); + } + + #[test] + fn cost_scales_linearly_with_context() { + let small = switch_cost_in_usd(10_000, 0.3, 0.8); + let large = switch_cost_in_usd(1_000_000, 0.3, 0.8); + assert!((large / small - 100.0).abs() < 1e-6); + } - let s1 = svc.get_cached_route("s1", None).await.unwrap(); - assert_eq!(s1.model_name, "model-a-updated"); - assert!(svc.get_cached_route("s2", None).await.is_some()); + #[test] + fn tiny_context_cost_is_negligible() { + // 2k-token chat: even an expensive candidate costs ~$0.009. + let cost = switch_cost_in_usd(2_000, 0.3, 5.0); + assert!(cost < 0.01); } } diff --git a/crates/brightstaff/src/session_cache/memory.rs b/crates/brightstaff/src/session_cache/memory.rs index 541857383..6d67a8b46 100644 --- a/crates/brightstaff/src/session_cache/memory.rs +++ b/crates/brightstaff/src/session_cache/memory.rs @@ -9,9 +9,9 @@ use lru::LruCache; use tokio::sync::Mutex; use tracing::info; -use super::{CachedRoute, SessionCache}; +use super::{SessionBinding, SessionCache}; -type CacheStore = Mutex>; +type CacheStore = Mutex>; pub struct MemorySessionCache { store: Arc, @@ -38,6 +38,8 @@ impl MemorySessionCache { async fn evict_expired(store: &CacheStore) { let mut cache = store.lock().await; + // The TTL is only a GC bound: drop bindings once they've outlived the window + // in which they could plausibly still be warm. let expired: Vec = cache .iter() .filter(|(_, (_, inserted_at, ttl))| inserted_at.elapsed() >= *ttl) @@ -59,21 +61,21 @@ impl MemorySessionCache { #[async_trait] impl SessionCache for MemorySessionCache { - async fn get(&self, key: &str) -> Option { + async fn get(&self, key: &str) -> Option { let mut cache = self.store.lock().await; - if let Some((route, inserted_at, ttl)) = cache.get(key) { + if let Some((binding, inserted_at, ttl)) = cache.get(key) { if inserted_at.elapsed() < *ttl { - return Some(route.clone()); + return Some(binding.clone()); } } None } - async fn put(&self, key: &str, route: CachedRoute, ttl: Duration) { + async fn put(&self, key: &str, binding: SessionBinding, ttl: Duration) { self.store .lock() .await - .put(key.to_string(), (route, Instant::now(), ttl)); + .put(key.to_string(), (binding, Instant::now(), ttl)); } async fn remove(&self, key: &str) { diff --git a/crates/brightstaff/src/session_cache/mod.rs b/crates/brightstaff/src/session_cache/mod.rs index 4cf0dda55..4bb913e32 100644 --- a/crates/brightstaff/src/session_cache/mod.rs +++ b/crates/brightstaff/src/session_cache/mod.rs @@ -1,4 +1,5 @@ use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; use async_trait::async_trait; use common::configuration::Configuration; @@ -8,21 +9,74 @@ use tracing::{debug, info}; pub mod memory; pub mod redis; +/// A conversation's binding to a model, plus the state the session router needs to +/// reason about cache warmth and switch affordability across turns. +/// +/// Warmth is no longer derived from the cache's own expiry — the entry is kept alive +/// as a plain KV value (subject only to a GC bound) and the router decides warmth from +/// [`SessionBinding::last_used`] against the provider's cache window. This is what lets +/// the decision path reason about warmth without ever seeing a provider response. #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] -pub struct CachedRoute { - pub model_name: String, +pub struct SessionBinding { + /// Provider-qualified model this session is anchored to (e.g. `openai/gpt-4o`). + pub anchor_model: String, + #[serde(default, skip_serializing_if = "Option::is_none")] pub route_name: Option, + /// Hash of the stable prompt prefix (system + tools) observed when the binding was + /// stored. Used to detect prefix drift: if a later request's prefix hash differs, + /// the provider cache is already lost so a switch is free. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub prefix_hash: Option, + /// When this session was last dispatched. Warmth = `now - last_used` compared + /// against the provider's idle/hard cache window. + #[serde(default = "SystemTime::now", with = "epoch_secs")] + pub last_used: SystemTime, + /// Best estimate of the cacheable context size (input tokens) — the tokens a switch + /// would have to re-ingest at the uncached rate. Refined from real usage on the + /// full-proxy path; the tokenizer estimate on the decision path. + #[serde(default)] + pub cached_tokens: u64, + /// Remaining cumulative switch budget (USD) for this session. Seeded on the first + /// warm binding and depleted by paid switches; free switches can credit it back. + #[serde(default)] + pub switch_budget_usd: f64, + /// Number of model switches taken during this warm session (observability). + #[serde(default)] + pub switches: u32, +} + +/// Serde helper: persist `SystemTime` as whole epoch seconds so the Redis wire format +/// is stable and compact (the default `SystemTime` representation is version-fragile). +mod epoch_secs { + use super::{Duration, SystemTime, UNIX_EPOCH}; + use serde::{Deserialize, Deserializer, Serializer}; + + pub fn serialize(t: &SystemTime, s: S) -> Result { + let secs = t + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0); + s.serialize_u64(secs) + } + + pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result { + let secs = u64::deserialize(d)?; + Ok(UNIX_EPOCH + Duration::from_secs(secs)) + } } #[async_trait] pub trait SessionCache: Send + Sync { - /// Look up a cached routing decision by key. - async fn get(&self, key: &str) -> Option; + /// Look up a session binding by key. `None` when absent or GC-evicted. Warmth is + /// the caller's concern (time since `last_used`), not the cache's. + async fn get(&self, key: &str) -> Option; - /// Store a routing decision in the session cache with the given TTL. - async fn put(&self, key: &str, route: CachedRoute, ttl: Duration); + /// Store a session binding with the given GC TTL. The TTL is only a memory bound + /// (keep the entry around at least as long as it could plausibly be warm); it does + /// not define warmth. + async fn put(&self, key: &str, binding: SessionBinding, ttl: Duration); - /// Remove a cached routing decision by key. + /// Remove a session binding by key. async fn remove(&self, key: &str); } diff --git a/crates/brightstaff/src/session_cache/redis.rs b/crates/brightstaff/src/session_cache/redis.rs index 29630acc3..d905965d4 100644 --- a/crates/brightstaff/src/session_cache/redis.rs +++ b/crates/brightstaff/src/session_cache/redis.rs @@ -4,7 +4,7 @@ use async_trait::async_trait; use redis::aio::MultiplexedConnection; use redis::AsyncCommands; -use super::{CachedRoute, SessionCache}; +use super::{SessionBinding, SessionCache}; const KEY_PREFIX: &str = "plano:affinity:"; @@ -26,18 +26,20 @@ impl RedisSessionCache { #[async_trait] impl SessionCache for RedisSessionCache { - async fn get(&self, key: &str) -> Option { + async fn get(&self, key: &str) -> Option { let mut conn = self.conn.clone(); let value: Option = conn.get(Self::make_key(key)).await.ok()?; value.and_then(|v| serde_json::from_str(&v).ok()) } - async fn put(&self, key: &str, route: CachedRoute, ttl: Duration) { + async fn put(&self, key: &str, binding: SessionBinding, ttl: Duration) { let mut conn = self.conn.clone(); - let Ok(json) = serde_json::to_string(&route) else { + // The Redis TTL is only a GC bound; warmth is decided by the router from + // `binding.last_used`, not by expiry here. + let ttl_secs = ttl.as_secs().max(1); + let Ok(json) = serde_json::to_string(&binding) else { return; }; - let ttl_secs = ttl.as_secs().max(1); let _: Result<(), _> = conn.set_ex(Self::make_key(key), json, ttl_secs).await; } diff --git a/crates/brightstaff/src/streaming.rs b/crates/brightstaff/src/streaming.rs index 4c532f303..f3016ae0c 100644 --- a/crates/brightstaff/src/streaming.rs +++ b/crates/brightstaff/src/streaming.rs @@ -22,10 +22,14 @@ const STREAM_BUFFER_SIZE: usize = 16; const USAGE_BUFFER_MAX: usize = 2 * 1024 * 1024; use crate::metrics as bs_metrics; use crate::metrics::labels as metric_labels; +use crate::router::orchestrator::OrchestratorService; +use crate::session_cache::SessionBinding; use crate::signals::otel::emit_signals_to_span; use crate::signals::{SignalAnalyzer, FLAG_MARKER}; use crate::tracing::{llm, set_service_name}; use hermesllm::apis::openai::Message; +use std::sync::Arc; +use std::time::{Duration, SystemTime}; /// Parsed usage + resolved-model details from a provider response. #[derive(Debug, Default, Clone)] @@ -187,6 +191,29 @@ pub struct LlmMetricsCtx { pub upstream_status: u16, } +/// Response-side session-update context: refreshes the session binding once the real +/// response is in hand, so `last_used` and the context-size estimate (`cached_tokens`) +/// reflect the turn that just completed. The routing decision (model, budget, switches) +/// was already made and persisted on the request side by [`super::handlers::llm::session_router`]; +/// this only refines the fields that need the response. +pub struct SessionUpdateCtx { + pub orchestrator: Arc, + pub session_id: String, + pub tenant_id: Option, + /// Provider-qualified model this request actually ran on (the router's final pick). + pub anchor_model: String, + pub route_name: Option, + pub prefix_hash: Option, + /// Remaining switch budget from the routing decision — preserved across the refresh. + pub switch_budget_usd: f64, + /// Cumulative switch count from the routing decision — preserved across the refresh. + pub switches: u32, + /// Context-size estimate to fall back to when the response carries no usage block. + pub est_context_tokens: u64, + /// GC bound to store the refreshed binding with. + pub gc_ttl: Duration, +} + /// A processor that tracks streaming metrics pub struct ObservableStreamProcessor { service_name: String, @@ -201,6 +228,7 @@ pub struct ObservableStreamProcessor { /// from the buffer (they still pass through to the client). response_buffer: Vec, llm_metrics: Option, + session_update: Option, metrics_recorded: bool, } @@ -237,6 +265,7 @@ impl ObservableStreamProcessor { messages, response_buffer: Vec::new(), llm_metrics: None, + session_update: None, metrics_recorded: false, } } @@ -247,6 +276,58 @@ impl ObservableStreamProcessor { self.llm_metrics = Some(ctx); self } + + /// Attach session-update context so the processor refreshes `last_used` and the + /// context-size estimate from the real response once usage is known. + pub fn with_session_update(mut self, ctx: SessionUpdateCtx) -> Self { + self.session_update = Some(ctx); + self + } + + /// Refresh the session binding from the response so warmth (`last_used`) and the + /// context-size estimate (`cached_tokens`) reflect the completed turn. + fn handle_session_update(&mut self, usage: &ExtractedUsage) { + let Some(update) = self.session_update.take() else { + return; + }; + let SessionUpdateCtx { + orchestrator, + session_id, + tenant_id, + anchor_model, + route_name, + prefix_hash, + switch_budget_usd, + switches, + est_context_tokens, + gc_ttl, + } = update; + + // Prefer the real prompt-token count (the tokens a future switch would re-read) + // over the request-side estimate; fall back to the estimate when absent. + let cached_tokens = usage + .prompt_tokens + .filter(|&p| p > 0) + .map(|p| p as u64) + .unwrap_or(est_context_tokens); + + bs_metrics::record_session_pin_event(metric_labels::PIN_EVENT_REFRESH); + let binding = SessionBinding { + anchor_model, + route_name, + prefix_hash, + last_used: SystemTime::now(), + cached_tokens, + switch_budget_usd, + switches, + }; + // Fire-and-forget: binding bookkeeping must not delay stream completion. + tokio::spawn(async move { + orchestrator + .store_binding(&session_id, tenant_id.as_deref(), binding, Some(gc_ttl)) + .await; + }); + } } impl StreamProcessor for ObservableStreamProcessor { @@ -357,11 +438,47 @@ impl StreamProcessor for ObservableStreamProcessor { v.max(0) as u64, ); } + // Prompt-cache token counters + hit/miss baseline (cache-blindness is + // invisible without these: a reroute that burns a warm cache shows up + // here as a miss with zero cache_read tokens). + if let Some(v) = usage.cached_input_tokens { + bs_metrics::record_llm_tokens( + &ctx.provider, + &ctx.model, + metric_labels::TOKEN_KIND_CACHE_READ, + v.max(0) as u64, + ); + } + if let Some(v) = usage.cache_creation_tokens { + bs_metrics::record_llm_tokens( + &ctx.provider, + &ctx.model, + metric_labels::TOKEN_KIND_CACHE_WRITE, + v.max(0) as u64, + ); + } + if usage.prompt_tokens.is_some() { + let cache_active = usage.cached_input_tokens.unwrap_or(0) > 0 + || usage.cache_creation_tokens.unwrap_or(0) > 0; + bs_metrics::record_prompt_cache_outcome( + &ctx.provider, + &ctx.model, + if cache_active { + metric_labels::PROMPT_CACHE_HIT + } else { + metric_labels::PROMPT_CACHE_MISS + }, + ); + } if usage.prompt_tokens.is_none() && usage.completion_tokens.is_none() { bs_metrics::record_llm_tokens_usage_missing(&ctx.provider, &ctx.model); } self.metrics_recorded = true; } + + // Session-binding refresh: update `last_used` + the context-size estimate from + // the completed turn (the routing decision itself was made request-side). + self.handle_session_update(&usage); // Release the buffered bytes early; nothing downstream needs them. self.response_buffer.clear(); self.response_buffer.shrink_to_fit(); diff --git a/crates/brightstaff/src/tracing/constants.rs b/crates/brightstaff/src/tracing/constants.rs index f88cf99cb..46ab32959 100644 --- a/crates/brightstaff/src/tracing/constants.rs +++ b/crates/brightstaff/src/tracing/constants.rs @@ -150,6 +150,37 @@ pub mod plano { /// fields (e.g. PostHog). Sourced from the configured /// `tracing.exporters[].distinct_id_header`. Absent for anonymous calls. pub const DISTINCT_ID: &str = "plano.distinct_id"; + + /// Whether the session's provider cache was inferred warm at decision time + /// (from the idle gap vs. the provider's cache window). + pub const CACHE_WARM: &str = "plano.cache.warm"; + + /// How long (ms) since the session was last used — the idle gap warmth is measured + /// against. + pub const CACHE_IDLE_MS: &str = "plano.cache.idle_ms"; + + /// Remaining cumulative switch budget (USD) for the session after this decision. + pub const SESSION_BUDGET_REMAINING_IN_USD: &str = "plano.session.budget_remaining_in_usd"; + + /// Cumulative number of model switches taken during this warm session. + pub const SESSION_SWITCHES: &str = "plano.session.switches"; + + /// Actual input-cost (USD) of the proposed model switch — computed from input-token + /// pricing only, output-token cost deliberately excluded. Negative when the candidate + /// is outright cheaper than staying on the warm anchor. + pub const SWITCH_COST_IN_USD: &str = "plano.switch.cost_in_usd"; + + /// The switch budget (USD) available when the switch cost was evaluated. + pub const SWITCH_THRESHOLD_IN_USD: &str = "plano.switch.threshold_in_usd"; + + /// Switch outcome: "allowed" or "retained". + pub const SWITCH_DECISION: &str = "plano.switch.decision"; + + /// The route (`provider/model`, plus route name when routed) the routing-budget gate + /// *would* have selected had the switch been allowed. Recorded only on a `retained` + /// decision when `routing.routing_budget.record_counterfactual` is enabled. + /// Telemetry only — the counterfactual model is never dispatched. + pub const SWITCH_COUNTERFACTUAL_ROUTE: &str = "plano.switch.counterfactual_route"; } // ============================================================================= diff --git a/crates/common/src/configuration.rs b/crates/common/src/configuration.rs index ceea57b32..1f03debec 100644 --- a/crates/common/src/configuration.rs +++ b/crates/common/src/configuration.rs @@ -33,6 +33,10 @@ pub struct Routing { pub session_ttl_seconds: Option, pub session_max_entries: Option, pub session_cache: Option, + /// Cost gate on model switching within a session. Independent of prompt caching: + /// this is a routing decision that applies whenever it is configured, whether or + /// not `prompt_caching` is enabled. Presence of this block turns it on. + pub routing_budget: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -133,7 +137,7 @@ pub enum StateStorageType { } #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "lowercase")] +#[serde(rename_all = "snake_case")] pub enum SelectionPreference { Cheapest, Fastest, @@ -218,6 +222,10 @@ pub struct Configuration { pub model_aliases: Option>, pub overrides: Option, pub routing: Option, + /// Automatic provider prompt caching. Disabled by default; opt in globally with + /// `prompt_caching: { enabled: true }`. Applies across the entire Plano instance + /// and never changes which model routing selects. + pub prompt_caching: Option, pub system_prompt: Option, pub prompt_guards: Option, pub prompt_targets: Option>, @@ -244,6 +252,180 @@ pub struct Overrides { pub disable_signals: Option, } +/// Automatic prompt caching, configured once for the whole Plano instance. +/// +/// Prompt caching keeps a multi-turn conversation's stable prefix warm in the +/// upstream provider's cache. It never influences which model routing selects — it +/// only (a) auto-injects provider cache-control markers where supported and +/// (b) derives an implicit session key from the stable prompt prefix so follow-up +/// turns reuse the same warm cache. An explicit `X-Model-Affinity` header always wins. +/// +/// Disabled by default; opt in with `enabled: true`. The remaining knobs are optional +/// tuning that only take effect while caching is enabled. +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct PromptCaching { + /// Master switch. Defaults to `false` (opt-in). + #[serde(default)] + pub enabled: bool, + /// Derive an implicit session key from the stable prompt prefix so caches survive + /// across turns without client changes. Defaults to `true` when caching is enabled. + pub session_affinity: Option, + /// Auto-inject provider cache-control markers (e.g. Anthropic `cache_control`). + /// Defaults to `true` when caching is enabled. + pub inject_cache_control: Option, + /// Minimum estimated prefix tokens before a cache breakpoint is injected. + pub min_prefix_tokens: Option, + /// Session pin TTL; falls back to `routing.session_ttl_seconds` when unset. + pub session_ttl_seconds: Option, +} + +/// A cumulative per-session budget (USD) governing when routing may switch models. +/// +/// This is a routing concern, not a caching one: it applies whenever configured, +/// regardless of whether `prompt_caching` is enabled. The default posture is to stick +/// to the model a session is warm on. When routing proposes a *different* model while +/// that session's provider cache is plausibly still warm, switching forces the +/// candidate to re-ingest the whole context at its uncached input rate. That +/// input-token cost — +/// `context_tokens x (candidate_uncached_input_rate - anchor_cached_input_rate)` — +/// is drawn down from the session's remaining budget: a paid switch is allowed only +/// while budget remains, and a switch that is outright cheaper (negative cost) can +/// credit the budget back. Requires a cost source in `model_metrics_sources` so +/// per-model rates are available. Presence of the block turns it on. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)] +pub struct RoutingBudget { + /// Initial budget (USD) granted to a session. `0.0` means "never pay to switch" + /// (only outright-cheaper switches are ever allowed); larger values buy more + /// quality-driven switches over the session. + pub seed_usd: f64, + /// Re-seed the budget back to `seed_usd` when a session goes cold and re-binds + /// (a fresh warm episode). Defaults to `true`. + #[serde(default = "default_true")] + pub replenish_on_rebind: bool, + /// Credit the budget back when a switch is outright cheaper (negative cost) — the + /// candidate's uncached rate undercuts the anchor's cached rate. Defaults to `true`. + #[serde(default = "default_true")] + pub credit_negative: bool, + /// Fallback used to estimate a model's cached input rate when the pricing feed + /// doesn't publish one: `cached_rate = input_rate * cache_read_discount`. A + /// pricing detail, not a cost policy. Defaults to 0.1 (cached reads at 10% of + /// input, typical for Anthropic-style caches). + pub cache_read_discount: Option, + /// When true, a vetoed switch records the route the gate *would* have taken had + /// the switch been allowed, as the `plano.switch.counterfactual_route` span + /// attribute. Telemetry only — the counterfactual model is never dispatched. + /// Useful for evals/benchmarks that want to quantify the road not taken. + /// Defaults to `false`. + #[serde(default)] + pub record_counterfactual: bool, +} + +fn default_true() -> bool { + true +} + +/// Fully-resolved routing-budget settings (present only when configured and valid). +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct EffectiveRoutingBudget { + /// Initial per-session switch budget (USD). + pub seed_usd: f64, + /// Re-seed the budget on cold->warm re-bind. + pub replenish_on_rebind: bool, + /// Credit negative-cost switches back to the budget. + pub credit_negative: bool, + pub cache_read_discount: f64, + /// Emit `plano.switch.counterfactual_route` on vetoed switches. Telemetry only. + pub record_counterfactual: bool, +} + +pub const DEFAULT_CACHE_READ_DISCOUNT: f64 = 0.1; + +impl RoutingBudget { + /// Resolve to effective settings, validating the seed and cache-read discount. + pub fn resolve(&self) -> Result { + if !self.seed_usd.is_finite() || self.seed_usd < 0.0 { + return Err(format!( + "routing.routing_budget.seed_usd: must be a non-negative number, got {}", + self.seed_usd + )); + } + let cache_read_discount = self + .cache_read_discount + .unwrap_or(DEFAULT_CACHE_READ_DISCOUNT); + if !(0.0..=1.0).contains(&cache_read_discount) { + return Err(format!( + "routing.routing_budget.cache_read_discount: must be between 0.0 and 1.0, got {cache_read_discount}" + )); + } + Ok(EffectiveRoutingBudget { + seed_usd: self.seed_usd, + replenish_on_rebind: self.replenish_on_rebind, + credit_negative: self.credit_negative, + cache_read_discount, + record_counterfactual: self.record_counterfactual, + }) + } +} + +impl EffectiveRoutingBudget { + /// Resolve from an optional config block; `None` means the gate is off. + pub fn from_config(config: Option<&RoutingBudget>) -> Result, String> { + config.map(RoutingBudget::resolve).transpose() + } +} + +/// Fully-resolved, instance-wide prompt-caching settings. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct EffectivePromptCaching { + pub enabled: bool, + pub session_affinity: bool, + pub inject_cache_control: bool, + pub min_prefix_tokens: u32, + /// Pin TTL override; `None` uses `routing.session_ttl_seconds`. + pub session_ttl_seconds: Option, +} + +pub const DEFAULT_MIN_PREFIX_TOKENS: u32 = 1024; + +impl Default for EffectivePromptCaching { + fn default() -> Self { + EffectivePromptCaching { + enabled: false, + session_affinity: false, + inject_cache_control: false, + min_prefix_tokens: DEFAULT_MIN_PREFIX_TOKENS, + session_ttl_seconds: None, + } + } +} + +impl PromptCaching { + /// Resolve the instance-wide effective settings. When caching is disabled every + /// sub-feature is off, regardless of the individual knobs. + pub fn resolve(&self) -> Result { + if !self.enabled { + return Ok(EffectivePromptCaching::default()); + } + Ok(EffectivePromptCaching { + enabled: true, + session_affinity: self.session_affinity.unwrap_or(true), + inject_cache_control: self.inject_cache_control.unwrap_or(true), + min_prefix_tokens: self.min_prefix_tokens.unwrap_or(DEFAULT_MIN_PREFIX_TOKENS), + session_ttl_seconds: self.session_ttl_seconds, + }) + } +} + +impl EffectivePromptCaching { + /// Resolve from an optional config block; `None` means caching is off. + pub fn from_config(config: Option<&PromptCaching>) -> Result { + config + .map(PromptCaching::resolve) + .transpose() + .map(Option::unwrap_or_default) + } +} + #[derive(Debug, Clone, Serialize, Deserialize, Default)] pub struct Tracing { pub sampling_rate: Option, @@ -699,7 +881,10 @@ mod test { use pretty_assertions::assert_eq; use std::fs; - use super::{IntoModels, LlmProvider, LlmProviderType}; + use super::{ + EffectivePromptCaching, EffectiveRoutingBudget, IntoModels, LlmProvider, LlmProviderType, + PromptCaching, RoutingBudget, DEFAULT_CACHE_READ_DISCOUNT, DEFAULT_MIN_PREFIX_TOKENS, + }; use crate::api::open_ai::ToolType; #[test] @@ -901,6 +1086,131 @@ disable_signals: false assert_eq!(overrides.disable_signals, None); } + #[test] + fn test_prompt_caching_disabled_by_default() { + // Absent config → everything off. + let effective = EffectivePromptCaching::from_config(None).unwrap(); + assert!(!effective.enabled); + assert!(!effective.session_affinity); + assert!(!effective.inject_cache_control); + + // Present but not enabled → still off. + let cfg: PromptCaching = serde_yaml::from_str("enabled: false").unwrap(); + let effective = cfg.resolve().unwrap(); + assert!(!effective.enabled); + assert!(!effective.session_affinity); + assert!(!effective.inject_cache_control); + } + + #[test] + fn test_prompt_caching_enabled_defaults() { + // A bare `enabled: true` turns everything on with sensible defaults. + let cfg: PromptCaching = serde_yaml::from_str("enabled: true").unwrap(); + let effective = cfg.resolve().unwrap(); + assert!(effective.enabled); + assert!(effective.session_affinity); + assert!(effective.inject_cache_control); + assert_eq!(effective.min_prefix_tokens, DEFAULT_MIN_PREFIX_TOKENS); + assert_eq!(effective.session_ttl_seconds, None); + } + + #[test] + fn test_prompt_caching_optional_knobs() { + let yaml = r#" +enabled: true +session_affinity: false +inject_cache_control: false +min_prefix_tokens: 2048 +session_ttl_seconds: 3600 +"#; + let cfg: PromptCaching = serde_yaml::from_str(yaml).unwrap(); + let effective = cfg.resolve().unwrap(); + assert!(effective.enabled); + assert!(!effective.session_affinity); + assert!(!effective.inject_cache_control); + assert_eq!(effective.min_prefix_tokens, 2048); + assert_eq!(effective.session_ttl_seconds, Some(3600)); + } + + #[test] + fn test_prompt_caching_knobs_ignored_when_disabled() { + // Knobs only take effect while caching is enabled. + let yaml = r#" +enabled: false +session_affinity: true +inject_cache_control: true +"#; + let cfg: PromptCaching = serde_yaml::from_str(yaml).unwrap(); + let effective = cfg.resolve().unwrap(); + assert!(!effective.enabled); + assert!(!effective.session_affinity); + assert!(!effective.inject_cache_control); + } + + #[test] + fn test_routing_budget_parses() { + let yaml = r#" +seed_usd: 0.50 +"#; + let cfg: RoutingBudget = serde_yaml::from_str(yaml).unwrap(); + let budget = cfg.resolve().unwrap(); + assert_eq!(budget.seed_usd, 0.50); + // Replenish + credit-negative default on. + assert!(budget.replenish_on_rebind); + assert!(budget.credit_negative); + assert_eq!(budget.cache_read_discount, DEFAULT_CACHE_READ_DISCOUNT); + // Counterfactual recording is opt-in; off unless requested. + assert!(!budget.record_counterfactual); + } + + #[test] + fn test_routing_budget_record_counterfactual_parses() { + let yaml = r#" +seed_usd: 0.50 +record_counterfactual: true +"#; + let cfg: RoutingBudget = serde_yaml::from_str(yaml).unwrap(); + let budget = cfg.resolve().unwrap(); + assert!(budget.record_counterfactual); + } + + #[test] + fn test_routing_budget_flags_parse() { + let yaml = r#" +seed_usd: 1.0 +replenish_on_rebind: false +credit_negative: false +cache_read_discount: 0.25 +"#; + let cfg: RoutingBudget = serde_yaml::from_str(yaml).unwrap(); + let budget = cfg.resolve().unwrap(); + assert_eq!(budget.seed_usd, 1.0); + assert!(!budget.replenish_on_rebind); + assert!(!budget.credit_negative); + assert_eq!(budget.cache_read_discount, 0.25); + } + + #[test] + fn test_routing_budget_absent_is_off() { + // No block configured → gate is off. + assert!(EffectiveRoutingBudget::from_config(None).unwrap().is_none()); + } + + #[test] + fn test_routing_budget_invalid_values_rejected() { + let negative: RoutingBudget = serde_yaml::from_str("seed_usd: -1.0").unwrap(); + assert!(negative.resolve().is_err()); + + let bad_discount: RoutingBudget = serde_yaml::from_str( + r#" +seed_usd: 0.50 +cache_read_discount: 1.5 +"#, + ) + .unwrap(); + assert!(bad_discount.resolve().is_err()); + } + #[test] fn test_tracing_posthog_exporter_deserialize() { let yaml = r#" diff --git a/crates/common/src/consts.rs b/crates/common/src/consts.rs index c99639ad4..cf882217d 100644 --- a/crates/common/src/consts.rs +++ b/crates/common/src/consts.rs @@ -23,6 +23,12 @@ pub const X_ARCH_FC_MODEL_RESPONSE: &str = "x-arch-fc-model-response"; pub const ARCH_FC_MODEL_NAME: &str = "Arch-Function"; pub const REQUEST_ID_HEADER: &str = "x-request-id"; pub const MODEL_AFFINITY_HEADER: &str = "x-model-affinity"; +/// Per-request prompt-caching control. `off` disables implicit session affinity and +/// cache-control injection for that single request. +pub const PLANO_CACHE_HEADER: &str = "x-plano-cache"; +/// Hash of the stable prompt prefix, forwarded upstream so self-hosted multi-replica +/// backends can do KV-aware (consistent-hash) replica routing at the LB/Envoy layer. +pub const PLANO_PREFIX_HASH_HEADER: &str = "x-plano-prefix-hash"; pub const ENVOY_ORIGINAL_PATH_HEADER: &str = "x-envoy-original-path"; pub const TRACE_PARENT_HEADER: &str = "traceparent"; pub const ARCH_INTERNAL_CLUSTER_NAME: &str = "arch_internal"; diff --git a/crates/hermesllm/src/apis/anthropic.rs b/crates/hermesllm/src/apis/anthropic.rs index cfde591db..024905302 100644 --- a/crates/hermesllm/src/apis/anthropic.rs +++ b/crates/hermesllm/src/apis/anthropic.rs @@ -285,6 +285,7 @@ pub struct MessagesTool { pub name: String, pub description: Option, pub input_schema: Value, + pub cache_control: Option, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] @@ -305,6 +306,155 @@ pub struct MessagesToolChoice { pub disable_parallel_tool_use: Option, } +/// Rough chars-per-token heuristic used to threshold-guard cache-marker injection. +/// Injection below the provider's minimum cacheable prefix is a provider-side no-op, +/// so a cheap estimate is sufficient here. +const CHARS_PER_TOKEN: usize = 4; + +fn block_cache_control(block: &MessagesContentBlock) -> Option<&MessagesCacheControl> { + match block { + MessagesContentBlock::Text { cache_control, .. } + | MessagesContentBlock::Thinking { cache_control, .. } + | MessagesContentBlock::ToolUse { cache_control, .. } + | MessagesContentBlock::ToolResult { cache_control, .. } => cache_control.as_ref(), + _ => None, + } +} + +/// Set `cache_control: ephemeral` on the last block that supports it. +/// Returns true if a marker was placed. +fn mark_last_cacheable_block(blocks: &mut [MessagesContentBlock]) -> bool { + for block in blocks.iter_mut().rev() { + match block { + MessagesContentBlock::Text { cache_control, .. } + | MessagesContentBlock::Thinking { cache_control, .. } + | MessagesContentBlock::ToolUse { cache_control, .. } + | MessagesContentBlock::ToolResult { cache_control, .. } => { + *cache_control = Some(MessagesCacheControl::Ephemeral); + return true; + } + _ => continue, + } + } + false +} + +impl MessagesRequest { + /// True when the client already supplied any `cache_control` marker + /// (on system blocks, tools, or message content blocks). + pub fn has_cache_markers(&self) -> bool { + let system_marked = match &self.system { + Some(MessagesSystemPrompt::Blocks(blocks)) => { + blocks.iter().any(|b| block_cache_control(b).is_some()) + } + _ => false, + }; + let tools_marked = self + .tools + .as_ref() + .is_some_and(|tools| tools.iter().any(|t| t.cache_control.is_some())); + let messages_marked = self.messages.iter().any(|m| match &m.content { + MessagesMessageContent::Blocks(blocks) => { + blocks.iter().any(|b| block_cache_control(b).is_some()) + } + MessagesMessageContent::Single(_) => false, + }); + system_marked || tools_marked || messages_marked + } + + /// Estimated token length of the stable prompt prefix (tools + system), which is + /// what precedes conversation history in Anthropic's prompt ordering. + pub fn estimated_prefix_tokens(&self) -> u32 { + let mut chars = 0usize; + if let Some(tools) = &self.tools { + for tool in tools { + chars += tool.name.len(); + chars += tool.description.as_deref().map_or(0, str::len); + chars += tool.input_schema.to_string().len(); + } + } + match &self.system { + Some(MessagesSystemPrompt::Single(text)) => chars += text.len(), + Some(MessagesSystemPrompt::Blocks(blocks)) => { + chars += blocks.extract_text().len(); + } + None => {} + } + (chars / CHARS_PER_TOKEN) as u32 + } + + /// Auto-inject ephemeral cache breakpoints for providers that require explicit + /// markers (Anthropic-shaped requests): + /// + /// 1. at the end of the system prompt (covering the fully-stable tools + system + /// prefix), falling back to the last tool when there is no system prompt, and + /// 2. a rolling breakpoint on the last content block of the final message, so each + /// turn's cache write becomes the next turn's cache read as history grows. + /// + /// Idempotent: a request that already carries any client-supplied `cache_control` + /// marker is left untouched. Threshold-guarded: no-op when the estimated stable + /// prefix is below `min_prefix_tokens` (injection below the provider's minimum + /// cacheable prefix is wasted bytes). + /// + /// Returns true if any marker was injected. + pub fn inject_cache_breakpoints(&mut self, min_prefix_tokens: u32) -> bool { + if self.has_cache_markers() { + return false; + } + if self.estimated_prefix_tokens() < min_prefix_tokens { + return false; + } + + let mut injected = false; + + // Breakpoint 1: end of the stable prefix. + match self.system.take() { + Some(MessagesSystemPrompt::Single(text)) => { + self.system = Some(MessagesSystemPrompt::Blocks(vec![ + MessagesContentBlock::Text { + text, + cache_control: Some(MessagesCacheControl::Ephemeral), + }, + ])); + injected = true; + } + Some(MessagesSystemPrompt::Blocks(mut blocks)) => { + injected |= mark_last_cacheable_block(&mut blocks); + self.system = Some(MessagesSystemPrompt::Blocks(blocks)); + } + None => { + if let Some(last_tool) = self.tools.as_mut().and_then(|t| t.last_mut()) { + last_tool.cache_control = Some(MessagesCacheControl::Ephemeral); + injected = true; + } + } + } + + // Breakpoint 2: rolling tail of conversation history. + if let Some(last_msg) = self.messages.last_mut() { + let content = std::mem::replace( + &mut last_msg.content, + MessagesMessageContent::Single(String::new()), + ); + last_msg.content = match content { + MessagesMessageContent::Single(text) => { + injected = true; + MessagesMessageContent::Blocks(vec![MessagesContentBlock::Text { + text, + cache_control: Some(MessagesCacheControl::Ephemeral), + }]) + } + MessagesMessageContent::Blocks(mut blocks) => { + injected |= mark_last_cacheable_block(&mut blocks); + MessagesMessageContent::Blocks(blocks) + } + }; + } + + injected + } +} + #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)] #[serde(rename_all = "snake_case")] pub enum MessagesStopReason { @@ -680,6 +830,134 @@ mod tests { use super::*; use serde_json::json; + fn cache_test_request(system: Option) -> MessagesRequest { + MessagesRequest { + model: "claude-sonnet-4".to_string(), + messages: vec![ + MessagesMessage { + role: MessagesRole::User, + content: MessagesMessageContent::Single("first user turn".to_string()), + }, + MessagesMessage { + role: MessagesRole::Assistant, + content: MessagesMessageContent::Single("assistant reply".to_string()), + }, + MessagesMessage { + role: MessagesRole::User, + content: MessagesMessageContent::Single("second user turn".to_string()), + }, + ], + max_tokens: 100, + container: None, + mcp_servers: None, + system, + metadata: None, + service_tier: None, + thinking: None, + temperature: None, + top_p: None, + top_k: None, + stream: None, + stop_sequences: None, + tools: None, + tool_choice: None, + } + } + + #[test] + fn test_inject_cache_breakpoints_marks_system_and_tail() { + let long_system = "x".repeat(8192); // ~2048 estimated tokens + let mut req = cache_test_request(Some(MessagesSystemPrompt::Single(long_system))); + + assert!(req.inject_cache_breakpoints(1024)); + + // System converted to a marked block. + match &req.system { + Some(MessagesSystemPrompt::Blocks(blocks)) => { + assert_eq!(blocks.len(), 1); + assert!(matches!( + &blocks[0], + MessagesContentBlock::Text { + cache_control: Some(MessagesCacheControl::Ephemeral), + .. + } + )); + } + other => panic!("expected marked system blocks, got {:?}", other), + } + + // Rolling breakpoint on the last message. + match &req.messages.last().unwrap().content { + MessagesMessageContent::Blocks(blocks) => { + assert!(matches!( + &blocks[0], + MessagesContentBlock::Text { + cache_control: Some(MessagesCacheControl::Ephemeral), + .. + } + )); + } + other => panic!("expected marked tail blocks, got {:?}", other), + } + } + + #[test] + fn test_inject_cache_breakpoints_below_threshold_is_noop() { + let mut req = cache_test_request(Some(MessagesSystemPrompt::Single( + "short system".to_string(), + ))); + assert!(!req.inject_cache_breakpoints(1024)); + assert!(matches!(req.system, Some(MessagesSystemPrompt::Single(_)))); + } + + #[test] + fn test_inject_cache_breakpoints_respects_client_markers() { + let long_system = "x".repeat(8192); + let mut req = cache_test_request(Some(MessagesSystemPrompt::Blocks(vec![ + MessagesContentBlock::Text { + text: long_system, + cache_control: Some(MessagesCacheControl::Ephemeral), + }, + ]))); + // Client already placed a marker: injection must be a no-op. + assert!(!req.inject_cache_breakpoints(1024)); + // Tail message untouched. + assert!(matches!( + req.messages.last().unwrap().content, + MessagesMessageContent::Single(_) + )); + } + + #[test] + fn test_inject_cache_breakpoints_marks_last_tool_when_no_system() { + let mut req = cache_test_request(None); + req.tools = Some(vec![MessagesTool { + name: "big_tool".to_string(), + description: Some("d".repeat(8192)), + input_schema: json!({"type": "object"}), + cache_control: None, + }]); + + assert!(req.inject_cache_breakpoints(1024)); + assert_eq!( + req.tools.as_ref().unwrap().last().unwrap().cache_control, + Some(MessagesCacheControl::Ephemeral) + ); + } + + #[test] + fn test_tool_cache_control_roundtrip() { + // Client-supplied cache_control on tools must survive serde passthrough. + let tool_json = json!({ + "name": "get_weather", + "input_schema": {"type": "object"}, + "cache_control": {"type": "ephemeral"} + }); + let tool: MessagesTool = serde_json::from_value(tool_json.clone()).unwrap(); + assert_eq!(tool.cache_control, Some(MessagesCacheControl::Ephemeral)); + assert_eq!(serde_json::to_value(&tool).unwrap(), tool_json); + } + #[test] fn test_anthropic_required_fields() { // Create a JSON object with only required fields diff --git a/crates/hermesllm/src/apis/openai.rs b/crates/hermesllm/src/apis/openai.rs index 8e66f0ad6..cf87d8ea9 100644 --- a/crates/hermesllm/src/apis/openai.rs +++ b/crates/hermesllm/src/apis/openai.rs @@ -163,6 +163,112 @@ impl ChatCompletionsRequest { self.store = None; } } + + /// True when any message content part already carries a `cache_control` marker + /// (client-supplied or previously injected). + pub fn has_cache_control_markers(&self) -> bool { + self.messages.iter().any(|m| match &m.content { + Some(MessageContent::Parts(parts)) => parts.iter().any(|p| { + matches!( + p, + ContentPart::Text { + cache_control: Some(_), + .. + } + ) + }), + _ => false, + }) + } + + /// Estimated token length of the stable prompt prefix — the system/developer + /// message(s) plus tool definitions — using a chars/4 heuristic. Precision is + /// not required; this only gates whether marking the prefix is worthwhile. + pub fn estimated_cache_prefix_tokens(&self) -> u32 { + const CHARS_PER_TOKEN: usize = 4; + let mut chars = 0usize; + for m in &self.messages { + if matches!(m.role, Role::System | Role::Developer) { + if let Some(content) = &m.content { + chars += content.extract_text().len(); + } + } + } + if let Some(tools) = &self.tools { + for t in tools { + chars += t.function.name.len(); + chars += t.function.description.as_deref().map_or(0, str::len); + chars += t.function.parameters.to_string().len(); + } + } + (chars / CHARS_PER_TOKEN) as u32 + } + + /// Auto-inject a single ephemeral `cache_control` breakpoint at the end of the + /// stable prompt prefix, for OpenAI-compatible gateways that proxy Anthropic-family + /// models (DigitalOcean, OpenRouter). The marker is attached to the last text + /// content part of the system/developer message (falling back to the first message + /// when there is no system prompt), normalizing a plain-string content into a + /// one-element content-part array as needed. + /// + /// Idempotent: a request that already carries any `cache_control` marker is left + /// untouched. Threshold-guarded: no-op when the estimated stable prefix is below + /// `min_prefix_tokens`. Returns true if a marker was injected. + pub fn inject_cache_control(&mut self, ttl: Option, min_prefix_tokens: u32) -> bool { + if self.has_cache_control_markers() { + return false; + } + if self.estimated_cache_prefix_tokens() < min_prefix_tokens { + return false; + } + + let target_idx = self + .messages + .iter() + .position(|m| matches!(m.role, Role::System | Role::Developer)) + .or(if self.messages.is_empty() { + None + } else { + Some(0) + }); + + let Some(idx) = target_idx else { + return false; + }; + let marker = CacheControl { + kind: "ephemeral".to_string(), + ttl, + }; + attach_cache_control(&mut self.messages[idx], marker) + } +} + +/// Attach a `cache_control` marker to the last text content part of `message`, +/// normalizing string content into a one-element parts array. Returns false when +/// there is no text part to mark (e.g. an image-only message). +fn attach_cache_control(message: &mut Message, marker: CacheControl) -> bool { + let mut parts = match message.content.take() { + Some(MessageContent::Text(text)) => vec![ContentPart::Text { + text, + cache_control: None, + }], + Some(MessageContent::Parts(parts)) => parts, + None => return false, + }; + + let marked = if let Some(ContentPart::Text { cache_control, .. }) = parts + .iter_mut() + .rev() + .find(|p| matches!(p, ContentPart::Text { .. })) + { + *cache_control = Some(marker); + true + } else { + false + }; + + message.content = Some(MessageContent::Parts(parts)); + marked } /// True when the upstream model id is Moonshot's Kimi Code endpoint model. @@ -277,7 +383,7 @@ impl ExtractText for Vec { fn extract_text(&self) -> String { self.iter() .filter_map(|part| match part { - ContentPart::Text { text } => Some(text.as_str()), + ContentPart::Text { text, .. } => Some(text.as_str()), _ => None, }) .collect::>() @@ -296,11 +402,32 @@ impl Display for MessageContent { #[serde(tag = "type")] pub enum ContentPart { #[serde(rename = "text")] - Text { text: String }, + Text { + text: String, + /// Prompt-cache breakpoint for OpenAI-compatible gateways that proxy + /// Anthropic-family models (e.g. DigitalOcean, OpenRouter). Round-trips so + /// client-supplied markers survive deserialization and are respected by the + /// idempotent injector. + #[serde(skip_serializing_if = "Option::is_none")] + cache_control: Option, + }, #[serde(rename = "image_url")] ImageUrl { image_url: ImageUrl }, } +/// Prompt-cache control marker carried on an OpenAI content part. Mirrors the +/// Anthropic `cache_control` object as accepted by OpenAI-compatible gateways that +/// front Anthropic models. +#[skip_serializing_none] +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)] +pub struct CacheControl { + #[serde(rename = "type")] + pub kind: String, // "ephemeral" + /// Cache lifetime hint: "5m" | "1h" (DigitalOcean / OpenRouter). Omitted for + /// plain Anthropic ephemeral caching (defaults to 5 minutes upstream). + pub ttl: Option, +} + /// Image URL configuration for vision capabilities #[skip_serializing_none] #[derive(Serialize, Deserialize, Debug, Clone)] @@ -445,11 +572,35 @@ pub struct Usage { pub total_tokens: u32, pub prompt_tokens_details: Option, pub completion_tokens_details: Option, + /// Anthropic-style cache-read counter emitted by OpenAI-compatible gateways that + /// front Anthropic models (e.g. DigitalOcean), which do *not* populate + /// `prompt_tokens_details.cached_tokens`. Captured here so it can be surfaced to + /// OpenAI clients via [`Usage::normalize_cache_tokens`]. + pub cache_read_input_tokens: Option, + /// Anthropic-style cache-write counter (see `cache_read_input_tokens`). + pub cache_creation_input_tokens: Option, +} + +impl Usage { + /// Fold gateway-specific Anthropic-style `cache_read_input_tokens` into the + /// OpenAI-standard `prompt_tokens_details.cached_tokens` so OpenAI-compatible + /// clients (and downstream cost accounting) observe the cache hit. No-op when the + /// standard field is already populated or no cache-read counter is present. + pub fn normalize_cache_tokens(&mut self) { + if let Some(read) = self.cache_read_input_tokens { + let details = self + .prompt_tokens_details + .get_or_insert_with(Default::default); + if details.cached_tokens.is_none() { + details.cached_tokens = Some(read); + } + } + } } /// Detailed breakdown of prompt tokens #[skip_serializing_none] -#[derive(Serialize, Deserialize, Debug, Clone)] +#[derive(Serialize, Deserialize, Debug, Clone, Default)] pub struct PromptTokensDetails { pub cached_tokens: Option, pub audio_tokens: Option, @@ -632,7 +783,15 @@ impl TokenUsage for Usage { fn cached_input_tokens(&self) -> Option { self.prompt_tokens_details .as_ref() - .and_then(|d| d.cached_tokens.map(|t| t as usize)) + .and_then(|d| d.cached_tokens) + // Gateways fronting Anthropic (e.g. DigitalOcean) report cache reads here + // instead of prompt_tokens_details.cached_tokens. + .or(self.cache_read_input_tokens) + .map(|t| t as usize) + } + + fn cache_creation_tokens(&self) -> Option { + self.cache_creation_input_tokens.map(|t| t as usize) } fn reasoning_tokens(&self) -> Option { @@ -666,7 +825,7 @@ impl ProviderRequest for ChatCompletionsRequest { MessageContent::Parts(parts) => parts .iter() .map(|part| match part { - ContentPart::Text { text } => text.clone(), + ContentPart::Text { text, .. } => text.clone(), ContentPart::ImageUrl { .. } => "[Image]".to_string(), }) .collect::>() @@ -796,6 +955,137 @@ mod tests { use super::*; use serde_json::json; + /// Build a request with a long system prompt (well over the token threshold) plus + /// a short user turn. + fn request_with_system(system: &str, user: &str) -> ChatCompletionsRequest { + ChatCompletionsRequest { + model: "anthropic/claude-3.5-sonnet".to_string(), + messages: vec![ + Message { + role: Role::System, + content: Some(MessageContent::Text(system.to_string())), + name: None, + tool_calls: None, + tool_call_id: None, + }, + Message { + role: Role::User, + content: Some(MessageContent::Text(user.to_string())), + name: None, + tool_calls: None, + tool_call_id: None, + }, + ], + ..Default::default() + } + } + + fn system_cache_control(req: &ChatCompletionsRequest) -> Option<&CacheControl> { + match &req.messages[0].content { + Some(MessageContent::Parts(parts)) => parts.iter().find_map(|p| match p { + ContentPart::Text { cache_control, .. } => cache_control.as_ref(), + _ => None, + }), + _ => None, + } + } + + #[test] + fn inject_cache_control_normalizes_string_content_and_marks_prefix() { + let mut req = request_with_system(&"x".repeat(8000), "hi"); + let injected = req.inject_cache_control(None, 1024); + assert!(injected); + // The plain-string system content was normalized into a one-element parts array. + let marker = system_cache_control(&req).expect("system content part should be marked"); + assert_eq!(marker.kind, "ephemeral"); + assert_eq!(marker.ttl, None); + // The user message is untouched. + assert!(matches!( + req.messages[1].content, + Some(MessageContent::Text(_)) + )); + } + + #[test] + fn inject_cache_control_is_idempotent() { + let mut req = request_with_system(&"x".repeat(8000), "hi"); + assert!(req.inject_cache_control(Some("1h".to_string()), 1024)); + // A second pass must be a no-op because a marker already exists. + assert!(!req.inject_cache_control(None, 1024)); + assert_eq!( + system_cache_control(&req).and_then(|c| c.ttl.clone()), + Some("1h".to_string()) + ); + } + + #[test] + fn inject_cache_control_respects_min_prefix_threshold() { + // A tiny prefix (below the threshold) is not worth marking. + let mut req = request_with_system("short system", "hi"); + assert!(!req.inject_cache_control(None, 1024)); + assert!(system_cache_control(&req).is_none()); + } + + #[test] + fn client_supplied_cache_control_round_trips_and_blocks_injection() { + // A client that already set cache_control must survive deserialization and be + // respected (no double injection). + let raw = json!({ + "model": "anthropic/claude-3.5-sonnet", + "messages": [ + { + "role": "system", + "content": [ + {"type": "text", "text": "big stable prefix", + "cache_control": {"type": "ephemeral", "ttl": "5m"}} + ] + }, + {"role": "user", "content": "hi"} + ] + }); + let mut req: ChatCompletionsRequest = serde_json::from_value(raw).unwrap(); + assert!(req.has_cache_control_markers()); + // Injection is a no-op, and re-serialization preserves the client's marker. + assert!(!req.inject_cache_control(None, 0)); + let reserialized = serde_json::to_value(&req).unwrap(); + let marker = &reserialized["messages"][0]["content"][0]["cache_control"]; + assert_eq!(marker["type"], "ephemeral"); + assert_eq!(marker["ttl"], "5m"); + } + + #[test] + fn cache_control_omitted_when_absent() { + // A plain text part must not emit a null cache_control field. + let part = ContentPart::Text { + text: "hello".to_string(), + cache_control: None, + }; + let v = serde_json::to_value(&part).unwrap(); + assert!(v.get("cache_control").is_none()); + } + + #[test] + fn usage_normalizes_gateway_cache_read_into_cached_tokens() { + // DigitalOcean-style usage: Anthropic cache_read with no prompt_tokens_details. + let raw = json!({ + "prompt_tokens": 100, + "completion_tokens": 20, + "total_tokens": 120, + "cache_read_input_tokens": 80, + "cache_creation_input_tokens": 12 + }); + let mut usage: Usage = serde_json::from_value(raw).unwrap(); + // TokenUsage falls back to the gateway field. + assert_eq!(usage.cached_input_tokens(), Some(80)); + assert_eq!(usage.cache_creation_tokens(), Some(12)); + // Normalization surfaces it as the OpenAI-standard cached_tokens. + usage.normalize_cache_tokens(); + assert_eq!( + usage.prompt_tokens_details.and_then(|d| d.cached_tokens), + Some(80) + ); + } + #[test] fn test_required_fields() { // Create a JSON object with only required fields @@ -995,7 +1285,7 @@ mod tests { assert_eq!(content_parts.len(), 2); // Validate text content part - if let ContentPart::Text { text } = &content_parts[0] { + if let ContentPart::Text { text, .. } = &content_parts[0] { assert_eq!(text, "What can you see in this image and what's the weather like in the location shown?"); } else { panic!("Expected text content part"); diff --git a/crates/hermesllm/src/lib.rs b/crates/hermesllm/src/lib.rs index 3b9611e00..cb867fa1e 100644 --- a/crates/hermesllm/src/lib.rs +++ b/crates/hermesllm/src/lib.rs @@ -6,10 +6,14 @@ pub mod clients; pub mod providers; pub mod transforms; // Re-export important types and traits +pub use apis::openai::CacheControl; pub use apis::streaming_shapes::amazon_bedrock_binary_frame::BedrockBinaryFrameDecoder; pub use apis::streaming_shapes::sse::{SseEvent, SseStreamIter}; pub use aws_smithy_eventstream::frame::DecodedFrame; -pub use providers::id::ProviderId; +pub use providers::id::{ + cache_marker_strategy, provider_cache_capability, CacheMarkerStrategy, ProviderCacheCapability, + ProviderId, +}; pub use providers::request::{ProviderRequest, ProviderRequestError, ProviderRequestType}; pub use providers::response::{ ProviderResponse, ProviderResponseError, ProviderResponseType, TokenUsage, diff --git a/crates/hermesllm/src/providers/id.rs b/crates/hermesllm/src/providers/id.rs index aaf1c16ed..7e030368a 100644 --- a/crates/hermesllm/src/providers/id.rs +++ b/crates/hermesllm/src/providers/id.rs @@ -4,6 +4,7 @@ use serde::Deserialize; use std::collections::HashMap; use std::fmt::Display; use std::sync::OnceLock; +use std::time::Duration; static PROVIDER_MODELS_YAML: &str = include_str!(concat!( env!("CARGO_MANIFEST_DIR"), @@ -92,6 +93,192 @@ impl TryFrom<&str> for ProviderId { } } +/// How Plano should mark a request for prompt caching, resolved from the *combination* +/// of gateway provider, the underlying model family, and the upstream API shape — not +/// the gateway alone. This is what lets DigitalOcean-Anthropic and OpenRouter-Anthropic +/// (both OpenAI-compatible chat completions fronting Anthropic models) cache through one +/// path, while OpenAI-family models stay correctly automatic and unimplemented backends +/// (Bedrock) are an honest `None` rather than a silent no-op. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum CacheMarkerStrategy { + /// No known prompt-caching support for this combination — do nothing. + None, + /// Provider caches stable prefixes automatically (OpenAI-family anywhere); no + /// request markers are needed. Plano only keeps the prefix byte-stable and pinned. + Automatic, + /// OpenAI-compatible chat completions fronting Anthropic-family models + /// (DigitalOcean, OpenRouter): attach `cache_control` to content parts. + OpenAiContentPartCacheControl { + /// Minimum cacheable prefix length in tokens; injecting below this is a no-op. + min_prefix_tokens: u32, + /// Optional cache lifetime hint ("5m" | "1h"). + ttl: Option, + }, + /// Native Anthropic Messages API (native `anthropic/*`, Vercel-Anthropic): inject + /// ephemeral breakpoints on the Anthropic-shaped request. + AnthropicMessagesBreakpoints { + /// Maximum number of cache breakpoints the provider accepts per request. + max_breakpoints: u8, + /// Minimum cacheable prefix length in tokens; injecting below this is a no-op. + min_prefix_tokens: u32, + }, + // BedrockCachePoint { .. } // left as explicit `None` until implemented. +} + +/// Coarse model family, inferred from the model id. Works across gateway naming +/// conventions (DigitalOcean's `anthropic-claude-…` dash form and OpenRouter's +/// `anthropic/claude-…` slash form) via substring matching. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ModelFamily { + Anthropic, + OpenAI, + Other, +} + +fn model_family(model_name: &str) -> ModelFamily { + let m = model_name.to_ascii_lowercase(); + if m.contains("claude") || m.contains("anthropic") { + ModelFamily::Anthropic + } else if m.contains("gpt") || m.contains("openai") || m.contains("chatgpt") { + ModelFamily::OpenAI + } else { + ModelFamily::Other + } +} + +/// Whether a gateway accepts Anthropic-style `cache_control` on OpenAI content parts +/// over its chat-completions endpoint. +fn accepts_openai_content_part_cache_control(provider: ProviderId) -> bool { + matches!(provider, ProviderId::DigitalOcean | ProviderId::OpenRouter) +} + +/// Whether a gateway/model relies on automatic prefix caching (no markers required) +/// over an OpenAI-compatible surface. +fn is_automatic_cache_provider(provider: ProviderId) -> bool { + matches!( + provider, + ProviderId::OpenAI + | ProviderId::AzureOpenAI + | ProviderId::ChatGPT + | ProviderId::Groq + | ProviderId::Deepseek + | ProviderId::Gemini + | ProviderId::Moonshotai + | ProviderId::XAI + | ProviderId::DigitalOcean + | ProviderId::OpenRouter + ) +} + +/// Resolve the cache-marking strategy for a `(gateway provider × underlying model × +/// upstream API)` combination. +/// +/// - `model_name` is the id *after* the gateway prefix (e.g. `anthropic-claude-3-5-sonnet` +/// for DigitalOcean, `anthropic/claude-3.5-sonnet` for OpenRouter). +pub fn cache_marker_strategy( + provider: ProviderId, + model_name: &str, + upstream_api: &SupportedUpstreamAPIs, +) -> CacheMarkerStrategy { + // Anthropic minimum cacheable prefix is ~1024 tokens (2048 for Haiku-class); + // callers may raise this via config. + const ANTHROPIC_MIN_PREFIX_TOKENS: u32 = 1024; + + match upstream_api { + // Native Anthropic Messages API — inject ephemeral breakpoints. + SupportedUpstreamAPIs::AnthropicMessagesAPI(_) => { + CacheMarkerStrategy::AnthropicMessagesBreakpoints { + max_breakpoints: 4, + min_prefix_tokens: ANTHROPIC_MIN_PREFIX_TOKENS, + } + } + // OpenAI-compatible chat completions — strategy depends on the model family. + SupportedUpstreamAPIs::OpenAIChatCompletions(_) => match model_family(model_name) { + ModelFamily::Anthropic if accepts_openai_content_part_cache_control(provider) => { + CacheMarkerStrategy::OpenAiContentPartCacheControl { + min_prefix_tokens: ANTHROPIC_MIN_PREFIX_TOKENS, + ttl: None, + } + } + // Anthropic-family behind a gateway that doesn't accept content-part + // cache_control over chat completions: no honest way to mark it. + ModelFamily::Anthropic => CacheMarkerStrategy::None, + ModelFamily::OpenAI => CacheMarkerStrategy::Automatic, + ModelFamily::Other if is_automatic_cache_provider(provider) => { + CacheMarkerStrategy::Automatic + } + ModelFamily::Other => CacheMarkerStrategy::None, + }, + // OpenAI Responses API — OpenAI-family automatic prefix caching. + SupportedUpstreamAPIs::OpenAIResponsesAPI(_) => CacheMarkerStrategy::Automatic, + // Bedrock cache points not yet implemented — honest None instead of a + // silent no-op. + SupportedUpstreamAPIs::AmazonBedrockConverse(_) + | SupportedUpstreamAPIs::AmazonBedrockConverseStream(_) => CacheMarkerStrategy::None, + } +} + +/// Provider prompt-cache retention behavior, used to decide whether a session's +/// upstream cache is still plausibly warm from the time since it was last used. +/// +/// This is deliberately time/behavior only — it says nothing about *how* to mark a +/// request for caching (that's [`CacheMarkerStrategy`]). Warmth is a function of the +/// idle gap vs the provider's cache window, so the session router can reason about +/// stickiness without ever seeing a provider response. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ProviderCacheCapability { + /// Sliding idle window: the cache stays warm as long as it is touched at least + /// this often. Anthropic's default ephemeral cache is 5 minutes. + pub idle_ttl: Duration, + /// Absolute ceiling on how long a cache entry can live regardless of activity. + /// Conservative default of 1h matches Anthropic's extended (1h) tier ceiling. + pub hard_ttl: Duration, + /// Whether the provider (as configured) actually retains caches out to the + /// extended window. Off by default — extended retention is opt-in per provider. + pub extended_retention: bool, + /// The extended idle window when `extended_retention` is enabled (e.g. 1h). + pub extended_ttl: Duration, +} + +impl Default for ProviderCacheCapability { + fn default() -> Self { + // Conservative, provider-agnostic defaults: a 5-minute sliding window capped + // at 1 hour, no extended retention. Anything unknown is treated as short-lived + // so the router doesn't over-stick to a cache that has likely gone cold. + ProviderCacheCapability { + idle_ttl: Duration::from_secs(5 * 60), + hard_ttl: Duration::from_secs(60 * 60), + extended_retention: false, + extended_ttl: Duration::from_secs(60 * 60), + } + } +} + +/// Resolve the prompt-cache retention window for a gateway provider. Data-driven so +/// tuning a provider's window needs no code changes at the call sites — only this +/// table. Unknown providers fall back to the conservative [`ProviderCacheCapability::default`]. +pub fn provider_cache_capability(provider: ProviderId) -> ProviderCacheCapability { + match provider { + // Anthropic-family caches (native or fronted): 5-minute sliding default, + // 1-hour hard ceiling. Extended (1h) retention is opt-in and left off here. + ProviderId::Anthropic + | ProviderId::DigitalOcean + | ProviderId::OpenRouter + | ProviderId::Vercel => ProviderCacheCapability::default(), + // OpenAI-family automatic prefix caching also lives on the order of minutes; + // the conservative default holds. + ProviderId::OpenAI + | ProviderId::AzureOpenAI + | ProviderId::ChatGPT + | ProviderId::Groq + | ProviderId::Deepseek + | ProviderId::Gemini + | ProviderId::Moonshotai + | ProviderId::XAI => ProviderCacheCapability::default(), + _ => ProviderCacheCapability::default(), + } +} + impl ProviderId { /// Get all available models for this provider /// Returns model names without the provider prefix (e.g., "gpt-4" not "openai/gpt-4") @@ -291,6 +478,100 @@ impl Display for ProviderId { #[cfg(test)] mod tests { use super::*; + use crate::apis::{AnthropicApi, OpenAIApi}; + + fn chat_completions() -> SupportedUpstreamAPIs { + SupportedUpstreamAPIs::OpenAIChatCompletions(OpenAIApi::ChatCompletions) + } + + fn anthropic_messages() -> SupportedUpstreamAPIs { + SupportedUpstreamAPIs::AnthropicMessagesAPI(AnthropicApi::Messages) + } + + #[test] + fn digitalocean_anthropic_uses_openai_content_part_markers() { + // DO fronts Anthropic over an OpenAI-compatible surface (dash-form model id). + let strategy = cache_marker_strategy( + ProviderId::DigitalOcean, + "anthropic-claude-3-5-sonnet", + &chat_completions(), + ); + assert!(matches!( + strategy, + CacheMarkerStrategy::OpenAiContentPartCacheControl { .. } + )); + } + + #[test] + fn openrouter_anthropic_uses_openai_content_part_markers() { + // OpenRouter uses slash-form model ids after the gateway prefix. + let strategy = cache_marker_strategy( + ProviderId::OpenRouter, + "anthropic/claude-3.5-sonnet", + &chat_completions(), + ); + assert!(matches!( + strategy, + CacheMarkerStrategy::OpenAiContentPartCacheControl { .. } + )); + } + + #[test] + fn openai_family_over_chat_completions_is_automatic() { + assert_eq!( + cache_marker_strategy( + ProviderId::DigitalOcean, + "openai-gpt-4o", + &chat_completions() + ), + CacheMarkerStrategy::Automatic + ); + assert_eq!( + cache_marker_strategy(ProviderId::OpenAI, "gpt-4o", &chat_completions()), + CacheMarkerStrategy::Automatic + ); + } + + #[test] + fn native_anthropic_uses_messages_breakpoints() { + let strategy = cache_marker_strategy( + ProviderId::Anthropic, + "claude-3-5-sonnet-20241022", + &anthropic_messages(), + ); + assert!(matches!( + strategy, + CacheMarkerStrategy::AnthropicMessagesBreakpoints { .. } + )); + } + + #[test] + fn anthropic_family_without_content_part_support_is_none() { + // An Anthropic-family model over chat completions on a gateway that does not + // accept content-part cache_control has no honest marking path. + assert_eq!( + cache_marker_strategy( + ProviderId::Vercel, + "anthropic/claude-3.5", + &chat_completions() + ), + CacheMarkerStrategy::None + ); + } + + #[test] + fn bedrock_is_honest_none() { + assert_eq!( + cache_marker_strategy( + ProviderId::AmazonBedrock, + "anthropic.claude-3-5-sonnet", + &SupportedUpstreamAPIs::AmazonBedrockConverse( + crate::apis::AmazonBedrockApi::Converse + ) + ), + CacheMarkerStrategy::None + ); + } #[test] fn test_models_loaded_from_yaml() { diff --git a/crates/hermesllm/src/transforms/lib.rs b/crates/hermesllm/src/transforms/lib.rs index 5308cc471..d936a7696 100644 --- a/crates/hermesllm/src/transforms/lib.rs +++ b/crates/hermesllm/src/transforms/lib.rs @@ -68,7 +68,10 @@ impl ContentUtils for Vec { for block in self { match block { MessagesContentBlock::Text { text, .. } => { - content_parts.push(ContentPart::Text { text: text.clone() }); + content_parts.push(ContentPart::Text { + text: text.clone(), + cache_control: None, + }); } MessagesContentBlock::Image { source } => { let url = convert_image_source_to_url(source); @@ -198,7 +201,7 @@ pub fn convert_openai_message_to_anthropic_content( Some(MessageContent::Parts(parts)) => { for part in parts { match part { - ContentPart::Text { text } => { + ContentPart::Text { text, .. } => { blocks.push(MessagesContentBlock::Text { text: text.clone(), cache_control: None, diff --git a/crates/hermesllm/src/transforms/request/from_anthropic.rs b/crates/hermesllm/src/transforms/request/from_anthropic.rs index 20442c59c..bb42ff988 100644 --- a/crates/hermesllm/src/transforms/request/from_anthropic.rs +++ b/crates/hermesllm/src/transforms/request/from_anthropic.rs @@ -324,7 +324,7 @@ fn build_openai_content( None } else if content_parts.len() == 1 && tool_calls.is_empty() { match &content_parts[0] { - ContentPart::Text { text } => Some(MessageContent::Text(text.clone())), + ContentPart::Text { text, .. } => Some(MessageContent::Text(text.clone())), _ => Some(MessageContent::Parts(content_parts)), } } else if content_parts.is_empty() { @@ -562,6 +562,7 @@ mod tests { }, "required": ["location"] }), + cache_control: None, }]), tool_choice: Some(MessagesToolChoice { kind: MessagesToolChoiceType::Tool, @@ -620,6 +621,7 @@ mod tests { "type": "object", "properties": {} }), + cache_control: None, }]), tool_choice: Some(MessagesToolChoice { kind: MessagesToolChoiceType::Auto, diff --git a/crates/hermesllm/src/transforms/request/from_openai.rs b/crates/hermesllm/src/transforms/request/from_openai.rs index 0514c039a..deca6cc6a 100644 --- a/crates/hermesllm/src/transforms/request/from_openai.rs +++ b/crates/hermesllm/src/transforms/request/from_openai.rs @@ -127,6 +127,7 @@ impl TryFrom for Vec { | InputContent::OutputText { text } => { Some(crate::apis::openai::ContentPart::Text { text: text.clone(), + cache_control: None, }) } InputContent::InputImage { image_url, .. } => { @@ -154,6 +155,7 @@ impl TryFrom for Vec { | InputContent::OutputText { text } => { Some(crate::apis::openai::ContentPart::Text { text: text.clone(), + cache_control: None, }) } InputContent::InputImage { image_url, .. } => { @@ -330,7 +332,7 @@ impl TryFrom for BedrockMessage { // Convert OpenAI content parts to Bedrock ContentBlocks for part in parts { match part { - crate::apis::openai::ContentPart::Text { text } => { + crate::apis::openai::ContentPart::Text { text, .. } => { if !text.is_empty() { content_blocks.push(ContentBlock::Text { text }); } @@ -882,6 +884,7 @@ fn convert_openai_tools(tools: Vec) -> Vec { name: tool.function.name, description: tool.function.description, input_schema: tool.function.parameters, + cache_control: None, }) .collect() } diff --git a/crates/hermesllm/src/transforms/response/to_openai.rs b/crates/hermesllm/src/transforms/response/to_openai.rs index c6d6f86dc..f2051e118 100644 --- a/crates/hermesllm/src/transforms/response/to_openai.rs +++ b/crates/hermesllm/src/transforms/response/to_openai.rs @@ -14,13 +14,18 @@ use crate::transforms::lib::*; // Usage Conversions impl From for Usage { fn from(val: MessagesUsage) -> Self { - Usage { + let mut usage = Usage { prompt_tokens: val.input_tokens, completion_tokens: val.output_tokens, total_tokens: val.input_tokens + val.output_tokens, prompt_tokens_details: None, completion_tokens_details: None, - } + cache_read_input_tokens: val.cache_read_input_tokens, + cache_creation_input_tokens: val.cache_creation_input_tokens, + }; + // Surface Anthropic cache reads to OpenAI clients as cached_tokens. + usage.normalize_cache_tokens(); + usage } } @@ -244,6 +249,7 @@ impl TryFrom for ChatCompletionsResponse { total_tokens: resp.usage.input_tokens + resp.usage.output_tokens, prompt_tokens_details: None, completion_tokens_details: None, + ..Default::default() }; Ok(ChatCompletionsResponse { @@ -312,6 +318,7 @@ impl TryFrom for ChatCompletionsResponse { total_tokens: resp.usage.total_tokens, prompt_tokens_details: None, completion_tokens_details: None, + ..Default::default() }; // Generate a response ID (using timestamp since Bedrock doesn't provide one) @@ -980,6 +987,7 @@ mod tests { total_tokens: 30, prompt_tokens_details: None, completion_tokens_details: None, + ..Default::default() }, system_fingerprint: None, service_tier: Some("default".to_string()), @@ -1061,6 +1069,7 @@ mod tests { total_tokens: 40, prompt_tokens_details: None, completion_tokens_details: None, + ..Default::default() }, system_fingerprint: None, service_tier: None, @@ -1134,6 +1143,7 @@ mod tests { total_tokens: 101, prompt_tokens_details: None, completion_tokens_details: None, + ..Default::default() }, system_fingerprint: Some("fp_7eeb46f068".to_string()), service_tier: Some("default".to_string()), diff --git a/crates/hermesllm/src/transforms/response_streaming/to_openai_streaming.rs b/crates/hermesllm/src/transforms/response_streaming/to_openai_streaming.rs index 4aa719afb..a45a84835 100644 --- a/crates/hermesllm/src/transforms/response_streaming/to_openai_streaming.rs +++ b/crates/hermesllm/src/transforms/response_streaming/to_openai_streaming.rs @@ -247,6 +247,7 @@ impl TryFrom for ChatCompletionsStreamResponse { total_tokens: metadata_event.usage.total_tokens, prompt_tokens_details: None, completion_tokens_details: None, + ..Default::default() }; Ok(create_openai_chunk( diff --git a/demos/llm_routing/routing_budget/GUIDE.md b/demos/llm_routing/routing_budget/GUIDE.md new file mode 100644 index 000000000..bec347b72 --- /dev/null +++ b/demos/llm_routing/routing_budget/GUIDE.md @@ -0,0 +1,251 @@ +# How-To: See Prompt Caching + Routing Budget in Action + +A hands-on guide for running Plano's automatic prompt caching and the routing budget locally, and for measuring the win in evals/benchmarks (e.g. on DigitalOcean models). + +There are two independent behaviors to observe: + +1. **Automatic prompt caching** — staying on one model keeps the stable prefix + warm, so per-turn input cost drops sharply across a multi-turn conversation. +2. **Routing budget** — the router still runs every turn, but when it proposes a + *different* model while the session's cache is plausibly warm, Plano only switches + while the session's cumulative budget covers the input-cost of abandoning that + cache. This is a routing concern and works whether or not prompt caching is on. + +--- + + + +## 1. Prerequisites + +- Plano CLI installed: `pip install planoai` (or `uv sync` from `cli/` for a dev build). +- Provider credentials as env vars, e.g.: + - `export DIGITALOCEAN_API_KEY=...` (DO SI) + - `export OPENAI_API_KEY=...`, `export ANTHROPIC_API_KEY=...` (if comparing) +- `curl` + `jq` for poking the endpoint. + +--- + + + +## 2. Configuration + +Start from `[config.yaml](config.yaml)` in this folder. The parts that matter: + +```yaml +# Per-model pricing is REQUIRED for the routing budget — the switch cost math needs +# each model's input and cached-input rates. +model_metrics_sources: + - type: cost + provider: models.dev # publishes real cache_read rates + refresh_interval: 86400 + +prompt_caching: + enabled: true # automatic caching + session affinity (separate concern) + +routing: + routing_budget: # no default — presence turns it on + seed_usd: 0.50 # cumulative budget for quality-driven switches + # replenish_on_rebind: true # re-seed when a cold session re-binds + # credit_negative: true # cheaper switches credit the budget back + # cache_read_discount: 0.1 # fallback when a feed omits cache_read +``` + +The routing budget lives under `routing` and is independent of prompt caching — it +applies whether or not `prompt_caching.enabled` is set. + + + +### DigitalOcean variant + +Address DO GenAI models with the `digitalocean/` prefix and point the cost feed +at the DO catalog (or keep `models.dev`, which publishes cached-read rates the +DO catalog doesn't): + +```yaml +model_providers: + - model: digitalocean/anthropic-claude-4.6-sonnet + access_key: $DIGITALOCEAN_API_KEY + default: true + - model: digitalocean/openai-gpt-4o + access_key: $DIGITALOCEAN_API_KEY + +model_metrics_sources: + - type: cost + provider: digitalocean # DO catalog + refresh_interval: 86400 +``` + +> The DO catalog does not publish a cached-read rate, so for DO-only setups the +> gate falls back to `input_rate × cache_read_discount`. For exact cached rates, +> add a `models.dev` cost source instead. + +--- + + + +## 3. Run it + +```bash +# From this directory. --with-tracing starts a local OTLP collector on :4317. +planoai up config.yaml --with-tracing + +# Tail logs (cache injections, pin events, switch decisions) +planoai logs --follow + +# Stop +planoai down +``` + +The model listener comes up on **:12000** (per `config.yaml`). + +--- + + + +## 4. See caching in action (single model) + +Send the same large system prompt across several turns. With caching enabled, +Plano derives an implicit session from the stable prefix and pins the model, so +turns 2+ read the prefix from the provider cache. + +```bash +curl -s localhost:12000/v1/chat/completions \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "digitalocean/anthropic-claude-4.6-sonnet", + "messages": [ + {"role": "system", "content": ""}, + {"role": "user", "content": "Scaffold the service"} + ] + }' | jq '.usage' +``` + +Watch `usage.prompt_tokens_details.cached_tokens` climb from 0 on turn 1 to +(nearly) the full prefix on later turns, and the billed cost fall accordingly — +this is exactly the ~4× per-turn drop in the caching-ON vs -OFF comparison. + +--- + + + +## 5. See the routing budget in action (model switch) + +The budget is consulted only when the router proposes a model that differs from +the session's warm anchor. Warmth is inferred from how long ago the session was +last used vs. the provider's cache window (no per-call cache-hit signal needed). +To observe it: + +- **Vetoed switch (paid, over budget):** with a warm session on an expensive model +and a large context, a switch to a pricier candidate costs more than the session's +remaining budget → Plano **retains** the anchor. +- **Paid switch (within budget):** the same switch while budget remains → Plano +**switches** and debits `switch_cost` from the session budget. +- **Free switch (cheaper candidate):** a candidate whose *uncached* input rate +undercuts the anchor's *cached* rate → switch cost ≤ 0 → Plano **switches** for free +(and, with `credit_negative`, credits the budget back). +- **Cold session:** the session went idle past the provider cache window → treated +as cold → the router's pick is dispatched with no budget penalty (and the budget +re-seeds on `replenish_on_rebind`). + +Each decision is emitted to metrics and traces (below) with a `reason` label +(`same_anchor | free | within_budget | over_budget | no_pricing`). + +--- + + + +## 6. Observability (for evals & benchmarks) + +**Prometheus metrics** — brightstaff exposes `/metrics` on **:9092** +(Envoy admin/stats on **:9901/stats**): + + +| Metric | What it tells you | +| ---------------------------------------------------------------------------------- | -------------------------------------------------------------------- | +| `brightstaff_session_switch_decisions_total{decision="allowed"|"retained",reason}` | How often the budget let a switch through vs. vetoed it, and why | +| `brightstaff_prompt_cache_requests_total{provider,model,outcome="hit"|"miss"}` | Real provider cache hit rate | +| `brightstaff_session_cache_events_total{outcome}` | Session binding lookups/stores | + + +```bash +curl -s localhost:9092/metrics | grep -E 'session_switch_decisions|prompt_cache_requests' +``` + +**Traces** — run with `--with-tracing` and inspect the routing span per request: + +- `plano.cache.warm` — whether the session's cache was considered warm this turn +- `plano.cache.idle_ms` — how long since the session was last used +- `plano.switch.cost_in_usd` — actual input-token cost of the proposed switch (output excluded) +- `plano.switch.threshold_in_usd` — budget remaining when the switch was evaluated +- `plano.switch.decision` — `allowed` or `retained` +- `plano.session.budget_remaining_in_usd` — switch budget left after this turn +- `plano.session.switches` — switches taken so far this session +- `plano.switch.counterfactual_route` — on a `retained` decision, the route the gate + *would* have taken had the switch been allowed (only when `record_counterfactual: true`) +- `plano.session_id`, `plano.route.name` + +**Grafana** — a ready dashboard + compose live in `config/grafana/` +(`docker compose up` there, using `prometheus_scrape.yaml`). + +--- + + + +## 7. A/B methodology (baseline vs treatment) + +The cleanest benchmark is same-workload, caching off vs on — the exact shape of +the caching-ON/OFF comparison: + +- **Baseline (no caching):** send requests with header `X-Plano-Cache: off` +(disables implicit pinning + marker injection per request), or run with +`prompt_caching.enabled: false`. +- **Treatment (caching on):** default config in this folder. + +Compare, over an identical multi-turn eval set: + +- total `prompt_tokens` billed at the uncached vs cached rate, +- `cached_tokens` ratio (cache hit rate), +- total USD cost, +- and — for routing-heavy workloads — `session_switch_decisions_total` and the +per-request `plano.switch.*` attributes to confirm switches happen only when +affordable. + +```bash +# Baseline call (caching bypassed) +curl -s localhost:12000/v1/chat/completions \ + -H 'Content-Type: application/json' \ + -H 'X-Plano-Cache: off' \ + -d '{ ... }' | jq '.usage' +``` + +--- + + + +## 8. Knobs to sweep + + +| Setting | Effect | +| ------------------------------------------------ | ------------------------------------------------------------------------------- | +| `routing.routing_budget.seed_usd` | Cumulative budget per session (higher = quality-first, more switching) | +| `routing.routing_budget.replenish_on_rebind` | Re-seed the budget when a cold session re-binds | +| `routing.routing_budget.credit_negative` | Credit the budget back on outright-cheaper switches | +| `routing.routing_budget.cache_read_discount` | Assumed cached rate when a feed omits `cache_read` (DO fallback) | +| `routing.routing_budget.record_counterfactual` | Emit `plano.switch.counterfactual_route` on vetoed switches (the road not taken)| +| `prompt_caching.session_ttl_seconds` | Session binding GC lifetime | +| `prompt_caching.min_prefix_tokens` | Minimum stable-prefix size before markers are injected | +| Header `X-Model-Affinity: ` | Explicit session key (overrides the implicit prefix hash) | +| Header `X-Plano-Cache: off` | Per-request bypass for baseline runs | + + +--- + + + +## Notes + +- Caching **never** changes which model routing selects — the router still makes +the quality call; the budget only vetoes a switch that the session can't afford. +- The routing budget is independent of prompt caching (it lives under `routing`) and +is fully opt-in with **no baked-in budget**: configuring it without a `seed_usd` +(or without a cost source) fails startup with a clear message. diff --git a/demos/llm_routing/routing_budget/README.md b/demos/llm_routing/routing_budget/README.md new file mode 100644 index 000000000..2b24dac9b --- /dev/null +++ b/demos/llm_routing/routing_budget/README.md @@ -0,0 +1,76 @@ +# Routing Budget + +Preference-based routing with a cumulative per-session **routing budget** that +protects warm provider caches, plus automatic prompt caching. The budget is a +routing concern configured under `routing` — independent of prompt caching. + +## The problem + +Provider prompt caches are per-model. When intelligent routing moves a +conversation to a different model, the new model re-ingests the full context at +its **uncached** input rate. In input-heavy, append-only workloads (coding +agents especially), a nominally cheaper model can end up more expensive than +the cached rate you abandoned. + +## What the routing budget does + +The router runs every turn (routing stays cache-blind). When it proposes a model +that differs from the session's warm anchor, Plano computes the **actual +input-token cost** of abandoning the anchor's cache: + +``` +switch_cost_in_usd = context_tokens x (candidate_uncached_input - anchor_cached_input) / 1M +``` + +- **Switch cost <= 0** — the candidate's uncached rate undercuts the anchor's + cached rate. Losing the cache costs nothing; switch freely (and, with + `credit_negative`, credit the budget back). +- **Switch cost > 0** — drawn from the session's remaining budget. If the budget + covers it, the switch proceeds and the budget is debited; otherwise Plano + retains the anchor and its warm cache. + +Warmth is inferred from how long ago the session was last used vs. the +provider's cache window — no per-call cache-hit signal is required, so the same +decision works on both the full-proxy and `/routing` decision paths. + +The math is **input-only by design**: output-token cost is deliberately excluded, +because output length is unknowable before generation. Quality and cost stay +separate — the router still picks the best model; the budget only vetoes switches +the session can't afford. + +## Configuration + +See [config.yaml](config.yaml). Requirements: + +- a cost source in `model_metrics_sources` (per-model rates feed the switch cost math) +- a `routing.routing_budget` block — there is no default; presence turns it on and + startup fails without a `seed_usd` (or without a cost source) + +`routing.routing_budget` fields: + +| Field | Meaning | +|---|---| +| `seed_usd` | Cumulative budget (USD) per session. `0` = never pay to switch | +| `replenish_on_rebind` | Re-seed the budget when a cold session re-binds (default true) | +| `credit_negative` | Credit the budget back on outright-cheaper switches (default true) | +| `cache_read_discount` | Assumed cached rate when a feed omits `cache_read` (default 0.1) | +| `record_counterfactual` | Record the switch that was vetoed, as a trace attribute (default false) | + +Prompt caching (`prompt_caching.enabled`) is a separate, optional concern that keeps +the upstream cache warm and injects provider cache-control markers. + +## Observability + +Every decision is visible: + +- Metric: `brightstaff_session_switch_decisions_total{decision="allowed"|"retained",reason}` + (`reason` ∈ `same_anchor | free | within_budget | over_budget | no_pricing`) +- Span attributes: `plano.cache.warm`, `plano.cache.idle_ms`, + `plano.switch.cost_in_usd`, `plano.switch.threshold_in_usd`, `plano.switch.decision`, + `plano.session.budget_remaining_in_usd`, `plano.session.switches` + +## Run + +```bash +planoai up config.yaml +``` diff --git a/demos/llm_routing/routing_budget/config.yaml b/demos/llm_routing/routing_budget/config.yaml new file mode 100644 index 000000000..eb139e7ca --- /dev/null +++ b/demos/llm_routing/routing_budget/config.yaml @@ -0,0 +1,76 @@ +version: v0.3.0 + +listeners: + - type: model + name: model_listener + port: 12000 + +model_providers: + + - model: openai/gpt-4o-mini + access_key: $OPENAI_API_KEY + default: true + + - model: openai/gpt-4o + access_key: $OPENAI_API_KEY + routing_preferences: + - name: code understanding + description: understand and explain existing code snippets, functions, or libraries + + - model: anthropic/claude-sonnet-4-6 + access_key: $ANTHROPIC_API_KEY + routing_preferences: + - name: code generation + description: generating new code snippets, functions, or boilerplate based on user prompts or requirements + +# Per-model pricing is required for the routing budget: the switch-cost +# calculation needs each model's input and cached-input rates. +model_metrics_sources: + - type: cost + provider: models.dev + refresh_interval: 86400 + +# Automatic prompt caching (opt-in). Keeps a conversation pinned to the same +# model so the upstream provider's prompt cache stays warm across turns, and +# auto-injects cache-control markers where the provider needs them. This is a +# separate concern from the routing budget below. +prompt_caching: + enabled: true + +routing: + # Per-session cost gate on model switching. Independent of prompt caching: it + # applies whenever configured, and presence of this block turns it on. The + # default posture is to stick to the model a session is warm on. When routing + # proposes a *different* model while that session's provider cache is plausibly + # still warm (inferred from how long ago the session was last used vs. the + # provider's cache window), the actual input-token cost of abandoning the cache: + # + # switch_cost = context_tokens x (candidate_uncached_input - anchor_cached_input) + # + # (output cost deliberately excluded) is drawn from the session's budget. A paid + # switch is allowed only while budget remains; an outright-cheaper switch is free + # and credits the budget back. The budget is yours to define -- Plano never + # invents one, and startup fails without it (or without a cost source). + routing_budget: + # Cumulative budget (USD) for quality-driven switches over a session. 0 means + # "never pay to switch"; larger values buy more switches before sticking. + seed_usd: 0.50 + + # Re-seed the budget when a cold session re-binds (new warm episode). Default true. + # replenish_on_rebind: true + + # Credit the budget back on outright-cheaper (negative-cost) switches. Default true. + # credit_negative: true + + # Fallback used only when the pricing feed doesn't publish a cached-read rate for + # a model: cached_rate = input_rate x cache_read_discount. Default 0.1. + # cache_read_discount: 0.1 + + # Record the route the gate WOULD have taken when it vetoes a switch, as the + # `plano.switch.counterfactual_route` trace attribute. Telemetry only -- the + # counterfactual model is never dispatched. Handy for evals that want to + # measure the road not taken. Default false. + record_counterfactual: true + +tracing: + random_sampling: 100