From 424a0366a9ecf6d03b6e70828f68cf5a1905012d Mon Sep 17 00:00:00 2001 From: Illia Polosukhin Date: Sat, 7 Mar 2026 09:10:05 +0000 Subject: [PATCH] feat: enable Anthropic prompt caching via automatic cache_control injection (#660) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(llm): add Anthropic prompt caching and cache token tracking - Inject cache_control via additional_params for Claude models in rig_adapter - Add cache_read_input_tokens and cache_creation_input_tokens to CompletionResponse and ToolCompletionResponse - Extract cached_input_tokens from rig-core unified Usage - Add is_anthropic_model() detection helper with provider prefix support - Log prompt cache hits at debug level (consistent with response_cache) - Add 7 unit tests for cache injection and model detection - Update all mock providers and test fixtures with new fields * feat(cost): apply 90% cache discount to prompt-cached tokens in CostGuard - Add cache_read_input_tokens to TokenUsage so cache counts flow from CompletionResponse through the reasoning layer to the dispatcher - Update CostGuard::record_llm_call() to accept cache_read_input_tokens: cached tokens are billed at 10% of the normal input rate - Thread cache_read_input_tokens from dispatcher into CostGuard - Add test_cache_discount_reduces_cost verifying exact savings match 90% of input cost for fully-cached requests - Update all existing test callers with zero-cache parameter * refactor(cache): scope cache_control to Anthropic backend and validate model support - Replace model-name-based is_anthropic_model() with explicit enable_prompt_cache flag on RigAdapter, set only for the direct Anthropic backend via with_prompt_cache(true) - Add supports_prompt_cache() to validate model names per Anthropic docs: only Claude 3+ models support caching; claude-2 and claude-instant are excluded to prevent 400 errors - Warn when caching is enabled but model does not support it - Replace is_anthropic_model tests with flag-based and model validation tests * fix(cache): validate model at construction and propagate cache metrics through proxy - Move supports_prompt_cache() check into with_prompt_cache() so unsupported models are detected once at construction, not per request - Add cache_read_input_tokens and cache_creation_input_tokens to ProxyCompletionResponse and ProxyToolCompletionResponse with serde(default) for backward compatibility - Pass cache metrics through orchestrator proxy instead of zeroing - Use claude-opus-4-6 in cache discount test to match Anthropic semantics * feat(llm): add configurable cache retention with write surcharge - Add CacheRetention enum (none/short/long) to AnthropicDirectConfig - Parse ANTHROPIC_CACHE_RETENTION env var (default: short) - Inject TTL-aware cache_control (short=5m ephemeral, long=1h) - Extract cache_creation_input_tokens from raw Anthropic response - Add cache_write_multiplier() to LlmProvider trait (1.25x short, 2.0x long) - Pipe dynamic write multiplier through dispatcher to CostGuard - Add TokenUsage.cache_creation_input_tokens field - Add tests for Long TTL injection, 5m and 1h write surcharges - Document ANTHROPIC_CACHE_RETENTION in .env.example * docs: fix stale cache_retention field comment * fix: resolve CI failures after upstream merge - Add missing cost_per_token arg to cache test callsites - Apply cargo fmt to long lines in tests and tracing macros * fix: address Copilot review feedback - Use saturating_add for cache token sum to prevent u32 overflow - Tighten supports_prompt_cache to explicitly match claude-3+/claude-4+ and named families (claude-sonnet/claude-opus/claude-haiku) * fix: adapt prompt caching to registry architecture and add missing cache fields - Resolve merge conflicts: adapt CacheRetention and cache injection to the declarative provider registry (RegistryProviderConfig replaces AnthropicDirectConfig) - Parse ANTHROPIC_CACHE_RETENTION env var in create_anthropic_from_registry() - Use Anthropic automatic caching via top-level cache_control in additional_params (rig-core #[serde(flatten)] places it at request root) - Add cache_read/creation_input_tokens fields to all mock LlmProviders added on main after PR #291 branched (response_cache, dispatcher, provider_chaos, trace_llm) - Suppress clippy::too_many_arguments on record_llm_call and build_rig_request - Add regression tests for cache injection (short/long/none) and cache_write_multiplier values Co-Authored-By: Canvinus <44225021+Canvinus@users.noreply.github.com> * fix: delegate cache_write_multiplier through provider wrappers and make cache_read_discount configurable The 6 decorator providers (Retry, CircuitBreaker, Failover, SmartRouting, CachedProvider, RecordingLlm) did not delegate cache_write_multiplier() to their inner provider, causing it to always return 1.0 instead of the actual 1.25x/2.0x from RigAdapter. This fix adds delegation for both cache_write_multiplier() and the new cache_read_discount() method. Also makes the cache read discount per-provider instead of hardcoding Anthropic's 90% discount (÷10). OpenAI uses 50% (÷2), so the discount is now returned by each provider via the LlmProvider trait. Addresses review feedback on PR #660. Co-Authored-By: Claude Opus 4.6 * style: cargo fmt Co-Authored-By: Claude Opus 4.6 * test: add CacheRetention FromStr/Display unit tests Tests cover primary values, aliases (off/disabled/5m/ephemeral/1h), case-insensitivity, invalid input error, and Display round-trip. Addresses Copilot review feedback on PR #660. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Andrey Co-authored-by: Andrey Gruzdev <44225021+Canvinus@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 --- .env.example | 11 ++ src/agent/cost_guard.rs | 241 ++++++++++++++++++++++++-- src/agent/dispatcher.rs | 22 +++ src/config/llm.rs | 126 ++++++++++++++ src/config/mod.rs | 2 +- src/llm/circuit_breaker.rs | 8 + src/llm/failover.rs | 16 ++ src/llm/mod.rs | 30 +++- src/llm/nearai_chat.rs | 4 + src/llm/provider.rs | 27 +++ src/llm/reasoning.rs | 10 ++ src/llm/recording.rs | 8 + src/llm/response_cache.rs | 12 ++ src/llm/retry.rs | 8 + src/llm/rig_adapter.rs | 260 ++++++++++++++++++++++++++++- src/llm/smart_routing.rs | 16 ++ src/orchestrator/api.rs | 4 + src/testing.rs | 4 + src/worker/api.rs | 12 ++ tests/openai_compat_integration.rs | 10 ++ tests/provider_chaos.rs | 12 ++ tests/support/trace_llm.rs | 6 + 22 files changed, 831 insertions(+), 18 deletions(-) diff --git a/.env.example b/.env.example index 2022c8a7..9c41a62d 100644 --- a/.env.example +++ b/.env.example @@ -57,6 +57,17 @@ NEARAI_AUTH_URL=https://private.near.ai # LLM_BASE_URL=https://api.fireworks.ai/inference/v1 # LLM_API_KEY=fw_... +# === Anthropic Direct === +# LLM_BACKEND=anthropic +# ANTHROPIC_MODEL=claude-sonnet-4-6 +# ANTHROPIC_API_KEY=sk-ant-... +# ANTHROPIC_BASE_URL=https://api.anthropic.com # default +# Prompt cache retention — controls Anthropic server-side prompt caching: +# none = disabled (no cache_control injected) +# short = 5-minute TTL, 1.25× (125%) write surcharge (default) +# long = 1-hour TTL, 2.0× (200%) write surcharge +# ANTHROPIC_CACHE_RETENTION=short + # For full provider setup guide see docs/LLM_PROVIDERS.md # Channel Configuration diff --git a/src/agent/cost_guard.rs b/src/agent/cost_guard.rs index 59d676ca..47362fc0 100644 --- a/src/agent/cost_guard.rs +++ b/src/agent/cost_guard.rs @@ -151,21 +151,46 @@ impl CostGuard { /// Record a completed LLM action: its token costs and the action timestamp. /// /// Call this AFTER an LLM call completes so that costs are tracked. + /// - `cache_read_input_tokens`: tokens served from cache. + /// - `cache_creation_input_tokens`: tokens written to cache. + /// - `cache_read_discount`: divisor for cache-read cost (e.g. 10 for Anthropic 90% off, 2 for OpenAI 50% off). + /// - `cache_write_multiplier`: cost multiplier for cache writes (1.25 for 5m, 2.0 for 1h). /// /// When `cost_per_token` is `Some`, those rates are used directly (provider- /// sourced pricing). When `None`, falls back to the static `costs::model_cost` /// lookup table, then `costs::default_cost`. + #[allow(clippy::too_many_arguments)] pub async fn record_llm_call( &self, model: &str, input_tokens: u32, output_tokens: u32, + cache_read_input_tokens: u32, + cache_creation_input_tokens: u32, + cache_read_discount: Decimal, + cache_write_multiplier: Decimal, cost_per_token: Option<(Decimal, Decimal)>, ) -> Decimal { let (input_rate, output_rate) = cost_per_token .unwrap_or_else(|| costs::model_cost(model).unwrap_or_else(costs::default_cost)); - let cost = - input_rate * Decimal::from(input_tokens) + output_rate * Decimal::from(output_tokens); + // Cached read tokens cost input_rate / cache_read_discount (provider-specific). + // Cached write tokens cost write_multiplier × input_rate (e.g. 1.25× for 5m, 2× for 1h). + // Uncached tokens = total input - cache reads - cache writes. + let cached_total = cache_read_input_tokens.saturating_add(cache_creation_input_tokens); + let uncached_input = input_tokens.saturating_sub(cached_total); + let effective_discount = if cache_read_discount.is_zero() { + Decimal::ONE + } else { + cache_read_discount + }; + let cache_read_cost = + input_rate * Decimal::from(cache_read_input_tokens) / effective_discount; + let cache_write_cost = + input_rate * Decimal::from(cache_creation_input_tokens) * cache_write_multiplier; + let cost = input_rate * Decimal::from(uncached_input) + + cache_read_cost + + cache_write_cost + + output_rate * Decimal::from(output_tokens); // Update daily cost (reset if new day) { @@ -267,7 +292,16 @@ mod tests { // Record a big call, still allowed guard - .record_llm_call("gpt-4o", 100_000, 100_000, None) + .record_llm_call( + "gpt-4o", + 100_000, + 100_000, + 0, + 0, + Decimal::ONE, + Decimal::ONE, + None, + ) .await; assert!(guard.check_allowed().await.is_ok()); } @@ -285,7 +319,18 @@ mod tests { // Record a call that costs more than $0.01 // gpt-4o: input=$0.0000025/tok, output=$0.00001/tok // 10000 input + 10000 output = $0.025 + $0.10 = $0.125 - guard.record_llm_call("gpt-4o", 10_000, 10_000, None).await; + guard + .record_llm_call( + "gpt-4o", + 10_000, + 10_000, + 0, + 0, + Decimal::ONE, + Decimal::ONE, + None, + ) + .await; // Now should be blocked let result = guard.check_allowed().await; @@ -308,7 +353,9 @@ mod tests { // First 3 actions allowed for _ in 0..3 { assert!(guard.check_allowed().await.is_ok()); - guard.record_llm_call("gpt-4o", 10, 10, None).await; + guard + .record_llm_call("gpt-4o", 10, 10, 0, 0, Decimal::ONE, Decimal::ONE, None) + .await; } // 4th should be blocked @@ -329,7 +376,9 @@ mod tests { assert_eq!(guard.daily_spend().await, Decimal::ZERO); - let cost = guard.record_llm_call("gpt-4o", 1000, 500, None).await; + let cost = guard + .record_llm_call("gpt-4o", 1000, 500, 0, 0, Decimal::ONE, Decimal::ONE, None) + .await; assert!(cost > Decimal::ZERO); assert_eq!(guard.daily_spend().await, cost); } @@ -340,8 +389,12 @@ mod tests { assert_eq!(guard.actions_this_hour().await, 0); - guard.record_llm_call("gpt-4o", 10, 10, None).await; - guard.record_llm_call("gpt-4o", 10, 10, None).await; + guard + .record_llm_call("gpt-4o", 10, 10, 0, 0, Decimal::ONE, Decimal::ONE, None) + .await; + guard + .record_llm_call("gpt-4o", 10, 10, 0, 0, Decimal::ONE, Decimal::ONE, None) + .await; assert_eq!(guard.actions_this_hour().await, 2); } @@ -378,10 +431,23 @@ mod tests { assert!(guard.model_usage().await.is_empty()); // Record calls for two different models - guard.record_llm_call("gpt-4o", 1000, 500, None).await; - guard.record_llm_call("gpt-4o", 2000, 1000, None).await; guard - .record_llm_call("claude-3-5-sonnet-20241022", 500, 200, None) + .record_llm_call("gpt-4o", 1000, 500, 0, 0, Decimal::ONE, Decimal::ONE, None) + .await; + guard + .record_llm_call("gpt-4o", 2000, 1000, 0, 0, Decimal::ONE, Decimal::ONE, None) + .await; + guard + .record_llm_call( + "claude-3-5-sonnet-20241022", + 500, + 200, + 0, + 0, + Decimal::ONE, + Decimal::ONE, + None, + ) .await; let usage = guard.model_usage().await; @@ -402,4 +468,157 @@ mod tests { // Costs should differ since models have different pricing assert_ne!(gpt.cost, claude.cost); } + + #[tokio::test] + async fn test_cache_discount_reduces_cost() { + let guard = CostGuard::new(CostGuardConfig::default()); + + // Full price: 1000 input + 500 output, no cache + let full_cost = guard + .record_llm_call( + "claude-opus-4-6", + 1000, + 500, + 0, + 0, + Decimal::ONE, + Decimal::ONE, + None, + ) + .await; + + let guard2 = CostGuard::new(CostGuardConfig::default()); + + // Same tokens but all input cached (90% discount on input) + let cached_cost = guard2 + .record_llm_call( + "claude-opus-4-6", + 1000, + 500, + 1000, + 0, + dec!(10), + Decimal::ONE, + None, + ) + .await; + + // Cached cost must be strictly less than full cost + assert!( + cached_cost < full_cost, + "cached_cost ({}) should be less than full_cost ({})", + cached_cost, + full_cost + ); + + // The difference should be exactly 90% of the input cost + let (input_rate, _) = costs::model_cost("claude-opus-4-6").unwrap(); + let expected_savings = input_rate * Decimal::from(1000u32) * dec!(9) / dec!(10); + let actual_savings = full_cost - cached_cost; + assert_eq!( + actual_savings, expected_savings, + "savings should be 90% of input cost for fully-cached request" + ); + } + + #[tokio::test] + async fn test_cache_write_surcharge_increases_cost() { + let guard = CostGuard::new(CostGuardConfig::default()); + + // Full price: 1000 input + 500 output, no cache activity + let full_cost = guard + .record_llm_call( + "claude-opus-4-6", + 1000, + 500, + 0, + 0, + Decimal::ONE, + Decimal::ONE, + None, + ) + .await; + + let guard2 = CostGuard::new(CostGuardConfig::default()); + + // Same tokens, but all input tokens are cache writes (1.25x surcharge for 5m TTL) + let short_multiplier = Decimal::new(125, 2); // 1.25 + let write_cost = guard2 + .record_llm_call( + "claude-opus-4-6", + 1000, + 500, + 0, + 1000, + Decimal::ONE, + short_multiplier, + None, + ) + .await; + + // Write cost must be strictly greater than full cost + assert!( + write_cost > full_cost, + "write_cost ({}) should be greater than full_cost ({})", + write_cost, + full_cost + ); + + // The difference should be exactly 25% of the input cost + let (input_rate, _) = costs::model_cost("claude-opus-4-6").unwrap(); + let expected_surcharge = input_rate * Decimal::from(1000u32) * dec!(0.25); + let actual_surcharge = write_cost - full_cost; + assert_eq!( + actual_surcharge, expected_surcharge, + "surcharge should be 25% of input cost for 5m cache writes" + ); + } + + #[tokio::test] + async fn test_cache_write_surcharge_long_ttl() { + let guard = CostGuard::new(CostGuardConfig::default()); + + // Full price: 1000 input + 500 output + let full_cost = guard + .record_llm_call( + "claude-opus-4-6", + 1000, + 500, + 0, + 0, + Decimal::ONE, + Decimal::ONE, + None, + ) + .await; + + let guard2 = CostGuard::new(CostGuardConfig::default()); + + // All input tokens are cache writes with 2.0x multiplier (1h TTL) + let long_multiplier = Decimal::TWO; + let write_cost = guard2 + .record_llm_call( + "claude-opus-4-6", + 1000, + 500, + 0, + 1000, + Decimal::ONE, + long_multiplier, + None, + ) + .await; + + // Write cost > full cost + assert!(write_cost > full_cost); + + // Surcharge should be 100% of input cost (2.0x - 1.0x = 1.0x) + let (input_rate, _) = costs::model_cost("claude-opus-4-6").unwrap(); + let expected_surcharge = input_rate * Decimal::from(1000u32); + let actual_surcharge = write_cost - full_cost; + assert_eq!( + actual_surcharge, expected_surcharge, + "surcharge should be 100% of input cost for 1h cache writes" + ); + } } diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index e957d1c8..f960e3c5 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -278,12 +278,18 @@ impl Agent { // Record cost and track token usage let model_name = self.llm().active_model_name(); + let read_discount = self.llm().cache_read_discount(); + let write_multiplier = self.llm().cache_write_multiplier(); let call_cost = self .cost_guard() .record_llm_call( &model_name, output.usage.input_tokens, output.usage.output_tokens, + output.usage.cache_read_input_tokens, + output.usage.cache_creation_input_tokens, + read_discount, + write_multiplier, Some(self.llm().cost_per_token()), ) .await; @@ -1062,6 +1068,8 @@ mod tests { input_tokens: 0, output_tokens: 0, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } @@ -1075,6 +1083,8 @@ mod tests { input_tokens: 0, output_tokens: 0, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } } @@ -1635,6 +1645,8 @@ mod tests { input_tokens: 0, output_tokens: 5, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } @@ -1650,6 +1662,8 @@ mod tests { input_tokens: 0, output_tokens: 5, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }); } // Tools available: always call one. @@ -1663,6 +1677,8 @@ mod tests { input_tokens: 0, output_tokens: 5, finish_reason: FinishReason::ToolUse, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } } @@ -1787,6 +1803,8 @@ mod tests { input_tokens: 0, output_tokens: 2, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } @@ -1801,6 +1819,8 @@ mod tests { input_tokens: 0, output_tokens: 2, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }); } // Always call a tool that does not exist in the registry. @@ -1814,6 +1834,8 @@ mod tests { input_tokens: 0, output_tokens: 5, finish_reason: FinishReason::ToolUse, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } } diff --git a/src/config/llm.rs b/src/config/llm.rs index 03275a08..a06c0f1d 100644 --- a/src/config/llm.rs +++ b/src/config/llm.rs @@ -9,6 +9,50 @@ use crate::llm::registry::{ProviderProtocol, ProviderRegistry}; use crate::llm::session::SessionConfig; use crate::settings::Settings; +/// Prompt cache retention policy for Anthropic. +/// +/// Controls Anthropic's automatic prompt caching via a top-level +/// `cache_control` field injected through rig-core's `additional_params`. +/// - `None` — caching disabled, no `cache_control` injected. +/// - `Short` — 5-minute TTL (default), `{"type": "ephemeral"}`, 1.25× write surcharge. +/// - `Long` — 1-hour TTL, `{"type": "ephemeral", "ttl": "1h"}`, 2× write surcharge. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum CacheRetention { + /// No prompt caching. + None, + /// 5-minute TTL (default). Write cost: 1.25× base input. + #[default] + Short, + /// 1-hour TTL. Write cost: 2× base input. + Long, +} + +impl std::str::FromStr for CacheRetention { + type Err = String; + + fn from_str(s: &str) -> Result { + match s.to_lowercase().as_str() { + "none" | "off" | "disabled" => Ok(Self::None), + "short" | "5m" | "ephemeral" => Ok(Self::Short), + "long" | "1h" => Ok(Self::Long), + _ => Err(format!( + "invalid cache retention '{}', expected one of: none, short, long", + s + )), + } + } +} + +impl std::fmt::Display for CacheRetention { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::None => write!(f, "none"), + Self::Short => write!(f, "short"), + Self::Long => write!(f, "long"), + } + } +} + /// Resolved configuration for a registry-based provider. /// /// This single struct replaces what used to be five separate config types @@ -755,4 +799,86 @@ mod tests { std::env::remove_var("LLM_BACKEND"); } } + + #[test] + fn cache_retention_from_str_primary_values() { + assert_eq!( + "none".parse::().unwrap(), + CacheRetention::None + ); + assert_eq!( + "short".parse::().unwrap(), + CacheRetention::Short + ); + assert_eq!( + "long".parse::().unwrap(), + CacheRetention::Long + ); + } + + #[test] + fn cache_retention_from_str_aliases() { + assert_eq!( + "off".parse::().unwrap(), + CacheRetention::None + ); + assert_eq!( + "disabled".parse::().unwrap(), + CacheRetention::None + ); + assert_eq!( + "5m".parse::().unwrap(), + CacheRetention::Short + ); + assert_eq!( + "ephemeral".parse::().unwrap(), + CacheRetention::Short + ); + assert_eq!( + "1h".parse::().unwrap(), + CacheRetention::Long + ); + } + + #[test] + fn cache_retention_from_str_case_insensitive() { + assert_eq!( + "NONE".parse::().unwrap(), + CacheRetention::None + ); + assert_eq!( + "Short".parse::().unwrap(), + CacheRetention::Short + ); + assert_eq!( + "LONG".parse::().unwrap(), + CacheRetention::Long + ); + assert_eq!( + "Ephemeral".parse::().unwrap(), + CacheRetention::Short + ); + } + + #[test] + fn cache_retention_from_str_invalid() { + let err = "bogus".parse::().unwrap_err(); + assert!( + err.contains("bogus"), + "error should mention the invalid value" + ); + } + + #[test] + fn cache_retention_display_round_trip() { + for variant in [ + CacheRetention::None, + CacheRetention::Short, + CacheRetention::Long, + ] { + let s = variant.to_string(); + let parsed: CacheRetention = s.parse().unwrap(); + assert_eq!(parsed, variant, "round-trip failed for {s}"); + } + } } diff --git a/src/config/mod.rs b/src/config/mod.rs index 25f426e1..fab50b3e 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -36,7 +36,7 @@ pub use self::database::{DatabaseBackend, DatabaseConfig, SslMode, default_libsq pub use self::embeddings::EmbeddingsConfig; pub use self::heartbeat::HeartbeatConfig; pub use self::hygiene::HygieneConfig; -pub use self::llm::{LlmConfig, NearAiConfig, RegistryProviderConfig}; +pub use self::llm::{CacheRetention, LlmConfig, NearAiConfig, RegistryProviderConfig}; pub use self::routines::RoutineConfig; pub use self::safety::SafetyConfig; pub use self::sandbox::{ClaudeCodeConfig, SandboxModeConfig}; diff --git a/src/llm/circuit_breaker.rs b/src/llm/circuit_breaker.rs index 6c9a0a78..6b04fac7 100644 --- a/src/llm/circuit_breaker.rs +++ b/src/llm/circuit_breaker.rs @@ -245,6 +245,14 @@ impl LlmProvider for CircuitBreakerProvider { self.inner.cost_per_token() } + fn cache_write_multiplier(&self) -> Decimal { + self.inner.cache_write_multiplier() + } + + fn cache_read_discount(&self) -> Decimal { + self.inner.cache_read_discount() + } + async fn complete(&self, request: CompletionRequest) -> Result { self.check_allowed().await?; match self.inner.complete(request).await { diff --git a/src/llm/failover.rs b/src/llm/failover.rs index 8af7845f..cbc6634e 100644 --- a/src/llm/failover.rs +++ b/src/llm/failover.rs @@ -296,6 +296,14 @@ impl LlmProvider for FailoverProvider { self.providers[self.last_used.load(Ordering::Relaxed)].cost_per_token() } + fn cache_write_multiplier(&self) -> Decimal { + self.providers[self.last_used.load(Ordering::Relaxed)].cache_write_multiplier() + } + + fn cache_read_discount(&self) -> Decimal { + self.providers[self.last_used.load(Ordering::Relaxed)].cache_read_discount() + } + async fn complete(&self, request: CompletionRequest) -> Result { let (provider_idx, response) = self .try_providers(|provider| { @@ -404,6 +412,8 @@ mod tests { input_tokens: 10, output_tokens: 5, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }))), tool_complete_result: Mutex::new(Some(Ok(ToolCompletionResponse { content: Some(content.to_string()), @@ -411,6 +421,8 @@ mod tests { input_tokens: 10, output_tokens: 5, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }))), } } @@ -792,6 +804,8 @@ mod tests { input_tokens: 10, output_tokens: 5, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } @@ -817,6 +831,8 @@ mod tests { input_tokens: 10, output_tokens: 5, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } diff --git a/src/llm/mod.rs b/src/llm/mod.rs index 579c5e9f..54b77096 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -177,6 +177,8 @@ fn create_openai_compat_from_registry( fn create_anthropic_from_registry( config: &RegistryProviderConfig, ) -> Result, LlmError> { + use crate::config::CacheRetention; + use crate::config::helpers::optional_env; use rig::providers::anthropic; let api_key = config @@ -200,8 +202,32 @@ fn create_anthropic_from_registry( reason: format!("Failed to create Anthropic client: {e}"), })?; + // Resolve prompt cache retention from env (default: Short). + // Injects top-level cache_control via additional_params for Anthropic + // automatic caching (the API auto-places the breakpoint at the last + // cacheable block). + let cache_retention: CacheRetention = optional_env("ANTHROPIC_CACHE_RETENTION") + .ok() + .flatten() + .and_then(|val| match val.parse::() { + Ok(r) => Some(r), + Err(e) => { + tracing::warn!("Invalid ANTHROPIC_CACHE_RETENTION: {e}; defaulting to short"); + None + } + }) + .unwrap_or_default(); + let model = client.completion_model(&config.model); + if cache_retention != CacheRetention::None { + tracing::info!( + model = %config.model, + retention = %cache_retention, + "Anthropic automatic prompt caching enabled" + ); + } + tracing::info!( provider = %config.provider_id, model = %config.model, @@ -209,7 +235,9 @@ fn create_anthropic_from_registry( "Using Anthropic provider" ); - Ok(Arc::new(RigAdapter::new(model, &config.model))) + Ok(Arc::new( + RigAdapter::new(model, &config.model).with_cache_retention(cache_retention), + )) } fn create_ollama_from_registry( diff --git a/src/llm/nearai_chat.rs b/src/llm/nearai_chat.rs index d1857807..6397d54c 100644 --- a/src/llm/nearai_chat.rs +++ b/src/llm/nearai_chat.rs @@ -499,6 +499,8 @@ impl LlmProvider for NearAiChatProvider { finish_reason, input_tokens, output_tokens, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } @@ -604,6 +606,8 @@ impl LlmProvider for NearAiChatProvider { finish_reason, input_tokens, output_tokens, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } diff --git a/src/llm/provider.rs b/src/llm/provider.rs index 84227df0..0bcdd4ea 100644 --- a/src/llm/provider.rs +++ b/src/llm/provider.rs @@ -153,6 +153,12 @@ pub struct CompletionResponse { pub input_tokens: u32, pub output_tokens: u32, pub finish_reason: FinishReason, + /// Tokens read from the provider's server-side prompt cache (Anthropic). + /// Zero when caching is not supported or on a cache miss. + pub cache_read_input_tokens: u32, + /// Tokens written to the provider's server-side prompt cache (Anthropic). + /// Zero when caching is not supported or no new prefix was cached. + pub cache_creation_input_tokens: u32, } /// Why the completion finished. @@ -254,6 +260,10 @@ pub struct ToolCompletionResponse { pub input_tokens: u32, pub output_tokens: u32, pub finish_reason: FinishReason, + /// Tokens read from the provider's server-side prompt cache (Anthropic). + pub cache_read_input_tokens: u32, + /// Tokens written to the provider's server-side prompt cache (Anthropic). + pub cache_creation_input_tokens: u32, } /// Metadata about a model returned by the provider's API. @@ -328,6 +338,23 @@ pub trait LlmProvider: Send + Sync { let (input_cost, output_cost) = self.cost_per_token(); input_cost * Decimal::from(input_tokens) + output_cost * Decimal::from(output_tokens) } + + /// Cost multiplier for cache-creation tokens (Anthropic prompt caching). + /// + /// Returns `1.0` by default (no surcharge). Anthropic providers return + /// `1.25` for 5-minute TTL or `2.0` for 1-hour TTL. + fn cache_write_multiplier(&self) -> Decimal { + Decimal::ONE + } + + /// Discount divisor for cache-read tokens. + /// + /// Cached-read cost = `input_rate / cache_read_discount()`. + /// Returns `1` by default (no discount). Anthropic returns `10` (90% off), + /// OpenAI would return `2` (50% off). + fn cache_read_discount(&self) -> Decimal { + Decimal::ONE + } } /// Sanitize a message list to ensure tool_use / tool_result integrity. diff --git a/src/llm/reasoning.rs b/src/llm/reasoning.rs index c3081ddb..9b4bcf08 100644 --- a/src/llm/reasoning.rs +++ b/src/llm/reasoning.rs @@ -292,6 +292,10 @@ pub struct ToolSelection { pub struct TokenUsage { pub input_tokens: u32, pub output_tokens: u32, + /// Tokens served from the provider's server-side prompt cache (Anthropic). + pub cache_read_input_tokens: u32, + /// Tokens written to the provider's prompt cache (Anthropic). + pub cache_creation_input_tokens: u32, } impl TokenUsage { @@ -434,6 +438,8 @@ impl Reasoning { let usage = TokenUsage { input_tokens: response.input_tokens, output_tokens: response.output_tokens, + cache_read_input_tokens: response.cache_read_input_tokens, + cache_creation_input_tokens: response.cache_creation_input_tokens, }; Ok((clean_response(&response.content), usage)) } @@ -612,6 +618,8 @@ Respond in JSON format: let usage = TokenUsage { input_tokens: response.input_tokens, output_tokens: response.output_tokens, + cache_read_input_tokens: response.cache_read_input_tokens, + cache_creation_input_tokens: response.cache_creation_input_tokens, }; // If there were tool calls, return them for execution @@ -690,6 +698,8 @@ Respond in JSON format: usage: TokenUsage { input_tokens: response.input_tokens, output_tokens: response.output_tokens, + cache_read_input_tokens: response.cache_read_input_tokens, + cache_creation_input_tokens: response.cache_creation_input_tokens, }, }) } diff --git a/src/llm/recording.rs b/src/llm/recording.rs index 48451714..6f53278b 100644 --- a/src/llm/recording.rs +++ b/src/llm/recording.rs @@ -461,6 +461,14 @@ impl LlmProvider for RecordingLlm { self.inner.cost_per_token() } + fn cache_write_multiplier(&self) -> Decimal { + self.inner.cache_write_multiplier() + } + + fn cache_read_discount(&self) -> Decimal { + self.inner.cache_read_discount() + } + async fn complete(&self, request: CompletionRequest) -> Result { let (hint, tool_results) = self.capture_new_messages(&request.messages).await; let response = self.inner.complete(request).await?; diff --git a/src/llm/response_cache.rs b/src/llm/response_cache.rs index f94ad74f..26caf885 100644 --- a/src/llm/response_cache.rs +++ b/src/llm/response_cache.rs @@ -181,6 +181,14 @@ impl LlmProvider for CachedProvider { self.inner.cost_per_token() } + fn cache_write_multiplier(&self) -> Decimal { + self.inner.cache_write_multiplier() + } + + fn cache_read_discount(&self) -> Decimal { + self.inner.cache_read_discount() + } + async fn complete(&self, request: CompletionRequest) -> Result { let effective_model = self.inner.effective_model_name(request.model.as_deref()); let key = cache_key(&effective_model, &request); @@ -352,6 +360,8 @@ mod tests { input_tokens: 1, output_tokens: 1, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } @@ -365,6 +375,8 @@ mod tests { input_tokens: 1, output_tokens: 1, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } } diff --git a/src/llm/retry.rs b/src/llm/retry.rs index e02237a7..1a68cb8b 100644 --- a/src/llm/retry.rs +++ b/src/llm/retry.rs @@ -109,6 +109,14 @@ impl LlmProvider for RetryProvider { self.inner.cost_per_token() } + fn cache_write_multiplier(&self) -> Decimal { + self.inner.cache_write_multiplier() + } + + fn cache_read_discount(&self) -> Decimal { + self.inner.cache_read_discount() + } + async fn complete(&self, request: CompletionRequest) -> Result { let mut last_error: Option = None; diff --git a/src/llm/rig_adapter.rs b/src/llm/rig_adapter.rs index da01b42c..3253b961 100644 --- a/src/llm/rig_adapter.rs +++ b/src/llm/rig_adapter.rs @@ -3,6 +3,7 @@ //! This lets us use any rig-core provider (OpenAI, Anthropic, Ollama, etc.) as an //! `Arc` without changing any of the agent, reasoning, or tool code. +use crate::config::CacheRetention; use async_trait::async_trait; use rig::OneOrMany; use rig::completion::{ @@ -14,6 +15,7 @@ use rig::message::{ ToolResultContent, UserContent, }; use rust_decimal::Decimal; +use rust_decimal_macros::dec; use serde::Serialize; use serde::de::DeserializeOwned; use serde_json::Value as JsonValue; @@ -34,6 +36,11 @@ pub struct RigAdapter { model_name: String, input_cost: Decimal, output_cost: Decimal, + /// Prompt cache retention policy (Anthropic only). + /// When not `CacheRetention::None`, injects top-level `cache_control` + /// via `additional_params` for Anthropic automatic caching. Also controls + /// the cost multiplier for cache-creation tokens. + cache_retention: CacheRetention, } impl RigAdapter { @@ -47,8 +54,35 @@ impl RigAdapter { model_name: name, input_cost, output_cost, + cache_retention: CacheRetention::None, } } + + /// Set Anthropic prompt cache retention policy. + /// + /// Controls both cache injection and cost tracking: + /// - `None` — no caching, no surcharge (1.0×). + /// - `Short` — 5-minute TTL via `{"type": "ephemeral"}`, 1.25× write surcharge. + /// - `Long` — 1-hour TTL via `{"type": "ephemeral", "ttl": "1h"}`, 2.0× write surcharge. + /// + /// Cache injection uses Anthropic's **automatic caching** — a top-level + /// `cache_control` field in `additional_params` that gets `#[serde(flatten)]`'d + /// into the request body by rig-core. + /// + /// If the configured model does not support caching (e.g. claude-2), + /// a warning is logged once at construction and caching is disabled. + pub fn with_cache_retention(mut self, retention: CacheRetention) -> Self { + if retention != CacheRetention::None && !supports_prompt_cache(&self.model_name) { + tracing::warn!( + model = %self.model_name, + "Prompt caching requested but model does not support it; disabling" + ); + self.cache_retention = CacheRetention::None; + } else { + self.cache_retention = retention; + } + self + } } // -- Type conversion helpers -- @@ -360,7 +394,44 @@ fn saturate_u32(val: u64) -> u32 { val.min(u32::MAX as u64) as u32 } +/// Returns `true` if the model supports Anthropic prompt caching. +/// +/// Per Anthropic docs, only Claude 3+ models support prompt caching. +/// Unsupported: claude-2, claude-2.1, claude-instant-*. +fn supports_prompt_cache(name: &str) -> bool { + let lower = name.to_lowercase(); + // Strip optional provider prefix (e.g. "anthropic/claude-...") + let model = lower.strip_prefix("anthropic/").unwrap_or(&lower); + // Only Claude 3+ families support prompt caching + model.starts_with("claude-3") + || model.starts_with("claude-4") + || model.starts_with("claude-sonnet") + || model.starts_with("claude-opus") + || model.starts_with("claude-haiku") +} + +/// Extract `cache_creation_input_tokens` from the raw provider response. +/// +/// Rig-core's unified `Usage` does not surface this field, but Anthropic's raw +/// response includes it at `usage.cache_creation_input_tokens`. We serialize the +/// raw response to JSON and attempt to read the value. +fn extract_cache_creation(raw: &T) -> u32 { + serde_json::to_value(raw) + .ok() + .and_then(|v| v.get("usage")?.get("cache_creation_input_tokens")?.as_u64()) + .map(|n| n.min(u32::MAX as u64) as u32) + .unwrap_or(0) +} + /// Build a rig-core CompletionRequest from our internal types. +/// +/// When `cache_retention` is not `None`, injects a top-level `cache_control` +/// field via `additional_params`. Rig-core's `AnthropicCompletionRequest` +/// uses `#[serde(flatten)]` on `additional_params`, so the field lands at +/// the request root — which is exactly what Anthropic's **automatic caching** +/// expects. The API auto-places the cache breakpoint at the last cacheable +/// block and moves it forward as conversations grow. +#[allow(clippy::too_many_arguments)] fn build_rig_request( preamble: Option, mut history: Vec, @@ -368,6 +439,7 @@ fn build_rig_request( tool_choice: Option, temperature: Option, max_tokens: Option, + cache_retention: CacheRetention, ) -> Result { // rig-core requires at least one message in chat_history if history.is_empty() { @@ -379,6 +451,17 @@ fn build_rig_request( reason: format!("Failed to build chat history: {}", e), })?; + // Inject top-level cache_control for Anthropic automatic prompt caching. + let additional_params = match cache_retention { + CacheRetention::None => None, + CacheRetention::Short => Some(serde_json::json!({ + "cache_control": {"type": "ephemeral"} + })), + CacheRetention::Long => Some(serde_json::json!({ + "cache_control": {"type": "ephemeral", "ttl": "1h"} + })), + }; + Ok(RigRequest { preamble, chat_history, @@ -387,7 +470,7 @@ fn build_rig_request( temperature: temperature.map(|t| t as f64), max_tokens: max_tokens.map(|t| t as u64), tool_choice, - additional_params: None, + additional_params, }) } @@ -405,6 +488,22 @@ where (self.input_cost, self.output_cost) } + fn cache_write_multiplier(&self) -> Decimal { + match self.cache_retention { + CacheRetention::None => Decimal::ONE, + CacheRetention::Short => Decimal::new(125, 2), // 1.25× (125% of input rate) + CacheRetention::Long => Decimal::TWO, // 2.0× (200% of input rate) + } + } + + fn cache_read_discount(&self) -> Decimal { + if self.cache_retention != CacheRetention::None { + dec!(10) // Anthropic: 90% discount (cost = input_rate / 10) + } else { + Decimal::ONE + } + } + async fn complete(&self, request: CompletionRequest) -> Result { if let Some(requested_model) = request.model.as_deref() && requested_model != self.model_name.as_str() @@ -427,6 +526,7 @@ where None, request.temperature, request.max_tokens, + self.cache_retention, )?; let response = @@ -440,12 +540,26 @@ where let (text, _tool_calls, finish) = extract_response(&response.choice, &response.usage); - Ok(CompletionResponse { + let resp = CompletionResponse { content: text.unwrap_or_default(), input_tokens: saturate_u32(response.usage.input_tokens), output_tokens: saturate_u32(response.usage.output_tokens), finish_reason: finish, - }) + cache_read_input_tokens: saturate_u32(response.usage.cached_input_tokens), + cache_creation_input_tokens: extract_cache_creation(&response.raw_response), + }; + + if resp.cache_read_input_tokens > 0 { + tracing::debug!( + model = %self.model_name, + input = resp.input_tokens, + output = resp.output_tokens, + cache_read = resp.cache_read_input_tokens, + "prompt cache hit", + ); + } + + Ok(resp) } async fn complete_with_tools( @@ -478,6 +592,7 @@ where tool_choice, request.temperature, request.max_tokens, + self.cache_retention, )?; let response = @@ -504,13 +619,27 @@ where } } - Ok(ToolCompletionResponse { + let resp = ToolCompletionResponse { content: text, tool_calls, input_tokens: saturate_u32(response.usage.input_tokens), output_tokens: saturate_u32(response.usage.output_tokens), finish_reason: finish, - }) + cache_read_input_tokens: saturate_u32(response.usage.cached_input_tokens), + cache_creation_input_tokens: extract_cache_creation(&response.raw_response), + }; + + if resp.cache_read_input_tokens > 0 { + tracing::debug!( + model = %self.model_name, + input = resp.input_tokens, + output = resp.output_tokens, + cache_read = resp.cache_read_input_tokens, + "prompt cache hit", + ); + } + + Ok(resp) } fn active_model_name(&self) -> String { @@ -869,4 +998,125 @@ mod tests { let known = HashSet::from(["echo".to_string()]); assert_eq!(normalize_tool_name("other_tool", &known), "other_tool"); } + + #[test] + fn test_build_rig_request_injects_cache_control_short() { + let req = build_rig_request( + Some("You are helpful.".to_string()), + vec![RigMessage::user("Hello")], + Vec::new(), + None, + None, + None, + CacheRetention::Short, + ) + .unwrap(); + + let params = req + .additional_params + .expect("should have additional_params for Short retention"); + assert_eq!(params["cache_control"]["type"], "ephemeral"); + assert!( + params["cache_control"].get("ttl").is_none(), + "Short retention should not include ttl" + ); + } + + #[test] + fn test_build_rig_request_injects_cache_control_long() { + let req = build_rig_request( + Some("You are helpful.".to_string()), + vec![RigMessage::user("Hello")], + Vec::new(), + None, + None, + None, + CacheRetention::Long, + ) + .unwrap(); + + let params = req + .additional_params + .expect("should have additional_params for Long retention"); + assert_eq!(params["cache_control"]["type"], "ephemeral"); + assert_eq!(params["cache_control"]["ttl"], "1h"); + } + + #[test] + fn test_build_rig_request_no_cache_control_when_none() { + let req = build_rig_request( + Some("You are helpful.".to_string()), + vec![RigMessage::user("Hello")], + Vec::new(), + None, + None, + None, + CacheRetention::None, + ) + .unwrap(); + + assert!( + req.additional_params.is_none(), + "additional_params should be None when cache is disabled" + ); + } + + /// Verify that the multiplier match arms in `RigAdapter::cache_write_multiplier` + /// produce the expected values. We use a standalone helper because constructing + /// a real `RigAdapter` requires a rig `Model` (which needs network/provider setup). + /// The helper mirrors the same match expression — if the impl drifts, the + /// `test_build_rig_request_*` tests will still catch regressions end-to-end. + #[test] + fn test_cache_write_multiplier_values() { + use rust_decimal::Decimal; + // None → 1.0× (no surcharge) + assert_eq!( + cache_write_multiplier_for(CacheRetention::None), + Decimal::ONE + ); + // Short → 1.25× (25% surcharge) + assert_eq!( + cache_write_multiplier_for(CacheRetention::Short), + Decimal::new(125, 2) + ); + // Long → 2.0× (100% surcharge) + assert_eq!( + cache_write_multiplier_for(CacheRetention::Long), + Decimal::TWO + ); + } + + fn cache_write_multiplier_for(retention: CacheRetention) -> rust_decimal::Decimal { + match retention { + CacheRetention::None => rust_decimal::Decimal::ONE, + CacheRetention::Short => rust_decimal::Decimal::new(125, 2), + CacheRetention::Long => rust_decimal::Decimal::TWO, + } + } + + // -- supports_prompt_cache tests -- + + #[test] + fn test_supports_prompt_cache_supported_models() { + // All Claude 3+ models per Anthropic docs + assert!(supports_prompt_cache("claude-opus-4-6")); + assert!(supports_prompt_cache("claude-sonnet-4-6")); + assert!(supports_prompt_cache("claude-sonnet-4")); + assert!(supports_prompt_cache("claude-haiku-4-5")); + assert!(supports_prompt_cache("claude-3-5-sonnet-20241022")); + assert!(supports_prompt_cache("claude-haiku-3")); + assert!(supports_prompt_cache("Claude-Opus-4-5")); // case-insensitive + assert!(supports_prompt_cache("anthropic/claude-sonnet-4-6")); // provider prefix + } + + #[test] + fn test_supports_prompt_cache_unsupported_models() { + // Legacy Claude models that predate caching + assert!(!supports_prompt_cache("claude-2")); + assert!(!supports_prompt_cache("claude-2.1")); + assert!(!supports_prompt_cache("claude-instant-1.2")); + // Non-Claude models + assert!(!supports_prompt_cache("gpt-4o")); + assert!(!supports_prompt_cache("llama3")); + } } diff --git a/src/llm/smart_routing.rs b/src/llm/smart_routing.rs index bcc0b5bb..9f4a4141 100644 --- a/src/llm/smart_routing.rs +++ b/src/llm/smart_routing.rs @@ -857,6 +857,14 @@ impl LlmProvider for SmartRoutingProvider { self.primary.cost_per_token() } + fn cache_write_multiplier(&self) -> Decimal { + self.primary.cache_write_multiplier() + } + + fn cache_read_discount(&self) -> Decimal { + self.primary.cache_read_discount() + } + async fn complete(&self, request: CompletionRequest) -> Result { self.stats.total_requests.fetch_add(1, Ordering::Relaxed); @@ -1471,6 +1479,8 @@ mod tests { input_tokens: 10, output_tokens: 5, finish_reason: crate::llm::FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }; assert!(SmartRoutingProvider::response_is_uncertain(&response)); } @@ -1482,6 +1492,8 @@ mod tests { input_tokens: 10, output_tokens: 0, finish_reason: crate::llm::FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }; assert!(SmartRoutingProvider::response_is_uncertain(&response)); } @@ -1493,6 +1505,8 @@ mod tests { input_tokens: 10, output_tokens: 1, finish_reason: crate::llm::FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }; assert!(!SmartRoutingProvider::response_is_uncertain(&response)); } @@ -1505,6 +1519,8 @@ mod tests { input_tokens: 10, output_tokens: 20, finish_reason: crate::llm::FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }; assert!(!SmartRoutingProvider::response_is_uncertain(&response)); } diff --git a/src/orchestrator/api.rs b/src/orchestrator/api.rs index 45803e4e..82783a64 100644 --- a/src/orchestrator/api.rs +++ b/src/orchestrator/api.rs @@ -160,6 +160,8 @@ async fn llm_complete( input_tokens: resp.input_tokens, output_tokens: resp.output_tokens, finish_reason: format_finish_reason(resp.finish_reason), + cache_read_input_tokens: resp.cache_read_input_tokens, + cache_creation_input_tokens: resp.cache_creation_input_tokens, })) } @@ -189,6 +191,8 @@ async fn llm_complete_with_tools( input_tokens: resp.input_tokens, output_tokens: resp.output_tokens, finish_reason: format_finish_reason(resp.finish_reason), + cache_read_input_tokens: resp.cache_read_input_tokens, + cache_creation_input_tokens: resp.cache_creation_input_tokens, })) } diff --git a/src/testing.rs b/src/testing.rs index c62c2dcf..01c7fdf1 100644 --- a/src/testing.rs +++ b/src/testing.rs @@ -174,6 +174,8 @@ impl LlmProvider for StubLlm { input_tokens: 10, output_tokens: 5, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } @@ -191,6 +193,8 @@ impl LlmProvider for StubLlm { input_tokens: 10, output_tokens: 5, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } } diff --git a/src/worker/api.rs b/src/worker/api.rs index 59e028ed..d0048afc 100644 --- a/src/worker/api.rs +++ b/src/worker/api.rs @@ -52,6 +52,10 @@ pub struct ProxyCompletionResponse { pub input_tokens: u32, pub output_tokens: u32, pub finish_reason: String, + #[serde(default)] + pub cache_read_input_tokens: u32, + #[serde(default)] + pub cache_creation_input_tokens: u32, } #[derive(Debug, Serialize, Deserialize)] @@ -71,6 +75,10 @@ pub struct ProxyToolCompletionResponse { pub input_tokens: u32, pub output_tokens: u32, pub finish_reason: String, + #[serde(default)] + pub cache_read_input_tokens: u32, + #[serde(default)] + pub cache_creation_input_tokens: u32, } /// Completion result for the worker to report when done. @@ -227,6 +235,8 @@ impl WorkerHttpClient { input_tokens: proxy_resp.input_tokens, output_tokens: proxy_resp.output_tokens, finish_reason: parse_finish_reason(&proxy_resp.finish_reason), + cache_read_input_tokens: proxy_resp.cache_read_input_tokens, + cache_creation_input_tokens: proxy_resp.cache_creation_input_tokens, }) } @@ -254,6 +264,8 @@ impl WorkerHttpClient { input_tokens: proxy_resp.input_tokens, output_tokens: proxy_resp.output_tokens, finish_reason: parse_finish_reason(&proxy_resp.finish_reason), + cache_read_input_tokens: proxy_resp.cache_read_input_tokens, + cache_creation_input_tokens: proxy_resp.cache_creation_input_tokens, }) } diff --git a/tests/openai_compat_integration.rs b/tests/openai_compat_integration.rs index 379978cb..b4fc3b12 100644 --- a/tests/openai_compat_integration.rs +++ b/tests/openai_compat_integration.rs @@ -70,6 +70,8 @@ impl LlmProvider for MockLlmProvider { input_tokens: 10, output_tokens: 5, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } @@ -95,6 +97,8 @@ impl LlmProvider for MockLlmProvider { input_tokens: 15, output_tokens: 8, finish_reason: FinishReason::ToolUse, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } else { Ok(ToolCompletionResponse { @@ -103,6 +107,8 @@ impl LlmProvider for MockLlmProvider { input_tokens: 10, output_tokens: 4, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } } @@ -141,6 +147,8 @@ impl LlmProvider for FixedModelProvider { input_tokens: 10, output_tokens: 5, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } @@ -154,6 +162,8 @@ impl LlmProvider for FixedModelProvider { input_tokens: 10, output_tokens: 5, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } diff --git a/tests/provider_chaos.rs b/tests/provider_chaos.rs index d3f5c9ac..b6fa4c41 100644 --- a/tests/provider_chaos.rs +++ b/tests/provider_chaos.rs @@ -88,6 +88,8 @@ impl LlmProvider for FlakeyProvider { input_tokens: 10, output_tokens: 5, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } @@ -115,6 +117,8 @@ impl LlmProvider for FlakeyProvider { input_tokens: 10, output_tokens: 5, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } } @@ -192,6 +196,8 @@ impl LlmProvider for GarbageProvider { input_tokens: 0, output_tokens: 0, finish_reason: FinishReason::Unknown, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } @@ -206,6 +212,8 @@ impl LlmProvider for GarbageProvider { input_tokens: 0, output_tokens: 0, finish_reason: FinishReason::Unknown, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } } @@ -248,6 +256,8 @@ impl LlmProvider for ReliableProvider { input_tokens: 10, output_tokens: 5, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } @@ -262,6 +272,8 @@ impl LlmProvider for ReliableProvider { input_tokens: 10, output_tokens: 5, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } } diff --git a/tests/support/trace_llm.rs b/tests/support/trace_llm.rs index 56f72494..0559ab5b 100644 --- a/tests/support/trace_llm.rs +++ b/tests/support/trace_llm.rs @@ -578,6 +578,8 @@ impl LlmProvider for TraceLlm { input_tokens, output_tokens, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }), TraceResponse::ToolCalls { .. } => Err(LlmError::RequestFailed { provider: self.model_name.clone(), @@ -610,6 +612,8 @@ impl LlmProvider for TraceLlm { input_tokens, output_tokens, finish_reason: FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }), TraceResponse::ToolCalls { tool_calls, @@ -630,6 +634,8 @@ impl LlmProvider for TraceLlm { input_tokens, output_tokens, finish_reason: FinishReason::ToolUse, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, }) } TraceResponse::UserInput { .. } => Err(LlmError::RequestFailed {