From 043afc6ebead629fda9f3e8bff5a026bc51197d8 Mon Sep 17 00:00:00 2001 From: Zhongxuan Wang Date: Fri, 10 Jul 2026 16:30:22 -0700 Subject: [PATCH 1/2] feat(response-cache): opt-in exact-match LLM response cache MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An opt-in feature of the adaptive plugin (a response_cache config section, not a new plugin kind): managed LLM calls are keyed by a SHA-256 fingerprint of the normalized request and repeats are served from the store — instant, free, reproducible. Buffered and streaming calls share one keyspace; a streamed miss is teed and stored as its aggregated response, a hit replays provider-native chunks. Keying auto-detects the provider surface from the request shape and trusts the decode only where it is faithful: known-lossy shapes and decodes that fail to round-trip fall back to raw-body fingerprinting, which can only cost a miss. Stateful calls (Responses persistence, conversations, containers) and nondeterministic calls under the safety toggle bypass entirely. Stored answers must be complete: non-null errors, non-final statuses, truncated or lossily-collected streams, and upstream failures are never cached — the CLI gateway relays failed upstream replies to the client verbatim while keeping them invisible to the execution chain. The in-memory store is bounded by an honest resident-size budget with oldest-first eviction; the Redis backend runs under hard deadlines and re-checks entry expiry. Everything fails open: any cache error falls through to a live call. Same config surface in Rust, Python, Node, and Go; doctor reports configuration and backend reachability; hit/miss/ bypass marks carry fingerprints and savings, never bodies. Signed-off-by: Zhongxuan Wang --- .gitignore | 11 + crates/adaptive/Cargo.toml | 10 +- crates/adaptive/src/config.rs | 82 ++ crates/adaptive/src/lib.rs | 7 +- crates/adaptive/src/plugin_component.rs | 66 + crates/adaptive/src/response_cache/config.rs | 64 + .../adaptive/src/response_cache/intercept.rs | 464 +++++++ crates/adaptive/src/response_cache/key.rs | 919 +++++++++++++ crates/adaptive/src/response_cache/mark.rs | 204 +++ crates/adaptive/src/response_cache/mod.rs | 25 + crates/adaptive/src/response_cache/replay.rs | 325 +++++ crates/adaptive/src/response_cache/store.rs | 697 ++++++++++ crates/adaptive/src/runtime/features.rs | 53 +- crates/adaptive/src/runtime/validation.rs | 151 ++- .../response_cache_benchmark_tests.rs | 571 +++++++++ .../tests/integration/response_cache_tests.rs | 1142 +++++++++++++++++ crates/adaptive/tests/unit/config_tests.rs | 16 + crates/cli/src/doctor.rs | 86 +- crates/cli/src/gateway.rs | 72 +- crates/cli/tests/coverage/doctor_tests.rs | 100 ++ crates/cli/tests/coverage/server_tests.rs | 250 ++++ crates/node/adaptive.d.ts | 28 + crates/node/adaptive.js | 30 + go/nemo_relay/adaptive.go | 76 ++ go/nemo_relay/adaptive/adaptive.go | 21 + go/nemo_relay/adaptive_runtime_test.go | 66 + go/nemo_relay/plugin.go | 6 +- python/nemo_relay/adaptive.py | 52 + python/nemo_relay/adaptive.pyi | 35 + python/tests/test_adaptive_config.py | 41 + 30 files changed, 5647 insertions(+), 23 deletions(-) create mode 100644 crates/adaptive/src/response_cache/config.rs create mode 100644 crates/adaptive/src/response_cache/intercept.rs create mode 100644 crates/adaptive/src/response_cache/key.rs create mode 100644 crates/adaptive/src/response_cache/mark.rs create mode 100644 crates/adaptive/src/response_cache/mod.rs create mode 100644 crates/adaptive/src/response_cache/replay.rs create mode 100644 crates/adaptive/src/response_cache/store.rs create mode 100644 crates/adaptive/tests/integration/response_cache_benchmark_tests.rs create mode 100644 crates/adaptive/tests/integration/response_cache_tests.rs diff --git a/.gitignore b/.gitignore index 525a9c40a..d2d6f263b 100644 --- a/.gitignore +++ b/.gitignore @@ -44,6 +44,8 @@ integrations/openclaw/.test-dist/ go/nemo_relay/coverage.out go_junit_report.xml *.dSYM/ +# Runtime observability capture files (written by examples/tests) +nemo-relay-events-*.jsonl # OS — macOS .DS_Store @@ -85,3 +87,12 @@ CHANGELOG.md # Relay /.nemo-relay + +# Local response-cache scratch files (proposals/, demos, benchmarks, quickstarts), +# kept on disk for local use only, never committed. +/proposals/ +crates/adaptive/examples/response_cache_*.rs +python/examples/response_cache_*.py + +# Local-only response-cache notes (kept local, never committed) +RESPONSE_CACHE*.md diff --git a/crates/adaptive/Cargo.toml b/crates/adaptive/Cargo.toml index c59bf27e4..6b5391803 100644 --- a/crates/adaptive/Cargo.toml +++ b/crates/adaptive/Cargo.toml @@ -22,6 +22,7 @@ chrono = { version = "0.4", features = ["serde"] } thiserror = "2" tdigest = { version = "0.2", features = ["use_serde"] } tokio = { version = "1", features = ["rt", "macros", "sync", "time"] } +tokio-stream = { version = "0.1", default-features = false, features = ["sync"] } regex = "1" serde_json_canonicalizer = "0.3" sha2 = "0.11" @@ -33,7 +34,6 @@ redis-backend = ["redis"] [dev-dependencies] tokio = { version = "1", default-features = false, features = ["rt", "macros", "sync", "time", "test-util", "rt-multi-thread"] } -tokio-stream = { version = "0.1", default-features = false } [[test]] name = "redis_integration" @@ -46,3 +46,11 @@ path = "tests/integration/runtime_integration_tests.rs" [[test]] name = "acg_module_surface" path = "tests/integration/acg_module_surface_tests.rs" + +[[test]] +name = "response_cache_integration" +path = "tests/integration/response_cache_tests.rs" + +[[test]] +name = "response_cache_benchmark_integration" +path = "tests/integration/response_cache_benchmark_tests.rs" diff --git a/crates/adaptive/src/config.rs b/crates/adaptive/src/config.rs index c24bc4e50..a26637264 100644 --- a/crates/adaptive/src/config.rs +++ b/crates/adaptive/src/config.rs @@ -7,6 +7,8 @@ use nemo_relay::plugin::ConfigPolicy; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value as Json}; +use crate::response_cache::config::{BackendConfig, KEY_STRATEGY_EXACT_REQUEST}; + /// Canonical config document for the adaptive plugin component. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AdaptiveConfig { @@ -32,6 +34,10 @@ pub struct AdaptiveConfig { /// Adaptive Cache Governor settings. #[serde(default, skip_serializing_if = "Option::is_none")] pub acg: Option, + /// Opt-in LLM response cache (exact-match). When present, the + /// adaptive plugin installs the response-cache execution intercept(s). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub response_cache: Option, /// Adaptive-local unsupported-config policy. #[serde(default)] pub policy: ConfigPolicy, @@ -47,6 +53,7 @@ impl Default for AdaptiveConfig { adaptive_hints: None, tool_parallelism: None, acg: None, + response_cache: None, policy: ConfigPolicy::default(), } } @@ -184,6 +191,55 @@ impl Default for AcgComponentConfig { } } +/// Configuration for the adaptive plugin's `response_cache` feature +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default)] +pub struct ResponseCacheConfig { + /// How long a stored answer stays reusable, in seconds. + pub ttl_seconds: u64, + /// Namespace folded into every key to separate environments/tenants. + pub namespace: String, + /// Execution-intercept priority; lower runs first/outermost (default `50`). + pub priority: i32, + /// Probability in `[0.0, 1.0]` of skipping the cache and running live. + pub bypass_rate: f64, + /// Cache nondeterministic requests too. Set `false` to cache only + /// requests explicitly pinned deterministic (`temperature` = 0) — absent + /// or unreadable temperatures count as nondeterministic. + pub cache_nondeterministic: bool, + /// Key strategy. Only [`KEY_STRATEGY_EXACT_REQUEST`] is supported. + pub key_strategy: String, + /// Request headers (case-insensitive) folded into the key; never auth headers. + pub header_allowlist: Vec, + /// Extra top-level request-body keys to drop from the key, beyond the noise defaults. + pub skip_keys: Vec, + /// Storage backend selection. + pub backend: BackendConfig, +} + +impl Default for ResponseCacheConfig { + fn default() -> Self { + Self { + ttl_seconds: 3600, + namespace: String::new(), + priority: 50, + bypass_rate: 0.0, + cache_nondeterministic: true, + key_strategy: KEY_STRATEGY_EXACT_REQUEST.to_string(), + header_allowlist: Vec::new(), + skip_keys: Vec::new(), + backend: BackendConfig::default(), + } + } +} + +impl ResponseCacheConfig { + /// TTL as a [`std::time::Duration`]. + pub fn ttl(&self) -> std::time::Duration { + std::time::Duration::from_secs(self.ttl_seconds) + } +} + fn default_adaptive_config_version() -> u32 { 1 } @@ -254,6 +310,13 @@ nemo_relay::editor_config! { nested: AcgComponentConfig, default: AcgComponentConfig, }, + response_cache => { + label: "response_cache", + kind: Section, + optional: true, + nested: ResponseCacheConfig, + default: ResponseCacheConfig, + }, policy => { label: "policy", kind: Section, @@ -326,6 +389,25 @@ nemo_relay::editor_config! { } } +nemo_relay::editor_config! { + impl ResponseCacheConfig { + ttl_seconds => { label: "ttl_seconds", kind: Integer }, + namespace => { label: "namespace", kind: String }, + priority => { label: "priority", kind: Integer }, + bypass_rate => { label: "bypass_rate", kind: Float }, + cache_nondeterministic => { label: "cache_nondeterministic", kind: Boolean }, + key_strategy => { label: "key_strategy", kind: String }, + header_allowlist => { label: "header_allowlist", kind: Json }, + skip_keys => { label: "skip_keys", kind: Json }, + backend => { + label: "backend", + kind: Section, + nested: BackendConfig, + default: BackendConfig, + }, + } +} + nemo_relay::editor_config! { impl crate::acg::stability::StabilityThresholds { stable_threshold => { label: "stable_threshold", kind: Float }, diff --git a/crates/adaptive/src/lib.rs b/crates/adaptive/src/lib.rs index ebe78d534..c1ca8b44c 100644 --- a/crates/adaptive/src/lib.rs +++ b/crates/adaptive/src/lib.rs @@ -30,6 +30,8 @@ pub mod learner; pub mod plugin_component; #[cfg(feature = "redis-backend")] pub mod redis; +/// Opt-in LLM response cache (exact-match). +pub mod response_cache; mod runtime; /// Storage backends and backend traits for adaptive state persistence. pub mod storage; @@ -41,8 +43,8 @@ pub mod trie; pub mod types; pub use config::{ - AcgComponentConfig, AdaptiveConfig, AdaptiveHintsComponentConfig, BackendSpec, StateConfig, - TelemetryComponentConfig, ToolParallelismComponentConfig, + AcgComponentConfig, AdaptiveConfig, AdaptiveHintsComponentConfig, BackendSpec, + ResponseCacheConfig, StateConfig, TelemetryComponentConfig, ToolParallelismComponentConfig, }; pub use context_helpers::{ LATENCY_SENSITIVITY_POINTER, extract_scope_path, read_manual_latency_sensitivity, @@ -51,6 +53,7 @@ pub use context_helpers::{ pub use error::{AdaptiveError, Result}; #[cfg(feature = "redis-backend")] pub use redis::RedisBackend; +pub use response_cache::RESPONSE_CACHE_MARK; pub use runtime::features::AdaptiveRuntime; pub use storage::erased::AnyBackend; pub use storage::memory::InMemoryBackend; diff --git a/crates/adaptive/src/plugin_component.rs b/crates/adaptive/src/plugin_component.rs index 62ccfd4b6..cc7507ff9 100644 --- a/crates/adaptive/src/plugin_component.rs +++ b/crates/adaptive/src/plugin_component.rs @@ -182,6 +182,7 @@ fn validate_adaptive_plugin_config(plugin_config: &Map) -> Vec) -> Vec, + policy: &ConfigPolicy, + backend_kind: &str, + backend_config: &Map, +) { + let known_fields: &[&str] = match backend_kind { + "in_memory" => &["max_bytes", "max_entries"], + "redis" => &["url", "key_prefix"], + _ => return, + }; + validate_unknown_fields( + diagnostics, + policy, + Some(format!("response_cache.backend.{backend_kind}")), + backend_config, + known_fields, + ); +} + fn validate_backend_config_fields( diagnostics: &mut Vec, policy: &ConfigPolicy, diff --git a/crates/adaptive/src/response_cache/config.rs b/crates/adaptive/src/response_cache/config.rs new file mode 100644 index 000000000..e8dbbdad1 --- /dev/null +++ b/crates/adaptive/src/response_cache/config.rs @@ -0,0 +1,64 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Backend selection for the adaptive plugin's `response_cache` feature. +//! +//! The `response_cache` section struct ([`crate::config::ResponseCacheConfig`]) +//! lives in [`crate::config`] alongside the other `AdaptiveConfig` sections +//! (`acg`, `adaptive_hints`, `tool_parallelism`). This module keeps the +//! response-cache-specific backend config and the key-strategy constant next to +//! the key/store code that consumes them. + +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value as Json}; + +/// Exact-request key strategy identifier. +pub const KEY_STRATEGY_EXACT_REQUEST: &str = "exact_request"; + +/// Default in-memory byte budget: 256 MiB. +pub const DEFAULT_MAX_BYTES: usize = 256 * 1024 * 1024; + +/// Backend selection mirroring the adaptive [`crate::config::BackendSpec`] shape. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default)] +pub struct BackendConfig { + /// Backend kind: `"in_memory"` or `"redis"` (needs the `redis-backend` feature). + pub kind: String, + /// Backend-specific options (in_memory: `max_bytes`/`max_entries`; redis: `url`/`key_prefix`). + pub config: Map, +} + +impl Default for BackendConfig { + fn default() -> Self { + Self { + kind: "in_memory".to_string(), + config: Map::new(), + } + } +} + +impl BackendConfig { + /// In-memory total-bytes budget before oldest-first eviction. + pub fn max_bytes(&self) -> usize { + self.config + .get("max_bytes") + .and_then(Json::as_u64) + .map(|value| value as usize) + .unwrap_or(DEFAULT_MAX_BYTES) + } + + /// Optional in-memory entry-count cap (on top of the byte budget). + pub fn max_entries(&self) -> Option { + self.config + .get("max_entries") + .and_then(Json::as_u64) + .map(|value| value as usize) + } +} + +nemo_relay::editor_config! { + impl BackendConfig { + kind => { label: "kind", kind: Enum, values: ["in_memory", "redis"] }, + config => { label: "config", kind: Json }, + } +} diff --git a/crates/adaptive/src/response_cache/intercept.rs b/crates/adaptive/src/response_cache/intercept.rs new file mode 100644 index 000000000..e3443cd59 --- /dev/null +++ b/crates/adaptive/src/response_cache/intercept.rs @@ -0,0 +1,464 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The buffered and streaming LLM execution intercepts: cache decisions, +//! the streaming tee, and the storage rules. + +use std::sync::Arc; + +use nemo_relay::api::llm::LlmRequest; +use nemo_relay::api::runtime::{ + LlmExecutionFn, LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionFn, + LlmStreamExecutionNextFn, +}; +use nemo_relay::codec::anthropic::AnthropicMessagesStreamingCodec; +use nemo_relay::codec::openai_chat::OpenAIChatStreamingCodec; +use nemo_relay::codec::openai_responses::OpenAIResponsesStreamingCodec; +use nemo_relay::codec::resolve::{ProviderSurface, detect_request_surface_with_hint}; +use nemo_relay::codec::streaming::StreamingCodec; +use nemo_relay::error::Result as FlowResult; +use serde_json::Value as Json; +use std::cell::Cell; +use std::time::{SystemTime, UNIX_EPOCH}; +use tokio_stream::StreamExt; +use tokio_stream::wrappers::ReceiverStream; + +use crate::config::ResponseCacheConfig; +use crate::response_cache::key::{KeyOutcome, build_cache_key}; +use crate::response_cache::mark::{CacheMark, emit_cache_mark, savings_from}; +use crate::response_cache::replay::replay_aggregate; +use crate::response_cache::store::{CacheEntry, CacheStore, now_unix_ms}; + +/// Bounded channel capacity for the streaming tee: it forwards live chunks to +/// the consumer while accumulating them, applying backpressure when the consumer +/// is slow. +const STREAM_TEE_CHANNEL_CAP: usize = 64; + +/// Builds the buffered LLM execution intercept for the response cache. +/// +/// Called from the adaptive runtime when the `response_cache` section is present. +pub(crate) fn make_intercept( + store: Arc, + config: Arc, +) -> LlmExecutionFn { + Arc::new( + move |provider: &str, request: LlmRequest, next: LlmExecutionNextFn| { + // Clone the shared handles before the async block so the intercept + // never holds a lock across an await (the remote-backend pattern). + let store = Arc::clone(&store); + let config = Arc::clone(&config); + let provider = provider.to_string(); + Box::pin(run_cache(provider, request, next, store, config)) + }, + ) +} + +/// Builds the streaming LLM execution intercept for the response cache. +/// +/// On a miss it tees the live stream — forwarding chunks while feeding them to +/// the codec — and on natural completion stores the codec-assembled +/// **aggregate** response (the same shape a buffered call stores), so buffered +/// and streaming entries share one key. On a hit it replays that aggregate as +/// provider-native chunks. The codec is inferred from the request surface; only +/// a request whose surface can't be inferred runs live (uncached). +pub(crate) fn make_stream_intercept( + store: Arc, + config: Arc, +) -> LlmStreamExecutionFn { + Arc::new( + move |provider: &str, request: LlmRequest, next: LlmStreamExecutionNextFn| { + let store = Arc::clone(&store); + let config = Arc::clone(&config); + let provider = provider.to_string(); + Box::pin(run_cache_stream(provider, request, next, store, config)) + }, + ) +} + +/// Core get-or-miss logic. Separated into a free `async fn` so the future type +/// is explicit and `Send`. +async fn run_cache( + provider: String, + request: LlmRequest, + next: LlmExecutionNextFn, + store: Arc, + config: Arc, +) -> FlowResult { + let backend = store.backend_kind(); + + // Decision marks are emitted before `next()` (like the runtime's start + // event) so every decision is recorded even when the provider then errors. + let key = match build_cache_key(&provider, &request, &config) { + KeyOutcome::Key(key) => key, + KeyOutcome::Bypass(reason) => { + // Not cacheable (e.g. stateful Responses call): run live. + emit_cache_mark(CacheMark::new("bypass", backend).reason(reason)); + return next(request).await; + } + }; + + // Capture the model before `request` is moved into the provider call. + let model = request_model(&request); + + // Sampled bypass: re-run live to catch drift, refreshing the stored answer. + if should_bypass(config.bypass_rate) { + emit_cache_mark( + CacheMark::new("bypass", backend) + .reason("sampled") + .key_hash(&key), + ); + let response = next(request).await?; + maybe_store(&store, &config, &key, &provider, model, &response).await; + return Ok(response); + } + + match store.get(&key).await { + Ok(Some(entry)) => { + let age_ms = now_unix_ms().saturating_sub(entry.created_unix_ms); + let (saved_tokens, saved_cost) = savings_from(&entry); + emit_cache_mark( + CacheMark::new("hit", backend) + .key_hash(&key) + .age_ms(age_ms) + .ttl_ms(config.ttl().as_millis() as u64) + .savings(saved_tokens, saved_cost), + ); + // A reuse must be shape-identical to a live call (usage intact); + // savings are reported on the mark, never by mutating the body. + Ok(entry.response.clone()) + } + Ok(None) => { + emit_cache_mark( + CacheMark::new("miss", backend) + .key_hash(&key) + .ttl_ms(config.ttl().as_millis() as u64), + ); + let response = next(request).await?; + maybe_store(&store, &config, &key, &provider, model, &response).await; + Ok(response) + } + Err(_) => { + // Cache read failed: fail open as a live call and do not store. + emit_cache_mark( + CacheMark::new("miss", backend) + .reason("store_error") + .key_hash(&key), + ); + next(request).await + } + } +} + +/// Streaming counterpart of [`run_cache`]. Assembles the streamed chunks into a +/// single aggregate response (via the configured codec) and stores **that** — the +/// same shape a buffered call stores — so buffered and streaming entries share +/// one key. Replays the stored aggregate on a hit. Aggregation needs a codec, +/// inferred from the request surface; only an unrecognized surface runs live +/// (uncached). Fails open. +async fn run_cache_stream( + provider: String, + request: LlmRequest, + next: LlmStreamExecutionNextFn, + store: Arc, + config: Arc, +) -> FlowResult { + let backend = store.backend_kind(); + + // Assembling streamed chunks into a stored response needs a streaming codec, + // inferred via the shared request-surface detector (the observability/ACG + // decode path), so gateway traffic caches with zero configuration. The + // detector is hinted with the provider name so a system-less Anthropic + // request is not misread as OpenAI Chat (an ambiguity core documents); the + // inference is guarded against a mistake at store time (see + // `tee_and_aggregate`). + let codec = match detect_request_surface_with_hint(&request.content, Some(&provider)) { + Some(surface) => streaming_codec_for_surface(surface), + None => { + emit_cache_mark(CacheMark::new("bypass", backend).reason("stream_no_codec")); + return next(request).await; + } + }; + + // As in `run_cache`, the decision mark is emitted before `next()`. + let key = match build_cache_key(&provider, &request, &config) { + KeyOutcome::Key(key) => key, + KeyOutcome::Bypass(reason) => { + emit_cache_mark(CacheMark::new("bypass", backend).reason(reason)); + return next(request).await; + } + }; + + let model = request_model(&request); + + // Sampled bypass: run live (and re-aggregate to refresh the stored answer). + if should_bypass(config.bypass_rate) { + emit_cache_mark( + CacheMark::new("bypass", backend) + .reason("sampled") + .key_hash(&key), + ); + let live = next(request).await?; + return Ok(tee_and_aggregate( + live, codec, store, config, key, provider, model, + )); + } + + match store.get(&key).await { + Ok(Some(entry)) => { + let age_ms = now_unix_ms().saturating_sub(entry.created_unix_ms); + let (saved_tokens, saved_cost) = savings_from(&entry); + emit_cache_mark( + CacheMark::new("hit", backend) + .key_hash(&key) + .age_ms(age_ms) + .ttl_ms(config.ttl().as_millis() as u64) + .savings(saved_tokens, saved_cost), + ); + // Replay the stored aggregate as provider-native chunks. + Ok(replay_aggregate(entry.response.clone())) + } + Ok(None) => { + emit_cache_mark( + CacheMark::new("miss", backend) + .key_hash(&key) + .ttl_ms(config.ttl().as_millis() as u64), + ); + let live = next(request).await?; + Ok(tee_and_aggregate( + live, codec, store, config, key, provider, model, + )) + } + Err(_) => { + // Cache read failed: fail open as a live stream and do not store. + emit_cache_mark( + CacheMark::new("miss", backend) + .reason("store_error") + .key_hash(&key), + ); + next(request).await + } + } +} + +/// Maps a detected provider surface to its streaming codec. Lets streaming reuse +/// the shared surface detector (the observability/ACG decode path) so a request +/// keys and caches without an explicit `codec` config. +pub(crate) fn streaming_codec_for_surface(surface: ProviderSurface) -> Box { + match surface { + ProviderSurface::OpenAIChat => Box::new(OpenAIChatStreamingCodec::new()), + ProviderSurface::OpenAIResponses => Box::new(OpenAIResponsesStreamingCodec::new()), + ProviderSurface::AnthropicMessages => Box::new(AnthropicMessagesStreamingCodec::new()), + } +} + +/// Tees a live stream: forwards each chunk to the consumer while feeding it to the +/// codec's collector, and on natural, error-free completion stores the +/// codec-assembled **aggregate**. An upstream error, a chunk the codec rejects, or +/// a dropped consumer caches nothing. A content-empty aggregate is also not +/// stored — the codec was inferred from the request surface, so an empty result +/// signals a mis-inference (e.g. a system-less Anthropic request read as OpenAI +/// Chat, whose collector silently drops the foreign chunks), and caching it +/// would serve a wrong empty response on the repeat. +fn tee_and_aggregate( + live: LlmJsonStream, + codec: Box, + store: Arc, + config: Arc, + key: String, + provider: String, + model: Option, +) -> LlmJsonStream { + let (tx, rx) = tokio::sync::mpsc::channel::>(STREAM_TEE_CHANNEL_CAP); + tokio::spawn(async move { + let mut collect = codec.collector(); + let mut live = live; + let mut collector_failed = false; + let mut saw_terminal = false; + while let Some(item) = live.next().await { + match &item { + Ok(chunk) => { + // In-band provider errors never surface as stream-level Err. + if chunk_is_inband_error(chunk) || collect(chunk.clone()).is_err() { + collector_failed = true; + } + saw_terminal |= chunk_is_terminal(chunk); + } + Err(_) => { + // Upstream error is a failed call: forward it, never cache. + let _ = tx.send(item).await; + return; + } + } + // Forward to the consumer; a send error means it was dropped. + if tx.send(item).await.is_err() { + return; + } + } + // Store only protocol-complete streams: every collector finalizes a + // clean truncation as a well-formed partial. + if !collector_failed && saw_terminal { + let aggregate = codec.finalizer()(); + // Empty = mis-inferred surface; lossy = unfaithful replay. + if aggregate_has_no_content(&aggregate) || aggregate_replay_lossy(&aggregate) { + return; + } + maybe_store(&store, &config, &key, &provider, model, &aggregate).await; + } + }); + Box::pin(ReceiverStream::new(rx)) +} + +/// Streamed content the collectors assemble lossily: thinking blocks lose +/// their deltas/signature; refusal-only chat choices lose the refusal text. +fn aggregate_replay_lossy(aggregate: &Json) -> bool { + if let Some(content) = aggregate.get("content").and_then(Json::as_array) + && content.iter().any(|block| { + matches!( + block.get("type").and_then(Json::as_str), + Some("thinking" | "redacted_thinking") + ) + }) + { + return true; + } + if let Some(choices) = aggregate.get("choices").and_then(Json::as_array) + && !choices.is_empty() + && choices.iter().all(|choice| { + let message = choice.get("message"); + let content_empty = message + .and_then(|message| message.get("content")) + .is_none_or(|content| { + content.is_null() || content.as_str().is_some_and(str::is_empty) + }); + let no_tool_calls = message + .and_then(|message| message.get("tool_calls")) + .and_then(Json::as_array) + .is_none_or(|calls| calls.is_empty()); + content_empty && no_tool_calls + }) + { + return true; + } + false +} + +/// A chunk that marks the stream complete (`response.incomplete` is +/// deliberately excluded: a capped answer must not replay as "the" answer). +fn chunk_is_terminal(chunk: &Json) -> bool { + if let Some(choices) = chunk.get("choices").and_then(Json::as_array) + && choices.iter().any(|choice| { + choice + .get("finish_reason") + .is_some_and(|reason| !reason.is_null()) + }) + { + return true; + } + matches!( + chunk.get("type").and_then(Json::as_str), + Some("message_stop" | "response.completed") + ) +} + +/// Provider-native in-band error chunk; a false positive only skips a store. +fn chunk_is_inband_error(chunk: &Json) -> bool { + if chunk.get("error").is_some_and(|error| !error.is_null()) { + return true; + } + matches!( + chunk.get("type").and_then(Json::as_str), + Some("error" | "response.failed") + ) +} + +/// True when a finalized streaming aggregate carries no response content in any +/// known provider shape (OpenAI Chat `choices`, OpenAI Responses `output`, +/// Anthropic `content`) — used to reject a mis-inferred codec's empty output. +fn aggregate_has_no_content(aggregate: &Json) -> bool { + let empty_array = |key: &str| { + aggregate + .get(key) + .and_then(Json::as_array) + .is_none_or(|items| items.is_empty()) + }; + empty_array("choices") && empty_array("output") && empty_array("content") +} + +async fn maybe_store( + store: &Arc, + config: &ResponseCacheConfig, + key: &str, + provider: &str, + model: Option, + response: &Json, +) { + // Failed calls are never cached (out of scope for v1). + if is_error_response(response) { + return; + } + let entry = CacheEntry::new( + response.clone(), + config.ttl(), + key.to_string(), + model, + Some(provider.to_string()), + ); + // Fail open: a store error must never break the live call. + let _ = store.set(key, entry, config.ttl()).await; +} + +fn request_model(request: &LlmRequest) -> Option { + request + .content + .get("model") + .and_then(Json::as_str) + .map(str::to_string) +} + +/// Non-null `error` or non-final `status` = not a complete, replayable +/// answer. Must tolerate `error: null` — real Responses success bodies carry it. +fn is_error_response(response: &Json) -> bool { + let Some(object) = response.as_object() else { + return false; + }; + if object.get("error").is_some_and(|error| !error.is_null()) { + return true; + } + matches!( + object.get("status").and_then(Json::as_str), + Some("failed" | "cancelled" | "canceled" | "incomplete" | "in_progress" | "queued") + ) +} + +thread_local! { + static RNG_STATE: Cell = Cell::new(rng_seed()); +} + +fn rng_seed() -> u64 { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|delta| delta.as_nanos() as u64) + .unwrap_or(0); + (nanos ^ 0x9E37_79B9_7F4A_7C15) | 1 +} + +fn next_unit_f64() -> f64 { + RNG_STATE.with(|cell| { + let mut x = cell.get(); + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + cell.set(x); + (x >> 11) as f64 / ((1u64 << 53) as f64) + }) +} + +pub(crate) fn should_bypass(rate: f64) -> bool { + if rate <= 0.0 { + false + } else if rate >= 1.0 { + true + } else { + next_unit_f64() < rate + } +} diff --git a/crates/adaptive/src/response_cache/key.rs b/crates/adaptive/src/response_cache/key.rs new file mode 100644 index 000000000..a116f3123 --- /dev/null +++ b/crates/adaptive/src/response_cache/key.rs @@ -0,0 +1,919 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Exact-match cache-key derivation. +//! +//! The LLM execution intercept receives the *raw* [`LlmRequest`] (the decoded +//! `AnnotatedLlmRequest` is only handed to *request* intercepts), so the plugin +//! decodes the request itself and keys on the normalized form: provider-shaped +//! differences that mean the same thing collapse to one key. The provider +//! surface is auto-detected from the request shape (hinted by the provider +//! name); there is nothing to configure. Only when detection or decode fails is +//! the raw body fingerprinted (the fallback). Either way RFC +//! 8785 canonicalization removes field-order/whitespace noise, an always-on +//! skip-list drops volatile/identity fields, tool-call IDs are normalized, and +//! only allowlisted headers fold in. + +use nemo_relay::api::llm::LlmRequest; +use nemo_relay::codec::anthropic::AnthropicMessagesCodec; +use nemo_relay::codec::openai_chat::OpenAIChatCodec; +use nemo_relay::codec::openai_responses::OpenAIResponsesCodec; +use nemo_relay::codec::request::AnnotatedLlmRequest; +use nemo_relay::codec::resolve::{ProviderSurface, detect_request_surface_with_hint}; +use nemo_relay::codec::traits::LlmCodec; +use serde_json::{Map, Value as Json, json}; +use sha2::{Digest, Sha256}; + +use crate::config::ResponseCacheConfig; +use crate::response_cache::store::CACHE_SCHEMA_VERSION; + +/// Top-level request-body keys that never affect the answer and are always +/// dropped before fingerprinting (IDs, routing, bookkeeping, streaming flag). +pub const DEFAULT_SKIP_KEYS: &[&str] = &["stream", "user", "metadata", "service_tier", "store"]; + +/// Result of deriving a cache key for a request. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum KeyOutcome { + /// A usable exact-match key fingerprint (`"sha256:…"`). + Key(String), + /// The request is intentionally not cacheable; the reason is a short, + /// stable label suitable for telemetry. + Bypass(&'static str), +} + +/// Derives the cache key for a request, or decides it must bypass the cache. +/// +/// The request is decoded to the normalized `AnnotatedLlmRequest` and keyed on +/// that (so provider-shaped differences that mean the same thing collapse to +/// one key); the surface is auto-detected from the request shape. Only when +/// detection or decode fails is the raw request body fingerprinted. Either way the answer-determining fields are +/// kept, noise is dropped, tool-call IDs are normalized, and only allowlisted +/// headers fold in. +pub fn build_cache_key( + provider: &str, + request: &LlmRequest, + config: &ResponseCacheConfig, +) -> KeyOutcome { + // Unparseable bodies arrive as null; they would all share one key. + if request.content.is_null() { + return KeyOutcome::Bypass("unparseable_body"); + } + // Cacheability gates run on the RAW request, so they are correct regardless + // of which codec (if any) decodes the body — a chat codec may park `store` + // in `extra` rather than the typed field, so we must not rely on the decode. + if let Some(object) = request.content.as_object() { + // Any present, non-`false` `store` opts into server-side persistence — + // bypass even a malformed non-boolean rather than risk caching a stateful + // call (whose result is otherwise keyed with `store` stripped). + if object + .get("store") + .is_some_and(|value| !matches!(value, Json::Bool(false) | Json::Null)) + { + return KeyOutcome::Bypass("stateful_store"); + } + if object.contains_key("previous_response_id") { + return KeyOutcome::Bypass("stateful_previous_response_id"); + } + // Server-side conversation state the key cannot see. + if object.contains_key("conversation") || object.contains_key("container") { + return KeyOutcome::Bypass("stateful_conversation"); + } + // Responses persists by default; only an explicit opt-out is stateless. + if (object.contains_key("input") || object.contains_key("instructions")) + && !object + .get("store") + .is_some_and(|store| store == &Json::Bool(false)) + { + return KeyOutcome::Bypass("stateful_store"); + } + } + // Toggle off = explicit temperature 0 only; absent defaults to sampling. + if !config.cache_nondeterministic + && request_temperature(&request.content).is_none_or(|temperature| temperature > 0.0) + { + return KeyOutcome::Bypass("nondeterministic_temperature"); + } + + // Body to fingerprint: the decoded/normalized form when a surface resolves + // and decode succeeds, otherwise the raw request body. + let (mut body, effective_codec) = resolved_body(provider, request); + + if let Some(object) = body.as_object_mut() { + // `AnnotatedLlmRequest.extra` is `#[serde(flatten)]`, so provider fields + // also land at top level — one skip pass covers both. + for key in DEFAULT_SKIP_KEYS { + object.remove(*key); + } + for key in &config.skip_keys { + object.remove(key); + } + normalize_tool_call_ids(object); + } + + let headers = allowlisted_headers(&request.headers, &config.header_allowlist); + + let key_doc = json!({ + "v": CACHE_SCHEMA_VERSION, + "ns": config.namespace, + "provider": provider, + "strategy": config.key_strategy, + "codec": effective_codec, + "body": body, + "headers": headers, + }); + + match fingerprint(&key_doc) { + Some(key) => KeyOutcome::Key(key), + None => KeyOutcome::Bypass("canonicalization_failed"), + } +} + +/// Canonicalizes straight into SHA-256; byte-identical output to +/// `sha256_hex(&canonicalize_value(doc)?)` (proven by test). +fn fingerprint(doc: &T) -> Option { + let mut hasher = Sha256::new(); + serde_json_canonicalizer::to_writer(doc, &mut HashWriter(&mut hasher)).ok()?; + let digest = hasher.finalize(); + let mut out = String::with_capacity(7 + 64); + out.push_str("sha256:"); + for byte in digest { + out.push(char::from(HEX[(byte >> 4) as usize])); + out.push(char::from(HEX[(byte & 0x0f) as usize])); + } + Some(out) +} + +const HEX: [u8; 16] = *b"0123456789abcdef"; + +/// Feeds canonical bytes straight into the hasher (`digest` 0.11 no longer +/// implements `io::Write` for hashers). +struct HashWriter<'a>(&'a mut Sha256); + +impl std::io::Write for HashWriter<'_> { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.0.update(buf); + Ok(buf.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +/// The body to fingerprint plus the codec that actually produced it. +/// +/// The surface is auto-detected from the request shape, hinted by the provider +/// name — the same detector the streaming path uses. `(raw clone, None)` when +/// nothing detects or decodes. +fn resolved_body(provider: &str, request: &LlmRequest) -> (Json, Option<&'static str>) { + let surface = detect_request_surface_with_hint(&request.content, Some(provider)); + let decoded = surface.and_then(|surface| { + let annotated = decode_surface(surface, request)?; + let body = serde_json::to_value(&annotated).ok()?; + Some((body, codec_name(surface))) + }); + match decoded { + Some((body, name)) => (body, Some(name)), + None => (request.content.clone(), None), + } +} + +/// Canonical name for a surface, recorded in the key document. +fn codec_name(surface: ProviderSurface) -> &'static str { + match surface { + ProviderSurface::OpenAIChat => "openai_chat", + ProviderSurface::OpenAIResponses => "openai_responses", + ProviderSurface::AnthropicMessages => "anthropic_messages", + } +} + +/// Decodes the raw request via the surface's codec. Returns `None` on a decode +/// failure (the caller falls back to raw-body fingerprinting). +fn decode_surface(surface: ProviderSurface, request: &LlmRequest) -> Option { + // Known-lossy shapes stay raw-keyed: a fallback only ever costs a miss. + if lossy_request_shape(surface, &request.content) { + return None; + } + let annotated = match surface { + ProviderSurface::OpenAIChat => OpenAIChatCodec.decode(request).ok(), + ProviderSurface::OpenAIResponses => OpenAIResponsesCodec.decode(request).ok(), + ProviderSurface::AnthropicMessages => AnthropicMessagesCodec.decode(request).ok(), + }?; + // The chat/anthropic codecs parse `messages` with `unwrap_or_default()`, so a + // message carrying content the normalized type cannot represent (Anthropic + // `tool_use`/`tool_result`/`image`, OpenAI `input_audio`/`file`) collapses the + // WHOLE array to empty — silently dropping answer-affecting content from the + // key, so two such requests would collide. If the raw request carried messages + // but the decode produced none, reject the decode so the caller falls back to + // raw-body fingerprinting, which keeps distinct requests distinct (lower + // normalization, but never a wrong reuse). + if annotated.messages.is_empty() && raw_has_messages(request) { + return None; + } + // The closed tool types drop unmodeled fields (`function.strict`, + // server-tool settings); trust the decode only if it round-trips. + if let Some(raw_tools) = request.content.get("tools") { + let decoded_tools = serde_json::to_value(&annotated.tools).ok()?; + if decoded_tools != *raw_tools { + return None; + } + } + // Same round-trip rule for `tool_choice`. + if let Some(raw_tool_choice) = request.content.get("tool_choice") { + let decoded_tool_choice = serde_json::to_value(&annotated.tool_choice).ok()?; + if decoded_tool_choice != *raw_tool_choice { + return None; + } + } + // Closed message types drop `function_call`/`refusal`; legitimately + // restructured shapes (synthesized system) just key raw — a dedup loss only. + if let Some(raw_conversation) = request + .content + .get("messages") + .or_else(|| request.content.get("input")) + { + let decoded_messages = serde_json::to_value(&annotated.messages).ok()?; + if decoded_messages != *raw_conversation { + return None; + } + } + Some(annotated) +} + +/// Raw shapes the built-in codecs are known to decode lossily — answer-affecting +/// fields the normalized types drop or merge. +fn lossy_request_shape(surface: ProviderSurface, content: &Json) -> bool { + let Some(object) = content.as_object() else { + return false; + }; + // Shapes the decode silently drops: a raw `params` object overwrites the + // typed field via the flattened extra; wrong-typed scalars vanish entirely. + if object.contains_key("params") { + return true; + } + let non_number = |key: &str| { + object + .get(key) + .is_some_and(|value| value.as_f64().is_none()) + }; + let non_u64 = |key: &str| { + object + .get(key) + .is_some_and(|value| value.as_u64().is_none()) + }; + if non_number("temperature") + || non_number("top_p") + || non_u64("max_tokens") + || non_u64("max_completion_tokens") + || non_u64("max_output_tokens") + { + return true; + } + match surface { + // `stop` is modeled only as an array of strings — any other present + // form (a bare string, a mixed array, an object) is dropped on decode. + // And when both token caps are present the decode keeps one. + ProviderSurface::OpenAIChat => { + object + .get("stop") + .is_some_and(|stop| !is_string_array(stop)) + || (object.contains_key("max_tokens") + && object.contains_key("max_completion_tokens")) + } + // Same for `stop_sequences`; and system content blocks are flattened + // to their text on decode, so any non-text block, non-string text, or + // block metadata beyond the provider cache hint is lost. + ProviderSurface::AnthropicMessages => { + object + .get("stop_sequences") + .is_some_and(|stop| !is_string_array(stop)) + || object + .get("system") + .and_then(Json::as_array) + .is_some_and(|blocks| blocks.iter().any(lossy_system_block)) + } + ProviderSurface::OpenAIResponses => false, + } +} + +/// Whether `value` is an array whose items are all strings — the only `stop` / +/// `stop_sequences` form the normalized types keep faithfully. +fn is_string_array(value: &Json) -> bool { + value + .as_array() + .is_some_and(|items| items.iter().all(Json::is_string)) +} + +/// Whether an Anthropic `system` content block carries anything the decode's +/// text-flattening would lose. +fn lossy_system_block(block: &Json) -> bool { + let Some(fields) = block.as_object() else { + return true; + }; + fields.get("type").and_then(Json::as_str) != Some("text") + || !fields.get("text").is_some_and(Json::is_string) + || fields + .keys() + .any(|key| !matches!(key.as_str(), "type" | "text" | "cache_control")) +} + +/// Whether the raw request body carries a non-empty `messages` (chat) or `input` +/// (Responses) array — the content a decode is expected to preserve. +fn raw_has_messages(request: &LlmRequest) -> bool { + let non_empty_array = |key: &str| { + request + .content + .get(key) + .and_then(Json::as_array) + .is_some_and(|items| !items.is_empty()) + }; + non_empty_array("messages") || non_empty_array("input") +} + +fn request_temperature(content: &Json) -> Option { + content + .get("temperature") + .and_then(Json::as_f64) + .or_else(|| { + content + .pointer("/params/temperature") + .and_then(Json::as_f64) + }) +} + +/// Keeps only allowlisted headers, matched case-insensitively and emitted with +/// lowercased names so header-case noise does not change the key. +fn allowlisted_headers(headers: &Map, allowlist: &[String]) -> Map { + let mut kept = Map::new(); + for name in allowlist { + for (header_name, value) in headers { + if header_name.eq_ignore_ascii_case(name) { + kept.insert(header_name.to_ascii_lowercase(), value.clone()); + } + } + } + kept +} + +/// Renumbers tool-call IDs consistently so random IDs do not cause misses. +/// +/// Tool calls and their results carry random IDs; what matters is which result +/// pairs with which call, not the literal value. We map IDs to `tcid_0`, +/// `tcid_1`, … in first-seen order across `messages`, rewriting both +/// `tool_calls[].id` and `tool_call_id`. Shapes without these keys are left +/// untouched (this only affects hit-rate, never correctness). +fn normalize_tool_call_ids(body: &mut Map) { + let Some(messages) = body.get_mut("messages").and_then(Json::as_array_mut) else { + return; + }; + + let mut mapping: Map = Map::new(); + for message in messages.iter_mut() { + let Some(object) = message.as_object_mut() else { + continue; + }; + if let Some(tool_calls) = object.get_mut("tool_calls").and_then(Json::as_array_mut) { + for call in tool_calls.iter_mut() { + if let Some(id_value) = call.get_mut("id") { + rewrite_id(id_value, &mut mapping); + } + } + } + if let Some(id_value) = object.get_mut("tool_call_id") { + rewrite_id(id_value, &mut mapping); + } + } +} + +/// Rewrites a single JSON string id in place to its stable, first-seen mapping. +fn rewrite_id(id_value: &mut Json, mapping: &mut Map) { + let Some(original) = id_value.as_str().map(str::to_string) else { + return; + }; + let stable = match mapping.get(&original) { + Some(Json::String(existing)) => existing.clone(), + _ => { + let stable = format!("tcid_{}", mapping.len()); + mapping.insert(original, Json::String(stable.clone())); + stable + } + }; + *id_value = Json::String(stable); +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::acg::canonicalize::{canonicalize_value, sha256_hex}; + + #[test] + fn fingerprint_matches_canonicalize_then_hash() { + // `fingerprint` streams canonical bytes into the hasher; it must produce + // byte-for-byte the same digest as materializing the canonical string + // first, or every existing key would silently change. + let doc = json!({ + "z": {"nested": [3, 1.5, -0.0, 1e21]}, + "a": "höllo \u{1F600} wörld", + "codec": null, + "headers": {}, + "body": {"messages": [{"role": "user", "content": "hi"}]} + }); + assert_eq!( + fingerprint(&doc).unwrap(), + sha256_hex(&canonicalize_value(&doc).unwrap()), + ); + } + + fn request(content: Json) -> LlmRequest { + LlmRequest { + headers: Map::new(), + content, + } + } + + fn key_of(provider: &str, request: &LlmRequest, config: &ResponseCacheConfig) -> String { + match build_cache_key(provider, request, config) { + KeyOutcome::Key(key) => key, + other => panic!("expected a key, got {other:?}"), + } + } + + #[test] + fn field_order_and_whitespace_do_not_change_the_key() { + let config = ResponseCacheConfig::default(); + let first = request( + json!({"model": "m", "messages": [{"role": "user", "content": "hi"}], "tool_choice": "auto"}), + ); + let second = request( + json!({"tool_choice": "auto", "messages": [{"content": "hi", "role": "user"}], "model": "m"}), + ); + assert_eq!( + key_of("openai", &first, &config), + key_of("openai", &second, &config) + ); + } + + #[test] + fn skiplist_fields_do_not_change_the_key() { + let config = ResponseCacheConfig::default(); + let base = request(json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]})); + let noisy = request(json!({ + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "stream": true, + "user": "abc", + "metadata": {"trace": "xyz"} + })); + assert_eq!( + key_of("openai", &base, &config), + key_of("openai", &noisy, &config) + ); + } + + #[test] + fn namespace_and_provider_separate_keys() { + let request = + request(json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]})); + let ns_a = ResponseCacheConfig { + namespace: "a".to_string(), + ..ResponseCacheConfig::default() + }; + let ns_b = ResponseCacheConfig { + namespace: "b".to_string(), + ..ResponseCacheConfig::default() + }; + assert_ne!( + key_of("openai", &request, &ns_a), + key_of("openai", &request, &ns_b) + ); + // Same namespace, different provider/family also separates. + assert_ne!( + key_of("openai", &request, &ns_a), + key_of("anthropic", &request, &ns_a) + ); + } + + #[test] + fn random_tool_call_ids_are_normalized_to_one_key() { + let config = ResponseCacheConfig::default(); + let make = |call_id: &str| { + request(json!({ + "model": "m", + "messages": [ + {"role": "user", "content": "weather?"}, + {"role": "assistant", "tool_calls": [{"id": call_id, "type": "function", "function": {"name": "get_weather", "arguments": "{}"}}]}, + {"role": "tool", "tool_call_id": call_id, "content": "sunny"} + ] + })) + }; + assert_eq!( + key_of("openai", &make("call_RANDOM_1"), &config), + key_of("openai", &make("call_RANDOM_2"), &config), + "random tool-call ids must not change the key" + ); + } + + #[test] + fn raw_params_objects_do_not_collide_with_typed_caps() { + // A top-level `params` object lands in the flattened `extra` and + // overwrites the typed field on serialization — the token caps would + // vanish from the key. + let config = ResponseCacheConfig::default(); + let make = |cap: u64| { + request(json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "write"}], + "max_tokens": cap, + "params": {"vendor": "x"} + })) + }; + assert_ne!( + key_of("openai", &make(1), &config), + key_of("openai", &make(100), &config), + "a raw params object must not erase the token cap from the key" + ); + } + + #[test] + fn wrong_typed_generation_scalars_do_not_collide() { + let config = ResponseCacheConfig::default(); + let plain = request(json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}] + })); + // A string temperature is dropped by the typed extraction and excluded + // from `extra`; it must not key like the temperature-less request. + let string_temp = request(json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "temperature": "0.9" + })); + assert_ne!( + key_of("openai", &plain, &config), + key_of("openai", &string_temp, &config), + "a wrong-typed temperature must separate keys" + ); + // Same for a float token cap (as_u64 yields None). + let float_cap = request(json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 100.0 + })); + assert_ne!( + key_of("openai", &plain, &config), + key_of("openai", &float_cap, &config), + "a float token cap must separate keys" + ); + } + + #[test] + fn unmodeled_tool_choice_fields_do_not_collide() { + let config = ResponseCacheConfig::default(); + let make = |strict: bool| { + request(json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "call it"}], + "tool_choice": {"type": "function", + "function": {"name": "docs_lookup", "strict": strict}} + })) + }; + assert_ne!( + key_of("openai", &make(true), &config), + key_of("openai", &make(false), &config), + "unmodeled tool_choice fields must separate keys" + ); + } + + #[test] + fn stateful_conversation_and_default_store_bypass() { + let config = ResponseCacheConfig::default(); + // Server-side conversation state. + let with_conversation = request(json!({ + "model": "gpt-4o", "input": "summarize", "store": false, + "conversation": "conv_1" + })); + assert_eq!( + build_cache_key("openai", &with_conversation, &config), + KeyOutcome::Bypass("stateful_conversation") + ); + // Responses persists by default: no explicit `store: false` opt-out + // means the call is stateful. + let default_store = request(json!({"model": "gpt-4o", "input": "hello"})); + assert_eq!( + build_cache_key("openai", &default_store, &config), + KeyOutcome::Bypass("stateful_store") + ); + let opted_out = request(json!({"model": "gpt-4o", "input": "hello", "store": false})); + assert!(matches!( + build_cache_key("openai", &opted_out, &config), + KeyOutcome::Key(_) + )); + } + + #[test] + fn toggle_off_caches_only_explicitly_deterministic_requests() { + let skip = ResponseCacheConfig { + cache_nondeterministic: false, + ..ResponseCacheConfig::default() + }; + // Absent temperature: providers default to positive sampling. + let absent = request(json!({"model": "m", "messages": []})); + assert_eq!( + build_cache_key("openai", &absent, &skip), + KeyOutcome::Bypass("nondeterministic_temperature") + ); + // Explicitly pinned deterministic stays cacheable. + let pinned = request(json!({"model": "m", "messages": [], "temperature": 0.0})); + assert!(matches!( + build_cache_key("openai", &pinned, &skip), + KeyOutcome::Key(_) + )); + } + + #[test] + fn null_bodies_bypass_the_cache() { + // The gateway parses unparseable upstream bodies to `null`; every such + // request would share one key, so they are never cacheable. + assert_eq!( + build_cache_key( + "openai", + &request(Json::Null), + &ResponseCacheConfig::default() + ), + KeyOutcome::Bypass("unparseable_body") + ); + } + + #[test] + fn stateful_responses_calls_bypass() { + let config = ResponseCacheConfig::default(); + let with_store = request(json!({"model": "m", "messages": [], "store": true})); + assert_eq!( + build_cache_key("openai", &with_store, &config), + KeyOutcome::Bypass("stateful_store") + ); + let with_prev = + request(json!({"model": "m", "messages": [], "previous_response_id": "resp_1"})); + assert_eq!( + build_cache_key("openai", &with_prev, &config), + KeyOutcome::Bypass("stateful_previous_response_id") + ); + // A truthy non-boolean `store` must still bypass (it is otherwise stripped + // from the key), while `store: false` stays cacheable. + let with_truthy = request(json!({"model": "m", "messages": [], "store": "true"})); + assert_eq!( + build_cache_key("openai", &with_truthy, &config), + KeyOutcome::Bypass("stateful_store") + ); + let not_stored = request(json!({"model": "m", "messages": [], "store": false})); + assert!(matches!( + build_cache_key("openai", ¬_stored, &config), + KeyOutcome::Key(_) + )); + // The gate runs pre-decode, so a cleanly decodable body bypasses too. + let decodable = request(json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "store": true + })); + assert_eq!( + build_cache_key("openai", &decodable, &config), + KeyOutcome::Bypass("stateful_store") + ); + } + + #[test] + fn nondeterministic_calls_bypass_only_when_disabled() { + let request = request(json!({"model": "m", "messages": [], "temperature": 0.7})); + let skip = ResponseCacheConfig { + cache_nondeterministic: false, + ..ResponseCacheConfig::default() + }; + assert_eq!( + build_cache_key("openai", &request, &skip), + KeyOutcome::Bypass("nondeterministic_temperature") + ); + // Default keeps caching temperature > 0 calls. + assert!(matches!( + build_cache_key("openai", &request, &ResponseCacheConfig::default()), + KeyOutcome::Key(_) + )); + } + + #[test] + fn chat_shaped_requests_key_on_the_detected_decode() { + // A request the OpenAI-chat codec can decode: detection must pick the + // chat surface and the keyed body must be the decode, not the raw body + // — a silent decode regression cannot hide behind identical keys. + let request = request(json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "temperature": 0.0 + })); + let (body, effective_codec) = resolved_body("openai", &request); + assert_eq!(effective_codec, Some("openai_chat")); + assert_ne!( + body, request.content, + "the keyed body must be the normalized decode, not the raw body" + ); + } + + #[test] + fn undetectable_shape_falls_back_to_raw_keying() { + // No `messages`/`input`/`system` top-level key: no surface detects, so + // the raw body is fingerprinted — still a usable, stable key. + let config = ResponseCacheConfig::default(); + let request = request(json!({"model": "m", "prompt": "hi"})); + let (body, effective_codec) = resolved_body("openai", &request); + assert_eq!(effective_codec, None, "nothing must detect this shape"); + assert_eq!(body, request.content, "the raw body is keyed as-is"); + let first = key_of("openai", &request, &config); + let second = key_of("openai", &request, &config); + assert_eq!(first, second, "raw-fallback keys must be stable"); + } + + #[test] + fn dual_token_caps_do_not_collide() { + // With both caps present the decode keeps one; requests differing in + // the other must not merge. + let config = ResponseCacheConfig::default(); + let low = request(json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "write a story"}], + "max_completion_tokens": 100, + "max_tokens": 1 + })); + let high = request(json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "write a story"}], + "max_completion_tokens": 100, + "max_tokens": 9999 + })); + assert_ne!( + key_of("openai", &low, &config), + key_of("openai", &high, &config), + "requests carrying both token caps must key on both" + ); + } + + #[test] + fn anthropic_system_block_metadata_does_not_collide() { + // System content blocks are flattened to their text on decode; block + // fields beyond the provider cache hint must not vanish from the key. + let config = ResponseCacheConfig::default(); + let make = |marker: u64| { + request(json!({ + "model": "claude-sonnet-4", + "max_tokens": 64, + "system": [{"type": "text", "text": "Use policy X", "priority": marker}], + "messages": [{"role": "user", "content": "Answer"}] + })) + }; + assert_ne!( + key_of("anthropic", &make(1), &config), + key_of("anthropic", &make(2), &config), + "system-block metadata must separate keys" + ); + } + + #[test] + fn unmodeled_tool_fields_do_not_collide() { + // `FunctionDefinition` has no unknown-field catch-all, so a field like + // OpenAI's `function.strict` is silently dropped on decode; requests + // differing in it must not merge. + let config = ResponseCacheConfig::default(); + let make = |strict: bool| { + request(json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "look it up"}], + "tools": [{"type": "function", "function": { + "name": "docs_lookup", + "description": "Look up docs", + "parameters": {"type": "object", "properties": {}}, + "strict": strict + }}] + })) + }; + assert_ne!( + key_of("openai", &make(true), &config), + key_of("openai", &make(false), &config), + "unmodeled tool fields must separate keys" + ); + } + + #[test] + fn cleanly_modeled_tools_still_key_on_the_decode() { + // The tools round-trip guard must not disable normalized keying for + // tools the normalized types represent fully. + let request = request(json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "look it up"}], + "tools": [{"type": "function", "function": { + "name": "docs_lookup", + "description": "Look up docs", + "parameters": {"type": "object", "properties": {}} + }}] + })); + let (body, effective_codec) = resolved_body("openai", &request); + assert_eq!( + effective_codec, + Some("openai_chat"), + "cleanly modeled tools must keep the decode" + ); + // This minimal fixture round-trips byte-identically — which is exactly + // what the tools guard verifies before trusting the decode. + assert_eq!(body["tools"], request.content["tools"]); + } + + #[test] + fn system_less_anthropic_body_uses_the_provider_hint() { + // An Anthropic request without a top-level `system` is shape-identical + // to OpenAI Chat; the provider-name hint must resolve it. + let request = request(json!({ + "model": "claude-sonnet-4", + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}] + })); + let (_, hinted) = resolved_body("anthropic", &request); + assert_eq!( + hinted, + Some("anthropic_messages"), + "the hint must detect the anthropic surface for a system-less body" + ); + let (_, unhinted) = resolved_body("openai", &request); + assert_eq!( + unhinted, + Some("openai_chat"), + "without the anthropic hint the same shape reads as chat" + ); + } + + #[test] + fn unmodeled_message_fields_do_not_collide() { + // The normalized message types are closed, so an assistant field they + // do not model — the deprecated `function_call`, `refusal` — decodes + // to nothing; conversations differing only there must not merge. + let config = ResponseCacheConfig::default(); + let make = |arguments: &str| { + request(json!({ + "model": "gpt-4o", + "messages": [ + {"role": "assistant", "content": null, + "function_call": {"name": "lookup", "arguments": arguments}}, + {"role": "user", "content": "Continue"} + ] + })) + }; + assert_ne!( + key_of("openai", &make("{\"q\":\"alpha\"}"), &config), + key_of("openai", &make("{\"q\":\"beta\"}"), &config), + "unmodeled message fields must separate keys" + ); + } + + #[test] + fn non_array_stop_forms_do_not_collide_with_a_stopless_request() { + // Only an array of strings decodes faithfully; every other `stop` + // form is silently dropped and must stay raw-keyed. + let config = ResponseCacheConfig::default(); + let without = request(json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "count: one END two"}] + })); + for stop in [json!("END"), json!(["END", 7]), json!({}), json!(7)] { + let with_stop = request(json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "count: one END two"}], + "stop": stop + })); + assert_ne!( + key_of("openai", &with_stop, &config), + key_of("openai", &without, &config), + "a malformed stop ({stop}) must not share a key with a stopless request" + ); + } + } + + #[test] + fn null_text_system_block_does_not_collide_with_no_system() { + // A `text: null` block decodes to no system prompt at all; it must + // not share a key with a request that genuinely has no system. + let config = ResponseCacheConfig::default(); + let malformed = request(json!({ + "model": "claude-sonnet-4", + "max_tokens": 64, + "system": [{"type": "text", "text": null}], + "messages": [{"role": "user", "content": "Answer"}] + })); + let clean = request(json!({ + "model": "claude-sonnet-4", + "max_tokens": 64, + "messages": [{"role": "user", "content": "Answer"}] + })); + assert_ne!( + key_of("anthropic", &malformed, &config), + key_of("anthropic", &clean, &config), + "a null-text system block must not key like an absent system" + ); + } +} diff --git a/crates/adaptive/src/response_cache/mark.rs b/crates/adaptive/src/response_cache/mark.rs new file mode 100644 index 000000000..5f01c723f --- /dev/null +++ b/crates/adaptive/src/response_cache/mark.rs @@ -0,0 +1,204 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The `response_cache` mark: every cache decision, its metadata, and the +//! savings a reuse reports. + +use nemo_relay::api::scope::{EmitMarkEventParams, event}; +use nemo_relay::codec::model_pricing::estimate_cost_for_provider; +use nemo_relay::codec::response::Usage; +use serde_json::{Map, Value as Json, json}; + +use crate::response_cache::store::CacheEntry; + +/// Mark-event name emitted on every cache decision. +pub const RESPONSE_CACHE_MARK: &str = "response_cache"; + +/// Pulls saved token count and cost out of a stored entry (the aggregate +/// response object — buffered and streaming both store this shape). +/// +/// Cost precedence mirrors the runtime's own `attach_estimated_cost`: a +/// provider-reported cost on the body wins; otherwise the cost is derived from +/// the token counts via the active pricing resolver (so the dollar figure works +/// for the common case where providers report tokens, not cost — provided a +/// pricing catalog is configured for the model). `saved_tokens` is available +/// whenever the response carries usage. +pub(crate) fn savings_from(entry: &CacheEntry) -> (Option, Option) { + let response = &entry.response; + let usage = response + .get("usage") + .or_else(|| response.get("token_usage")); + let tokens = usage.and_then(|usage| { + usage + .get("total_tokens") + .and_then(Json::as_u64) + .or_else(|| { + let prompt = usage.get("prompt_tokens").and_then(Json::as_u64)?; + let completion = usage.get("completion_tokens").and_then(Json::as_u64)?; + Some(prompt.saturating_add(completion)) + }) + .or_else(|| { + // Anthropic Messages usage uses input_tokens / output_tokens. + let input = usage.get("input_tokens").and_then(Json::as_u64)?; + let output = usage.get("output_tokens").and_then(Json::as_u64)?; + Some(input.saturating_add(output)) + }) + }); + // 1. Provider-reported cost on the body wins. + let body_cost = usage + .and_then(|usage| { + usage + .get("cost_usd") + .and_then(Json::as_f64) + .or_else(|| usage.pointer("/cost/total").and_then(Json::as_f64)) + }) + .or_else(|| response.get("cost_usd").and_then(Json::as_f64)); + // 2. Otherwise price the token counts with the active resolver. + let cost = body_cost.or_else(|| { + let model = entry.model_name.as_deref()?; + let usage: Usage = serde_json::from_value(usage?.clone()).ok()?; + estimate_cost_for_provider(entry.provider_name.as_deref(), model, &usage) + .and_then(|estimate| estimate.total) + }); + (tokens, cost) +} + +pub(crate) struct CacheMark<'a> { + status: &'a str, + reason: Option<&'a str>, + backend: &'a str, + key_hash: Option<&'a str>, + age_ms: Option, + ttl_ms: Option, + saved_tokens: Option, + saved_cost_usd: Option, +} + +impl<'a> CacheMark<'a> { + pub(crate) fn new(status: &'a str, backend: &'a str) -> Self { + Self { + status, + reason: None, + backend, + key_hash: None, + age_ms: None, + ttl_ms: None, + saved_tokens: None, + saved_cost_usd: None, + } + } + + pub(crate) fn reason(mut self, reason: &'a str) -> Self { + self.reason = Some(reason); + self + } + + pub(crate) fn key_hash(mut self, key_hash: &'a str) -> Self { + self.key_hash = Some(key_hash); + self + } + + pub(crate) fn age_ms(mut self, age_ms: u64) -> Self { + self.age_ms = Some(age_ms); + self + } + + pub(crate) fn ttl_ms(mut self, ttl_ms: u64) -> Self { + self.ttl_ms = Some(ttl_ms); + self + } + + pub(crate) fn savings(mut self, tokens: Option, cost: Option) -> Self { + self.saved_tokens = tokens; + self.saved_cost_usd = cost; + self + } +} + +/// Emits the `response_cache` mark. Only the key fingerprint is ever recorded — +/// never raw prompts, answers, or credentials. Fails open. +pub(crate) fn emit_cache_mark(mark: CacheMark<'_>) { + let mut metadata = Map::new(); + metadata.insert( + "nemo_relay.response_cache.surface".to_string(), + json!("llm"), + ); + metadata.insert( + "nemo_relay.response_cache.backend".to_string(), + json!(mark.backend), + ); + if let Some(reason) = mark.reason { + metadata.insert( + "nemo_relay.response_cache.reason".to_string(), + json!(reason), + ); + } + if let Some(key_hash) = mark.key_hash { + metadata.insert( + "nemo_relay.response_cache.key_hash".to_string(), + json!(key_hash), + ); + } + if let Some(age_ms) = mark.age_ms { + metadata.insert( + "nemo_relay.response_cache.age_ms".to_string(), + json!(age_ms), + ); + } + if let Some(ttl_ms) = mark.ttl_ms { + metadata.insert( + "nemo_relay.response_cache.ttl_ms".to_string(), + json!(ttl_ms), + ); + } + if let Some(saved_tokens) = mark.saved_tokens { + metadata.insert( + "nemo_relay.response_cache.saved_tokens".to_string(), + json!(saved_tokens), + ); + } + if let Some(saved_cost_usd) = mark.saved_cost_usd { + metadata.insert( + "nemo_relay.response_cache.saved_cost_usd".to_string(), + json!(saved_cost_usd), + ); + } + + let _ = event( + EmitMarkEventParams::builder() + .name(RESPONSE_CACHE_MARK) + .data(json!({ "status": mark.status })) + .metadata(Json::Object(metadata)) + .build(), + ); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn savings_from_counts_anthropic_input_output_tokens() { + // Anthropic Messages reports usage as input_tokens/output_tokens, not the + // OpenAI total_tokens / prompt_tokens+completion_tokens. Savings must still + // be counted, or every Claude hit reports zero avoided tokens. + let entry = CacheEntry { + response: json!({ + "content": [{"type": "text", "text": "hi"}], + "usage": {"input_tokens": 100, "output_tokens": 25} + }), + created_unix_ms: 0, + expires_unix_ms: 0, + schema_version: 1, + key_hash: "sha256:x".to_string(), + model_name: Some("claude-x".to_string()), + provider_name: Some("anthropic.messages".to_string()), + }; + let (tokens, _cost) = savings_from(&entry); + assert_eq!( + tokens, + Some(125), + "anthropic input+output tokens must be counted for savings" + ); + } +} diff --git a/crates/adaptive/src/response_cache/mod.rs b/crates/adaptive/src/response_cache/mod.rs new file mode 100644 index 000000000..e1bfaa62f --- /dev/null +++ b/crates/adaptive/src/response_cache/mod.rs @@ -0,0 +1,25 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Opt-in LLM response cache (exact-match): a feature of the adaptive plugin, +//! configured through [`crate::config::AdaptiveConfig::response_cache`]. +//! +//! [`intercept`] holds the execution intercepts and storage rules, [`key`] the +//! cache-key derivation, [`store`] the backends, [`replay`] the streaming +//! replay, and [`mark`] the observability surface. + +pub mod config; +pub(crate) mod intercept; +pub mod key; +pub(crate) mod mark; +pub(crate) mod replay; +pub mod store; + +pub use crate::config::ResponseCacheConfig; +pub use crate::response_cache::config::{BackendConfig, KEY_STRATEGY_EXACT_REQUEST}; +pub(crate) use crate::response_cache::intercept::{make_intercept, make_stream_intercept}; +pub use crate::response_cache::mark::RESPONSE_CACHE_MARK; +pub(crate) use crate::response_cache::store::build_store; +pub use crate::response_cache::store::{ + CacheEntry, CacheStore, InMemoryCacheStore, check_backend_health, +}; diff --git a/crates/adaptive/src/response_cache/replay.rs b/crates/adaptive/src/response_cache/replay.rs new file mode 100644 index 000000000..501b488cb --- /dev/null +++ b/crates/adaptive/src/response_cache/replay.rs @@ -0,0 +1,325 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Replay of a stored aggregate as provider-native streaming chunks. + +use nemo_relay::api::runtime::LlmJsonStream; +use nemo_relay::error::FlowError; +use serde_json::{Map, Value as Json, json}; + +/// Replays a stored aggregate response as a stream of **provider-native chunks**. +/// +/// A strict streaming client parses only its provider's wire chunks (Anthropic +/// `message_start → content_block_delta → … → message_stop`, OpenAI Chat +/// `chat.completion.chunk` deltas, OpenAI Responses lifecycle events) — replaying +/// the aggregate as one frame breaks such clients even though the body is correct. +/// The surface is detected from the stored aggregate's own shape (the codec +/// finalizer's output, or a buffered body of the same shape), so no codec handle +/// is needed at the call sites. An unrecognized shape falls back to the original +/// single-frame replay rather than dropping the answer. +pub(crate) fn replay_aggregate(response: Json) -> LlmJsonStream { + let chunks = synthesize_replay_chunks(&response).unwrap_or_else(|| vec![response]); + Box::pin(tokio_stream::iter( + chunks.into_iter().map(Ok::), + )) +} + +/// Detects the aggregate's provider surface from its shape and synthesizes the +/// native chunk sequence for it. `None` when the shape is not recognized. +/// +/// Detection prefers the explicit discriminators (`type: "message"`, +/// `object: "chat.completion"` / `"response"`) but also accepts the structural +/// shape (`choices[].message`, `output[]` + `status`) — buffered bodies from +/// OpenAI-compatible providers sometimes omit `object`, and a buffered entry can +/// be served to a streaming caller. +fn synthesize_replay_chunks(aggregate: &Json) -> Option> { + let object = aggregate.get("object").and_then(Json::as_str); + if aggregate.get("type").and_then(Json::as_str) == Some("message") { + return Some(synthesize_anthropic_chunks(aggregate)); + } + let chat_shaped = object == Some("chat.completion") + || aggregate + .pointer("/choices/0/message") + .is_some_and(Json::is_object); + if chat_shaped { + return Some(synthesize_chat_chunks(aggregate)); + } + let responses_shaped = object == Some("response") + || (aggregate.get("output").is_some_and(Json::is_array) + && aggregate.get("status").is_some()); + if responses_shaped { + return Some(synthesize_responses_chunks(aggregate)); + } + None +} + +/// Anthropic Messages: `message_start`, then per content block a +/// `content_block_start`/delta/`content_block_stop` triple (text and `tool_use` +/// stream their payload as a native delta; other block types ship complete at +/// start, as the live API does), then `message_delta` carrying the stored usage +/// verbatim (the collector replaces usage wholesale, so the reassembled aggregate +/// round-trips), then `message_stop`. +fn synthesize_anthropic_chunks(aggregate: &Json) -> Vec { + let mut start_message = Map::new(); + start_message.insert("type".to_string(), json!("message")); + for key in ["id", "role", "model"] { + if let Some(value) = aggregate.get(key) { + start_message.insert(key.to_string(), value.clone()); + } + } + start_message.insert("content".to_string(), json!([])); + start_message.insert("stop_reason".to_string(), Json::Null); + start_message.insert("stop_sequence".to_string(), Json::Null); + let input_tokens = aggregate + .pointer("/usage/input_tokens") + .cloned() + .unwrap_or(json!(0)); + start_message.insert( + "usage".to_string(), + json!({"input_tokens": input_tokens, "output_tokens": 0}), + ); + let mut chunks = vec![json!({"type": "message_start", "message": start_message})]; + + let blocks = aggregate + .get("content") + .and_then(Json::as_array) + .cloned() + .unwrap_or_default(); + for (index, block) in blocks.iter().enumerate() { + let block_type = block.get("type").and_then(Json::as_str).unwrap_or(""); + match block_type { + "text" => { + let mut skeleton = block.clone(); + let text = skeleton + .as_object_mut() + .and_then(|map| map.insert("text".to_string(), json!(""))) + .and_then(|old| old.as_str().map(str::to_string)) + .unwrap_or_default(); + chunks.push(json!({"type": "content_block_start", "index": index, + "content_block": skeleton})); + chunks.push(json!({"type": "content_block_delta", "index": index, + "delta": {"type": "text_delta", "text": text}})); + } + "tool_use" => { + let mut skeleton = block.clone(); + let input = skeleton + .as_object_mut() + .and_then(|map| map.insert("input".to_string(), json!({}))) + .unwrap_or(json!({})); + let partial = serde_json::to_string(&input).unwrap_or_else(|_| "{}".to_string()); + chunks.push(json!({"type": "content_block_start", "index": index, + "content_block": skeleton})); + chunks.push(json!({"type": "content_block_delta", "index": index, + "delta": {"type": "input_json_delta", "partial_json": partial}})); + } + // Thinking, server_tool_use, and other block types ship complete at + // start (the collector keeps the skeleton verbatim). + _ => { + chunks.push(json!({"type": "content_block_start", "index": index, + "content_block": block.clone()})); + } + } + chunks.push(json!({"type": "content_block_stop", "index": index})); + } + + let mut delta = Map::new(); + if let Some(reason) = aggregate.get("stop_reason") { + delta.insert("stop_reason".to_string(), reason.clone()); + } + if let Some(sequence) = aggregate.get("stop_sequence") { + delta.insert("stop_sequence".to_string(), sequence.clone()); + } + let mut message_delta = Map::new(); + message_delta.insert("type".to_string(), json!("message_delta")); + message_delta.insert("delta".to_string(), Json::Object(delta)); + if let Some(usage) = aggregate.get("usage") { + message_delta.insert("usage".to_string(), usage.clone()); + } + chunks.push(Json::Object(message_delta)); + chunks.push(json!({"type": "message_stop"})); + chunks +} + +/// OpenAI Chat: per choice a role delta, a content delta (when the message has +/// text), a `tool_calls` delta (full arguments in one fragment — spec-valid), and +/// a finish chunk; then a final usage-bearing chunk with empty `choices` (the +/// `include_usage` wire shape). Top-level id/created/model ride on every chunk. +fn synthesize_chat_chunks(aggregate: &Json) -> Vec { + let base = |choices: Json| -> Json { + let mut chunk = Map::new(); + for key in ["id", "created", "model"] { + if let Some(value) = aggregate.get(key) { + chunk.insert(key.to_string(), value.clone()); + } + } + chunk.insert("object".to_string(), json!("chat.completion.chunk")); + chunk.insert("choices".to_string(), choices); + Json::Object(chunk) + }; + let mut chunks = Vec::new(); + let choices = aggregate + .get("choices") + .and_then(Json::as_array) + .cloned() + .unwrap_or_default(); + for (position, choice) in choices.iter().enumerate() { + let index = choice + .get("index") + .and_then(Json::as_u64) + .unwrap_or(position as u64); + let message = choice.get("message").cloned().unwrap_or(json!({})); + if let Some(role) = message.get("role") { + chunks.push(base( + json!([{"index": index, "delta": {"role": role}, "finish_reason": null}]), + )); + } + if let Some(content) = message.get("content").and_then(Json::as_str) + && !content.is_empty() + { + chunks.push(base( + json!([{"index": index, "delta": {"content": content}, "finish_reason": null}]), + )); + } + if let Some(tool_calls) = message.get("tool_calls").and_then(Json::as_array) { + let deltas: Vec = tool_calls + .iter() + .enumerate() + .map(|(call_index, call)| { + let mut delta = call.clone(); + if let Some(map) = delta.as_object_mut() { + map.entry("index".to_string()) + .or_insert(json!(call_index as u64)); + } + delta + }) + .collect(); + chunks.push(base( + json!([{"index": index, "delta": {"tool_calls": deltas}, "finish_reason": null}]), + )); + } + let finish = choice.get("finish_reason").cloned().unwrap_or(Json::Null); + chunks.push(base( + json!([{"index": index, "delta": {}, "finish_reason": finish}]), + )); + } + if let Some(usage) = aggregate.get("usage") { + let mut usage_chunk = base(json!([])); + if let Some(map) = usage_chunk.as_object_mut() { + map.insert("usage".to_string(), usage.clone()); + } + chunks.push(usage_chunk); + } + chunks +} + +/// OpenAI Responses: a `response.created` snapshot, one `response.output_item.done` +/// per output item, and a `response.completed` carrying the full stored aggregate +/// (the collector keeps the last snapshot wholesale, so reassembly is exact). +fn synthesize_responses_chunks(aggregate: &Json) -> Vec { + let mut created = Map::new(); + for key in ["id", "object", "model"] { + if let Some(value) = aggregate.get(key) { + created.insert(key.to_string(), value.clone()); + } + } + created.insert("status".to_string(), json!("in_progress")); + let mut chunks = vec![json!({"type": "response.created", "response": created})]; + if let Some(items) = aggregate.get("output").and_then(Json::as_array) { + for (index, item) in items.iter().enumerate() { + chunks.push(json!({"type": "response.output_item.done", + "output_index": index, "item": item.clone()})); + } + } + chunks.push(json!({"type": "response.completed", "response": aggregate.clone()})); + chunks +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::response_cache::intercept::streaming_codec_for_surface; + use nemo_relay::codec::resolve::ProviderSurface; + + /// The replay synthesizer must emit chunks the provider's own streaming codec + /// reassembles into EXACTLY the stored aggregate — the property that makes a + /// replayed hit indistinguishable from a live stream to a strict client. + #[test] + fn replay_chunks_roundtrip_through_the_codecs() { + let anthropic = json!({"id": "msg_1", "type": "message", "role": "assistant", + "model": "m", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "tool_use", "id": "t1", "name": "get", "input": {"q": "x"}} + ], + "stop_reason": "tool_use", + "usage": {"input_tokens": 9, "output_tokens": 3}}); + let chat = json!({"id": "c1", "object": "chat.completion", "created": 1, + "model": "m", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop"}], + "usage": {"prompt_tokens": 9, "completion_tokens": 3, "total_tokens": 12}}); + let responses = json!({"id": "r1", "object": "response", "status": "completed", + "model": "m", + "output": [{"type": "message", "role": "assistant", + "content": [{"type": "output_text", "text": "hello"}]}], + "usage": {"input_tokens": 9, "output_tokens": 3, "total_tokens": 12}}); + for (aggregate, surface, codec_name) in [ + ( + anthropic, + ProviderSurface::AnthropicMessages, + "anthropic_messages", + ), + (chat, ProviderSurface::OpenAIChat, "openai_chat"), + ( + responses, + ProviderSurface::OpenAIResponses, + "openai_responses", + ), + ] { + let codec = streaming_codec_for_surface(surface); + let chunks = + synthesize_replay_chunks(&aggregate).expect("aggregate shape must be recognized"); + assert!( + chunks.len() > 1, + "{codec_name}: a native replay must be a chunk sequence, not one frame" + ); + let mut collect = codec.collector(); + for chunk in &chunks { + collect(chunk.clone()).expect("codec must accept its own native chunk shape"); + } + let reassembled = codec.finalizer()(); + assert_eq!( + reassembled, aggregate, + "{codec_name}: replayed chunks must reassemble to the stored aggregate" + ); + } + } + + /// A chat replay must carry tool calls in the streaming delta shape a client + /// can accumulate (full arguments in one fragment is spec-valid). + #[test] + fn chat_replay_streams_tool_calls_as_deltas() { + let aggregate = json!({"id": "c1", "object": "chat.completion", + "choices": [{"index": 0, "message": {"role": "assistant", "content": null, + "tool_calls": [{"id": "call1", "type": "function", + "function": {"name": "f", "arguments": "{\"a\":1}"}}]}, + "finish_reason": "tool_calls"}]}); + let chunks = synthesize_replay_chunks(&aggregate).expect("chat shape"); + let tool_delta = chunks + .iter() + .find(|chunk| chunk.pointer("/choices/0/delta/tool_calls").is_some()) + .expect("a tool_calls delta chunk must be synthesized"); + assert_eq!( + tool_delta.pointer("/choices/0/delta/tool_calls/0/function/arguments"), + Some(&json!("{\"a\":1}")) + ); + } + + /// An unknown aggregate shape must fall back to the single-frame replay, + /// never drop the answer. + #[test] + fn replay_of_an_unknown_shape_falls_back_to_one_frame() { + assert!(synthesize_replay_chunks(&json!({"weird": true})).is_none()); + assert!(synthesize_replay_chunks(&json!("bare string")).is_none()); + } +} diff --git a/crates/adaptive/src/response_cache/store.rs b/crates/adaptive/src/response_cache/store.rs new file mode 100644 index 000000000..021fb3684 --- /dev/null +++ b/crates/adaptive/src/response_cache/store.rs @@ -0,0 +1,697 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! General-purpose cache store with expiry for the response cache. +//! +//! This is intentionally separate from the adaptive learning-data storage +//! ([`crate::storage`]), which is shaped for specific learning records rather +//! than a "store any value with a TTL" cache. The trait keeps the same +//! object-safe, boxed-future style so a Redis backend can drop in later behind +//! the existing `redis-backend` feature. + +use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use serde::{Deserialize, Serialize}; +use serde_json::Value as Json; + +use crate::config::ResponseCacheConfig; +use crate::error::{AdaptiveError, Result}; + +/// Boxed, `Send` future returned by [`CacheStore`] operations. +pub type BoxCacheFuture<'a, T> = Pin> + Send + 'a>>; + +/// Current on-disk/over-the-wire schema version for a stored entry. +/// +/// Bump this to invalidate every previously stored entry when the entry shape +/// or the key derivation changes in an incompatible way. +pub const CACHE_SCHEMA_VERSION: u32 = 1; + +/// Wall-clock milliseconds since the Unix epoch. +pub fn now_unix_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|delta| delta.as_millis() as u64) + .unwrap_or(0) +} + +/// One cached response together with the metadata needed to expire it and +/// report savings on a reuse. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CacheEntry { + /// The saved answer, returned **unchanged** on a hit (usage intact, so a + /// reuse is shape-identical to a live call). Its token counts are read to + /// report savings on the cache-hit mark — the response is never mutated. + pub response: Json, + /// When the entry was stored. + pub created_unix_ms: u64, + /// When the entry expires (`created + ttl`). + pub expires_unix_ms: u64, + /// Schema version that produced this entry. + pub schema_version: u32, + /// The full cache key fingerprint (`"sha256:…"`). + pub key_hash: String, + /// Optional model name recorded for diagnostics. + pub model_name: Option, + /// Optional provider/family name recorded for diagnostics. + pub provider_name: Option, +} + +impl CacheEntry { + /// Builds an entry that expires `ttl` from now. + pub fn new( + response: Json, + ttl: Duration, + key_hash: String, + model_name: Option, + provider_name: Option, + ) -> Self { + let created = now_unix_ms(); + Self { + response, + created_unix_ms: created, + expires_unix_ms: created + .saturating_add(u64::try_from(ttl.as_millis()).unwrap_or(u64::MAX)), + schema_version: CACHE_SCHEMA_VERSION, + key_hash, + model_name, + provider_name, + } + } + + /// Whether the entry is expired relative to `now_ms`. + pub fn is_expired(&self, now_ms: u64) -> bool { + now_ms >= self.expires_unix_ms + } +} + +/// A general "store a value with an expiry" cache. +/// +/// Implementations must be cheap to clone behind an [`std::sync::Arc`] and must +/// not block: the in-memory backend locks only for the duration of a +/// synchronous map operation and never holds the lock across an `await`. +pub trait CacheStore: Send + Sync + 'static { + /// Looks up an entry. Returns `Ok(None)` when the key is absent or expired. + /// Entries come back shared behind an [`Arc`] so a lookup never deep-clones + /// the stored response (the in-memory backend holds its lock O(1)). + fn get<'a>(&'a self, key: &'a str) -> BoxCacheFuture<'a, Option>>; + /// Stores an entry under `key`. `ttl` is supplied for backends with native + /// expiry (e.g. Redis `SETEX`); the in-memory backend uses + /// [`CacheEntry::expires_unix_ms`]. + fn set<'a>(&'a self, key: &'a str, entry: CacheEntry, ttl: Duration) -> BoxCacheFuture<'a, ()>; + /// Removes an entry if present. + fn delete<'a>(&'a self, key: &'a str) -> BoxCacheFuture<'a, ()>; + /// Cheap reachability check surfaced by `doctor`/validation. + fn health<'a>(&'a self) -> BoxCacheFuture<'a, ()>; + /// Backend kind label for telemetry (e.g. `"in_memory"`). + fn backend_kind(&self) -> &'static str; +} + +/// One stored entry plus the byte size it accounts for. The entry is shared +/// behind an [`Arc`] so `get` hands out a reference instead of a deep clone. +struct StoredEntry { + entry: Arc, + size: usize, +} + +/// Locked map plus a running byte total, guarded together so the budget never +/// drifts from the contents. +#[derive(Default)] +struct Inner { + map: HashMap, + total_bytes: usize, +} + +/// Process-local cache backed by a locked map. +/// +/// Drops entries lazily when they are read after expiry, and is bounded by a +/// **total-bytes** budget (`max_bytes`) so a few large completions cannot OOM +/// the process; an optional `max_entries` adds a count cap. Eviction is +/// oldest-first by `created_unix_ms`. This is the single-process path; a shared +/// deployment uses the Redis-backed [`CacheStore`]. +pub struct InMemoryCacheStore { + inner: Mutex, + max_bytes: usize, + max_entries: Option, +} + +/// Approximate resident size: response bytes + metadata strings (key hash is +/// held twice) + fixed slot overhead. An OOM guard, not exact accounting. +fn entry_size(entry: &CacheEntry) -> usize { + const ENTRY_OVERHEAD: usize = 256; + serde_json::to_vec(&entry.response) + .map(|bytes| bytes.len()) + .unwrap_or(0) + + entry.key_hash.len() * 2 + + entry.model_name.as_deref().map_or(0, str::len) + + entry.provider_name.as_deref().map_or(0, str::len) + + ENTRY_OVERHEAD +} + +impl InMemoryCacheStore { + /// Creates a store bounded by `max_bytes` (minimum 1) and an optional + /// `max_entries` count cap. + pub fn new(max_bytes: usize, max_entries: Option) -> Self { + Self { + inner: Mutex::new(Inner::default()), + max_bytes: max_bytes.max(1), + max_entries: max_entries.map(|count| count.max(1)), + } + } + + /// Current number of stored entries (including not-yet-reaped expired ones). + pub fn len(&self) -> usize { + self.inner.lock().map(|guard| guard.map.len()).unwrap_or(0) + } + + /// Whether the store currently holds no entries. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Current total accounted bytes. + pub fn total_bytes(&self) -> usize { + self.inner + .lock() + .map(|guard| guard.total_bytes) + .unwrap_or(0) + } +} + +/// Removes the oldest entry (smallest `created_unix_ms`), updating the byte +/// total. Returns whether an entry was removed. +fn evict_oldest(inner: &mut Inner) -> bool { + let Some(oldest_key) = inner + .map + .iter() + .min_by_key(|(_, stored)| stored.entry.created_unix_ms) + .map(|(key, _)| key.clone()) + else { + return false; + }; + if let Some(stored) = inner.map.remove(&oldest_key) { + inner.total_bytes = inner.total_bytes.saturating_sub(stored.size); + } + true +} + +impl CacheStore for InMemoryCacheStore { + fn get<'a>(&'a self, key: &'a str) -> BoxCacheFuture<'a, Option>> { + Box::pin(async move { + let now = now_unix_ms(); + let mut guard = self + .inner + .lock() + .map_err(|err| AdaptiveError::Internal(format!("cache lock poisoned: {err}")))?; + match guard.map.get(key) { + Some(stored) if !stored.entry.is_expired(now) => { + Ok(Some(Arc::clone(&stored.entry))) + } + Some(_) => { + if let Some(stored) = guard.map.remove(key) { + guard.total_bytes = guard.total_bytes.saturating_sub(stored.size); + } + Ok(None) + } + None => Ok(None), + } + }) + } + + fn set<'a>( + &'a self, + key: &'a str, + entry: CacheEntry, + _ttl: Duration, + ) -> BoxCacheFuture<'a, ()> { + Box::pin(async move { + let size = entry_size(&entry); + let mut guard = self + .inner + .lock() + .map_err(|err| AdaptiveError::Internal(format!("cache lock poisoned: {err}")))?; + // Never store past the whole budget, but still drop the stale + // entry — a fresher answer exists even if it cannot be stored. + if size > self.max_bytes { + if let Some(previous) = guard.map.remove(key) { + guard.total_bytes = guard.total_bytes.saturating_sub(previous.size); + } + return Ok(()); + } + + // Replacing an existing key frees its old size first. + if let Some(previous) = guard.map.remove(key) { + guard.total_bytes = guard.total_bytes.saturating_sub(previous.size); + } + + // Evict oldest-first until the new entry fits the byte budget and + // the optional count cap (always keep room for the entry itself). + while !guard.map.is_empty() + && (guard.total_bytes + size > self.max_bytes + || self + .max_entries + .is_some_and(|cap| guard.map.len() + 1 > cap)) + && evict_oldest(&mut guard) + {} + + guard.total_bytes += size; + guard.map.insert( + key.to_string(), + StoredEntry { + entry: Arc::new(entry), + size, + }, + ); + Ok(()) + }) + } + + fn delete<'a>(&'a self, key: &'a str) -> BoxCacheFuture<'a, ()> { + Box::pin(async move { + let mut guard = self + .inner + .lock() + .map_err(|err| AdaptiveError::Internal(format!("cache lock poisoned: {err}")))?; + if let Some(stored) = guard.map.remove(key) { + guard.total_bytes = guard.total_bytes.saturating_sub(stored.size); + } + Ok(()) + }) + } + + fn health<'a>(&'a self) -> BoxCacheFuture<'a, ()> { + Box::pin(async move { + // Touch the lock so a poisoned mutex surfaces as unhealthy. + self.inner + .lock() + .map(|_| ()) + .map_err(|err| AdaptiveError::Internal(format!("cache lock poisoned: {err}"))) + }) + } + + fn backend_kind(&self) -> &'static str { + "in_memory" + } +} + +/// Redis-backed [`CacheStore`] for sharing one cache across processes / a team. +/// +/// Uses Redis's native expiry (`SET … PX `) so entries drop themselves, +/// and a configurable key prefix. The connection is a [`ConnectionManager`] +/// (cheap to clone, auto-reconnecting), cloned before each await so the store +/// never blocks. Mirrors the connection pattern of [`crate::redis::RedisBackend`]. +#[cfg(feature = "redis-backend")] +pub struct RedisCacheStore { + conn: redis::aio::ConnectionManager, + key_prefix: String, +} + +#[cfg(feature = "redis-backend")] +impl RedisCacheStore { + /// Connects to Redis and returns a new store. + /// + /// # Errors + /// Returns [`AdaptiveError::Storage`] when the client or connection fails. + pub async fn new(url: &str, key_prefix: impl Into) -> Result { + let client = redis::Client::open(url) + .map_err(|err| AdaptiveError::Storage(format!("redis client: {err}")))?; + let conn = client + .get_connection_manager() + .await + .map_err(|err| AdaptiveError::Storage(format!("redis connection: {err}")))?; + Ok(Self { + conn, + key_prefix: key_prefix.into(), + }) + } + + fn full_key(&self, key: &str) -> String { + format!("{}{}", self.key_prefix, key) + } +} + +/// A hung Redis peer must degrade to fail-open, not block the request. +#[cfg(feature = "redis-backend")] +const REDIS_OP_TIMEOUT: Duration = Duration::from_secs(2); + +#[cfg(feature = "redis-backend")] +async fn with_redis_deadline( + operation: &'static str, + future: impl Future>, +) -> Result { + match tokio::time::timeout(REDIS_OP_TIMEOUT, future).await { + Ok(result) => result, + Err(_) => Err(AdaptiveError::Storage(format!( + "redis {operation}: timed out after {REDIS_OP_TIMEOUT:?}" + ))), + } +} + +#[cfg(feature = "redis-backend")] +impl CacheStore for RedisCacheStore { + fn get<'a>(&'a self, key: &'a str) -> BoxCacheFuture<'a, Option>> { + Box::pin(with_redis_deadline("GET", async move { + let mut conn = self.conn.clone(); + let full_key = self.full_key(key); + let maybe_json: Option = redis::cmd("GET") + .arg(&full_key) + .query_async(&mut conn) + .await + .map_err(|err| AdaptiveError::Storage(format!("redis GET: {err}")))?; + // Redis already dropped the key if it expired, so any value present + // is live. + match maybe_json { + Some(json) => { + let entry: CacheEntry = + serde_json::from_str(&json).map_err(AdaptiveError::Serialization)?; + // The entry's own stamp is authoritative over Redis PX. + if entry.is_expired(now_unix_ms()) { + return Ok(None); + } + Ok(Some(Arc::new(entry))) + } + None => Ok(None), + } + })) + } + + fn set<'a>(&'a self, key: &'a str, entry: CacheEntry, ttl: Duration) -> BoxCacheFuture<'a, ()> { + Box::pin(with_redis_deadline("SET", async move { + let mut conn = self.conn.clone(); + let full_key = self.full_key(key); + let json = serde_json::to_string(&entry).map_err(AdaptiveError::Serialization)?; + let ttl_ms = u64::try_from(ttl.as_millis()).unwrap_or(u64::MAX).max(1); // PX requires a positive TTL + redis::cmd("SET") + .arg(&full_key) + .arg(json) + .arg("PX") + .arg(ttl_ms) + .exec_async(&mut conn) + .await + .map_err(|err| AdaptiveError::Storage(format!("redis SET: {err}")))?; + Ok(()) + })) + } + + fn delete<'a>(&'a self, key: &'a str) -> BoxCacheFuture<'a, ()> { + Box::pin(with_redis_deadline("DEL", async move { + let mut conn = self.conn.clone(); + let full_key = self.full_key(key); + redis::cmd("DEL") + .arg(&full_key) + .exec_async(&mut conn) + .await + .map_err(|err| AdaptiveError::Storage(format!("redis DEL: {err}")))?; + Ok(()) + })) + } + + fn health<'a>(&'a self) -> BoxCacheFuture<'a, ()> { + Box::pin(with_redis_deadline("PING", async move { + let mut conn = self.conn.clone(); + redis::cmd("PING") + .exec_async(&mut conn) + .await + .map_err(|err| AdaptiveError::Storage(format!("redis PING: {err}")))?; + Ok(()) + })) + } + + fn backend_kind(&self) -> &'static str { + "redis" + } +} + +/// Builds the configured backend and runs its reachability check. +/// +/// Intended for `nemo-relay doctor` / health surfaces: it constructs a +/// throwaway store from the response-cache config and calls +/// [`CacheStore::health`], returning the backend kind label on success or a +/// human-readable error. +pub async fn check_backend_health( + config: &ResponseCacheConfig, +) -> std::result::Result { + let store = build_store(config).await.map_err(|err| err.to_string())?; + store.health().await.map_err(|err| err.to_string())?; + Ok(store.backend_kind().to_string()) +} + +/// Builds the response-cache backend from config. +/// +/// Returns the boxed [`CacheStore`] used by the intercept and the `doctor` +/// health check. Redis support is gated behind the `redis-backend` feature. +pub(crate) async fn build_store(config: &ResponseCacheConfig) -> Result> { + match config.backend.kind.as_str() { + "in_memory" => Ok(Arc::new(InMemoryCacheStore::new( + config.backend.max_bytes(), + config.backend.max_entries(), + ))), + "redis" => build_redis_store(config).await, + other => Err(AdaptiveError::InvalidConfig(format!( + "response_cache: unknown backend kind '{other}'" + ))), + } +} + +#[cfg(feature = "redis-backend")] +async fn build_redis_store(config: &ResponseCacheConfig) -> Result> { + use crate::response_cache::store::RedisCacheStore; + + let url = config + .backend + .config + .get("url") + .and_then(Json::as_str) + .ok_or_else(|| { + AdaptiveError::InvalidConfig( + "response_cache: redis backend requires backend.config.url".to_string(), + ) + })?; + let key_prefix = config + .backend + .config + .get("key_prefix") + .and_then(Json::as_str) + .unwrap_or("nemo-relay:llm-cache:"); + let store = RedisCacheStore::new(url, key_prefix) + .await + .map_err(|err| AdaptiveError::Storage(format!("response_cache: {err}")))?; + Ok(Arc::new(store)) +} + +#[cfg(not(feature = "redis-backend"))] +async fn build_redis_store(_config: &ResponseCacheConfig) -> Result> { + Err(AdaptiveError::InvalidConfig( + "response_cache: backend.kind = \"redis\" requires building with the 'redis-backend' \ + feature" + .to_string(), + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn entry(key: &str, created: u64, expires: u64) -> CacheEntry { + CacheEntry { + response: json!({ "answer": key }), + created_unix_ms: created, + expires_unix_ms: expires, + schema_version: CACHE_SCHEMA_VERSION, + key_hash: key.to_string(), + model_name: None, + provider_name: None, + } + } + + const BIG: usize = 1 << 20; // 1 MiB — never evicts in these tests + + #[test] + fn ttl_arithmetic_is_milliseconds() { + // Pin the seconds->milliseconds conversion: a units regression would + // expire entries 1000x early or late. + let entry = CacheEntry::new( + json!({"ok": true}), + Duration::from_secs(60), + "sha256:t".to_string(), + None, + None, + ); + assert_eq!(entry.expires_unix_ms - entry.created_unix_ms, 60_000); + } + + #[tokio::test] + async fn same_key_replacement_swaps_content_and_accounting() { + let store = InMemoryCacheStore::new(BIG, None); + let small = entry("a", 100, u64::MAX); + let small_size = entry_size(&small); + store.set("k", small, Duration::MAX).await.unwrap(); + + let big = CacheEntry { + response: json!({ "answer": "a much longer replacement body".repeat(4) }), + ..entry("a", 200, u64::MAX) + }; + let big_size = entry_size(&big); + assert_ne!(small_size, big_size); + store.set("k", big, Duration::MAX).await.unwrap(); + + let got = store.get("k").await.unwrap().expect("entry present"); + assert_eq!( + got.response["answer"], + json!("a much longer replacement body".repeat(4)), + "a same-key set must serve the replacement content" + ); + assert_eq!( + store.total_bytes(), + big_size, + "replacement must swap the accounted size, not add to it" + ); + } + + #[tokio::test] + async fn an_oversized_refresh_drops_the_stale_entry() { + let small = entry("a", 100, u64::MAX); + let budget = entry_size(&small); + let store = InMemoryCacheStore::new(budget, None); + store.set("k", small, Duration::MAX).await.unwrap(); + + // A fresher answer too large to store must still invalidate the stale + // one — otherwise the old answer keeps serving after a newer one exists. + let oversized = CacheEntry::new( + json!({ "blob": "x".repeat(budget * 4 + 512) }), + Duration::MAX, + "sha256:o".to_string(), + None, + None, + ); + store.set("k", oversized, Duration::MAX).await.unwrap(); + assert!( + store.get("k").await.unwrap().is_none(), + "the stale entry must not outlive a fresher, unstorable answer" + ); + assert_eq!(store.total_bytes(), 0); + } + + #[tokio::test] + async fn set_get_delete_roundtrip() { + let store = InMemoryCacheStore::new(BIG, None); + let fresh = CacheEntry::new( + json!({"ok": true}), + Duration::from_secs(60), + "sha256:a".to_string(), + Some("gpt-4o".to_string()), + Some("openai".to_string()), + ); + store + .set("a", fresh, Duration::from_secs(60)) + .await + .unwrap(); + assert!(store.get("a").await.unwrap().is_some()); + store.delete("a").await.unwrap(); + assert!(store.get("a").await.unwrap().is_none()); + assert_eq!(store.total_bytes(), 0, "delete must reclaim bytes"); + } + + #[tokio::test] + async fn expired_entry_reads_as_absent_and_is_reaped() { + let store = InMemoryCacheStore::new(BIG, None); + // expires_unix_ms = 1 (1ms after the epoch) is firmly in the past. + store + .set("k", entry("k", 0, 1), Duration::from_secs(0)) + .await + .unwrap(); + assert!( + store.get("k").await.unwrap().is_none(), + "expired entry must read as absent" + ); + assert!(store.is_empty(), "reading an expired entry should reap it"); + assert_eq!(store.total_bytes(), 0, "reaping must reclaim bytes"); + } + + #[tokio::test] + async fn eviction_drops_the_oldest_entry_when_over_the_byte_budget() { + // Budget holds exactly two of these entries; the third forces eviction. + let one = entry("a", 100, u64::MAX); + let size = entry_size(&one); + let store = InMemoryCacheStore::new(size * 2, None); + + store.set("a", one, Duration::MAX).await.unwrap(); + store + .set("b", entry("b", 200, u64::MAX), Duration::MAX) + .await + .unwrap(); + // Third insert exceeds max_bytes -> evict oldest (created = 100 -> "a"). + store + .set("c", entry("c", 300, u64::MAX), Duration::MAX) + .await + .unwrap(); + + assert!( + store.get("a").await.unwrap().is_none(), + "oldest entry should be evicted once over the byte budget" + ); + assert!(store.get("b").await.unwrap().is_some()); + assert!(store.get("c").await.unwrap().is_some()); + assert!( + store.total_bytes() <= size * 2, + "must stay within the budget" + ); + } + + #[tokio::test] + async fn eviction_honors_the_optional_entry_count_cap() { + let store = InMemoryCacheStore::new(BIG, Some(2)); + store + .set("a", entry("a", 100, u64::MAX), Duration::MAX) + .await + .unwrap(); + store + .set("b", entry("b", 200, u64::MAX), Duration::MAX) + .await + .unwrap(); + store + .set("c", entry("c", 300, u64::MAX), Duration::MAX) + .await + .unwrap(); + + assert!( + store.get("a").await.unwrap().is_none(), + "count cap evicts oldest" + ); + assert!(store.get("b").await.unwrap().is_some()); + assert!(store.get("c").await.unwrap().is_some()); + } + + #[tokio::test] + async fn an_entry_larger_than_the_budget_is_not_cached_and_keeps_existing_entries() { + let small = entry("a", 100, u64::MAX); + let budget = entry_size(&small); + let store = InMemoryCacheStore::new(budget, None); + store.set("a", small, Duration::MAX).await.unwrap(); + + // An entry whose size alone exceeds max_bytes must be skipped — not stored + // (breaching the cap) and not flushing the cache to make room it can't use. + let oversized = CacheEntry::new( + json!({ "blob": "x".repeat(budget * 10 + 100) }), + Duration::MAX, + "b".to_string(), + None, + None, + ); + store.set("b", oversized, Duration::MAX).await.unwrap(); + + assert!( + store.get("b").await.unwrap().is_none(), + "an oversized entry must not be cached" + ); + assert!( + store.get("a").await.unwrap().is_some(), + "an oversized set must not flush existing entries" + ); + assert!(store.total_bytes() <= budget, "the byte budget must hold"); + } +} diff --git a/crates/adaptive/src/runtime/features.rs b/crates/adaptive/src/runtime/features.rs index af2553a9e..bcbf93289 100644 --- a/crates/adaptive/src/runtime/features.rs +++ b/crates/adaptive/src/runtime/features.rs @@ -34,14 +34,15 @@ use crate::acg_learner::AcgLearner; use crate::adaptive_hints_intercept::AdaptiveHintsIntercept; use crate::cache_diagnostics::{self, CacheDiagnosticsTracker}; use crate::config::{ - AcgComponentConfig, AdaptiveConfig, AdaptiveHintsComponentConfig, TelemetryComponentConfig, - ToolParallelismComponentConfig, + AcgComponentConfig, AdaptiveConfig, AdaptiveHintsComponentConfig, ResponseCacheConfig, + TelemetryComponentConfig, ToolParallelismComponentConfig, }; use crate::context_helpers::resolve_agent_id; use crate::error::{AdaptiveError, Result}; use crate::intercepts::create_tool_execution_intercept_with_mode; use crate::learner::latency::LatencySensitivityLearner; use crate::learner::traits::Learner; +use crate::response_cache::{build_store, make_intercept, make_stream_intercept}; use crate::runtime::backend::build_backend; use crate::runtime::validation::validate_config; use crate::storage::traits::StorageBackendDyn; @@ -473,6 +474,11 @@ impl AdaptiveRuntime { self.runtime_id, ))); } + // The response cache is independent of the learning-state backend: it has + // its own CacheStore and installs a buffered LLM execution intercept. + if let Some(config) = self.config.response_cache.clone() { + pending.push(Box::new(ResponseCacheFeature::new(config, self.runtime_id))); + } pending } @@ -767,6 +773,49 @@ impl AdaptiveFeature for AcgFeature { } } +struct ResponseCacheFeature { + name: String, + stream_name: String, + priority: i32, + config: ResponseCacheConfig, +} + +impl ResponseCacheFeature { + fn new(config: ResponseCacheConfig, runtime_id: Uuid) -> Self { + Self { + name: format!("adaptive_{runtime_id}_response_cache_llm_execution"), + stream_name: format!("adaptive_{runtime_id}_response_cache_llm_stream_execution"), + priority: config.priority, + config, + } + } +} + +impl AdaptiveFeature for ResponseCacheFeature { + fn register<'a>( + &'a mut self, + ctx: &'a mut RegistrationContext<'_>, + ) -> Pin> + Send + 'a>> { + Box::pin(async move { + // Build the backend once, shared by both intercepts. A Redis backend + // unreachable at startup fails registration (fail-fast); the + // intercepts themselves fail open at runtime. + let store = build_store(&self.config).await?; + let config = Arc::new(self.config.clone()); + ctx.register_llm_execution_intercept( + &self.name, + self.priority, + make_intercept(store.clone(), config.clone()), + )?; + ctx.register_llm_stream_execution_intercept( + &self.stream_name, + self.priority, + make_stream_intercept(store, config), + ) + }) + } +} + fn build_learners( agent_id: &str, learners: &[String], diff --git a/crates/adaptive/src/runtime/validation.rs b/crates/adaptive/src/runtime/validation.rs index 8d0393791..27d021919 100644 --- a/crates/adaptive/src/runtime/validation.rs +++ b/crates/adaptive/src/runtime/validation.rs @@ -4,8 +4,10 @@ use nemo_relay::plugin::{ ConfigDiagnostic, ConfigPolicy, ConfigReport, DiagnosticLevel, UnsupportedBehavior, }; +use serde_json::Value as Json; -use crate::config::{AdaptiveConfig, BackendSpec}; +use crate::config::{AdaptiveConfig, BackendSpec, ResponseCacheConfig}; +use crate::response_cache::config::KEY_STRATEGY_EXACT_REQUEST; pub fn validate_config(config: &AdaptiveConfig) -> ConfigReport { let mut report = ConfigReport::default(); @@ -81,9 +83,156 @@ pub fn validate_config(config: &AdaptiveConfig) -> ConfigReport { ); } + if let Some(response_cache) = &config.response_cache { + validate_response_cache(&mut report, response_cache); + } + report } +/// Validates the adaptive plugin's `response_cache` section. +/// +/// These are hard errors (not policy-driven): an invalid response-cache config +/// fails adaptive runtime construction rather than silently disabling the cache, +/// so a misconfiguration is caught at startup instead of producing surprising +/// runtime behavior. +fn validate_response_cache(report: &mut ConfigReport, config: &ResponseCacheConfig) { + if config.ttl_seconds == 0 { + report.diagnostics.push(response_cache_error( + "response_cache.invalid_ttl", + Some("ttl_seconds"), + "ttl_seconds must be greater than 0".to_string(), + )); + } + if !(0.0..=1.0).contains(&config.bypass_rate) { + report.diagnostics.push(response_cache_error( + "response_cache.invalid_bypass_rate", + Some("bypass_rate"), + "bypass_rate must be in [0.0, 1.0]".to_string(), + )); + } + if config.key_strategy != KEY_STRATEGY_EXACT_REQUEST { + report.diagnostics.push(response_cache_error( + "response_cache.unsupported_key_strategy", + Some("key_strategy"), + format!("unsupported key_strategy; only \"{KEY_STRATEGY_EXACT_REQUEST}\" is supported"), + )); + } + // Dropping an answer-determining field merges requests that differ there. + const RESERVED_SKIP_KEYS: &[&str] = &[ + "messages", + "input", + "instructions", + "model", + "tools", + "tool_choice", + ]; + for key in &config.skip_keys { + if RESERVED_SKIP_KEYS.contains(&key.as_str()) { + report.diagnostics.push(response_cache_error( + "response_cache.reserved_skip_key", + Some("skip_keys"), + format!("skip_keys must not drop the answer-determining field '{key}'"), + )); + } + } + // Auth material must never enter the key or the stored entries. + const AUTH_HEADERS: &[&str] = &[ + "authorization", + "proxy-authorization", + "cookie", + "set-cookie", + "x-api-key", + "api-key", + "x-goog-api-key", + ]; + for name in &config.header_allowlist { + if AUTH_HEADERS.contains(&name.to_ascii_lowercase().as_str()) { + report.diagnostics.push(response_cache_error( + "response_cache.auth_header_allowlisted", + Some("header_allowlist"), + format!("'{name}' is an auth header and must never be folded into cache keys"), + )); + } + } + match config.backend.kind.as_str() { + "in_memory" => { + // Mistyped or zero caps silently disable the cache. + for option in ["max_bytes", "max_entries"] { + if let Some(value) = config.backend.config.get(option) + && value.as_u64().is_none_or(|value| value == 0) + { + report.diagnostics.push(response_cache_error( + "response_cache.invalid_backend_option", + Some("backend.config"), + format!("{option} must be a positive integer"), + )); + } + } + } + "redis" => { + if !cfg!(feature = "redis-backend") { + report.diagnostics.push(response_cache_error( + "response_cache.backend_unavailable", + Some("backend.kind"), + "redis backend requires building with the 'redis-backend' feature".to_string(), + )); + } else if config + .backend + .config + .get("url") + .and_then(Json::as_str) + .is_none() + { + report.diagnostics.push(response_cache_error( + "response_cache.missing_redis_url", + Some("backend.config.url"), + "redis backend requires backend.config.url".to_string(), + )); + } + // A shared redis cache with no namespace mixes every caller's + // responses together — caution against cross-tenant reuse. + if config.namespace.is_empty() { + report.diagnostics.push(response_cache_warning( + "response_cache.shared_empty_namespace", + Some("namespace"), + "redis backend with an empty namespace shares cached responses across all \ + callers; set a namespace per environment/tenant" + .to_string(), + )); + } + } + other => report.diagnostics.push(response_cache_error( + "response_cache.unknown_backend", + Some("backend.kind"), + format!("unknown backend kind '{other}'"), + )), + } +} + +fn response_cache_error(code: &str, field: Option<&str>, message: String) -> ConfigDiagnostic { + response_cache_diag(DiagnosticLevel::Error, code, field, message) +} + +fn response_cache_warning(code: &str, field: Option<&str>, message: String) -> ConfigDiagnostic { + response_cache_diag(DiagnosticLevel::Warning, code, field, message) +} + +fn response_cache_diag( + level: DiagnosticLevel, + code: &str, + field: Option<&str>, + message: String, +) -> ConfigDiagnostic { + ConfigDiagnostic { + level, + code: code.to_string(), + component: Some("response_cache".to_string()), + field: field.map(str::to_string), + message, + } +} + fn validate_backend(report: &mut ConfigReport, policy: &ConfigPolicy, backend: &BackendSpec) { let kind = backend.kind.as_str(); match kind { diff --git a/crates/adaptive/tests/integration/response_cache_benchmark_tests.rs b/crates/adaptive/tests/integration/response_cache_benchmark_tests.rs new file mode 100644 index 000000000..febc56302 --- /dev/null +++ b/crates/adaptive/tests/integration/response_cache_benchmark_tests.rs @@ -0,0 +1,571 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Integration tests for the cache-ON-vs-OFF savings benchmark +//! (`crates/adaptive`). +//! +//! These validate the accounting the benchmark relies on directly against the +//! real runtime: the `response_cache` marks partition every managed call into +//! `hits` / `misses` / `bypasses`; the actual provider call count equals +//! `total_requests - hits`; and the savings identity +//! +//! ```text +//! baseline_tokens == served_tokens + saved_tokens +//! ``` +//! +//! holds, where `baseline_tokens` is the all-live cost (`total_requests * +//! per_call_total_tokens`), `served_tokens` is what the provider actually +//! served (accumulated only when the stub runs), and `saved_tokens` is summed +//! from the hit marks. Because the provider stub runs exactly on misses and +//! bypasses (never on hits), `served_tokens` auto-excludes reused answers. + +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex as StdMutex}; + +use nemo_relay::api::event::Event; +use nemo_relay::api::llm::{LlmCallExecuteParams, LlmRequest, llm_call_execute}; +use nemo_relay::api::runtime::{LlmExecutionNextFn, NemoRelayContextState, global_context}; +use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; +use nemo_relay::plugin::{PluginConfig, clear_plugin_configuration, initialize_plugins}; +use nemo_relay_adaptive::plugin_component::{ComponentSpec, register_adaptive_component}; +use nemo_relay_adaptive::{AdaptiveConfig, ResponseCacheConfig}; +use serde_json::{Value as Json, json}; +use tokio::sync::Mutex; + +static TEST_MUTEX: Mutex<()> = Mutex::const_new(()); + +/// The constant per-call token cost every response reports. Keeping this fixed +/// is what makes the savings identity hold: served-per-call ≡ saved-per-call. +const PER_CALL_TOTAL_TOKENS: u64 = 1280; + +fn reset_global() { + let _ = clear_plugin_configuration(); + let ctx = global_context(); + let mut state = ctx.write().unwrap(); + *state = NemoRelayContextState::new(); +} + +/// A provider stub that (a) counts how many times it actually runs and (b) +/// accumulates the `usage.total_tokens` it serves. The token add lives INSIDE +/// the closure, which the cache skips on a hit — so `served` naturally excludes +/// reused answers. `total_tokens` is read from the body itself, so the body is +/// the single source of truth for both what is served and what a hit reports as +/// saved. +fn counting_provider( + calls: Arc, + served_tokens: Arc, + body: Json, +) -> LlmExecutionNextFn { + Arc::new(move |_req: LlmRequest| { + let calls = Arc::clone(&calls); + let served_tokens = Arc::clone(&served_tokens); + let body = body.clone(); + Box::pin(async move { + calls.fetch_add(1, Ordering::SeqCst); + let tokens = body + .pointer("/usage/total_tokens") + .and_then(Json::as_u64) + .unwrap_or(0); + served_tokens.fetch_add(tokens, Ordering::SeqCst); + Ok(body) + }) + }) +} + +/// The canned response body, carrying deterministic token usage and cost. Its +/// `usage.total_tokens` is both what the stub accumulates as served and what a +/// hit mark reports as saved. +fn sample_body() -> Json { + json!({ + "id": "resp_abc", + "model": "gpt-4o", + "choices": [{"message": {"role": "assistant", "content": "The answer is 42."}}], + "usage": { + "prompt_tokens": 1200, + "completion_tokens": 80, + "total_tokens": PER_CALL_TOTAL_TOKENS, + "cost_usd": 0.0123 + } + }) +} + +fn chat_request(prompt: &str) -> LlmRequest { + LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": prompt}], + "temperature": 0.0 + }), + } +} + +/// The response-cache config the benchmark tests use unless a test needs +/// otherwise: exact-match keying under a dedicated namespace, no sampled +/// bypass; keys are built from the auto-detected decode of each request. +fn bench_config() -> ResponseCacheConfig { + ResponseCacheConfig { + ttl_seconds: 3600, + namespace: "bench_test".into(), + bypass_rate: 0.0, + ..Default::default() + } +} + +/// Activates the response cache as a section of the adaptive plugin component +/// (its only supported configuration surface — there is no standalone kind). +async fn activate_cache(config: ResponseCacheConfig) { + register_adaptive_component().unwrap(); + let adaptive = AdaptiveConfig { + response_cache: Some(config), + ..AdaptiveConfig::default() + }; + let report = initialize_plugins(PluginConfig { + components: vec![ComponentSpec::new(adaptive).into()], + ..PluginConfig::default() + }) + .await + .unwrap(); + assert!( + report.diagnostics.is_empty(), + "unexpected diagnostics: {:?}", + report.diagnostics + ); +} + +async fn call(provider: &LlmExecutionNextFn, request: LlmRequest) -> Json { + llm_call_execute( + LlmCallExecuteParams::builder() + .name("openai") + .request(request) + .func(provider.clone()) + .model_name("gpt-4o") + .build(), + ) + .await + .unwrap() +} + +/// Tally of the `response_cache` marks observed during a workload, plus the +/// tokens the hit marks reported as saved. This mirrors exactly what the +/// benchmark accumulates from the mark stream. +#[derive(Debug, Default, Clone, Copy)] +struct CacheStats { + hits: usize, + misses: usize, + bypasses: usize, + saved_tokens: u64, +} + +/// Registers a subscriber on the `response_cache` mark that folds every mark +/// into the shared [`CacheStats`]. Registration happens BEFORE the workload +/// because marks fire during the managed calls. +fn register_stats_subscriber(name: &str, stats: Arc>) { + register_subscriber( + name, + Arc::new(move |event: &Event| { + if event.name() != "response_cache" { + return; + } + let status = event + .data() + .and_then(|data| data.get("status")) + .and_then(Json::as_str); + let mut stats = stats.lock().unwrap(); + match status { + Some("hit") => { + stats.hits += 1; + let saved = event + .metadata() + .and_then(|metadata| metadata.get("nemo_relay.response_cache.saved_tokens")) + .and_then(Json::as_u64) + .unwrap_or(0); + stats.saved_tokens += saved; + } + Some("miss") => stats.misses += 1, + Some("bypass") => stats.bypasses += 1, + _ => {} + } + }), + ) + .unwrap(); +} + +/// A repeat workload: the first `distinct` prompts are unique, then the next +/// `repeats` prompts re-issue earlier ones (index `i % distinct`). Every request +/// uses the same provider (so served tokens accumulate on live calls only) and +/// the same shared stats subscriber. +async fn run_repeat_workload(provider: &LlmExecutionNextFn, distinct: usize, repeats: usize) { + for i in 0..distinct { + call(provider, chat_request(&format!("prompt #{i}"))).await; + } + for i in 0..repeats { + let reused = i % distinct; + call(provider, chat_request(&format!("prompt #{reused}"))).await; + } +} + +#[tokio::test] +async fn repeat_workload_yields_expected_hits_and_savings() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(bench_config()).await; + + let stats = Arc::new(StdMutex::new(CacheStats::default())); + register_stats_subscriber("bench_repeat_savings", Arc::clone(&stats)); + + let calls = Arc::new(AtomicUsize::new(0)); + let served_tokens = Arc::new(AtomicU64::new(0)); + let provider = counting_provider( + Arc::clone(&calls), + Arc::clone(&served_tokens), + sample_body(), + ); + + // 4 distinct prompts, then 4 exact repeats of them -> exactly K = 4 hits. + let distinct = 4; + let repeats = 4; + let total_requests = distinct + repeats; + run_repeat_workload(&provider, distinct, repeats).await; + flush_subscribers().unwrap(); + + let stats = *stats.lock().unwrap(); + assert_eq!( + stats.hits, repeats, + "each of the K repeats must be served from cache as a hit" + ); + assert_eq!( + calls.load(Ordering::SeqCst), + total_requests - repeats, + "the provider must run exactly total_requests - hits times" + ); + // Supporting invariants make a failure point at the cause, not just a total. + assert_eq!( + stats.hits + stats.misses, + total_requests, + "every request is either a hit or a miss (no bypasses here)" + ); + assert_eq!( + calls.load(Ordering::SeqCst), + total_requests - stats.hits, + "provider_calls_on == total_requests - hits" + ); + assert!( + stats.saved_tokens > 0, + "hits must report a positive saved-token total" + ); + assert_eq!( + stats.saved_tokens, + (repeats as u64) * PER_CALL_TOTAL_TOKENS, + "saved_tokens == hits * per-call total_tokens" + ); + + deregister_subscriber("bench_repeat_savings").unwrap(); +} + +#[tokio::test] +async fn accounting_identity_baseline_equals_served_plus_saved() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(bench_config()).await; + + let stats = Arc::new(StdMutex::new(CacheStats::default())); + register_stats_subscriber("bench_identity", Arc::clone(&stats)); + + let calls = Arc::new(AtomicUsize::new(0)); + let served_tokens = Arc::new(AtomicU64::new(0)); + let provider = counting_provider( + Arc::clone(&calls), + Arc::clone(&served_tokens), + sample_body(), + ); + + // 3 distinct prompts + 5 repeats of them = 8 managed requests. + let distinct = 3; + let repeats = 5; + let total_requests = (distinct + repeats) as u64; + run_repeat_workload(&provider, distinct, repeats).await; + flush_subscribers().unwrap(); + + let stats = *stats.lock().unwrap(); + let served = served_tokens.load(Ordering::SeqCst); + let baseline = total_requests * PER_CALL_TOTAL_TOKENS; + + // Supporting invariants first, so a break localizes. + assert_eq!( + stats.hits + stats.misses, + total_requests as usize, + "every request is a hit or a miss" + ); + assert_eq!( + calls.load(Ordering::SeqCst), + total_requests as usize - stats.hits, + "provider ran on exactly the non-hit requests" + ); + assert_eq!( + served, + (total_requests - stats.hits as u64) * PER_CALL_TOTAL_TOKENS, + "served tokens accumulate only on live (non-hit) calls" + ); + assert_eq!( + stats.saved_tokens, + (stats.hits as u64) * PER_CALL_TOTAL_TOKENS, + "saved tokens accumulate only on hits" + ); + + // THE benchmark accounting identity: the all-live baseline equals what the + // provider actually served plus what the hits saved. + assert_eq!( + served + stats.saved_tokens, + baseline, + "baseline_tokens == served_tokens + saved_tokens" + ); + + deregister_subscriber("bench_identity").unwrap(); +} + +#[tokio::test] +async fn zero_repeat_workload_has_no_hits() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(bench_config()).await; + + let stats = Arc::new(StdMutex::new(CacheStats::default())); + register_stats_subscriber("bench_no_repeat", Arc::clone(&stats)); + + let calls = Arc::new(AtomicUsize::new(0)); + let served_tokens = Arc::new(AtomicU64::new(0)); + let provider = counting_provider( + Arc::clone(&calls), + Arc::clone(&served_tokens), + sample_body(), + ); + + // All-distinct workload: no prompt is ever repeated. + let distinct = 6; + run_repeat_workload(&provider, distinct, 0).await; + flush_subscribers().unwrap(); + + let stats = *stats.lock().unwrap(); + assert_eq!(stats.hits, 0, "distinct requests can never hit"); + assert_eq!(stats.saved_tokens, 0, "with no hits there is nothing saved"); + assert_eq!( + calls.load(Ordering::SeqCst), + distinct, + "the provider must run once per distinct request" + ); + // With no hits, served must equal the full baseline. + assert_eq!( + served_tokens.load(Ordering::SeqCst), + (distinct as u64) * PER_CALL_TOTAL_TOKENS, + "served tokens equal the all-live baseline when nothing is reused" + ); + + deregister_subscriber("bench_no_repeat").unwrap(); +} + +#[tokio::test] +async fn bypass_rate_one_disables_reuse() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + // bypass_rate = 1.0 forces every call live, even exact repeats. + activate_cache(ResponseCacheConfig { + bypass_rate: 1.0, + ..bench_config() + }) + .await; + + let stats = Arc::new(StdMutex::new(CacheStats::default())); + register_stats_subscriber("bench_bypass", Arc::clone(&stats)); + + let calls = Arc::new(AtomicUsize::new(0)); + let served_tokens = Arc::new(AtomicU64::new(0)); + let provider = counting_provider( + Arc::clone(&calls), + Arc::clone(&served_tokens), + sample_body(), + ); + + // A workload full of repeats that would otherwise hit. + let distinct = 2; + let repeats = 4; + let total_requests = distinct + repeats; + run_repeat_workload(&provider, distinct, repeats).await; + flush_subscribers().unwrap(); + + let stats = *stats.lock().unwrap(); + // Load-bearing: no reuse, and the provider ran for every request. + assert_eq!( + stats.hits, 0, + "bypass_rate = 1.0 must never serve a hit, even with repeats" + ); + assert_eq!( + calls.load(Ordering::SeqCst), + total_requests, + "every request must run live under full bypass" + ); + // The bypass status string is confirmed in the source, so assert it too. + assert_eq!( + stats.bypasses, total_requests, + "every request must be marked as a bypass" + ); + assert_eq!(stats.saved_tokens, 0, "a bypass saves nothing"); + + deregister_subscriber("bench_bypass").unwrap(); +} + +#[tokio::test] +async fn cosmetic_noise_still_hits_across_field_order_and_volatile_keys() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + // Cosmetic differences collapse to one key before the store is consulted: + // RFC 8785 canonicalization folds field order, and the always-on skip-list + // drops `user`. NOTE: both hold on the raw-body path too (they are not + // unique to the codec), so this asserts cosmetic invariance, not that the + // codec decode ran — that is covered by `key::chat_shaped_requests_key_on_the_detected_decode`. + activate_cache(bench_config()).await; + + let stats = Arc::new(StdMutex::new(CacheStats::default())); + register_stats_subscriber("bench_cosmetic", Arc::clone(&stats)); + + let calls = Arc::new(AtomicUsize::new(0)); + let served_tokens = Arc::new(AtomicU64::new(0)); + let provider = counting_provider( + Arc::clone(&calls), + Arc::clone(&served_tokens), + sample_body(), + ); + + // Semantically identical, but the fields are reordered and a per-call + // `user` id differs — neither must affect the key. + let first = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hello"}], + "temperature": 0.0, + "user": "user-A" + }), + }; + let second = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "user": "user-B", + "messages": [{"role": "user", "content": "hello"}], + "temperature": 0.0, + "model": "gpt-4o" + }), + }; + + call(&provider, first).await; + call(&provider, second).await; + flush_subscribers().unwrap(); + + let stats = *stats.lock().unwrap(); + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "the second, cosmetically-different request must be served from cache" + ); + assert_eq!( + stats.hits, 1, + "canonicalization + skip-list must reuse across field order / extra `user`" + ); + assert_eq!(stats.misses, 1, "only the first request is a miss"); + + deregister_subscriber("bench_cosmetic").unwrap(); +} + +#[tokio::test] +async fn determinism_same_workload_same_hits() { + let _guard = TEST_MUTEX.lock().await; + + // Runs one repeat workload against a fresh cache and returns the observed + // hit count. Each invocation resets the global state and re-activates the + // cache, so `initialize_plugins` builds a brand-new in-memory store. + async fn run_once(subscriber: &str) -> usize { + reset_global(); + activate_cache(bench_config()).await; + + let stats = Arc::new(StdMutex::new(CacheStats::default())); + register_stats_subscriber(subscriber, Arc::clone(&stats)); + + let calls = Arc::new(AtomicUsize::new(0)); + let served_tokens = Arc::new(AtomicU64::new(0)); + let provider = counting_provider( + Arc::clone(&calls), + Arc::clone(&served_tokens), + sample_body(), + ); + + let distinct = 3; + let repeats = 5; + run_repeat_workload(&provider, distinct, repeats).await; + flush_subscribers().unwrap(); + + let hits = stats.lock().unwrap().hits; + deregister_subscriber(subscriber).unwrap(); + hits + } + + let first_run = run_once("bench_determinism_a").await; + let second_run = run_once("bench_determinism_b").await; + + assert_eq!( + first_run, 5, + "the repeat workload should produce a stable, known hit count" + ); + assert_eq!( + first_run, second_run, + "the same workload against a fresh cache must yield identical hits" + ); +} + +#[tokio::test] +async fn warm_hits_stay_within_the_latency_budget() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + // No codec configured, so every warm hit also exercises surface + // auto-detection — the default configuration users actually run. + activate_cache(ResponseCacheConfig { + ttl_seconds: 3600, + namespace: "bench_latency".into(), + bypass_rate: 0.0, + ..Default::default() + }) + .await; + + let calls = Arc::new(AtomicUsize::new(0)); + let served = Arc::new(AtomicU64::new(0)); + let provider = counting_provider(Arc::clone(&calls), served, sample_body()); + let request = || chat_request("latency probe"); + + // One miss warms the entry; every timed iteration must then be a hit. + call(&provider, request()).await; + + const ITERATIONS: usize = 200; + let mut timings_us: Vec = Vec::with_capacity(ITERATIONS); + for _ in 0..ITERATIONS { + let started = std::time::Instant::now(); + call(&provider, request()).await; + timings_us.push(started.elapsed().as_micros()); + } + + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "every timed iteration must be served from cache" + ); + + // A warm hit is a key build plus a table lookup — around 0.2 ms measured. + // The 2 ms median budget is deliberately generous so CI machines never + // flake; at this fixture size it catches gross hit-path regressions — an + // accidental provider call, lock contention, a blowup that scales with the + // iteration count — not micro-costs like one extra clone of a small body. + timings_us.sort_unstable(); + let p50 = timings_us[ITERATIONS / 2]; + assert!( + p50 < 2_000, + "warm-hit p50 must stay under the 2 ms budget, got {p50} µs" + ); +} diff --git a/crates/adaptive/tests/integration/response_cache_tests.rs b/crates/adaptive/tests/integration/response_cache_tests.rs new file mode 100644 index 000000000..7aadd8469 --- /dev/null +++ b/crates/adaptive/tests/integration/response_cache_tests.rs @@ -0,0 +1,1142 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! End-to-end tests for the adaptive plugin's `response_cache` feature +//! (exact-match). + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex as StdMutex}; + +use nemo_relay::api::event::{Event, ScopeCategory}; +use nemo_relay::api::llm::{ + LlmCallExecuteParams, LlmRequest, LlmStreamCallExecuteParams, llm_call_execute, + llm_stream_call_execute, +}; +use nemo_relay::api::runtime::{ + LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, NemoRelayContextState, + global_context, +}; +use nemo_relay::api::scope::ScopeType; +use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; +use nemo_relay::error::FlowError; +use nemo_relay::plugin::{ + PluginConfig, clear_plugin_configuration, initialize_plugins, validate_plugin_config, +}; +use nemo_relay_adaptive::plugin_component::{ComponentSpec, register_adaptive_component}; +use nemo_relay_adaptive::{ + AcgComponentConfig, AdaptiveConfig, BackendSpec, ResponseCacheConfig, StateConfig, +}; +use serde_json::{Value as Json, json}; +use tokio::sync::Mutex; +use tokio_stream::StreamExt; + +static TEST_MUTEX: Mutex<()> = Mutex::const_new(()); + +fn reset_global() { + let _ = clear_plugin_configuration(); + let ctx = global_context(); + let mut state = ctx.write().unwrap(); + *state = NemoRelayContextState::new(); +} + +/// A provider stub that counts how many times it actually runs and returns a +/// canned body carrying token usage and cost. +fn counting_provider(calls: Arc, body: Json) -> LlmExecutionNextFn { + Arc::new(move |_req: LlmRequest| { + let calls = Arc::clone(&calls); + let body = body.clone(); + Box::pin(async move { + calls.fetch_add(1, Ordering::SeqCst); + Ok(body) + }) + }) +} + +/// A provider stub that counts runs and always returns a transport-style error. +fn erroring_provider(calls: Arc) -> LlmExecutionNextFn { + Arc::new(move |_req: LlmRequest| { + let calls = Arc::clone(&calls); + Box::pin(async move { + calls.fetch_add(1, Ordering::SeqCst); + Err(FlowError::Internal("provider boom".to_string())) + }) + }) +} + +fn sample_body() -> Json { + json!({ + "id": "resp_abc", + "model": "gpt-4o", + "choices": [{"message": {"role": "assistant", "content": "The answer is 42."}}], + "usage": {"prompt_tokens": 1200, "completion_tokens": 80, "total_tokens": 1280, "cost_usd": 0.0123} + }) +} + +fn chat_request(prompt: &str) -> LlmRequest { + LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": prompt}], + "temperature": 0.0 + }), + } +} + +/// Activates the response cache as a section of the adaptive plugin component +/// (its only supported configuration surface — there is no standalone kind). +async fn activate_cache(config: ResponseCacheConfig) { + register_adaptive_component().unwrap(); + let adaptive = AdaptiveConfig { + response_cache: Some(config), + ..AdaptiveConfig::default() + }; + let report = initialize_plugins(PluginConfig { + components: vec![ComponentSpec::new(adaptive).into()], + ..PluginConfig::default() + }) + .await + .unwrap(); + assert!( + report.diagnostics.is_empty(), + "unexpected diagnostics: {:?}", + report.diagnostics + ); +} + +async fn call(provider: &LlmExecutionNextFn, request: LlmRequest) -> Json { + llm_call_execute( + LlmCallExecuteParams::builder() + .name("openai") + .request(request) + .func(provider.clone()) + .model_name("gpt-4o") + .build(), + ) + .await + .unwrap() +} + +/// A streaming provider stub that counts how many times it actually runs and +/// yields a fixed sequence of chunks. +fn counting_stream_provider( + calls: Arc, + chunks: Vec, +) -> LlmStreamExecutionNextFn { + Arc::new(move |_req: LlmRequest| { + let calls = Arc::clone(&calls); + let chunks = chunks.clone(); + Box::pin(async move { + calls.fetch_add(1, Ordering::SeqCst); + Ok(Box::pin(tokio_stream::iter( + chunks.into_iter().map(Ok::), + )) as LlmJsonStream) + }) + }) +} + +/// Runs one managed streaming call and drains it, returning the chunks the +/// caller observed (after the cache intercept and the managed pipeline). +async fn stream_call(provider: &LlmStreamExecutionNextFn, request: LlmRequest) -> Vec { + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("openai") + .request(request) + .func(provider.clone()) + .collector(Box::new(|_chunk| Ok(()))) + .finalizer(Box::new(|| json!({"done": true}))) + .model_name("gpt-4o") + .build(), + ) + .await + .unwrap(); + let mut collected = Vec::new(); + while let Some(item) = stream.next().await { + collected.push(item.unwrap()); + } + collected +} + +#[tokio::test] +async fn exact_repeat_is_a_hit_that_skips_the_provider_and_returns_the_response_unchanged() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(ResponseCacheConfig::default()).await; + + let calls = Arc::new(AtomicUsize::new(0)); + let provider = counting_provider(Arc::clone(&calls), sample_body()); + + // First call: cache miss -> provider runs. + let first = call(&provider, chat_request("What is the answer?")).await; + // Second identical call: cache hit -> provider does NOT run. + let second = call(&provider, chat_request("What is the answer?")).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "provider should run once; the repeat must be served from cache" + ); + + // A reuse is shape-identical to a live call: the hit returns the stored + // response UNCHANGED, including usage — never a stripped/different shape. + assert_eq!( + first, second, + "a hit must return the stored response unchanged" + ); + assert!( + second.get("usage").is_some(), + "a hit must preserve usage (savings are reported on the mark, not by mutating the body)" + ); +} + +#[tokio::test] +async fn a_different_request_is_a_miss() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(ResponseCacheConfig::default()).await; + + let calls = Arc::new(AtomicUsize::new(0)); + let provider = counting_provider(Arc::clone(&calls), sample_body()); + + call(&provider, chat_request("Question one?")).await; + call(&provider, chat_request("A different question?")).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "distinct requests must each hit the provider" + ); +} + +#[tokio::test] +async fn cosmetic_noise_still_hits_the_cache() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(ResponseCacheConfig::default()).await; + + let calls = Arc::new(AtomicUsize::new(0)); + let provider = counting_provider(Arc::clone(&calls), sample_body()); + + // Same meaning, but with a streaming flag and a per-call user id that must + // not affect the key, plus reordered top-level fields. + let first = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hello"}], + "stream": false, + "user": "user-A" + }), + }; + let second = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "user": "user-B", + "stream": true, + "messages": [{"role": "user", "content": "hello"}], + "model": "gpt-4o" + }), + }; + + call(&provider, first).await; + call(&provider, second).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "field order, streaming flag, and per-call user id must not change the key" + ); +} + +#[tokio::test] +async fn bypass_rate_one_always_runs_live() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(ResponseCacheConfig { + bypass_rate: 1.0, + ..ResponseCacheConfig::default() + }) + .await; + + let calls = Arc::new(AtomicUsize::new(0)); + let provider = counting_provider(Arc::clone(&calls), sample_body()); + + call(&provider, chat_request("replay me")).await; + call(&provider, chat_request("replay me")).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "bypass_rate = 1.0 must always run live (never serve a hit)" + ); +} + +#[tokio::test] +async fn stateful_responses_calls_bypass_the_cache() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(ResponseCacheConfig::default()).await; + + let calls = Arc::new(AtomicUsize::new(0)); + let provider = counting_provider(Arc::clone(&calls), sample_body()); + + let stateful = || LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "stateful"}], + "store": true + }), + }; + + call(&provider, stateful()).await; + call(&provider, stateful()).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "stateful Responses calls (store=true) must never be cached" + ); +} + +#[tokio::test] +async fn invalid_config_is_rejected_by_validation() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + register_adaptive_component().unwrap(); + + let adaptive = AdaptiveConfig { + response_cache: Some(ResponseCacheConfig { + ttl_seconds: 0, + bypass_rate: 2.0, + ..ResponseCacheConfig::default() + }), + ..AdaptiveConfig::default() + }; + let report = validate_plugin_config(&PluginConfig { + components: vec![ComponentSpec::new(adaptive).into()], + ..PluginConfig::default() + }); + + assert!( + report + .diagnostics + .iter() + .any(|diagnostic| diagnostic.code == "response_cache.invalid_ttl"), + "ttl_seconds = 0 must produce a diagnostic" + ); + assert!( + report + .diagnostics + .iter() + .any(|diagnostic| diagnostic.code == "response_cache.invalid_bypass_rate"), + "bypass_rate out of range must produce a diagnostic" + ); +} + +#[tokio::test] +async fn unknown_and_unavailable_backends_are_rejected_by_validation() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + register_adaptive_component().unwrap(); + + let validate = |kind: &str| { + let mut config = ResponseCacheConfig::default(); + config.backend.kind = kind.to_string(); + validate_plugin_config(&PluginConfig { + components: vec![ + ComponentSpec::new(AdaptiveConfig { + response_cache: Some(config), + ..AdaptiveConfig::default() + }) + .into(), + ], + ..PluginConfig::default() + }) + }; + + // An unrecognized backend kind is always rejected. + let unknown = validate("sqlite"); + assert!( + unknown + .diagnostics + .iter() + .any(|diagnostic| diagnostic.code == "response_cache.unknown_backend"), + "an unknown backend kind must be rejected: {:?}", + unknown.diagnostics + ); + + // Without the redis-backend feature, selecting redis is rejected as unavailable. + #[cfg(not(feature = "redis-backend"))] + { + let redis = validate("redis"); + assert!( + redis + .diagnostics + .iter() + .any(|diagnostic| diagnostic.code == "response_cache.backend_unavailable"), + "redis without the redis-backend feature must be rejected: {:?}", + redis.diagnostics + ); + } +} + +#[tokio::test] +async fn hit_preserves_usage_on_the_end_event_and_reports_savings_on_the_mark() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(ResponseCacheConfig::default()).await; + + // Capture every event so we can inspect the LLM end events + cache marks. + let captured = Arc::new(StdMutex::new(Vec::::new())); + let sink = Arc::clone(&captured); + register_subscriber( + "response_cache_event_capture", + Arc::new(move |event: &Event| { + sink.lock().unwrap().push(event.clone()); + }), + ) + .unwrap(); + + let calls = Arc::new(AtomicUsize::new(0)); + let provider = counting_provider(Arc::clone(&calls), sample_body()); + + call(&provider, chat_request("end-event check")).await; // miss + call(&provider, chat_request("end-event check")).await; // hit + flush_subscribers().unwrap(); + + let events = captured.lock().unwrap(); + let end_events: Vec<&Event> = events + .iter() + .filter(|event| { + event.scope_type() == Some(ScopeType::Llm) + && event.scope_category() == Some(ScopeCategory::End) + }) + .collect(); + assert_eq!(end_events.len(), 2, "expected one LLM end event per call"); + + // Transparency: a hit is shape-identical to a live call, so BOTH end events + // carry usage. The cache never mutates the response to fake $0. + for end_event in &end_events { + let output = end_event.output().expect("llm end event has output"); + assert!( + output.get("usage").is_some(), + "both miss and hit end events must carry usage (no stripping)" + ); + } + + // Savings ride on the `response_cache` hit mark instead. + let hit_mark = events + .iter() + .find(|event| { + event.name() == "response_cache" + && event + .data() + .and_then(|data| data.get("status")) + .and_then(Json::as_str) + == Some("hit") + }) + .expect("a response_cache hit mark should be emitted"); + let saved_tokens = hit_mark + .metadata() + .and_then(|metadata| metadata.get("nemo_relay.response_cache.saved_tokens")) + .and_then(Json::as_u64); + assert_eq!( + saved_tokens, + Some(1280), + "the hit mark must report the avoided token count" + ); + let ttl_ms = hit_mark + .metadata() + .and_then(|metadata| metadata.get("nemo_relay.response_cache.ttl_ms")) + .and_then(Json::as_u64); + assert!( + ttl_ms.is_some_and(|ms| ms > 0), + "the hit mark must carry the configured ttl_ms" + ); + + drop(events); + deregister_subscriber("response_cache_event_capture").unwrap(); +} + +#[tokio::test] +async fn errors_are_not_cached() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(ResponseCacheConfig::default()).await; + + let calls = Arc::new(AtomicUsize::new(0)); + let error_body = json!({"error": {"message": "rate limited", "type": "rate_limit"}}); + let provider = counting_provider(Arc::clone(&calls), error_body); + + call(&provider, chat_request("will error")).await; + call(&provider, chat_request("will error")).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "error-shaped responses must never be cached" + ); +} + +#[tokio::test] +async fn error_null_success_bodies_are_cached() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(ResponseCacheConfig::default()).await; + + // Real OpenAI-Responses success bodies carry `"error": null`; a key-presence + // check would silently disable caching for that whole surface. + let calls = Arc::new(AtomicUsize::new(0)); + let body = json!({ + "id": "resp_1", "object": "response", "status": "completed", "error": null, + "model": "gpt-4o", + "output": [{"type": "message", "role": "assistant", + "content": [{"type": "output_text", "text": "hello"}]}], + "usage": {"input_tokens": 9, "output_tokens": 3, "total_tokens": 12} + }); + let provider = counting_provider(Arc::clone(&calls), body); + + call(&provider, chat_request("responses success")).await; + call(&provider, chat_request("responses success")).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "a success body with error: null must be cached and hit on repeat" + ); +} + +#[tokio::test] +async fn non_final_status_bodies_are_not_cached() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(ResponseCacheConfig::default()).await; + + // An incomplete/failed body is not a complete, replayable answer. + let calls = Arc::new(AtomicUsize::new(0)); + let body = json!({ + "id": "resp_1", "object": "response", "status": "incomplete", "error": null, + "model": "gpt-4o", + "output": [{"type": "message", "role": "assistant", + "content": [{"type": "output_text", "text": "partial"}]}] + }); + let provider = counting_provider(Arc::clone(&calls), body); + + call(&provider, chat_request("incomplete run")).await; + call(&provider, chat_request("incomplete run")).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "a non-final status body must never be stored" + ); +} + +#[tokio::test] +async fn miss_mark_is_emitted_even_when_the_provider_errors() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(ResponseCacheConfig::default()).await; + + let captured = Arc::new(StdMutex::new(Vec::::new())); + let sink = Arc::clone(&captured); + register_subscriber( + "response_cache_error_capture", + Arc::new(move |event: &Event| sink.lock().unwrap().push(event.clone())), + ) + .unwrap(); + + let calls = Arc::new(AtomicUsize::new(0)); + let provider = erroring_provider(Arc::clone(&calls)); + + // The managed call surfaces the provider error to the caller... + let result = llm_call_execute( + LlmCallExecuteParams::builder() + .name("openai") + .request(chat_request("will error")) + .func(provider.clone()) + .model_name("gpt-4o") + .build(), + ) + .await; + assert!(result.is_err(), "the provider error must surface"); + flush_subscribers().unwrap(); + + // ...but the cache still recorded the miss decision (mark emitted before the call). + let events = captured.lock().unwrap(); + let miss = events.iter().any(|event| { + event.name() == "response_cache" + && event + .data() + .and_then(|data| data.get("status")) + .and_then(Json::as_str) + == Some("miss") + }); + assert!( + miss, + "a miss mark must be emitted even when the provider errors" + ); + drop(events); + deregister_subscriber("response_cache_error_capture").unwrap(); +} + +/// The text a strict streaming client would accumulate from replayed +/// provider-native chunks (OpenAI chat `delta.content` fragments and Anthropic +/// `text_delta` fragments). +fn replayed_text(chunks: &[Json]) -> String { + chunks + .iter() + .filter_map(|chunk| { + chunk + .pointer("/choices/0/delta/content") + .and_then(Json::as_str) + .or_else(|| chunk.pointer("/delta/text").and_then(Json::as_str)) + }) + .collect() +} + +/// OpenAI-chat streaming deltas the codec can assemble into one response. +fn openai_chat_stream_chunks() -> Vec { + vec![ + json!({"id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1_700_000_000, "model": "gpt-4o", "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": null}]}), + json!({"id": "chatcmpl-1", "object": "chat.completion.chunk", "choices": [{"index": 0, "delta": {"content": "The answer "}, "finish_reason": null}]}), + json!({"id": "chatcmpl-1", "object": "chat.completion.chunk", "choices": [{"index": 0, "delta": {"content": "is 42."}, "finish_reason": null}]}), + json!({"id": "chatcmpl-1", "object": "chat.completion.chunk", "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14}}), + ] +} + +#[tokio::test] +async fn streaming_repeat_is_a_hit_that_skips_the_provider_and_replays_the_aggregate() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + // Streaming aggregation infers its codec from the request surface. + activate_cache(ResponseCacheConfig::default()).await; + + let calls = Arc::new(AtomicUsize::new(0)); + let chunks = openai_chat_stream_chunks(); + let provider = counting_stream_provider(Arc::clone(&calls), chunks.clone()); + + // First stream: miss -> provider runs; chunks pass through and are aggregated + // into a stored response. + let first = stream_call(&provider, chat_request("stream me")).await; + // Second identical stream: hit -> provider does NOT run; the stored aggregate + // is replayed as provider-native chunks. + let second = stream_call(&provider, chat_request("stream me")).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "provider should stream once; the repeat must be served from cache" + ); + assert_eq!( + first, chunks, + "a miss must pass the live chunks through unchanged" + ); + // The hit replays PROVIDER-NATIVE chunks a strict streaming client can parse + // (not one aggregate-shaped frame), and they carry the stored answer + usage. + assert!( + second.len() > 1, + "a hit must replay a native chunk sequence, got {second:?}" + ); + assert_eq!(second[0]["object"], json!("chat.completion.chunk")); + assert_eq!(replayed_text(&second), "The answer is 42."); + let usage_chunk = second + .iter() + .rev() + .find(|chunk| chunk.get("usage").is_some()) + .expect("the replay must carry a usage chunk"); + assert_eq!(usage_chunk["usage"]["total_tokens"], json!(14)); +} + +#[tokio::test] +async fn streaming_unrecognized_shape_without_a_codec_runs_live() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + // No codec, and a body the surface detector does not recognize (no + // messages/system/input) -> nothing to aggregate -> each call runs live. + activate_cache(ResponseCacheConfig::default()).await; + + let calls = Arc::new(AtomicUsize::new(0)); + let provider = counting_stream_provider(Arc::clone(&calls), openai_chat_stream_chunks()); + let unrecognized = || LlmRequest { + headers: serde_json::Map::new(), + content: json!({"model": "mystery", "prompt": "hi"}), + }; + + stream_call(&provider, unrecognized()).await; + stream_call(&provider, unrecognized()).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "an unrecognized-shape stream with no codec must still bypass and run live" + ); +} + +/// A minimal Anthropic Messages SSE sequence the anthropic codec assembles into +/// `content: [{type:text, text:"Hello, world."}]`. +fn anthropic_stream_chunks() -> Vec { + vec![ + json!({"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-haiku-4-5", "content": [], "usage": {"input_tokens": 10, "output_tokens": 0}}}), + json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello, "}}), + json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "world."}}), + json!({"type": "content_block_stop", "index": 0}), + json!({"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"input_tokens": 10, "output_tokens": 5}}), + json!({"type": "message_stop"}), + ] +} + +/// Like [`stream_call`], but with a caller-chosen provider name (the value the +/// gateway derives from the route, e.g. `"anthropic.messages"`). +async fn stream_call_named( + name: &str, + provider: &LlmStreamExecutionNextFn, + request: LlmRequest, +) -> Vec { + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name(name.to_string()) + .request(request) + .func(provider.clone()) + .collector(Box::new(|_chunk| Ok(()))) + .finalizer(Box::new(|| json!({"done": true}))) + .model_name("claude-haiku-4-5") + .build(), + ) + .await + .unwrap(); + let mut collected = Vec::new(); + while let Some(item) = stream.next().await { + collected.push(item.unwrap()); + } + collected +} + +#[tokio::test] +async fn no_codec_mis_inferred_empty_aggregate_is_not_cached() { + // The correctness guard: a chat-shaped request infers the openai_chat codec, + // but the provider streams NON-chat chunks (the misdetection case — e.g. a + // system-less Anthropic body read as chat). The chat collector drops them and + // finalizes an empty aggregate; the guard must refuse to cache it, so the + // repeat runs live instead of serving a wrong empty response. + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(ResponseCacheConfig::default()).await; // codec = None + + let calls = Arc::new(AtomicUsize::new(0)); + let foreign_chunks = vec![ + json!({"type": "content_block_delta", "delta": {"text": "The answer "}}), + json!({"type": "content_block_delta", "delta": {"text": "is 42."}}), + ]; + let provider = counting_stream_provider(Arc::clone(&calls), foreign_chunks); + + stream_call(&provider, chat_request("mis-infer")).await; // empty aggregate -> NOT stored + stream_call(&provider, chat_request("mis-infer")).await; // must run live again + + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "a mis-inferred empty aggregate must not be cached; the repeat runs live" + ); +} + +/// An Anthropic Messages SSE sequence that dies mid-answer with an in-band +/// `error` event (never a stream-level `Err`). +fn anthropic_stream_chunks_with_inband_error() -> Vec { + vec![ + json!({"type": "message_start", "message": {"id": "msg_x", "type": "message", + "role": "assistant", "content": [], "usage": {"input_tokens": 9, "output_tokens": 0}}}), + json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial ans"}}), + json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}), + ] +} + +#[tokio::test] +async fn streaming_inband_error_is_not_cached() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(ResponseCacheConfig::default()).await; + + let calls = Arc::new(AtomicUsize::new(0)); + let provider = counting_stream_provider( + Arc::clone(&calls), + anthropic_stream_chunks_with_inband_error(), + ); + + // A stream that carries an in-band `error` event must NOT have its partial + // aggregate cached: a later identical request would otherwise replay a + // truncated answer as if it were complete. + stream_call_named( + "anthropic.messages", + &provider, + chat_request("erroring stream"), + ) + .await; + stream_call_named( + "anthropic.messages", + &provider, + chat_request("erroring stream"), + ) + .await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "an erroring stream must not populate the cache; the repeat runs live" + ); +} + +#[tokio::test] +async fn thinking_streams_are_not_cached() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(ResponseCacheConfig::default()).await; + + // The anthropic collector drops thinking/signature deltas, so a stored + // thinking aggregate would replay a gutted, unsigned block. + let chunks = vec![ + json!({"type": "message_start", "message": {"id": "msg_t", "type": "message", + "role": "assistant", "content": [], "usage": {"input_tokens": 9, "output_tokens": 0}}}), + json!({"type": "content_block_start", "index": 0, + "content_block": {"type": "thinking", "thinking": ""}}), + json!({"type": "content_block_delta", "index": 0, + "delta": {"type": "thinking_delta", "thinking": "step by step"}}), + json!({"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}), + json!({"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "answer"}}), + json!({"type": "message_stop"}), + ]; + let calls = Arc::new(AtomicUsize::new(0)); + let provider = counting_stream_provider(Arc::clone(&calls), chunks); + + stream_call_named("anthropic.messages", &provider, chat_request("think hard")).await; + stream_call_named("anthropic.messages", &provider, chat_request("think hard")).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "a thinking stream must not be stored (its replay would be gutted)" + ); +} + +#[tokio::test] +async fn refusal_only_streams_are_not_cached() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(ResponseCacheConfig::default()).await; + + // The chat collector has no refusal handling: a refusal-only stream + // aggregates to a choice with no content and no tool calls, and a replay + // would look empty where the live stream carried the refusal text. + let chunks = vec![ + json!({"id": "chatcmpl-r", "object": "chat.completion.chunk", "created": 1_700_000_000, + "model": "gpt-4o", "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": null}]}), + json!({"id": "chatcmpl-r", "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {"refusal": "I cannot help with that."}, "finish_reason": null}]}), + json!({"id": "chatcmpl-r", "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}), + ]; + let calls = Arc::new(AtomicUsize::new(0)); + let provider = counting_stream_provider(Arc::clone(&calls), chunks); + + stream_call(&provider, chat_request("refused")).await; + stream_call(&provider, chat_request("refused")).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "a refusal-only stream must not be stored (its replay would be empty)" + ); +} + +#[tokio::test] +async fn misconfigured_keys_and_headers_are_rejected_by_validation() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + register_adaptive_component().unwrap(); + + let validate = |config: ResponseCacheConfig| { + let adaptive = AdaptiveConfig { + response_cache: Some(config), + ..AdaptiveConfig::default() + }; + validate_plugin_config(&PluginConfig { + components: vec![ComponentSpec::new(adaptive).into()], + ..PluginConfig::default() + }) + }; + + let report = validate(ResponseCacheConfig { + skip_keys: vec!["messages".to_string()], + ..ResponseCacheConfig::default() + }); + assert!( + report + .diagnostics + .iter() + .any(|d| d.code == "response_cache.reserved_skip_key"), + "dropping an answer-determining field must be rejected: {:?}", + report.diagnostics + ); + + let report = validate(ResponseCacheConfig { + header_allowlist: vec!["Authorization".to_string()], + ..ResponseCacheConfig::default() + }); + assert!( + report + .diagnostics + .iter() + .any(|d| d.code == "response_cache.auth_header_allowlisted"), + "auth headers must never be allowlisted into keys: {:?}", + report.diagnostics + ); + + let mut zeroed = ResponseCacheConfig::default(); + zeroed + .backend + .config + .insert("max_bytes".to_string(), json!(0)); + let report = validate(zeroed); + assert!( + report + .diagnostics + .iter() + .any(|d| d.code == "response_cache.invalid_backend_option"), + "max_bytes: 0 silently disables the cache and must be rejected: {:?}", + report.diagnostics + ); + + let mut mistyped = ResponseCacheConfig::default(); + mistyped + .backend + .config + .insert("max_entries".to_string(), json!("10")); + let report = validate(mistyped); + assert!( + report + .diagnostics + .iter() + .any(|d| d.code == "response_cache.invalid_backend_option"), + "a wrong-typed cap silently vanishes and must be rejected: {:?}", + report.diagnostics + ); +} + +#[tokio::test] +async fn truncated_stream_without_a_terminal_event_is_not_cached() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(ResponseCacheConfig::default()).await; + + // A clean end-of-file after N deltas (proxy idle-timeout, upstream abort + // without a finish event) is a truncated answer: every collector finalizes + // it as a well-formed partial, so only the missing terminal event reveals + // the truncation. It must not be stored and replayed as complete. + let truncated: Vec = openai_chat_stream_chunks() + .into_iter() + .filter(|chunk| chunk["choices"][0]["finish_reason"].is_null()) + .collect(); + let calls = Arc::new(AtomicUsize::new(0)); + let provider = counting_stream_provider(Arc::clone(&calls), truncated); + + stream_call(&provider, chat_request("cut off")).await; + stream_call(&provider, chat_request("cut off")).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "a stream with no terminal event must not populate the cache" + ); +} + +#[tokio::test] +async fn no_codec_system_less_anthropic_uses_provider_hint_and_caches() { + // A system-less Anthropic body is shape-ambiguous (reads as OpenAI Chat by + // shape alone). The provider-name hint disambiguates it, so the correct codec + // assembles a non-empty aggregate and the repeat is a hit. + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(ResponseCacheConfig::default()).await; // codec = None + + let calls = Arc::new(AtomicUsize::new(0)); + let provider = counting_stream_provider(Arc::clone(&calls), anthropic_stream_chunks()); + let body = || LlmRequest { + headers: serde_json::Map::new(), + content: json!({"model": "claude-haiku-4-5", "messages": [{"role": "user", "content": "hi"}]}), + }; + + stream_call_named("anthropic.messages", &provider, body()).await; // miss -> stores + let second = stream_call_named("anthropic.messages", &provider, body()).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "the anthropic provider hint must let the repeat hit the cache" + ); + assert_eq!( + second[0]["type"], + json!("message_start"), + "an anthropic replay must open with a native message_start event" + ); + assert_eq!( + replayed_text(&second), + "Hello, world.", + "the replayed native chunks carry the assembled anthropic answer" + ); +} + +#[tokio::test] +async fn no_codec_buffered_and_streaming_share_one_store_entry() { + // The interchange the inference change newly enables: with codec unset, + // streaming now participates in the SAME store as buffered (before, it + // bypassed). Both derive the same key from the auto-detected decode, so an + // entry stored by one is served to + // the other. Assert cross-replay is a hit (not a re-run) in BOTH directions. + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(ResponseCacheConfig::default()).await; // codec = None + + // Direction A: streaming miss stores the inferred-codec aggregate; a buffered + // repeat of the same request is served from it (buffered provider skipped). + let s_calls = Arc::new(AtomicUsize::new(0)); + let s_provider = counting_stream_provider(Arc::clone(&s_calls), openai_chat_stream_chunks()); + stream_call(&s_provider, chat_request("shared-A")).await; + let b_calls = Arc::new(AtomicUsize::new(0)); + let b_provider = counting_provider(Arc::clone(&b_calls), sample_body()); + let buffered_hit = call(&b_provider, chat_request("shared-A")).await; + assert_eq!( + b_calls.load(Ordering::SeqCst), + 0, + "a buffered repeat must hit the streamed entry, not run the provider" + ); + assert_eq!( + buffered_hit["choices"][0]["message"]["content"], + json!("The answer is 42.") + ); + + // Direction B: buffered miss stores the raw response; a streaming repeat of a + // (different) request is served from it, replayed as native chunks. + let b2 = Arc::new(AtomicUsize::new(0)); + let b2_provider = counting_provider(Arc::clone(&b2), sample_body()); + call(&b2_provider, chat_request("shared-B")).await; + let s2 = Arc::new(AtomicUsize::new(0)); + let s2_provider = counting_stream_provider(Arc::clone(&s2), openai_chat_stream_chunks()); + let streamed = stream_call(&s2_provider, chat_request("shared-B")).await; + assert_eq!( + s2.load(Ordering::SeqCst), + 0, + "a streaming repeat must hit the buffered entry, not run the provider" + ); + assert!( + streamed.len() > 1, + "the buffered entry is replayed as native streaming chunks" + ); + assert_eq!(replayed_text(&streamed), "The answer is 42."); +} + +#[tokio::test] +async fn cache_coexists_with_acg_execution_intercept() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + register_adaptive_component().unwrap(); + + // Configure BOTH the ACG execution intercept (needs a state backend) and the + // response cache on one adaptive component. The cache takes a lower priority + // so it sits outermost and keys on the original (pre-ACG) request. + let adaptive = AdaptiveConfig { + state: Some(StateConfig { + backend: BackendSpec::in_memory(), + }), + acg: Some(AcgComponentConfig::default()), + response_cache: Some(ResponseCacheConfig { + priority: 40, + ..ResponseCacheConfig::default() + }), + ..AdaptiveConfig::default() + }; + let report = initialize_plugins(PluginConfig { + components: vec![ComponentSpec::new(adaptive).into()], + ..PluginConfig::default() + }) + .await + .unwrap(); + assert!( + report.diagnostics.is_empty(), + "unexpected diagnostics: {:?}", + report.diagnostics + ); + + let calls = Arc::new(AtomicUsize::new(0)); + let provider = counting_provider(Arc::clone(&calls), sample_body()); + call(&provider, chat_request("co-configured")).await; + call(&provider, chat_request("co-configured")).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "the cache must still hit when co-configured with the ACG execution intercept" + ); +} + +/// Shared-cache test: a second, independent store instance (a stand-in for +/// another process / teammate) sees entries written by the first. Requires a +/// local Redis and `NEMO_RELAY_RUN_REDIS_TESTS=1`; otherwise it skips. Compile +/// with `--features redis-backend`. +#[cfg(feature = "redis-backend")] +#[tokio::test] +async fn redis_backend_shares_entries_across_store_instances() { + use std::time::Duration; + + use nemo_relay_adaptive::response_cache::store::{CacheEntry, CacheStore, RedisCacheStore}; + + let enabled = std::env::var("NEMO_RELAY_RUN_REDIS_TESTS") + .ok() + .is_some_and(|value| { + let value = value.trim(); + !value.is_empty() && value != "0" && !value.eq_ignore_ascii_case("false") + }); + if !enabled { + eprintln!( + "skipping redis shared-cache test: set NEMO_RELAY_RUN_REDIS_TESTS=1 with a local Redis" + ); + return; + } + + let url = "redis://127.0.0.1:6379"; + let prefix = "nemo-relay:test-response-cache:"; + let writer = RedisCacheStore::new(url, prefix) + .await + .expect("connect redis (writer)"); + let reader = RedisCacheStore::new(url, prefix) + .await + .expect("connect redis (reader)"); + + let key = "sha256:shared-cache-test-key"; + let _ = writer.delete(key).await; + + let entry = CacheEntry::new( + json!({"answer": "shared"}), + Duration::from_secs(60), + key.to_string(), + Some("gpt-4o".to_string()), + Some("openai".to_string()), + ); + writer + .set(key, entry, Duration::from_secs(60)) + .await + .expect("set"); + + // The independent reader instance sees the writer's entry — team sharing. + let got = reader.get(key).await.expect("get"); + assert!( + got.is_some(), + "a second store instance should see the shared entry" + ); + assert_eq!(got.unwrap().response["answer"], json!("shared")); + + writer.delete(key).await.expect("delete"); + assert!(reader.get(key).await.expect("get").is_none()); +} diff --git a/crates/adaptive/tests/unit/config_tests.rs b/crates/adaptive/tests/unit/config_tests.rs index 22f8ed6f6..b77034fe4 100644 --- a/crates/adaptive/tests/unit/config_tests.rs +++ b/crates/adaptive/tests/unit/config_tests.rs @@ -90,6 +90,7 @@ fn test_adaptive_editor_schema_covers_canonical_options() { "adaptive_hints", "tool_parallelism", "acg", + "response_cache", "policy", ] ); @@ -118,4 +119,19 @@ fn test_adaptive_editor_schema_covers_canonical_options() { .kind, EditorFieldKind::Integer ); + + let response_cache = schema.field("response_cache").unwrap().schema().unwrap(); + assert_eq!( + response_cache.field("ttl_seconds").unwrap().kind, + EditorFieldKind::Integer + ); + assert_eq!( + response_cache.field("bypass_rate").unwrap().kind, + EditorFieldKind::Float + ); + let response_cache_backend = response_cache.field("backend").unwrap().schema().unwrap(); + assert_eq!( + response_cache_backend.field("kind").unwrap().kind, + EditorFieldKind::Enum + ); } diff --git a/crates/cli/src/doctor.rs b/crates/cli/src/doctor.rs index b272d5388..becbac44c 100644 --- a/crates/cli/src/doctor.rs +++ b/crates/cli/src/doctor.rs @@ -18,7 +18,8 @@ use nemo_relay::api::event::{BaseEvent, Event, MarkEvent}; use nemo_relay::codec::model_pricing::{PricingCatalog, PricingConfig, PricingSourceConfig}; use nemo_relay::observability::plugin_component::OBSERVABILITY_PLUGIN_KIND; use nemo_relay::plugin::{DiagnosticLevel, PluginConfig, validate_plugin_config}; -use nemo_relay_adaptive::plugin_component::register_adaptive_component; +use nemo_relay_adaptive::plugin_component::{ADAPTIVE_PLUGIN_KIND, register_adaptive_component}; +use nemo_relay_adaptive::{ResponseCacheConfig, response_cache}; use nemo_relay_pii_redaction::component::register_pii_redaction_component; use serde::Serialize; use serde_json::{Value, json}; @@ -643,6 +644,13 @@ async fn collect_observability(gateway: &GatewayConfig) -> Vec { return checks; } let report = validate_plugin_config(&plugin_config); + let response_cache_invalid = report.diagnostics.iter().any(|diagnostic| { + diagnostic.level == DiagnosticLevel::Error + && diagnostic.component.as_deref().is_some_and(|component| { + // Backend diagnostics are tagged `response_cache.backend[.kind]`. + component == "response_cache" || component.starts_with("response_cache.") + }) + }); if report.diagnostics.is_empty() { checks.push(Check { name: "Plugin validation", @@ -673,10 +681,86 @@ async fn collect_observability(gateway: &GatewayConfig) -> Vec { }); } collect_pricing_component_checks(&mut checks, &plugin_config); + collect_response_cache_component_checks(&mut checks, &plugin_config, response_cache_invalid) + .await; checks } +async fn collect_response_cache_component_checks( + checks: &mut Vec, + plugin_config: &PluginConfig, + config_invalid: bool, +) { + // The response cache is a section of the adaptive component, not its own + // plugin kind: find the adaptive component and look for `response_cache`. + let Some(component) = plugin_config + .components + .iter() + .find(|component| component.kind == ADAPTIVE_PLUGIN_KIND) + else { + checks.push(Check { + name: "Response cache", + status: Status::Info, + details: "not configured".into(), + }); + return; + }; + let Some(response_cache_value) = component + .config + .get("response_cache") + .filter(|value| !value.is_null()) + else { + checks.push(Check { + name: "Response cache", + status: Status::Info, + details: "not configured".into(), + }); + return; + }; + if !component.enabled { + checks.push(Check { + name: "Response cache", + status: Status::Info, + details: "configured but disabled (adaptive plugin disabled)".into(), + }); + return; + } + // Don't probe the backend for a config validation already rejected — that + // would report a misleading "reachable" Pass next to the validation failure. + if config_invalid { + checks.push(Check { + name: "Response cache", + status: Status::Fail, + details: "config invalid (see plugin diagnostics above)".into(), + }); + return; + } + let config: ResponseCacheConfig = match serde_json::from_value(response_cache_value.clone()) { + Ok(config) => config, + Err(error) => { + checks.push(Check { + name: "Response cache", + status: Status::Fail, + details: format!("invalid response_cache config: {error}"), + }); + return; + } + }; + match response_cache::check_backend_health(&config).await { + Ok(backend) => checks.push(Check { + name: "Response cache", + status: Status::Pass, + details: format!("on; backend '{backend}' reachable"), + }), + Err(error) => checks.push(Check { + name: "Response cache", + status: Status::Fail, + details: format!("backend not reachable: {error}"), + }), + } +} + async fn collect_observability_component_checks(checks: &mut Vec, config: &Value) { for section in ["atof", "atif"] { if let Some(check) = observability_file_exporter_check(config, section) { diff --git a/crates/cli/src/gateway.rs b/crates/cli/src/gateway.rs index 87e0a34eb..07c0828bb 100644 --- a/crates/cli/src/gateway.rs +++ b/crates/cli/src/gateway.rs @@ -316,25 +316,52 @@ async fn run_managed_buffered( .await; match result { Ok(response_json) => { - state - .sessions - .record_gateway_response_hints(&session_id, owner_subagent_id, response_json) - .await; - state.sessions.finish_gateway_call(&session_id, false).await; - let (status, headers) = upstream_info + let (status, mut headers) = upstream_info .lock() .expect("upstream info lock poisoned") .take() .unwrap_or((StatusCode::OK, HeaderMap::new())); - let bytes = response_bytes + let captured = response_bytes .lock() .expect("response bytes lock poisoned") - .take() - .unwrap_or_default(); - build_response(status, headers, Body::from(bytes)) + .take(); + // A short-circuiting execution intercept (e.g. a response_cache hit) + // returns the response without an upstream call, so no upstream bytes + // were captured. Serialize the returned JSON so the client still gets + // the body — otherwise a cache hit responds 200 with an empty payload. + let body_bytes = match captured { + Some(bytes) => bytes, + None => { + if !headers.contains_key(http::header::CONTENT_TYPE) { + headers.insert( + http::header::CONTENT_TYPE, + HeaderValue::from_static("application/json"), + ); + } + Bytes::from(serde_json::to_vec(&response_json).unwrap_or_default()) + } + }; + state + .sessions + .record_gateway_response_hints(&session_id, owner_subagent_id, response_json) + .await; + state.sessions.finish_gateway_call(&session_id, false).await; + build_response(status, headers, Body::from(body_bytes)) } Err(error) => { state.sessions.finish_gateway_call(&session_id, false).await; + // Relay a captured upstream failure verbatim, like an unmanaged proxy. + let info = upstream_info + .lock() + .expect("upstream info lock poisoned") + .take(); + let captured = response_bytes + .lock() + .expect("response bytes lock poisoned") + .take(); + if let (Some((status, headers)), Some(bytes)) = (info, captured) { + return build_response(status, headers, Body::from(bytes)); + } Err(translate_runtime_error(error, &upstream_error)) } } @@ -394,12 +421,21 @@ fn build_buffered_func( return Err(FlowError::Internal(message)); } }; - let json = serde_json::from_slice::(&bytes) - .unwrap_or_else(|_| json!({ "body_bytes": bytes.len() })); + let parsed = serde_json::from_slice::(&bytes); *upstream_info.lock().expect("upstream info lock poisoned") = Some((status, response_headers)); *response_bytes.lock().expect("response bytes lock poisoned") = Some(bytes); - Ok(json) + // Failures flow back as Err so nothing downstream can cache them; + // the captured bytes still reach the client verbatim. + if !status.is_success() { + return Err(FlowError::Internal(format!("upstream returned {status}"))); + } + match parsed { + Ok(json) => Ok(json), + Err(_) => Err(FlowError::Internal( + "upstream returned a non-JSON body".to_string(), + )), + } }) }) } @@ -484,11 +520,19 @@ async fn run_managed_streaming( return Err(translate_runtime_error(error, &upstream_error)); } }; - let (status, headers) = upstream_info + let (status, mut headers) = upstream_info .lock() .expect("upstream info lock poisoned") .take() .unwrap_or((StatusCode::OK, HeaderMap::new())); + // A short-circuited replay captured no upstream headers; strict SSE + // clients require the content type. + if !headers.contains_key(http::header::CONTENT_TYPE) { + headers.insert( + http::header::CONTENT_TYPE, + HeaderValue::from_static("text/event-stream"), + ); + } let body = client_sse_body( json_stream, provider_route, diff --git a/crates/cli/tests/coverage/doctor_tests.rs b/crates/cli/tests/coverage/doctor_tests.rs index 8c7c90d9a..05cb6844a 100644 --- a/crates/cli/tests/coverage/doctor_tests.rs +++ b/crates/cli/tests/coverage/doctor_tests.rs @@ -987,6 +987,106 @@ async fn collect_observability_registers_adaptive_before_validation() { ); } +#[tokio::test] +async fn collect_observability_reports_response_cache_on_when_configured() { + // The response cache is a section of the adaptive component; doctor should + // find it and report the in-memory backend reachable. + let gateway = GatewayConfig { + plugin_config: Some(serde_json::json!({ + "version": 1, + "components": [ + { + "kind": "adaptive", + "enabled": true, + "config": { + "response_cache": { + "ttl_seconds": 3600, + "backend": { "kind": "in_memory" } + } + } + } + ] + })), + ..GatewayConfig::default() + }; + + let checks = collect_observability(&gateway).await; + + let cache = checks + .iter() + .find(|check| check.name == "Response cache") + .expect("a Response cache check should be present"); + assert_eq!(cache.status, Status::Pass, "checks: {checks:?}"); + assert!( + cache.details.contains("in_memory") && cache.details.contains("reachable"), + "details: {}", + cache.details + ); +} + +#[tokio::test] +async fn collect_observability_reports_response_cache_not_configured_without_section() { + // Adaptive present but no `response_cache` section -> reported as not configured. + let gateway = GatewayConfig { + plugin_config: Some(serde_json::json!({ + "version": 1, + "components": [ + { "kind": "adaptive", "enabled": true, "config": { "version": 1 } } + ] + })), + ..GatewayConfig::default() + }; + + let checks = collect_observability(&gateway).await; + + let cache = checks + .iter() + .find(|check| check.name == "Response cache") + .expect("a Response cache check should be present"); + assert_eq!(cache.status, Status::Info); + assert!( + cache.details.contains("not configured"), + "details: {}", + cache.details + ); +} + +#[tokio::test] +async fn collect_observability_reports_response_cache_fail_when_config_invalid() { + // An invalid section (ttl_seconds = 0) must not be reported as a reachable + // "Pass" alongside the validation failure. + let gateway = GatewayConfig { + plugin_config: Some(serde_json::json!({ + "version": 1, + "components": [ + { + "kind": "adaptive", + "enabled": true, + "config": { "response_cache": { "ttl_seconds": 0 } } + } + ] + })), + ..GatewayConfig::default() + }; + + let checks = collect_observability(&gateway).await; + + let cache = checks + .iter() + .find(|check| check.name == "Response cache") + .expect("a Response cache check should be present"); + assert_eq!(cache.status, Status::Fail); + assert!( + cache.details.contains("config invalid"), + "details: {}", + cache.details + ); + assert!( + !cache.details.contains("reachable"), + "must not claim reachable for an invalid config" + ); +} + #[tokio::test] async fn collect_observability_registers_pii_redaction_before_validation() { let gateway = GatewayConfig { diff --git a/crates/cli/tests/coverage/server_tests.rs b/crates/cli/tests/coverage/server_tests.rs index 25d7ac598..9cb0edc90 100644 --- a/crates/cli/tests/coverage/server_tests.rs +++ b/crates/cli/tests/coverage/server_tests.rs @@ -2680,3 +2680,253 @@ async fn spawn_anthropic_upstream() -> TestServer { handle, } } + +/// Spawns a minimal mock upstream that counts calls and returns a fixed +/// (status, content-type, body) for every POST. Returns its base URL and the +/// call counter. +async fn spawn_mock_upstream( + status: StatusCode, + content_type: &'static str, + body: &'static str, +) -> (String, std::sync::Arc) { + use axum::routing::any; + let calls = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let counter = calls.clone(); + let app = axum::Router::new().route( + "/{*path}", + any(move || { + let counter = counter.clone(); + async move { + counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + axum::response::Response::builder() + .status(status) + .header(http::header::CONTENT_TYPE, content_type) + .header("x-upstream-marker", "mock") + .body(axum::body::Body::from(body)) + .unwrap() + } + }), + ); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + (format!("http://{address}"), calls) +} + +/// Gateway config wired to `upstream` with the response cache enabled. +fn cache_gateway_config(upstream: &str) -> GatewayConfig { + let mut config = test_config(); + config.openai_base_url = upstream.into(); + config.anthropic_base_url = upstream.into(); + config.plugin_config = Some(json!({ + "version": 1, + "components": [{ + "kind": "adaptive", + "enabled": true, + "config": {"response_cache": { + "ttl_seconds": 3600, + "bypass_rate": 0.0, + "namespace": "gateway-test", + "backend": {"kind": "in_memory"} + }} + }] + })); + config +} + +const MOCK_CHAT_BODY: &str = r#"{"id":"chatcmpl-1","object":"chat.completion","created":1700000000,"model":"gpt-4o","choices":[{"index":0,"message":{"role":"assistant","content":"The answer is 42."},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":4,"total_tokens":14}}"#; + +#[tokio::test] +async fn gateway_serves_cached_hit_with_full_body_and_json_content_type() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let _ = nemo_relay::plugin::clear_plugin_configuration(); + + let (upstream, upstream_calls) = + spawn_mock_upstream(StatusCode::OK, "application/json", MOCK_CHAT_BODY).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let config = cache_gateway_config(&upstream); + let handle = + tokio::spawn(async move { serve_listener(listener, config, Some(shutdown_rx)).await }); + wait_for_gateway(&url).await; + + let client = test_http_client(); + let request_body = json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "What is the answer?"}], + "temperature": 0.0 + }); + let first = client + .post(format!("{url}/v1/chat/completions")) + .json(&request_body) + .send() + .await + .unwrap(); + assert_eq!(first.status(), StatusCode::OK); + let first_body: Value = first.json().await.unwrap(); + + let second = client + .post(format!("{url}/v1/chat/completions")) + .json(&request_body) + .send() + .await + .unwrap(); + assert_eq!(second.status(), StatusCode::OK); + assert_eq!( + second + .headers() + .get(http::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()), + Some("application/json"), + "a cached hit must still declare its JSON content type" + ); + let second_body: Value = second.json().await.unwrap(); + + assert_eq!( + upstream_calls.load(std::sync::atomic::Ordering::SeqCst), + 1, + "the repeat must be served from cache, not the upstream" + ); + assert_eq!( + second_body, first_body, + "the cached body must be byte-equivalent to the live one" + ); + assert_eq!( + second_body["choices"][0]["message"]["content"], + json!("The answer is 42.") + ); + + shutdown_tx.send(()).unwrap(); + handle.await.unwrap().unwrap(); + let _ = nemo_relay::plugin::clear_plugin_configuration(); +} + +#[tokio::test] +async fn gateway_relays_non_2xx_upstream_verbatim_and_never_caches_it() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let _ = nemo_relay::plugin::clear_plugin_configuration(); + + // A 503 whose body has NO top-level `error` key — the shape that used to be + // stored and replayed as a fabricated 200 for the full TTL. + let (upstream, upstream_calls) = spawn_mock_upstream( + StatusCode::SERVICE_UNAVAILABLE, + "application/json", + r#"{"message":"service temporarily unavailable"}"#, + ) + .await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let config = cache_gateway_config(&upstream); + let handle = + tokio::spawn(async move { serve_listener(listener, config, Some(shutdown_rx)).await }); + wait_for_gateway(&url).await; + + let client = test_http_client(); + let request_body = json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hello"}] + }); + for _ in 0..2 { + let response = client + .post(format!("{url}/v1/chat/completions")) + .json(&request_body) + .send() + .await + .unwrap(); + assert_eq!( + response.status(), + StatusCode::SERVICE_UNAVAILABLE, + "the upstream failure must be relayed with its real status" + ); + assert_eq!( + response + .headers() + .get("x-upstream-marker") + .and_then(|value| value.to_str().ok()), + Some("mock"), + "upstream headers must be relayed on failures" + ); + let body: Value = response.json().await.unwrap(); + assert_eq!(body, json!({"message": "service temporarily unavailable"})); + } + assert_eq!( + upstream_calls.load(std::sync::atomic::Ordering::SeqCst), + 2, + "an upstream failure must never be cached" + ); + + shutdown_tx.send(()).unwrap(); + handle.await.unwrap().unwrap(); + let _ = nemo_relay::plugin::clear_plugin_configuration(); +} + +const MOCK_CHAT_SSE: &str = "data: {\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"created\":1700000000,\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"The answer is 42.\"},\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"; + +#[tokio::test] +async fn gateway_streaming_hit_carries_event_stream_content_type() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let _ = nemo_relay::plugin::clear_plugin_configuration(); + + let (upstream, upstream_calls) = + spawn_mock_upstream(StatusCode::OK, "text/event-stream", MOCK_CHAT_SSE).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let config = cache_gateway_config(&upstream); + let handle = + tokio::spawn(async move { serve_listener(listener, config, Some(shutdown_rx)).await }); + wait_for_gateway(&url).await; + + let client = test_http_client(); + let request_body = json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "stream it"}], + "stream": true + }); + let first = client + .post(format!("{url}/v1/chat/completions")) + .json(&request_body) + .send() + .await + .unwrap(); + assert_eq!(first.status(), StatusCode::OK); + let _ = first.text().await.unwrap(); // drain so the tee stores the aggregate + + let second = client + .post(format!("{url}/v1/chat/completions")) + .json(&request_body) + .send() + .await + .unwrap(); + assert_eq!(second.status(), StatusCode::OK); + assert_eq!( + second + .headers() + .get(http::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()), + Some("text/event-stream"), + "a replayed stream must declare the SSE content type" + ); + let replayed = second.text().await.unwrap(); + assert!( + replayed.contains("The answer is 42."), + "the replay must carry the stored answer: {replayed}" + ); + assert!( + replayed.trim_end().ends_with("data: [DONE]"), + "a chat replay must terminate the SSE stream: {replayed}" + ); + assert_eq!( + upstream_calls.load(std::sync::atomic::Ordering::SeqCst), + 1, + "the streaming repeat must be served from cache" + ); + + shutdown_tx.send(()).unwrap(); + handle.await.unwrap().unwrap(); + let _ = nemo_relay::plugin::clear_plugin_configuration(); +} diff --git a/crates/node/adaptive.d.ts b/crates/node/adaptive.d.ts index a2ad4d30c..be4639491 100644 --- a/crates/node/adaptive.d.ts +++ b/crates/node/adaptive.d.ts @@ -52,6 +52,19 @@ export interface AcgConfig { stability_thresholds?: AcgStabilityThresholds; } +/** Opt-in LLM response cache (exact-match) settings. */ +export interface ResponseCacheConfig { + ttl_seconds?: number; + namespace?: string; + priority?: number; + bypass_rate?: number; + cache_nondeterministic?: boolean; + key_strategy?: string; + header_allowlist?: string[]; + skip_keys?: string[]; + backend?: BackendSpec; +} + /** Canonical config object for the top-level adaptive component. */ export interface Config { version?: number; @@ -61,6 +74,7 @@ export interface Config { adaptive_hints?: AdaptiveHintsConfig; tool_parallelism?: ToolParallelismConfig; acg?: AcgConfig; + response_cache?: ResponseCacheConfig; policy?: ConfigPolicy; } @@ -244,6 +258,20 @@ export declare function toolParallelismConfig(config?: ToolParallelismConfig): T * callers can override only the thresholds they need. */ export declare function acgConfig(config?: AcgConfig): AcgConfig; +/** + * Create response-cache settings with defaults applied. + * + * Merges caller-supplied overrides onto the opt-in LLM response-cache config + * shape (exact-match) used by the adaptive plugin. This is a section of + * the adaptive component, not a standalone plugin kind. + * + * @param config - Partial response-cache settings to override. + * @returns A normalized response-cache config object. + * @remarks The default backend is in-memory; pass a `backend` (e.g. + * `redisBackend(url)`) for a shared cache. `bypass_rate` defaults to `0.0` + * (always reuse / exact replay). + */ +export declare function responseCacheConfig(config?: ResponseCacheConfig): ResponseCacheConfig; /** * Wrap adaptive config as a top-level component. * diff --git a/crates/node/adaptive.js b/crates/node/adaptive.js index 8392f46a3..0f6448e82 100644 --- a/crates/node/adaptive.js +++ b/crates/node/adaptive.js @@ -147,6 +147,35 @@ function acgConfig(config = {}) { }; } +/** + * Create response-cache settings with defaults applied. + * + * Merges caller-supplied overrides onto the opt-in LLM response-cache config + * shape (exact-match) used by the adaptive plugin. This is a section of + * the adaptive component, not a standalone plugin kind. + * + * @param {object} [config={}] - Partial response-cache settings to override. + * @returns {object} A normalized response-cache config object. + * @remarks The default backend is in-memory; pass a `backend` (e.g. + * `redisBackend(url)`) for a shared cache. `bypass_rate` defaults to `0.0` + * (always reuse / exact replay). + */ +function responseCacheConfig(config = {}) { + const { backend, ...rest } = config; + return { + ttl_seconds: 3600, + namespace: '', + priority: 50, + bypass_rate: 0.0, + cache_nondeterministic: true, + key_strategy: 'exact_request', + header_allowlist: [], + skip_keys: [], + backend: backend ?? inMemoryBackend(), + ...rest, + }; +} + /** * Wrap adaptive config as a top-level plugin component. * @@ -205,6 +234,7 @@ module.exports = { adaptiveHintsConfig, toolParallelismConfig, acgConfig, + responseCacheConfig, ComponentSpec, validateConfig, buildCacheTelemetryEvent, diff --git a/go/nemo_relay/adaptive.go b/go/nemo_relay/adaptive.go index eb7ce81f5..f9073fdb4 100644 --- a/go/nemo_relay/adaptive.go +++ b/go/nemo_relay/adaptive.go @@ -17,6 +17,7 @@ type AdaptiveConfig struct { AdaptiveHints *AdaptiveHintsConfig `json:"adaptive_hints,omitempty"` ToolParallelism *ToolParallelismConfig `json:"tool_parallelism,omitempty"` Acg *AcgConfig `json:"acg,omitempty"` + ResponseCache *ResponseCacheConfig `json:"response_cache,omitempty"` Policy *ConfigPolicy `json:"policy,omitempty"` } @@ -66,6 +67,49 @@ type AcgConfig struct { StabilityThresholds *AcgStabilityThresholds `json:"stability_thresholds,omitempty"` } +// ResponseCacheConfig configures the opt-in LLM response cache: a section +// of the adaptive config (a sibling to acg/adaptive_hints/tool_parallelism), not a +// standalone plugin kind. The Rust core validates and installs it from the adaptive +// runtime; this struct only has to carry the section through to the FFI validator. +type ResponseCacheConfig struct { + // TTLSeconds is how long a stored answer stays reusable, in seconds (> 0). + // Always serialized: with omitempty an explicit 0 would vanish and Rust's + // default (3600) would silently replace it instead of failing validation. + TTLSeconds uint64 `json:"ttl_seconds"` + // Namespace is folded into every key to separate environments/tenants. + Namespace string `json:"namespace,omitempty"` + // Priority is the execution-intercept priority; lower runs first/outermost. + // Always serialized so an explicit 0 (outermost) stays expressible. + Priority int32 `json:"priority"` + // BypassRate is the probability in [0.0, 1.0] of skipping the cache and running live. + BypassRate float64 `json:"bypass_rate,omitempty"` + // CacheNondeterministic caches temperature>0 requests too (default true). + // + // Deliberately NOT omitempty: the Rust default is true, so omitting a false + // value would let the core re-apply true and silently re-enable caching of + // nondeterministic requests. Always serializing the field keeps an explicit + // false expressible — unlike the other zero-default fields here. + CacheNondeterministic bool `json:"cache_nondeterministic"` + // KeyStrategy is the key strategy. Only "exact_request" is supported. + KeyStrategy string `json:"key_strategy,omitempty"` + // HeaderAllowlist lists request headers folded into the key; never auth headers. + HeaderAllowlist []string `json:"header_allowlist,omitempty"` + // SkipKeys lists extra top-level request-body keys to drop from the key. + SkipKeys []string `json:"skip_keys,omitempty"` + // Backend selects the cache's own storage backend (distinct from the adaptive + // state backend). Defaults to in-memory when nil. + Backend *ResponseCacheBackendConfig `json:"backend,omitempty"` +} + +// ResponseCacheBackendConfig selects the response-cache backend kind and options. +type ResponseCacheBackendConfig struct { + // Kind is "in_memory" or "redis" (redis needs the redis-backend build feature). + Kind string `json:"kind"` + // Config holds backend-specific options (in_memory: max_bytes/max_entries; + // redis: url/key_prefix). + Config map[string]any `json:"config,omitempty"` +} + // AdaptiveComponentSpec wraps one adaptive config as a top-level plugin component. type AdaptiveComponentSpec struct { Enabled bool `json:"enabled,omitempty"` @@ -138,6 +182,38 @@ func NewAcgConfig() AcgConfig { } } +// NewResponseCacheConfig returns default response cache settings, mirroring +// the Rust ResponseCacheConfig defaults (caching on, exact-request keying). Backend +// is left nil so the core applies its in-memory default; set it for redis or to tune +// the in-memory budget. +func NewResponseCacheConfig() ResponseCacheConfig { + return ResponseCacheConfig{ + TTLSeconds: 3600, + Priority: 50, + CacheNondeterministic: true, + KeyStrategy: "exact_request", + } +} + +// NewInMemoryResponseCacheBackend returns an in-memory response-cache backend spec. +func NewInMemoryResponseCacheBackend() ResponseCacheBackendConfig { + return ResponseCacheBackendConfig{ + Kind: "in_memory", + Config: map[string]any{}, + } +} + +// NewRedisResponseCacheBackend returns a Redis response-cache backend spec. +func NewRedisResponseCacheBackend(url, keyPrefix string) ResponseCacheBackendConfig { + return ResponseCacheBackendConfig{ + Kind: "redis", + Config: map[string]any{ + "url": url, + "key_prefix": keyPrefix, + }, + } +} + // NewAdaptiveComponentSpec wraps adaptive config as an enabled top-level component. func NewAdaptiveComponentSpec(config AdaptiveConfig) AdaptiveComponentSpec { return AdaptiveComponentSpec{ diff --git a/go/nemo_relay/adaptive/adaptive.go b/go/nemo_relay/adaptive/adaptive.go index 493064225..e18877054 100644 --- a/go/nemo_relay/adaptive/adaptive.go +++ b/go/nemo_relay/adaptive/adaptive.go @@ -52,6 +52,12 @@ type AcgStabilityThresholds = nemo_relay.AcgStabilityThresholds // AcgConfig configures the adaptive cache governor. type AcgConfig = nemo_relay.AcgConfig +// ResponseCacheConfig configures the opt-in LLM response cache. +type ResponseCacheConfig = nemo_relay.ResponseCacheConfig + +// ResponseCacheBackendConfig selects the response-cache backend kind and options. +type ResponseCacheBackendConfig = nemo_relay.ResponseCacheBackendConfig + // CacheUsage is normalized LLM token usage for cache telemetry. type CacheUsage = nemo_relay.CacheUsage @@ -113,6 +119,21 @@ func NewAcgConfig() AcgConfig { return nemo_relay.NewAcgConfig() } +// NewResponseCacheConfig returns default response cache settings. +func NewResponseCacheConfig() ResponseCacheConfig { + return nemo_relay.NewResponseCacheConfig() +} + +// NewInMemoryResponseCacheBackend returns an in-memory response-cache backend spec. +func NewInMemoryResponseCacheBackend() ResponseCacheBackendConfig { + return nemo_relay.NewInMemoryResponseCacheBackend() +} + +// NewRedisResponseCacheBackend returns a Redis response-cache backend spec. +func NewRedisResponseCacheBackend(url, keyPrefix string) ResponseCacheBackendConfig { + return nemo_relay.NewRedisResponseCacheBackend(url, keyPrefix) +} + // NewComponentSpec wraps adaptive config as an enabled top-level adaptive component. func NewComponentSpec(config Config) ComponentSpec { return nemo_relay.NewAdaptiveComponentSpec(config) diff --git a/go/nemo_relay/adaptive_runtime_test.go b/go/nemo_relay/adaptive_runtime_test.go index 51db86a5d..d4f0cd153 100644 --- a/go/nemo_relay/adaptive_runtime_test.go +++ b/go/nemo_relay/adaptive_runtime_test.go @@ -169,3 +169,69 @@ func TestSetLatencySensitivityRejectsInvalidValue(t *testing.T) { t.Fatal("expected SetLatencySensitivity(0) to fail") } } + +func TestResponseCacheConfigReachesTypedSurface(t *testing.T) { + backend := NewInMemoryResponseCacheBackend() + rc := NewResponseCacheConfig() + rc.Namespace = "go-harness" + rc.CacheNondeterministic = false // must survive marshal (no omitempty) + rc.Backend = &backend + + config := NewAdaptiveConfig() + config.ResponseCache = &rc + + // 1. The typed AdaptiveConfig must carry response_cache through json.Marshal. + // The bug this guards is the struct silently DROPPING the section because it + // enumerates fields by name with no response_cache field and no catch-all. + payload, err := json.Marshal(config) + if err != nil { + t.Fatalf("marshal failed: %v", err) + } + var decoded map[string]any + if err := json.Unmarshal(payload, &decoded); err != nil { + t.Fatalf("unmarshal failed: %v", err) + } + section, ok := decoded["response_cache"].(map[string]any) + if !ok { + t.Fatalf("response_cache missing from marshaled config: %s", payload) + } + if section["namespace"] != "go-harness" { + t.Fatalf("response_cache fields not preserved: %#v", section) + } + if v, ok := section["cache_nondeterministic"].(bool); !ok || v { + t.Fatalf("explicit cache_nondeterministic=false was not preserved: %#v", section["cache_nondeterministic"]) + } + if b, ok := section["backend"].(map[string]any); !ok || b["kind"] != "in_memory" { + t.Fatalf("backend not preserved: %#v", section["backend"]) + } + + // 2. A valid section validates clean through the FFI -> Rust adaptive validator. + report, err := ValidateAdaptiveConfig(config) + if err != nil { + t.Fatalf("ValidateAdaptiveConfig failed: %v", err) + } + if len(report.Diagnostics) != 0 { + t.Fatalf("expected clean report, got %#v", report.Diagnostics) + } + + // 3. An invalid section produces a response_cache diagnostic. This proves the + // section is actually validated end-to-end (a dropped section would yield no + // diagnostic at all), not merely carried in the struct. + bad := NewResponseCacheConfig() + bad.BypassRate = 2.0 + badConfig := NewAdaptiveConfig() + badConfig.ResponseCache = &bad + badReport, err := ValidateAdaptiveConfig(badConfig) + if err != nil { + t.Fatalf("ValidateAdaptiveConfig (invalid bypass_rate) returned error: %v", err) + } + found := false + for _, d := range badReport.Diagnostics { + if d.Code == "response_cache.invalid_bypass_rate" { + found = true + } + } + if !found { + t.Fatalf("expected response_cache.invalid_bypass_rate diagnostic, got %#v", badReport.Diagnostics) + } +} diff --git a/go/nemo_relay/plugin.go b/go/nemo_relay/plugin.go index af39171b6..c35b70ea5 100644 --- a/go/nemo_relay/plugin.go +++ b/go/nemo_relay/plugin.go @@ -159,8 +159,10 @@ type ConfigReport struct { // PluginComponentSpec is one top-level plugin component. type PluginComponentSpec struct { - Kind string `json:"kind"` - Enabled bool `json:"enabled,omitempty"` + Kind string `json:"kind"` + // Always serialized: with omitempty an explicit Enabled=false is omitted + // and the Rust default (true) silently re-enables the component. + Enabled bool `json:"enabled"` Config map[string]any `json:"config,omitempty"` } diff --git a/python/nemo_relay/adaptive.py b/python/nemo_relay/adaptive.py index e4f375457..beb1a7b86 100644 --- a/python/nemo_relay/adaptive.py +++ b/python/nemo_relay/adaptive.py @@ -251,6 +251,54 @@ def to_dict(self) -> JsonObject: ) +@dataclass(slots=True) +class ResponseCacheConfig: + """Opt-in LLM response cache (exact-match) settings. + + This is a section of the adaptive component, not a standalone plugin kind. + When present, the adaptive plugin installs the response-cache execution + intercept that reuses an earlier answer for a repeated managed LLM call. + + Args: + ttl_seconds: How long a stored answer stays reusable, in seconds. + namespace: Namespace folded into every key to separate environments/tenants. + priority: Execution-intercept priority. Lower runs first/outermost. + bypass_rate: Probability in ``[0.0, 1.0]`` of skipping the cache and running live. + cache_nondeterministic: Cache nondeterministic requests too; ``False`` + caches only requests explicitly pinned deterministic (``temperature`` = 0). + key_strategy: Key strategy. Only ``"exact_request"`` is supported. + header_allowlist: Request headers folded into the key; never auth headers. + skip_keys: Extra top-level request-body keys to drop from the key. + backend: Cache storage backend (``in_memory`` or ``redis``). + """ + + ttl_seconds: int = 3600 + namespace: str = "" + priority: int = 50 + bypass_rate: float = 0.0 + cache_nondeterministic: bool = True + key_strategy: str = "exact_request" + header_allowlist: list[str] = field(default_factory=list) + skip_keys: list[str] = field(default_factory=list) + backend: BackendSpec = field(default_factory=BackendSpec.in_memory) + + def to_dict(self) -> JsonObject: + """Serialize this response-cache config to the canonical JSON object shape.""" + return _normalize_object( + { + "ttl_seconds": self.ttl_seconds, + "namespace": self.namespace, + "priority": self.priority, + "bypass_rate": self.bypass_rate, + "cache_nondeterministic": self.cache_nondeterministic, + "key_strategy": self.key_strategy, + "header_allowlist": self.header_allowlist, + "skip_keys": self.skip_keys, + "backend": _normalize(self.backend), + } + ) + + @dataclass(slots=True) class AdaptiveConfig: """Canonical config document for the top-level adaptive component. @@ -263,6 +311,7 @@ class AdaptiveConfig: adaptive_hints: Built-in LLM hint-injection settings. tool_parallelism: Built-in tool scheduling settings. acg: Adaptive Cache Governor settings. + response_cache: Opt-in LLM response cache settings. policy: Unsupported-config policy applied within the adaptive config. Behavior: @@ -277,6 +326,7 @@ class AdaptiveConfig: adaptive_hints: AdaptiveHintsConfig | None = None tool_parallelism: ToolParallelismConfig | None = None acg: AcgConfig | None = None + response_cache: ResponseCacheConfig | None = None policy: ConfigPolicy = field(default_factory=ConfigPolicy) def to_dict(self) -> JsonObject: @@ -289,6 +339,7 @@ def to_dict(self) -> JsonObject: "adaptive_hints": _normalize(self.adaptive_hints), "tool_parallelism": _normalize(self.tool_parallelism), "acg": _normalize(self.acg), + "response_cache": _normalize(self.response_cache), "policy": self.policy.to_dict(), } @@ -385,6 +436,7 @@ def set_latency_sensitivity(level: int) -> None: "ConfigPolicy", "ConfigReport", "ComponentSpec", + "ResponseCacheConfig", "StateConfig", "TelemetryConfig", "ToolParallelismConfig", diff --git a/python/nemo_relay/adaptive.pyi b/python/nemo_relay/adaptive.pyi index 892f7260b..b9e07cd9d 100644 --- a/python/nemo_relay/adaptive.pyi +++ b/python/nemo_relay/adaptive.pyi @@ -181,6 +181,39 @@ class AcgConfig: """Serialize this ACG config to the canonical JSON object shape.""" ... +@dataclass(slots=True) +class ResponseCacheConfig: + """Opt-in LLM response cache (exact-match) settings. + + A section of the adaptive component, not a standalone plugin kind. + + Args: + ttl_seconds: How long a stored answer stays reusable, in seconds. + namespace: Namespace folded into every key to separate environments/tenants. + priority: Execution-intercept priority. Lower runs first/outermost. + bypass_rate: Probability in ``[0.0, 1.0]`` of skipping the cache and running live. + cache_nondeterministic: Cache nondeterministic requests too; ``False`` + caches only requests explicitly pinned deterministic (``temperature`` = 0). + key_strategy: Key strategy. Only ``"exact_request"`` is supported. + header_allowlist: Request headers folded into the key. + skip_keys: Extra top-level request-body keys to drop from the key. + backend: Cache storage backend (``in_memory`` or ``redis``). + """ + + ttl_seconds: int = ... + namespace: str = ... + priority: int = ... + bypass_rate: float = ... + cache_nondeterministic: bool = ... + key_strategy: str = ... + header_allowlist: list[str] = ... + skip_keys: list[str] = ... + backend: BackendSpec = ... + + def to_dict(self) -> JsonObject: + """Serialize this response-cache config to the canonical JSON object shape.""" + ... + @dataclass(slots=True) class AdaptiveConfig: """Canonical config document for the top-level adaptive component. @@ -193,6 +226,7 @@ class AdaptiveConfig: adaptive_hints: Built-in adaptive request-hints configuration. tool_parallelism: Built-in adaptive tool-scheduling configuration. acg: Adaptive Cache Governor configuration. + response_cache: Opt-in LLM response cache configuration. policy: Policy for unsupported adaptive configuration. """ @@ -203,6 +237,7 @@ class AdaptiveConfig: adaptive_hints: AdaptiveHintsConfig | None = ... tool_parallelism: ToolParallelismConfig | None = ... acg: AcgConfig | None = ... + response_cache: ResponseCacheConfig | None = ... policy: ConfigPolicy = ... def to_dict(self) -> JsonObject: diff --git a/python/tests/test_adaptive_config.py b/python/tests/test_adaptive_config.py index eec629c4f..e1fc89145 100644 --- a/python/tests/test_adaptive_config.py +++ b/python/tests/test_adaptive_config.py @@ -15,6 +15,7 @@ BackendSpec, ComponentSpec, ConfigPolicy, + ResponseCacheConfig, StateConfig, TelemetryConfig, ToolParallelismConfig, @@ -156,6 +157,46 @@ def test_acg_config_allows_threshold_overrides(self): "min_observations_for_full_confidence": 12, } + def test_response_cache_config_serializes_with_defaults(self): + assert ResponseCacheConfig().to_dict() == { + "ttl_seconds": 3600, + "namespace": "", + "priority": 50, + "bypass_rate": 0.0, + "cache_nondeterministic": True, + "key_strategy": "exact_request", + "header_allowlist": [], + "skip_keys": [], + "backend": {"kind": "in_memory", "config": {}}, + } + + def test_response_cache_rides_the_adaptive_component(self): + component = ComponentSpec(AdaptiveConfig(response_cache=ResponseCacheConfig(namespace="dev"))).to_dict() + assert component["kind"] == "adaptive" + config = cast(dict[str, object], component["config"]) + response_cache = cast(dict[str, object], config["response_cache"]) + assert response_cache["namespace"] == "dev" + + def test_response_cache_clean_report(self): + report = plugin.validate( + plugin.PluginConfig( + components=[ComponentSpec(AdaptiveConfig(response_cache=ResponseCacheConfig(namespace="dev")))] + ) + ) + assert report["diagnostics"] == [] + + def test_invalid_response_cache_section_is_rejected(self): + report = plugin.validate( + plugin.PluginConfig( + components=[ + ComponentSpec(AdaptiveConfig(response_cache=ResponseCacheConfig(ttl_seconds=0, bypass_rate=2.0))) + ] + ) + ) + codes = {diag["code"] for diag in report["diagnostics"]} + assert "response_cache.invalid_ttl" in codes + assert "response_cache.invalid_bypass_rate" in codes + def test_canonical_cache_telemetry_helper_supports_openai_provider(self): event = adaptive_module.build_cache_telemetry_event( provider="openai", From 1e7cafa86e585abffc45b80792ff3c0d60d48ff3 Mon Sep 17 00:00:00 2001 From: Zhongxuan Wang Date: Fri, 10 Jul 2026 17:49:00 -0700 Subject: [PATCH 2/2] docs: document the opt-in response cache Adds a Response Cache page to the adaptive plugin docs: when to use it, a plugins.toml quickstart, per-language plugin configuration, the complete-answers-only storage rules with streaming aggregation and native replay, exact-request key normalization, the response_cache mark and doctor surfaces, a field reference, and the unredacted-at-rest caveat for shared Redis deployments. Links the page from the adaptive About and Configuration pages. Signed-off-by: Zhongxuan Wang --- docs/adaptive-plugin/about.mdx | 3 + docs/adaptive-plugin/configuration.mdx | 6 +- docs/adaptive-plugin/response-cache.mdx | 325 ++++++++++++++++++++++++ 3 files changed, 332 insertions(+), 2 deletions(-) create mode 100644 docs/adaptive-plugin/response-cache.mdx diff --git a/docs/adaptive-plugin/about.mdx b/docs/adaptive-plugin/about.mdx index b053c8dc9..9552096ce 100644 --- a/docs/adaptive-plugin/about.mdx +++ b/docs/adaptive-plugin/about.mdx @@ -22,6 +22,7 @@ The plugin can coordinate: - Adaptive hints injected into outgoing model requests. - Tool-parallelism observation or scheduling behavior. - Adaptive Cache Governor (ACG) prompt-cache planning. +- Opt-in response caching for repeated LLM calls. - Component-local validation policy. ## Use Adaptive When @@ -51,6 +52,8 @@ If instrumentation is not in place yet, start with cache planning accomplishes. - [Adaptive Hints](/adaptive-plugin/adaptive-hints) explains request hint injection and how downstream model paths can consume the hints. +- [Response Cache](/adaptive-plugin/response-cache) explains the opt-in LLM response cache: + turning it on, what gets cached, and how savings are reported. State, telemetry, tool parallelism, and policy are whole-plugin configuration areas. They are documented on [Adaptive Configuration](/adaptive-plugin/configuration) rather diff --git a/docs/adaptive-plugin/configuration.mdx b/docs/adaptive-plugin/configuration.mdx index 55621aef7..7c1cc1dfb 100644 --- a/docs/adaptive-plugin/configuration.mdx +++ b/docs/adaptive-plugin/configuration.mdx @@ -32,10 +32,12 @@ The top-level adaptive object contains: | `adaptive_hints` | Request hint-injection behavior. | | `tool_parallelism` | Tool scheduling observation or scheduling behavior. | | `acg` | Adaptive Cache Governor prompt-cache planning. | +| `response_cache` | Opt-in LLM response cache for repeated managed calls. | | `policy` | Adaptive-local handling for unknown fields and unsupported values. | -The requested area pages cover [Adaptive Cache Governor (ACG)](/adaptive-plugin/acg) and -[Adaptive Hints](/adaptive-plugin/adaptive-hints). State, telemetry, tool parallelism, and +The requested area pages cover [Adaptive Cache Governor (ACG)](/adaptive-plugin/acg), +[Adaptive Hints](/adaptive-plugin/adaptive-hints), and +[Response Cache](/adaptive-plugin/response-cache). State, telemetry, tool parallelism, and policy remain whole-plugin settings: - Use `state.backend.kind = "in_memory"` for local experiments. diff --git a/docs/adaptive-plugin/response-cache.mdx b/docs/adaptive-plugin/response-cache.mdx new file mode 100644 index 000000000..0fcb53fb5 --- /dev/null +++ b/docs/adaptive-plugin/response-cache.mdx @@ -0,0 +1,325 @@ +--- +title: "Response Cache" +description: "" +position: 5 +--- +{/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 */} + + +Use the response cache when the same LLM request is made more than once and the +repeat should be served from a store instead of paying for the call again. A +repeat of an identical request returns the earlier answer instantly, at no +cost, and byte-for-byte unchanged — usage fields intact, so a cached reuse is +shape-identical to a live call. + +The cache is an optional `response_cache` section of the +[Adaptive plugin](/adaptive-plugin/configuration), not a standalone plugin +kind. It is off until the section is present, applies to +[managed LLM calls](/instrument-applications/instrument-llm-call) without any +call-site changes, and fails open: any cache error falls back to a normal live +call. + +## When To Use It + +- Development and test loops that replay the same prompts — repeats are + instant, free, and reproducible. +- CI suites that exercise real models — only the first run of each distinct + request pays. +- Demos and workshops — stable answers and no cost spikes. +- A shared team cache — point the `redis` backend at one store and identical + requests hit across processes and machines. +- Gateway deployments — enable caching for an agent without touching its code. + +Reuse requires the request to match exactly after normalization, so prompts +that embed volatile content (timestamps, random IDs) will not repeat. A stored +answer is also frozen for its TTL — size `ttl_seconds` to how fresh answers +must be, and use `bypass_rate` to re-run a sample of cacheable calls live. + +## `plugins.toml` Example + +```toml +version = 1 + +[[components]] +kind = "adaptive" +enabled = true + +[components.config] +version = 1 + +[components.config.response_cache] +ttl_seconds = 3600 # how long a stored answer stays reusable +namespace = "dev-harness" # separates environments, tenants, and experiments +bypass_rate = 0.0 # 0.1 would re-run 10% of cacheable calls live + +[components.config.response_cache.backend] +kind = "in_memory" # or "redis" for a shared cache +``` + +With this configuration the first occurrence of a request runs the provider +and stores the answer; an identical repeat within the TTL is served from the +store and skips the provider entirely. The request's provider surface is +auto-detected from its shape, so there is nothing else to configure. For file +discovery and precedence rules, see +[Plugin Configuration Files](/build-plugins/plugin-configuration-files). + +## Plugin Configuration + +Use plugin configuration when the application should let NeMo Relay own the +cache lifecycle. + + + +```python +import nemo_relay + +adaptive_config = nemo_relay.adaptive.AdaptiveConfig( + response_cache=nemo_relay.adaptive.ResponseCacheConfig( + ttl_seconds=3600, + namespace="dev-harness", + ), +) + +plugin_config = nemo_relay.plugin.PluginConfig( + components=[nemo_relay.adaptive.ComponentSpec(adaptive_config)] +) + +report = nemo_relay.plugin.validate(plugin_config) +if any(diagnostic["level"] == "error" for diagnostic in report["diagnostics"]): + raise RuntimeError(report["diagnostics"]) + +await nemo_relay.plugin.initialize(plugin_config) + +# Managed calls need no changes. The first call runs the provider; +# an identical repeat is served from the cache. +request = nemo_relay.LLMRequest( + {}, {"model": "gpt-4o", "messages": [{"role": "user", "content": "What is NeMo Relay?"}]} +) +await nemo_relay.llm.execute("openai", request, call_model) # miss: runs call_model +await nemo_relay.llm.execute("openai", request, call_model) # hit: call_model is skipped +``` + + + +```js +const adaptive = require("nemo-relay-node/adaptive"); +const plugin = require("nemo-relay-node/plugin"); + +const adaptiveConfig = adaptive.defaultConfig(); +adaptiveConfig.response_cache = adaptive.responseCacheConfig({ + ttl_seconds: 3600, + namespace: "dev-harness", +}); + +const pluginConfig = plugin.defaultConfig(); +pluginConfig.components = [adaptive.ComponentSpec(adaptiveConfig)]; + +const report = plugin.validate(pluginConfig); +if (report.diagnostics.some((diagnostic) => diagnostic.level === "error")) { + throw new Error(JSON.stringify(report.diagnostics)); +} + +await plugin.initialize(pluginConfig); +// Managed LLM calls now serve identical repeats from the cache. +``` + + + +```rust +use nemo_relay::plugin::{initialize_plugins, validate_plugin_config, PluginConfig}; +use nemo_relay_adaptive::plugin_component::ComponentSpec; +use nemo_relay_adaptive::{AdaptiveConfig, ResponseCacheConfig}; + +let mut adaptive = AdaptiveConfig::default(); +adaptive.response_cache = Some(ResponseCacheConfig { + ttl_seconds: 3600, + namespace: "dev-harness".into(), + ..ResponseCacheConfig::default() +}); + +let mut plugin_config = PluginConfig::default(); +plugin_config.components.push(ComponentSpec::new(adaptive).into()); + +let report = validate_plugin_config(&plugin_config); +assert!(!report.has_errors()); + +let active = initialize_plugins(plugin_config).await?; +// Managed LLM calls now serve identical repeats from the cache. +``` + + + + + +NeMo Relay plugin configuration keys use `snake_case` regardless of language or +file format. The `backend` helpers differ slightly by binding: Python uses +`nemo_relay.adaptive.BackendSpec.in_memory()` / `BackendSpec.redis(url)`, and +Node.js uses `adaptive.inMemoryBackend()` / `adaptive.redisBackend(url)`. + + +## Manual API + +Use the manual runtime API when an integration needs to own the adaptive +lifecycle directly instead of activating the top-level plugin component. + + + +```python +import nemo_relay + +adaptive_config = nemo_relay.adaptive.AdaptiveConfig( + response_cache=nemo_relay.adaptive.ResponseCacheConfig(namespace="dev-harness"), +) + +runtime = nemo_relay.adaptive.AdaptiveRuntime(adaptive_config.to_dict()) +await runtime.register() +try: + # Run instrumented application work here. + runtime.wait_for_idle() +finally: + await runtime.shutdown() +``` + + + +The Node.js binding exposes the response cache through the adaptive plugin +component helpers. Use the Plugin Configuration example above when activating +the cache from Node.js. + + + +```rust +use nemo_relay_adaptive::{AdaptiveConfig, AdaptiveRuntime, ResponseCacheConfig}; + +let mut adaptive = AdaptiveConfig::default(); +adaptive.response_cache = Some(ResponseCacheConfig { + namespace: "dev-harness".into(), + ..ResponseCacheConfig::default() +}); + +let mut runtime = AdaptiveRuntime::new(adaptive).await?; +runtime.register().await?; + +// Run instrumented application work here. + +runtime.wait_for_idle(); +runtime.shutdown().await?; +``` + + + + +## What Gets Cached + +Only complete, replayable answers are stored: + +- A response with a non-null `error`, or a non-final `status` such as + `failed`, `cancelled`, `incomplete`, or `in_progress`, is never stored. A + failed upstream reply passes through to the caller untouched. +- Stateful requests bypass the cache entirely: a request that opts into + server-side persistence (a non-`false` `store` flag, or an OpenAI Responses + call without an explicit `store = false`), continues a stored interaction + (`previous_response_id`), or references server-side state (`conversation`, + `container`). Their answers depend on state the cache key cannot see. +- Streaming calls are cached too. On a miss the live chunks are forwarded to + the consumer while being assembled into one aggregate response — the same + shape a buffered call stores, so buffered and streaming calls share one + keyspace. On a hit the stored answer is replayed as provider-native chunks + (OpenAI Chat deltas, OpenAI Responses lifecycle events, or Anthropic + Messages events), so strict streaming clients parse it like a live stream. + Truncated or incomplete streams, and streams whose content cannot be + replayed faithfully (for example thinking blocks), are never stored; a + request whose surface cannot be inferred runs live, uncached. + +A hit returns the stored response unchanged. What the hit avoided — tokens and +estimated cost — is reported on the `response_cache` mark, never by editing +the response body. + + +By default the cache stores answers to nondeterministic requests as well: a +repeat serves one representative sample. Set `cache_nondeterministic = false` +to cache only requests explicitly pinned deterministic (`temperature = 0`). + + +## Cache Keys + +Two requests hit the same entry when they are the same request after +normalization, under the default `key_strategy = "exact_request"`: + +- The request is decoded to its normalized form and fingerprinted with + SHA-256, so provider-shaped differences that mean the same thing collapse to + one key. When a request cannot be decoded faithfully, the raw body is + fingerprinted instead — that fallback can only cost a miss, never a wrong + hit. +- Field order and whitespace never matter (RFC 8785 canonicalization). +- Volatile request fields are always dropped from the key: `stream`, `user`, + `metadata`, `service_tier`, and `store`. Add project-specific noise fields + with `skip_keys`. +- Tool-call IDs are normalized, so randomly generated per-call IDs do not + fragment the keyspace. +- Request headers stay out of the key unless explicitly named in + `header_allowlist`; auth headers are rejected outright. +- The `namespace` and the provider name are folded into every key, so + environments and providers never share entries. + +## Observability + +Every cache decision emits a `response_cache` mark with +`data.status` set to one of: + +| Status | Meaning | +|---|---| +| `hit` | Served the stored answer; the provider was skipped. | +| `miss` | No stored answer; the call ran live and, on a clean result, the answer was stored. | +| `bypass` | The request is not cacheable, or the `bypass_rate` sampler chose to run live. | + +Mark metadata rides under `nemo_relay.response_cache.*`: `backend`, +`key_hash` (the `sha256:…` fingerprint), `ttl_ms`, `age_ms` and +`saved_tokens` / `saved_cost_usd` on hits, and a `reason` on bypasses and +store-error misses (for example `sampled`, `stateful_store`, `store_error`, +`stream_no_codec`). Marks carry only fingerprints — never prompts, answers, or +credentials. + +`nemo-relay doctor` reports the cache state: `not configured` when the section +is absent, `configured but disabled (adaptive plugin disabled)` when the +adaptive component is off, `on; backend '' reachable` when healthy, and +a failure when the config is invalid or the backend is unreachable. + +## Fields + +| Field | Default | Notes | +|---|---|---| +| `ttl_seconds` | `3600` | How long a stored answer stays reusable, in seconds. | +| `namespace` | `""` | Folded into every key; separates environments, tenants, and experiments. | +| `priority` | `50` | LLM execution intercept priority. Lower values run earlier. | +| `bypass_rate` | `0.0` | Probability in `[0.0, 1.0]` of running a cacheable call live. A sampled call refreshes the stored answer, so this doubles as drift detection. | +| `cache_nondeterministic` | `true` | Set `false` to cache only requests pinned deterministic (`temperature = 0`). | +| `key_strategy` | `"exact_request"` | The only supported strategy: reuse requires the same normalized request. | +| `header_allowlist` | `[]` | Headers folded into the key (case-insensitive). Auth headers are rejected. | +| `skip_keys` | `[]` | Extra top-level body keys dropped from the key, beyond the built-in volatile set. | +| `backend.kind` | `"in_memory"` | `"in_memory"`, or `"redis"` (requires building with the `redis-backend` feature). | +| `backend.config.max_bytes` | 256 MiB | In-memory size budget; the oldest entries are evicted first. | +| `backend.config.max_entries` | unset | Optional entry-count cap on top of the byte budget. | +| `backend.config.url` | — | Redis connection URL. Required for the `redis` backend. | +| `backend.config.key_prefix` | `"nemo-relay:llm-cache:"` | Prefix for keys in Redis. | + + +Cached entries are stored unredacted. PII sanitize guardrails rewrite emitted +telemetry, never payloads, so the store holds full request and response +bodies. A shared Redis backend must be trusted, access-controlled, and +namespaced per environment and tenant. + + +## Common Validation Failures + +- `ttl_seconds` is `0`, or `bypass_rate` is outside `[0.0, 1.0]`. +- `key_strategy` is not `"exact_request"`. +- `skip_keys` names an answer-determining field such as `messages`, `model`, + `tools`, or `tool_choice`. +- `header_allowlist` names an auth header such as `authorization` or + `x-api-key`. +- `backend.kind = "redis"` without `backend.config.url`, or in a build without + the `redis-backend` feature. +- A `redis` backend with an empty `namespace` warns: every caller of the + shared store would exchange cached answers.