feat: enable Anthropic prompt caching via automatic cache_control injection (#660)

* 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 <[email protected]> * 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 <[email protected]> * style: cargo fmt Co-Authored-By: Claude Opus 4.6 <[email protected]> * 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 <[email protected]> --------- Co-authored-by: Andrey <[email protected]> Co-authored-by: Andrey Gruzdev <[email protected]> Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-08-27 08:00:17 +00:00 · 2026-03-07 09:10:05 +00:00
parent 633b234e44
commit 424a0366a9
22 changed files with 831 additions and 18 deletions
@@ -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"
+        );
+    }
 }
@@ -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,
            })
        }
    }