diff --git a/crates/adaptive/src/config.rs b/crates/adaptive/src/config.rs index d5db2445a..4ac432c36 100644 --- a/crates/adaptive/src/config.rs +++ b/crates/adaptive/src/config.rs @@ -7,7 +7,7 @@ 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}; +use crate::response_cache::config::{BackendConfig, KEY_STRATEGY_EXACT_REQUEST, ToolCacheConfig}; /// Canonical config document for the adaptive plugin component. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -34,8 +34,8 @@ 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). + /// Opt-in exact-match LLM response and tool-result cache. 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. @@ -191,7 +191,8 @@ impl Default for AcgComponentConfig { } } -/// Configuration for the adaptive plugin's `response_cache` feature +/// Configuration for the adaptive plugin's exact-match LLM response and +/// opt-in tool-result cache feature. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(default)] pub struct ResponseCacheConfig { @@ -217,6 +218,9 @@ pub struct ResponseCacheConfig { pub header_allowlist: Vec, /// Storage backend selection. pub backend: BackendConfig, + /// Opt-in tool-result cache configuration. + #[serde(skip_serializing_if = "Option::is_none")] + pub tools: Option, } impl Default for ResponseCacheConfig { @@ -230,6 +234,7 @@ impl Default for ResponseCacheConfig { key_strategy: KEY_STRATEGY_EXACT_REQUEST.to_string(), header_allowlist: Vec::new(), backend: BackendConfig::default(), + tools: None, } } } @@ -405,6 +410,13 @@ nemo_relay::editor_config! { nested: BackendConfig, default: BackendConfig, }, + tools => { + label: "tools", + kind: Section, + optional: true, + nested: ToolCacheConfig, + default: ToolCacheConfig, + }, } } diff --git a/crates/adaptive/src/lib.rs b/crates/adaptive/src/lib.rs index 24fd27cb1..237dfbe8e 100644 --- a/crates/adaptive/src/lib.rs +++ b/crates/adaptive/src/lib.rs @@ -34,7 +34,7 @@ pub mod learner; pub mod plugin_component; #[cfg(feature = "redis-backend")] pub mod redis; -/// Opt-in LLM response cache (exact-match). +/// Opt-in exact-match LLM response and tool-result cache. pub mod response_cache; mod runtime; /// Storage backends and backend traits for adaptive state persistence. @@ -58,6 +58,7 @@ pub use error::{AdaptiveError, Result}; #[cfg(feature = "redis-backend")] pub use redis::RedisBackend; pub use response_cache::RESPONSE_CACHE_MARK; +pub use response_cache::config::{ToolCacheConfig, ToolClass, ToolOverride}; 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 1d1b530a9..9e5ada654 100644 --- a/crates/adaptive/src/plugin_component.rs +++ b/crates/adaptive/src/plugin_component.rs @@ -338,6 +338,7 @@ fn validate_response_cache_section( "key_strategy", "header_allowlist", "backend", + "tools", ], ); if let Some(backend_json) = response_cache_json.get("backend").and_then(Json::as_object) { @@ -361,6 +362,9 @@ fn validate_response_cache_section( ); } } + if let Some(tools_json) = response_cache_json.get("tools").and_then(Json::as_object) { + validate_response_cache_tools_fields(diagnostics, policy, tools_json); + } } fn validate_response_cache_backend_config_fields( @@ -383,6 +387,78 @@ fn validate_response_cache_backend_config_fields( ); } +fn validate_response_cache_tools_fields( + diagnostics: &mut Vec, + policy: &ConfigPolicy, + tools_json: &Map, +) { + const CLASS_FIELDS: &[&str] = &[ + "cacheable", + "ttl_seconds", + "bypass_rate", + "arg_skip", + "members", + ]; + const OVERRIDE_FIELDS: &[&str] = &[ + "cacheable", + "ttl_seconds", + "bypass_rate", + "tool_version", + "arg_skip", + ]; + + validate_unknown_fields( + diagnostics, + policy, + Some("response_cache.tools".to_string()), + tools_json, + &[ + "enabled", + "priority", + "cache_errors", + "default", + "classes", + "overrides", + ], + ); + + if let Some(default_json) = tools_json.get("default").and_then(Json::as_object) { + validate_unknown_fields( + diagnostics, + policy, + Some("response_cache.tools.default".to_string()), + default_json, + CLASS_FIELDS, + ); + } + if let Some(classes_json) = tools_json.get("classes").and_then(Json::as_object) { + for (class_name, class_value) in classes_json { + if let Some(class_object) = class_value.as_object() { + validate_unknown_fields( + diagnostics, + policy, + Some(format!("response_cache.tools.classes.{class_name}")), + class_object, + CLASS_FIELDS, + ); + } + } + } + if let Some(overrides_json) = tools_json.get("overrides").and_then(Json::as_object) { + for (tool_name, override_value) in overrides_json { + if let Some(override_object) = override_value.as_object() { + validate_unknown_fields( + diagnostics, + policy, + Some(format!("response_cache.tools.overrides.{tool_name}")), + override_object, + OVERRIDE_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 index 42c8499a1..88757747b 100644 --- a/crates/adaptive/src/response_cache/config.rs +++ b/crates/adaptive/src/response_cache/config.rs @@ -9,6 +9,8 @@ //! response-cache-specific backend config and the key-strategy constant next to //! the key/store code that consumes them. +use std::collections::BTreeMap; + use serde::{Deserialize, Serialize}; use serde_json::{Map, Value as Json}; @@ -63,3 +65,112 @@ nemo_relay::editor_config! { config => { label: "config", kind: Json }, } } + +/// Opt-in tool-result cache configuration. +/// +/// Cache only tools that are read-only and stable for their TTL. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default)] +pub struct ToolCacheConfig { + /// Master switch; off by default. + pub enabled: bool, + /// Tool execution-intercept priority. The default keeps standard + /// priority-100 guardrails outside cache hits. + pub priority: i32, + /// Whether conventional in-band tool error results may be stored. + pub cache_errors: bool, + /// Policy for unclassified tools; not cacheable by default. + pub default: ToolClass, + /// Named tool classes. + pub classes: BTreeMap, + /// Per-tool refinements keyed by exact name or wildcard. + pub overrides: BTreeMap, +} + +impl Default for ToolCacheConfig { + fn default() -> Self { + Self { + enabled: false, + priority: 150, + cache_errors: false, + default: ToolClass::default(), + classes: BTreeMap::new(), + overrides: BTreeMap::new(), + } + } +} + +nemo_relay::editor_config! { + impl ToolCacheConfig { + enabled => { label: "enabled", kind: Boolean }, + priority => { label: "priority", kind: Integer }, + cache_errors => { label: "cache_errors", kind: Boolean }, + default => { + label: "default", + kind: Section, + nested: ToolClass, + default: ToolClass, + }, + classes => { label: "classes", kind: Json }, + overrides => { label: "overrides", kind: Json }, + } +} + +/// Policy shared by a class of tools. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default)] +pub struct ToolClass { + /// Whether class members may be served from cache. + pub cacheable: bool, + /// TTL in seconds; inherits the response-cache TTL when unset. + #[serde(skip_serializing_if = "Option::is_none")] + pub ttl_seconds: Option, + /// Live-rerun probability; inherits the response-cache rate when unset. + #[serde(skip_serializing_if = "Option::is_none")] + pub bypass_rate: Option, + /// Top-level argument keys dropped before keying. + pub arg_skip: Vec, + /// Exact tool names or `*` wildcard patterns in this class. + pub members: Vec, +} + +/// Per-tool refinement applied after class resolution. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default)] +pub struct ToolOverride { + /// Overrides the class cacheability. + #[serde(skip_serializing_if = "Option::is_none")] + pub cacheable: Option, + /// Overrides the class TTL. + #[serde(skip_serializing_if = "Option::is_none")] + pub ttl_seconds: Option, + /// Overrides the class bypass rate. + #[serde(skip_serializing_if = "Option::is_none")] + pub bypass_rate: Option, + /// Version string folded into the cache key. + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_version: Option, + /// Replaces the class argument skip list when set. + #[serde(skip_serializing_if = "Option::is_none")] + pub arg_skip: Option>, +} + +nemo_relay::editor_config! { + impl ToolClass { + cacheable => { label: "cacheable", kind: Boolean }, + ttl_seconds => { label: "ttl_seconds", kind: Integer, optional: true }, + bypass_rate => { label: "bypass_rate", kind: Float, optional: true }, + arg_skip => { label: "arg_skip", kind: Json }, + members => { label: "members", kind: Json }, + } +} + +nemo_relay::editor_config! { + impl ToolOverride { + cacheable => { label: "cacheable", kind: Boolean, optional: true }, + ttl_seconds => { label: "ttl_seconds", kind: Integer, optional: true }, + bypass_rate => { label: "bypass_rate", kind: Float, optional: true }, + tool_version => { label: "tool_version", kind: String, optional: true }, + arg_skip => { label: "arg_skip", kind: Json, optional: true }, + } +} diff --git a/crates/adaptive/src/response_cache/key.rs b/crates/adaptive/src/response_cache/key.rs index dbdb1726b..de751b949 100644 --- a/crates/adaptive/src/response_cache/key.rs +++ b/crates/adaptive/src/response_cache/key.rs @@ -14,6 +14,8 @@ //! skip-list drops volatile/identity fields, tool-call IDs are normalized, and //! only allowlisted headers plus Relay-owned routing partitions fold in. +use std::collections::BTreeSet; + use nemo_relay::api::llm::LlmRequest; use nemo_relay::codec::request::AnnotatedLlmRequest; use nemo_relay::codec::resolve::{ @@ -74,7 +76,8 @@ pub fn build_cache_key( normalize_tool_call_ids(object); } - let headers = cache_key_headers(&request.headers, &config.header_allowlist); + let header_allowlist = normalized_header_allowlist(&config.header_allowlist); + let headers = cache_key_headers(&request.headers, &header_allowlist); let key_doc = json!({ "v": CACHE_SCHEMA_VERSION, @@ -85,6 +88,7 @@ pub fn build_cache_key( "openai_chat_token_cap": chat_token_cap_spelling, "body": body, "headers": headers, + "header_allowlist": header_allowlist, }); if contains_unrepresentable_int(&key_doc) { return KeyOutcome::Bypass("unrepresentable_number"); @@ -211,6 +215,47 @@ impl std::io::Write for HashWriter<'_> { } } +/// Builds a tool-result key from its name, version, canonicalized arguments, +/// and the effective cache policies. +pub fn build_tool_cache_key( + namespace: &str, + tool_name: &str, + tool_version: Option<&str>, + args: &Json, + arg_skip: &[String], + cache_errors: bool, +) -> KeyOutcome { + let arg_skip = normalized_arg_skip(arg_skip); + let mut args = args.clone(); + if !arg_skip.is_empty() + && let Some(object) = args.as_object_mut() + { + for key in &arg_skip { + object.remove(key); + } + } + + if contains_unrepresentable_int(&args) { + return KeyOutcome::Bypass("unrepresentable_number"); + } + + let key_doc = json!({ + "v": CACHE_SCHEMA_VERSION, + "surface": "tool_result", + "ns": namespace, + "tool": tool_name, + "tool_version": tool_version, + "arg_skip": arg_skip, + "cache_errors": cache_errors, + "args": args, + }); + + match fingerprint(&key_doc) { + Some(key) => KeyOutcome::Key(key), + None => KeyOutcome::Bypass("canonicalization_failed"), + } +} + /// The body to fingerprint plus the codec that actually produced it. /// /// The surface is auto-detected from the request shape, hinted by the provider @@ -555,6 +600,26 @@ fn allowlisted_headers(headers: &Map, allowlist: &[String]) -> Map kept } +/// Normalizes case-insensitive header policy names before keying them. +fn normalized_header_allowlist(allowlist: &[String]) -> Vec { + allowlist + .iter() + .map(|name| name.to_ascii_lowercase()) + .collect::>() + .into_iter() + .collect() +} + +/// Normalizes the case-sensitive tool argument keys dropped before keying. +fn normalized_arg_skip(arg_skip: &[String]) -> Vec { + arg_skip + .iter() + .cloned() + .collect::>() + .into_iter() + .collect() +} + /// Builds the key's header partition from configured headers plus the /// Relay-owned Switchyard backend ID. fn cache_key_headers(headers: &Map, allowlist: &[String]) -> Map { diff --git a/crates/adaptive/src/response_cache/mark.rs b/crates/adaptive/src/response_cache/mark.rs index 2900ac1da..922c9d9e7 100644 --- a/crates/adaptive/src/response_cache/mark.rs +++ b/crates/adaptive/src/response_cache/mark.rs @@ -17,6 +17,24 @@ use crate::response_cache::store::CacheEntry; /// Mark-event name emitted on every cache decision. pub const RESPONSE_CACHE_MARK: &str = "response_cache"; +/// Execution surface that made a response-cache decision. +#[derive(Debug, Clone, Copy)] +pub(crate) enum CacheSurface { + /// An LLM response. + Llm, + /// A tool result. + Tool, +} + +impl CacheSurface { + const fn as_str(self) -> &'static str { + match self { + Self::Llm => "llm", + Self::Tool => "tool", + } + } +} + /// Pulls saved token count and cost out of a stored entry (the aggregate /// response object — buffered and streaming both store this shape). /// @@ -103,12 +121,14 @@ fn probed_savings(entry: &CacheEntry) -> (Option, Option) { pub(crate) struct CacheMark<'a> { status: &'a str, reason: Option<&'a str>, + surface: CacheSurface, backend: &'a str, key_hash: Option<&'a str>, age_ms: Option, ttl_ms: Option, saved_tokens: Option, saved_cost_usd: Option, + saved_invocations: Option, } impl<'a> CacheMark<'a> { @@ -116,12 +136,14 @@ impl<'a> CacheMark<'a> { Self { status, reason: None, + surface: CacheSurface::Llm, backend, key_hash: None, age_ms: None, ttl_ms: None, saved_tokens: None, saved_cost_usd: None, + saved_invocations: None, } } @@ -130,6 +152,12 @@ impl<'a> CacheMark<'a> { self } + /// Overrides the cache surface (defaults to [`CacheSurface::Llm`]). + pub(crate) fn surface(mut self, surface: CacheSurface) -> Self { + self.surface = surface; + self + } + pub(crate) fn key_hash(mut self, key_hash: &'a str) -> Self { self.key_hash = Some(key_hash); self @@ -150,6 +178,12 @@ impl<'a> CacheMark<'a> { self.saved_cost_usd = cost; self } + + /// Records the number of tool invocations a hit avoided (tool surface). + pub(crate) fn saved_invocations(mut self, invocations: u64) -> Self { + self.saved_invocations = Some(invocations); + self + } } /// Emits the `response_cache` mark. Only the key fingerprint is ever recorded — @@ -158,7 +192,7 @@ pub(crate) fn emit_cache_mark(mark: CacheMark<'_>) { let mut metadata = Map::new(); metadata.insert( "nemo_relay.response_cache.surface".to_string(), - json!("llm"), + json!(mark.surface.as_str()), ); metadata.insert( "nemo_relay.response_cache.backend".to_string(), @@ -200,6 +234,12 @@ pub(crate) fn emit_cache_mark(mark: CacheMark<'_>) { json!(saved_cost_usd), ); } + if let Some(saved_invocations) = mark.saved_invocations { + metadata.insert( + "nemo_relay.response_cache.saved_invocations".to_string(), + json!(saved_invocations), + ); + } let _ = event( EmitMarkEventParams::builder() diff --git a/crates/adaptive/src/response_cache/mod.rs b/crates/adaptive/src/response_cache/mod.rs index aac0bff7b..f8cb4b8b7 100644 --- a/crates/adaptive/src/response_cache/mod.rs +++ b/crates/adaptive/src/response_cache/mod.rs @@ -1,8 +1,11 @@ // 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`]. +//! Opt-in exact-match cache for LLM responses and tool results: a feature of +//! the adaptive plugin, configured through +//! [`crate::config::AdaptiveConfig::response_cache`]. +//! +//! The two surfaces share storage but use disjoint keys. //! //! [`intercept`] holds the execution intercepts and storage rules, [`key`] the //! cache-key derivation, [`store`] the backends, [`replay`] the streaming @@ -17,11 +20,15 @@ pub(crate) mod replay; /// health check; not part of the user-facing API. #[doc(hidden)] pub mod store; +pub(crate) mod tool; pub use crate::config::ResponseCacheConfig; -pub use crate::response_cache::config::{BackendConfig, KEY_STRATEGY_EXACT_REQUEST}; +pub use crate::response_cache::config::{ + BackendConfig, KEY_STRATEGY_EXACT_REQUEST, ToolCacheConfig, +}; 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; #[doc(hidden)] pub use crate::response_cache::store::check_backend_health; +pub(crate) use crate::response_cache::tool::make_tool_intercept; diff --git a/crates/adaptive/src/response_cache/tool.rs b/crates/adaptive/src/response_cache/tool.rs new file mode 100644 index 000000000..6855c41e6 --- /dev/null +++ b/crates/adaptive/src/response_cache/tool.rs @@ -0,0 +1,724 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Opt-in tool-result cache. +//! +//! A hit suppresses the real call, so caching is off by default and must be +//! enabled only for tools that are read-only and stable for the configured TTL. +//! Key and store failures fail open to the real call. + +use std::collections::HashSet; +use std::sync::Arc; +use std::time::Duration; + +use nemo_relay::api::runtime::{ToolExecutionFn, ToolExecutionNextFn}; +use nemo_relay::api::tool::ToolExecutionInterceptOutcome; +use nemo_relay::error::Result as FlowResult; +use serde_json::Value as Json; + +use crate::config::ResponseCacheConfig; +use crate::response_cache::config::{ToolCacheConfig, ToolClass, ToolOverride}; +use crate::response_cache::intercept::should_bypass; +use crate::response_cache::key::{KeyOutcome, build_tool_cache_key}; +use crate::response_cache::mark::{CacheMark, CacheSurface, emit_cache_mark}; +use crate::response_cache::store::{CacheEntry, CacheStore, now_unix_ms}; + +#[derive(Debug, Clone, PartialEq)] +struct ResolvedToolPolicy { + cacheable: bool, + ttl: Duration, + bypass_rate: f64, + arg_skip: Vec, + tool_version: Option, +} + +fn resolve_policy( + tool_name: &str, + response_cache: &ResponseCacheConfig, + tools: &ToolCacheConfig, +) -> ResolvedToolPolicy { + let class: &ToolClass = resolve_class(tool_name, tools).unwrap_or(&tools.default); + let over: Option<&ToolOverride> = resolve_override(tool_name, tools); + + let cacheable = over + .and_then(|over| over.cacheable) + .unwrap_or(class.cacheable); + + let ttl_seconds = over + .and_then(|over| over.ttl_seconds) + .or(class.ttl_seconds) + .unwrap_or(response_cache.ttl_seconds); + + let bypass_rate = over + .and_then(|over| over.bypass_rate) + .or(class.bypass_rate) + .unwrap_or(response_cache.bypass_rate); + + let arg_skip = match over.and_then(|over| over.arg_skip.clone()) { + Some(list) => list, + None => class.arg_skip.clone(), + }; + + let tool_version = over.and_then(|over| over.tool_version.clone()); + + ResolvedToolPolicy { + cacheable, + ttl: Duration::from_secs(ttl_seconds), + bypass_rate, + arg_skip, + tool_version, + } +} + +fn resolve_class<'a>(tool_name: &str, tools: &'a ToolCacheConfig) -> Option<&'a ToolClass> { + for class in tools.classes.values() { + if class + .members + .iter() + .any(|member| !member.contains('*') && member == tool_name) + { + return Some(class); + } + } + best_wildcard_match( + tools.classes.values().flat_map(|class| { + class + .members + .iter() + .map(move |member| (member.as_str(), class)) + }), + tool_name, + ) +} + +fn resolve_override<'a>(tool_name: &str, tools: &'a ToolCacheConfig) -> Option<&'a ToolOverride> { + if let Some(over) = tools.overrides.get(tool_name) { + return Some(over); + } + best_wildcard_match( + tools + .overrides + .iter() + .map(|(key, over)| (key.as_str(), over)), + tool_name, + ) +} + +fn best_wildcard_match<'a, T>( + candidates: impl Iterator, + name: &str, +) -> Option<&'a T> { + let mut best = None; + for (pattern, candidate) in candidates { + if !pattern.contains('*') || !wildcard_match(pattern, name) { + continue; + } + let rank = wildcard_rank(pattern); + if best.as_ref().is_none_or(|(_, current)| rank > *current) { + best = Some((candidate, rank)); + } + } + best.map(|(candidate, _)| candidate) +} + +type WildcardRank<'a> = (usize, std::cmp::Reverse, std::cmp::Reverse<&'a str>); + +/// Returns the deterministic specificity order for a wildcard pattern. +/// +/// Literal and wildcard counts are Unicode-character based. The final +/// lexicographic component only breaks otherwise equal ranks. +fn wildcard_rank(pattern: &str) -> WildcardRank<'_> { + let stars = pattern + .chars() + .filter(|character| *character == '*') + .count(); + ( + pattern.chars().count() - stars, + std::cmp::Reverse(stars), + std::cmp::Reverse(pattern), + ) +} + +/// Returns whether two `*` patterns can match at least one common tool name. +/// +/// This evaluates the product of the two wildcard automata, so it is exact for +/// this deliberately small pattern language without constructing a sample name. +pub(crate) fn wildcard_patterns_overlap(left: &str, right: &str) -> bool { + let left: Vec = left.chars().collect(); + let right: Vec = right.chars().collect(); + let mut pending = vec![(0, 0)]; + let mut visited = HashSet::new(); + + while let Some((left_index, right_index)) = pending.pop() { + if !visited.insert((left_index, right_index)) { + continue; + } + if left_index == left.len() && right_index == right.len() { + return true; + } + + if left.get(left_index) == Some(&'*') { + pending.push((left_index + 1, right_index)); + } + if right.get(right_index) == Some(&'*') { + pending.push((left_index, right_index + 1)); + } + + let (Some(left_character), Some(right_character)) = + (left.get(left_index), right.get(right_index)) + else { + continue; + }; + if *left_character == '*' || *right_character == '*' || left_character == right_character { + pending.push(( + left_index + usize::from(*left_character != '*'), + right_index + usize::from(*right_character != '*'), + )); + } + } + + false +} + +fn wildcard_match(pattern: &str, name: &str) -> bool { + if !pattern.contains('*') { + return pattern == name; + } + let segments: Vec<&str> = pattern.split('*').collect(); + let (first, rest) = segments + .split_first() + .expect("split always yields a segment"); + if !name.starts_with(first) { + return false; + } + let mut cursor = first.len(); + let (last, middles) = rest + .split_last() + .expect("a starred pattern splits into at least two segments"); + for segment in middles { + match name[cursor..].find(segment) { + Some(position) => cursor += position + segment.len(), + None => return false, + } + } + name.len() >= cursor + last.len() && name.ends_with(last) +} + +pub(crate) fn make_tool_intercept( + store: Arc, + response_cache: Arc, + tools: Arc, +) -> ToolExecutionFn { + Arc::new(move |name: &str, args: Json, next: ToolExecutionNextFn| { + let store = Arc::clone(&store); + let response_cache = Arc::clone(&response_cache); + let tools = Arc::clone(&tools); + let name = name.to_string(); + Box::pin(run_tool_cache( + name, + args, + next, + store, + response_cache, + tools, + )) + }) +} + +async fn run_tool_cache( + name: String, + args: Json, + next: ToolExecutionNextFn, + store: Arc, + response_cache: Arc, + tools: Arc, +) -> FlowResult { + let policy = resolve_policy(&name, &response_cache, &tools); + + if !policy.cacheable { + return next(args).await.map(Into::into); + } + + let backend = store.backend_kind(); + + let key = match build_tool_cache_key( + &response_cache.namespace, + &name, + policy.tool_version.as_deref(), + &args, + &policy.arg_skip, + tools.cache_errors, + ) { + KeyOutcome::Key(key) => key, + KeyOutcome::Bypass(reason) => { + emit_cache_mark( + CacheMark::new("bypass", backend) + .surface(CacheSurface::Tool) + .reason(reason), + ); + return next(args).await.map(Into::into); + } + }; + + if should_bypass(policy.bypass_rate) { + emit_cache_mark( + CacheMark::new("bypass", backend) + .surface(CacheSurface::Tool) + .reason("sampled") + .key_hash(&key), + ); + let result = next(args).await?; + store_tool_result(&store, &key, policy.ttl, &result, tools.cache_errors).await; + return Ok(result.into()); + } + + match store.get(&key).await { + Ok(Some(entry)) if !tools.cache_errors && is_error_shaped_tool_result(&entry.response) => { + // A prior Relay version could have stored a snake_case `is_error` + // result under this same policy. Never replay it after error + // caching is disabled; a successful live result replaces it. + emit_cache_mark( + CacheMark::new("bypass", backend) + .surface(CacheSurface::Tool) + .reason("cached_error") + .key_hash(&key), + ); + let result = next(args).await?; + store_tool_result(&store, &key, policy.ttl, &result, tools.cache_errors).await; + Ok(result.into()) + } + Ok(Some(entry)) => { + let age_ms = now_unix_ms().saturating_sub(entry.created_unix_ms); + emit_cache_mark( + CacheMark::new("hit", backend) + .surface(CacheSurface::Tool) + .key_hash(&key) + .age_ms(age_ms) + .ttl_ms(policy.ttl.as_millis() as u64) + .saved_invocations(1), + ); + Ok(entry.response.clone().into()) + } + Ok(None) => { + emit_cache_mark( + CacheMark::new("miss", backend) + .surface(CacheSurface::Tool) + .key_hash(&key) + .ttl_ms(policy.ttl.as_millis() as u64), + ); + let result = next(args).await?; + store_tool_result(&store, &key, policy.ttl, &result, tools.cache_errors).await; + Ok(result.into()) + } + Err(_) => { + emit_cache_mark( + CacheMark::new("miss", backend) + .surface(CacheSurface::Tool) + .reason("store_error") + .key_hash(&key), + ); + next(args).await.map(Into::into) + } + } +} + +async fn store_tool_result( + store: &Arc, + key: &str, + ttl: Duration, + result: &Json, + cache_errors: bool, +) { + if !cache_errors && is_error_shaped_tool_result(result) { + return; + } + let entry = CacheEntry::new(result.clone(), ttl, key.to_string(), None, None); + let _ = store.set(key, entry, ttl).await; +} + +/// A tool result has no universal provider envelope. Treat only the explicit, +/// widely used in-band error signals as failures by default; applications that +/// use these fields for stable data can opt into caching them. +fn is_error_shaped_tool_result(result: &Json) -> bool { + let Some(object) = result.as_object() else { + return false; + }; + object.get("error").is_some_and(|error| !error.is_null()) + || object.get("isError").and_then(Json::as_bool) == Some(true) + || object.get("is_error").and_then(Json::as_bool) == Some(true) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::response_cache::config::ToolClass; + use std::collections::BTreeMap; + + fn response_cache(ttl_seconds: u64, bypass_rate: f64) -> ResponseCacheConfig { + ResponseCacheConfig { + ttl_seconds, + bypass_rate, + ..ResponseCacheConfig::default() + } + } + + fn class(cacheable: bool, members: &[&str]) -> ToolClass { + ToolClass { + cacheable, + members: members.iter().map(|member| member.to_string()).collect(), + ..ToolClass::default() + } + } + + #[test] + fn unclassified_tool_falls_into_the_default_bucket_uncached() { + let tools = ToolCacheConfig::default(); + let policy = resolve_policy("anything", &response_cache(3600, 0.0), &tools); + assert!( + !policy.cacheable, + "an unknown tool must default to not cached" + ); + } + + #[test] + fn class_membership_makes_a_tool_cacheable() { + let mut classes = BTreeMap::new(); + classes.insert("read_only".to_string(), class(true, &["docs_lookup"])); + classes.insert( + "volatile".to_string(), + ToolClass { + cacheable: true, + ttl_seconds: Some(300), + bypass_rate: Some(0.2), + members: vec!["get_weather".to_string()], + ..ToolClass::default() + }, + ); + let tools = ToolCacheConfig { + classes, + ..ToolCacheConfig::default() + }; + let policy = resolve_policy("docs_lookup", &response_cache(3600, 0.0), &tools); + assert!(policy.cacheable); + assert_eq!(policy.ttl, Duration::from_secs(3600)); + assert_eq!(policy.bypass_rate, 0.0); + let policy = resolve_policy("get_weather", &response_cache(3600, 0.0), &tools); + assert_eq!(policy.ttl, Duration::from_secs(300)); + assert_eq!(policy.bypass_rate, 0.2); + } + + #[test] + fn per_tool_override_wins_over_its_class() { + let mut classes = BTreeMap::new(); + classes.insert( + "read_only".to_string(), + ToolClass { + cacheable: true, + arg_skip: vec!["request_id".to_string()], + members: vec!["docs_lookup".to_string()], + ..ToolClass::default() + }, + ); + let mut overrides = BTreeMap::new(); + overrides.insert( + "docs_lookup".to_string(), + ToolOverride { + cacheable: Some(false), + ttl_seconds: Some(30), + bypass_rate: Some(0.25), + tool_version: Some("v2".to_string()), + ..ToolOverride::default() + }, + ); + let tools = ToolCacheConfig { + classes, + overrides, + ..ToolCacheConfig::default() + }; + let policy = resolve_policy("docs_lookup", &response_cache(3600, 0.0), &tools); + assert!(!policy.cacheable, "override cacheable=false must win"); + assert_eq!(policy.ttl, Duration::from_secs(30)); + assert_eq!(policy.bypass_rate, 0.25); + assert_eq!(policy.tool_version.as_deref(), Some("v2")); + assert_eq!(policy.arg_skip, vec!["request_id".to_string()]); + } + + #[test] + fn override_arg_skip_replaces_the_class_list() { + let mut classes = BTreeMap::new(); + classes.insert( + "read_only".to_string(), + ToolClass { + cacheable: true, + arg_skip: vec!["session_id".to_string()], + members: vec!["lookup".to_string()], + ..ToolClass::default() + }, + ); + let mut overrides = BTreeMap::new(); + overrides.insert( + "lookup".to_string(), + ToolOverride { + arg_skip: Some(vec![]), + ..ToolOverride::default() + }, + ); + let tools = ToolCacheConfig { + classes, + overrides, + ..ToolCacheConfig::default() + }; + let policy = resolve_policy("lookup", &response_cache(3600, 0.0), &tools); + assert!( + policy.arg_skip.is_empty(), + "an override arg_skip (even empty) replaces the class list" + ); + } + + #[test] + fn default_bucket_can_be_flipped_on_for_broad_coverage() { + let tools = ToolCacheConfig { + default: ToolClass { + cacheable: true, + ttl_seconds: Some(60), + bypass_rate: Some(0.5), + ..ToolClass::default() + }, + ..ToolCacheConfig::default() + }; + let policy = resolve_policy("unknown_tool", &response_cache(3600, 0.0), &tools); + assert!( + policy.cacheable, + "default cacheable=true covers unknown tools" + ); + assert_eq!(policy.ttl, Duration::from_secs(60)); + assert_eq!(policy.bypass_rate, 0.5); + } + + #[test] + fn wildcard_match_table() { + let cases = [ + ("*", "", true), + ("*", "anything", true), + ("docs_*", "docs_lookup", true), + ("docs_*", "docs_", true), + ("docs_*", "doc_lookup", false), + ("*_price", "stock_price", true), + ("*_price", "price", false), + ("get_*_price", "get_stock_price", true), + ("get_*_price", "get_price", false), + ("a*a", "a", false), + ("a*a", "aa", true), + ("a*a", "aba", true), + ("a*b*c", "abc", true), + ("a*b*c", "acb", false), + ("Docs_*", "docs_lookup", false), // case-sensitive + ("abc*", "abc*", true), // no escaping: '*' matches itself via the span + ]; + for (pattern, name, expected) in cases { + assert_eq!( + wildcard_match(pattern, name), + expected, + "wildcard_match({pattern:?}, {name:?})" + ); + } + } + + #[test] + fn wildcard_overlap_table() { + let cases = [ + ("*_email", "send_*", true), + ("delete_*", "*_record", true), + ("docs_*", "send_*", false), + ("a*b*c", "a*c", true), + ("é*", "*é", true), + ("foo*", "bar*", false), + ]; + for (left, right, expected) in cases { + assert_eq!( + wildcard_patterns_overlap(left, right), + expected, + "wildcard_patterns_overlap({left:?}, {right:?})" + ); + } + } + + #[test] + fn wildcard_rank_counts_unicode_characters_not_utf8_bytes() { + assert_eq!(wildcard_rank("*é*").0, 1); + assert_eq!(wildcard_rank("*éé*").0, 2); + assert_eq!(wildcard_rank("*💡*").0, 1); + } + + #[test] + fn wildcard_member_classifies_a_matching_tool() { + let mut classes = BTreeMap::new(); + classes.insert("read_only".to_string(), class(true, &["docs_*"])); + let tools = ToolCacheConfig { + classes, + ..ToolCacheConfig::default() + }; + assert!(resolve_policy("docs_lookup", &response_cache(3600, 0.0), &tools).cacheable); + assert!( + !resolve_policy("send_email", &response_cache(3600, 0.0), &tools).cacheable, + "a non-matching tool still falls through to default" + ); + } + + #[test] + fn exact_member_beats_any_wildcard_match() { + let mut classes = BTreeMap::new(); + classes.insert("a_wildcards".to_string(), class(true, &["docs_*"])); + classes.insert("b_exact".to_string(), class(false, &["docs_lookup"])); + let tools = ToolCacheConfig { + classes, + ..ToolCacheConfig::default() + }; + let policy = resolve_policy("docs_lookup", &response_cache(3600, 0.0), &tools); + assert!( + !policy.cacheable, + "the exact member's class must win over a matching wildcard" + ); + } + + #[test] + fn most_specific_wildcard_wins() { + let mut classes = BTreeMap::new(); + classes.insert("a_catch_all".to_string(), class(false, &["*"])); + classes.insert("b_docs".to_string(), class(true, &["docs_*"])); + let tools = ToolCacheConfig { + classes, + ..ToolCacheConfig::default() + }; + assert!( + resolve_policy("docs_lookup", &response_cache(3600, 0.0), &tools).cacheable, + "the pattern with more literal characters must win" + ); + assert!(!resolve_policy("send_email", &response_cache(3600, 0.0), &tools).cacheable); + } + + #[test] + fn equal_literals_fewer_stars_then_smaller_pattern_break_ties() { + let mut classes = BTreeMap::new(); + classes.insert("two_stars".to_string(), class(false, &["a*b*"])); + classes.insert("one_star".to_string(), class(true, &["ab*"])); + let tools = ToolCacheConfig { + classes, + ..ToolCacheConfig::default() + }; + assert!( + resolve_policy("ab", &response_cache(3600, 0.0), &tools).cacheable, + "with equal literal counts the pattern with fewer stars must win" + ); + + let mut classes = BTreeMap::new(); + classes.insert("suffix".to_string(), class(false, &["*x"])); + classes.insert("prefix".to_string(), class(true, &["x*"])); + let tools = ToolCacheConfig { + classes, + ..ToolCacheConfig::default() + }; + assert!( + !resolve_policy("x", &response_cache(3600, 0.0), &tools).cacheable, + "'*x' sorts before 'x*', so the suffix class must win the tie" + ); + } + + #[test] + fn override_patterns_apply_with_exact_keys_winning() { + let mut classes = BTreeMap::new(); + classes.insert("read_only".to_string(), class(true, &["docs_*"])); + let mut overrides = BTreeMap::new(); + overrides.insert( + "docs_secret_*".to_string(), + ToolOverride { + cacheable: Some(false), + ..ToolOverride::default() + }, + ); + overrides.insert( + "docs_secret_audit".to_string(), + ToolOverride { + cacheable: Some(true), + ..ToolOverride::default() + }, + ); + let tools = ToolCacheConfig { + classes, + overrides, + ..ToolCacheConfig::default() + }; + let cacheable = + |name: &str| resolve_policy(name, &response_cache(3600, 0.0), &tools).cacheable; + assert!( + !cacheable("docs_secret_dump"), + "a pattern override must apply to the tools it matches" + ); + assert!( + cacheable("docs_secret_audit"), + "an exact override key must win over a matching pattern" + ); + assert!( + cacheable("docs_lookup"), + "tools no override matches keep their class policy" + ); + let mut overrides = BTreeMap::new(); + overrides.insert( + "docs_*".to_string(), + ToolOverride { + cacheable: Some(false), + ..ToolOverride::default() + }, + ); + let mut classes = BTreeMap::new(); + classes.insert("read_only".to_string(), class(true, &["docs_*"])); + let tools = ToolCacheConfig { + classes, + overrides, + ..ToolCacheConfig::default() + }; + assert!( + !resolve_policy("docs_*", &response_cache(3600, 0.0), &tools).cacheable, + "the literal name `docs_*` resolves its exact entry" + ); + assert!(!resolve_policy("docs_lookup", &response_cache(3600, 0.0), &tools).cacheable); + } + + #[test] + fn most_specific_override_pattern_wins() { + let mut classes = BTreeMap::new(); + classes.insert("read_only".to_string(), class(true, &["docs_*"])); + let mut overrides = BTreeMap::new(); + overrides.insert( + "docs_*".to_string(), + ToolOverride { + cacheable: Some(true), + ..ToolOverride::default() + }, + ); + overrides.insert( + "docs_secret_*".to_string(), + ToolOverride { + cacheable: Some(false), + ..ToolOverride::default() + }, + ); + let tools = ToolCacheConfig { + classes, + overrides, + ..ToolCacheConfig::default() + }; + assert!( + !resolve_policy("docs_secret_dump", &response_cache(3600, 0.0), &tools).cacheable, + "`docs_secret_*` (more literal characters) must beat `docs_*`" + ); + } +} + +#[cfg(test)] +#[path = "../../tests/unit/response_cache/tool_policy_tests.rs"] +mod policy_tests; + +#[cfg(test)] +#[path = "../../tests/unit/response_cache/tool_tests.rs"] +mod coverage_tests; diff --git a/crates/adaptive/src/runtime/features.rs b/crates/adaptive/src/runtime/features.rs index f83862458..1be023de0 100644 --- a/crates/adaptive/src/runtime/features.rs +++ b/crates/adaptive/src/runtime/features.rs @@ -42,7 +42,9 @@ 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::response_cache::{ + build_store, make_intercept, make_stream_intercept, make_tool_intercept, +}; use crate::runtime::backend::build_backend; use crate::runtime::validation::validate_config; use crate::storage::traits::StorageBackendDyn; @@ -484,7 +486,7 @@ impl AdaptiveRuntime { } // The response cache is independent of the learning-state backend: it has // its own CacheStore and installs buffered and streaming LLM execution - // intercepts. + // intercepts plus an opt-in tool execution intercept. if let Some(config) = self.config.response_cache.clone() { pending.push(Box::new(ResponseCacheFeature::new(config, self.runtime_id))); } @@ -785,6 +787,7 @@ impl AdaptiveFeature for AcgFeature { struct ResponseCacheFeature { name: String, stream_name: String, + tool_name: String, priority: i32, config: ResponseCacheConfig, } @@ -794,6 +797,7 @@ impl ResponseCacheFeature { Self { name: format!("adaptive_{runtime_id}_response_cache_llm_execution"), stream_name: format!("adaptive_{runtime_id}_response_cache_llm_stream_execution"), + tool_name: format!("adaptive_{runtime_id}_response_cache_tool_execution"), priority: config.priority, config, } @@ -806,9 +810,6 @@ impl AdaptiveFeature for ResponseCacheFeature { ctx: &'a mut RegistrationContext<'_>, ) -> Pin> + Send + 'a>> { Box::pin(async move { - // Build the backend once, shared by both intercepts. A Redis backend - // that is unreachable at startup disables this optional feature; - // the intercepts themselves fail open on later store errors. let store = match build_store(&self.config).await { Ok(store) => store, Err(AdaptiveError::Storage(error)) => { @@ -816,7 +817,7 @@ impl AdaptiveFeature for ResponseCacheFeature { target: "nemo_relay.runtime", event = "adaptive_response_cache_store_init_failed"; "Adaptive runtime could not initialize the optional response cache; \ - managed LLM calls will run live: {error}" + managed LLM and tool calls will run live: {error}" ); return Ok(()); } @@ -831,8 +832,17 @@ impl AdaptiveFeature for ResponseCacheFeature { ctx.register_llm_stream_execution_intercept( &self.stream_name, self.priority, - make_stream_intercept(store, config), - ) + make_stream_intercept(store.clone(), config.clone()), + )?; + if let Some(tools) = self.config.tools.clone().filter(|tools| tools.enabled) { + let priority = tools.priority; + ctx.register_tool_execution_intercept( + &self.tool_name, + priority, + make_tool_intercept(store, config, Arc::new(tools)), + )?; + } + Ok(()) }) } } diff --git a/crates/adaptive/src/runtime/validation.rs b/crates/adaptive/src/runtime/validation.rs index f5e4bb7a1..c1478bae4 100644 --- a/crates/adaptive/src/runtime/validation.rs +++ b/crates/adaptive/src/runtime/validation.rs @@ -1,13 +1,16 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use std::collections::HashMap; + use nemo_relay::plugin::{ ConfigDiagnostic, ConfigPolicy, ConfigReport, DiagnosticLevel, UnsupportedBehavior, }; use serde_json::Value as Json; use crate::config::{AdaptiveConfig, BackendSpec, ResponseCacheConfig}; -use crate::response_cache::config::KEY_STRATEGY_EXACT_REQUEST; +use crate::response_cache::config::{KEY_STRATEGY_EXACT_REQUEST, ToolCacheConfig}; +use crate::response_cache::tool::wildcard_patterns_overlap; pub fn validate_config(config: &AdaptiveConfig) -> ConfigReport { let mut report = ConfigReport::default(); @@ -199,12 +202,182 @@ fn validate_response_cache(report: &mut ConfigReport, config: &ResponseCacheConf format!("unknown backend kind '{other}'"), )), } + + if let Some(tools) = &config.tools { + validate_tool_cache(report, tools); + } +} + +fn validate_tool_cache(report: &mut ConfigReport, tools: &ToolCacheConfig) { + validate_tool_policy( + report, + "default", + tools.default.ttl_seconds, + tools.default.bypass_rate, + ); + if !tools.default.members.is_empty() { + report.diagnostics.push(response_cache_error( + "response_cache.tool_default_members", + Some("tools.default"), + "tools.default.members is never matched (the default bucket applies to every \ + unclassified tool); move these names into a named class" + .to_string(), + )); + } + + let mut owning_class: HashMap<&str, &str> = HashMap::new(); + for (class_name, class) in &tools.classes { + validate_tool_policy( + report, + &format!("classes.{class_name}"), + class.ttl_seconds, + class.bypass_rate, + ); + for member in &class.members { + if let Some(previous) = owning_class.get(member.as_str()) { + if *previous != class_name.as_str() { + report.diagnostics.push(response_cache_error( + "response_cache.tool_multiple_classes", + Some("tools.classes"), + format!( + "tool member '{member}' appears in multiple classes ('{previous}' and \ + '{class_name}'); a member — exact name or pattern — may appear in at \ + most one class" + ), + )); + } + } else { + owning_class.insert(member.as_str(), class_name.as_str()); + } + if class.cacheable && !member.is_empty() && member.chars().all(|c| c == '*') { + report.diagnostics.push(response_cache_warning( + "response_cache.tool_catch_all_member", + Some("tools.classes"), + format!( + "class '{class_name}' lists the catch-all member '{member}' with \ + cacheable = true, which caches every tool no other class claims; \ + prefer flipping default.cacheable on explicitly if broad coverage is \ + intended" + ), + )); + } + } + } + + validate_conflicting_tool_class_patterns(report, tools); + + for (tool_name, over) in &tools.overrides { + validate_tool_policy( + report, + &format!("overrides.{tool_name}"), + over.ttl_seconds, + over.bypass_rate, + ); + if over.cacheable == Some(true) + && !tool_name.is_empty() + && tool_name.chars().all(|c| c == '*') + { + report.diagnostics.push(response_cache_warning( + "response_cache.tool_catch_all_override", + Some("tools.overrides"), + format!( + "override '{tool_name}' sets cacheable = true for every tool; prefer \ + flipping default.cacheable on explicitly if broad coverage is intended" + ), + )); + } + } + + validate_conflicting_tool_override_patterns(report, tools); +} + +fn validate_conflicting_tool_class_patterns(report: &mut ConfigReport, tools: &ToolCacheConfig) { + for (index, (left_name, left_class)) in tools.classes.iter().enumerate() { + for (right_name, right_class) in tools.classes.iter().skip(index + 1) { + if left_class.cacheable == right_class.cacheable { + continue; + } + let conflict = left_class.members.iter().any(|left_member| { + left_member.contains('*') + && right_class.members.iter().any(|right_member| { + left_member != right_member + && right_member.contains('*') + && wildcard_patterns_overlap(left_member, right_member) + }) + }); + if conflict { + report.diagnostics.push(response_cache_error( + "response_cache.tool_conflicting_classes", + Some("tools.classes"), + format!( + "classes '{left_name}' and '{right_name}' contain overlapping wildcard \ + members with conflicting cacheable settings; split the patterns so one \ + policy applies to every tool" + ), + )); + } + } + } +} + +fn validate_conflicting_tool_override_patterns(report: &mut ConfigReport, tools: &ToolCacheConfig) { + for (index, (left_name, left_override)) in tools.overrides.iter().enumerate() { + for (right_name, right_override) in tools.overrides.iter().skip(index + 1) { + // An omitted value inherits from whichever class wins for the + // concrete tool name. Because overlapping patterns can select + // different classes, only identical declarations are safe. + let conflicting_cacheability = left_override.cacheable != right_override.cacheable; + if !conflicting_cacheability + || !left_name.contains('*') + || !right_name.contains('*') + || !wildcard_patterns_overlap(left_name, right_name) + { + continue; + } + report.diagnostics.push(response_cache_error( + "response_cache.tool_conflicting_overrides", + Some("tools.overrides"), + format!( + "overrides '{left_name}' and '{right_name}' overlap with conflicting \ + cacheable settings; split the patterns so one policy applies to every tool" + ), + )); + } + } +} + +fn validate_tool_policy( + report: &mut ConfigReport, + location: &str, + ttl_seconds: Option, + bypass_rate: Option, +) { + if ttl_seconds == Some(0) { + report.diagnostics.push(response_cache_error( + "response_cache.tool_invalid_ttl", + Some("tools"), + format!("tools.{location}.ttl_seconds must be greater than 0 when set"), + )); + } + if let Some(rate) = bypass_rate + && !(0.0..=1.0).contains(&rate) + { + report.diagnostics.push(response_cache_error( + "response_cache.tool_invalid_bypass_rate", + Some("tools"), + format!("tools.{location}.bypass_rate must be in [0.0, 1.0] when set"), + )); + } } 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, diff --git a/crates/adaptive/tests/integration/response_cache_benchmark_tests.rs b/crates/adaptive/tests/integration/response_cache_benchmark_tests.rs index b698a81da..49ad044f1 100644 --- a/crates/adaptive/tests/integration/response_cache_benchmark_tests.rs +++ b/crates/adaptive/tests/integration/response_cache_benchmark_tests.rs @@ -326,7 +326,6 @@ async fn reinitialized_cache_starts_empty() { would let the second run hit on the first run's distinct prompts" ); } - #[tokio::test] async fn warm_hits_stay_within_the_latency_budget() { let _guard = TEST_MUTEX.lock().await; diff --git a/crates/adaptive/tests/integration/response_cache_tests.rs b/crates/adaptive/tests/integration/response_cache_tests.rs index d16e8bf3f..a7365c7cf 100644 --- a/crates/adaptive/tests/integration/response_cache_tests.rs +++ b/crates/adaptive/tests/integration/response_cache_tests.rs @@ -16,19 +16,25 @@ use nemo_relay::api::llm::{ LlmCallExecuteParams, LlmRequest, LlmStreamCallExecuteParams, llm_call_execute, llm_stream_call_execute, }; +use nemo_relay::api::registry::{ + deregister_tool_execution_intercept, register_tool_execution_intercept, +}; use nemo_relay::api::runtime::{ LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, LlmStreamInner, - NemoRelayContextState, global_context, + NemoRelayContextState, ToolExecutionNextFn, global_context, }; use nemo_relay::api::scope::ScopeType; use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; +use nemo_relay::api::tool::{ToolCallExecuteParams, tool_call_execute}; use nemo_relay::error::FlowError; use nemo_relay::plugin::{ - PluginConfig, clear_plugin_configuration, initialize_plugins_exact, validate_plugin_config, + DiagnosticLevel, PluginConfig, clear_plugin_configuration, initialize_plugins_exact, + validate_plugin_config, }; use nemo_relay_adaptive::plugin_component::{ComponentSpec, register_adaptive_component}; use nemo_relay_adaptive::{ AcgComponentConfig, AdaptiveConfig, BackendSpec, ResponseCacheConfig, StateConfig, + ToolCacheConfig, ToolClass, ToolOverride, }; use serde_json::{Value as Json, json}; use tokio::sync::Mutex; @@ -549,6 +555,7 @@ async fn invalid_config_is_rejected_by_validation() { response_cache: Some(ResponseCacheConfig { ttl_seconds: 0, bypass_rate: 2.0, + key_strategy: "semantic".to_string(), namespace: "invalid-config-test".to_string(), ..ResponseCacheConfig::default() }), @@ -573,6 +580,13 @@ async fn invalid_config_is_rejected_by_validation() { .any(|diagnostic| diagnostic.code == "response_cache.invalid_bypass_rate"), "bypass_rate out of range must produce a diagnostic" ); + assert!( + report + .diagnostics + .iter() + .any(|diagnostic| diagnostic.code == "response_cache.unsupported_key_strategy"), + "an unsupported key strategy must produce a diagnostic" + ); } #[tokio::test] @@ -625,6 +639,81 @@ async fn unknown_and_unavailable_backends_are_rejected_by_validation() { } } +#[tokio::test] +async fn response_cache_validation_diagnostics_identify_the_invalid_setting() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + register_adaptive_component().unwrap(); + + let mut cache = ResponseCacheConfig { + namespace: "diagnostic-contract-test".to_string(), + key_strategy: "semantic".to_string(), + tools: Some(ToolCacheConfig { + enabled: true, + default: ToolClass { + bypass_rate: Some(-0.01), + ..ToolClass::default() + }, + ..ToolCacheConfig::default() + }), + ..ResponseCacheConfig::default() + }; + cache.backend.kind = "redis".to_string(); + + let report = validate_plugin_config(&PluginConfig { + components: vec![ + ComponentSpec::new(AdaptiveConfig { + response_cache: Some(cache), + ..AdaptiveConfig::default() + }) + .into(), + ], + ..PluginConfig::default() + }); + + assert!( + report.diagnostics.iter().any(|diagnostic| { + diagnostic.code == "response_cache.unsupported_key_strategy" + && diagnostic.level == DiagnosticLevel::Error + && diagnostic.component.as_deref() == Some("response_cache") + && diagnostic.field.as_deref() == Some("key_strategy") + }), + "an unsupported key strategy must identify its setting: {:?}", + report.diagnostics + ); + assert!( + report.diagnostics.iter().any(|diagnostic| { + diagnostic.code == "response_cache.tool_invalid_bypass_rate" + && diagnostic.level == DiagnosticLevel::Error + && diagnostic.field.as_deref() == Some("tools") + }), + "an invalid tool bypass rate must identify the tools section: {:?}", + report.diagnostics + ); + + #[cfg(not(feature = "redis-backend"))] + assert!( + report.diagnostics.iter().any(|diagnostic| { + diagnostic.code == "response_cache.backend_unavailable" + && diagnostic.level == DiagnosticLevel::Error + && diagnostic.field.as_deref() == Some("backend.kind") + }), + "redis must be rejected when its backend feature is not compiled: {:?}", + report.diagnostics + ); + + #[cfg(feature = "redis-backend")] + assert!( + report.diagnostics.iter().any(|diagnostic| { + diagnostic.code == "response_cache.missing_redis_url" + && diagnostic.level == DiagnosticLevel::Error + && diagnostic.field.as_deref() == Some("backend.config.url") + }), + "redis must identify a missing connection URL when its backend feature is compiled: {:?}", + report.diagnostics + ); +} + #[tokio::test] async fn hit_preserves_usage_on_the_end_event_and_reports_savings_on_the_mark() { let _guard = TEST_MUTEX.lock().await; @@ -1827,3 +1916,702 @@ async fn redis_backend_shares_entries_across_store_instances() { writer.delete(key).await.expect("delete"); assert!(reader.get(key).await.expect("get").is_none()); } + +fn counting_tool(calls: Arc, result: Json) -> ToolExecutionNextFn { + Arc::new(move |_args: Json| { + let calls = Arc::clone(&calls); + let result = result.clone(); + Box::pin(async move { + calls.fetch_add(1, Ordering::SeqCst); + Ok(result) + }) + }) +} + +async fn tool_call(name: &str, tool: &ToolExecutionNextFn, args: Json) -> Json { + tool_call_execute( + ToolCallExecuteParams::builder() + .name(name) + .args(args) + .func(tool.clone()) + .build(), + ) + .await + .unwrap() +} + +fn one_cacheable_class(members: &[&str]) -> ToolCacheConfig { + let mut classes = std::collections::BTreeMap::new(); + classes.insert( + "read_only".to_string(), + ToolClass { + cacheable: true, + members: members.iter().map(|member| member.to_string()).collect(), + ..ToolClass::default() + }, + ); + ToolCacheConfig { + enabled: true, + classes, + ..ToolCacheConfig::default() + } +} + +fn cache_with_tools(tools: ToolCacheConfig) -> ResponseCacheConfig { + ResponseCacheConfig { + namespace: "tool-cache-integration-test".to_string(), + tools: Some(tools), + ..ResponseCacheConfig::default() + } +} + +#[tokio::test] +async fn classified_tool_repeat_is_a_hit_that_skips_the_tool() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(cache_with_tools(one_cacheable_class(&["docs_lookup"]))).await; + + let calls = Arc::new(AtomicUsize::new(0)); + let tool = counting_tool(Arc::clone(&calls), json!({"doc": "the answer is 42"})); + + let first = tool_call("docs_lookup", &tool, json!({"q": "rust"})).await; + let second = tool_call("docs_lookup", &tool, json!({"q": "rust"})).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "a classified-cacheable tool must run once; the repeat is served from cache" + ); + assert_eq!(first, second, "a hit returns the stored result unchanged"); +} + +#[tokio::test] +async fn an_effectful_class_is_never_cached() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + let mut classes = std::collections::BTreeMap::new(); + classes.insert( + "effectful".to_string(), + ToolClass { + cacheable: false, + members: vec!["send_email".to_string()], + ..ToolClass::default() + }, + ); + activate_cache(cache_with_tools(ToolCacheConfig { + enabled: true, + classes, + ..ToolCacheConfig::default() + })) + .await; + + let calls = Arc::new(AtomicUsize::new(0)); + let tool = counting_tool(Arc::clone(&calls), json!({"sent": true})); + + tool_call("send_email", &tool, json!({"to": "a@b.c"})).await; + tool_call("send_email", &tool, json!({"to": "a@b.c"})).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "an effectful (cacheable=false) tool must run every time — a hit would skip the side effect" + ); +} + +#[tokio::test] +async fn disabled_tools_section_does_not_cache() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + let mut tools = one_cacheable_class(&["docs_lookup"]); + tools.enabled = false; + activate_cache(cache_with_tools(tools)).await; + + let calls = Arc::new(AtomicUsize::new(0)); + let tool = counting_tool(Arc::clone(&calls), json!({"doc": "x"})); + + tool_call("docs_lookup", &tool, json!({"q": "rust"})).await; + tool_call("docs_lookup", &tool, json!({"q": "rust"})).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "with tools.enabled = false the tool intercept is not installed" + ); +} + +#[tokio::test] +async fn conventional_error_shaped_tool_results_are_not_cached_by_default() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(cache_with_tools(one_cacheable_class(&["lookup"]))).await; + + let calls = Arc::new(AtomicUsize::new(0)); + let tool = counting_tool( + Arc::clone(&calls), + json!({"type": "tool_result", "is_error": true, "content": "not found"}), + ); + + tool_call("lookup", &tool, json!({"q": "missing"})).await; + tool_call("lookup", &tool, json!({"q": "missing"})).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "an Anthropic-style in-band tool error must run live again unless cache_errors is enabled" + ); +} + +#[tokio::test] +async fn conventional_error_shaped_tool_results_can_be_cached_when_opted_in() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + let mut tools = one_cacheable_class(&["lookup"]); + tools.cache_errors = true; + activate_cache(cache_with_tools(tools)).await; + + let calls = Arc::new(AtomicUsize::new(0)); + let tool = counting_tool( + Arc::clone(&calls), + json!({"type": "tool_result", "is_error": true, "content": "not found"}), + ); + + tool_call("lookup", &tool, json!({"q": "missing"})).await; + tool_call("lookup", &tool, json!({"q": "missing"})).await; + + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "cache_errors=true explicitly permits caching Anthropic-style in-band error results" + ); +} + +#[tokio::test] +async fn default_tool_cache_priority_keeps_standard_guardrails_on_hits() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + + let guardrail_runs = Arc::new(AtomicUsize::new(0)); + register_tool_execution_intercept( + "response_cache_standard_tool_guardrail_test", + 100, + Arc::new({ + let guardrail_runs = Arc::clone(&guardrail_runs); + move |_name, args, next| { + let guardrail_runs = Arc::clone(&guardrail_runs); + Box::pin(async move { + guardrail_runs.fetch_add(1, Ordering::SeqCst); + next(args).await.map(Into::into) + }) + } + }), + ) + .unwrap(); + activate_cache(cache_with_tools(one_cacheable_class(&["lookup"]))).await; + + let calls = Arc::new(AtomicUsize::new(0)); + let tool = counting_tool(Arc::clone(&calls), json!({"answer": "cached"})); + tool_call("lookup", &tool, json!({"q": "relay"})).await; + tool_call("lookup", &tool, json!({"q": "relay"})).await; + + deregister_tool_execution_intercept("response_cache_standard_tool_guardrail_test").unwrap(); + assert_eq!(calls.load(Ordering::SeqCst), 1, "the second call must hit"); + assert_eq!( + guardrail_runs.load(Ordering::SeqCst), + 2, + "the standard priority-100 guardrail must wrap and run on a cache hit" + ); +} + +#[tokio::test] +async fn tool_callback_errors_emit_misses_and_are_never_cached() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(cache_with_tools(one_cacheable_class(&["lookup"]))).await; + + let captured = Arc::new(StdMutex::new(Vec::::new())); + let sink = Arc::clone(&captured); + register_subscriber( + "response_cache_tool_callback_error_capture", + Arc::new(move |event: &Event| sink.lock().unwrap().push(event.clone())), + ) + .unwrap(); + + let calls = Arc::new(AtomicUsize::new(0)); + let tool: ToolExecutionNextFn = { + let calls = Arc::clone(&calls); + Arc::new(move |_args| { + let calls = Arc::clone(&calls); + Box::pin(async move { + calls.fetch_add(1, Ordering::SeqCst); + Err(FlowError::Internal("tool unavailable".to_string())) + }) + }) + }; + + for _ in 0..2 { + let error = tool_call_execute( + ToolCallExecuteParams::builder() + .name("lookup") + .args(json!({"q": "missing"})) + .func(tool.clone()) + .build(), + ) + .await + .expect_err("a tool callback error must reach the caller"); + assert!(matches!(error, FlowError::Internal(message) if message == "tool unavailable")); + } + assert_eq!(calls.load(Ordering::SeqCst), 2); + + flush_subscribers().unwrap(); + let misses = captured + .lock() + .unwrap() + .iter() + .filter(|event| { + event.name() == "response_cache" + && event + .data() + .and_then(|data| data.get("status")) + .and_then(Json::as_str) + == Some("miss") + && event + .metadata() + .and_then(|metadata| metadata.get("nemo_relay.response_cache.surface")) + .and_then(Json::as_str) + == Some("tool") + }) + .count(); + assert_eq!(misses, 2, "each failed call must still report a cache miss"); + deregister_subscriber("response_cache_tool_callback_error_capture").unwrap(); +} + +#[tokio::test] +async fn execution_intercepts_outside_the_cache_run_on_hits() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + + let outer_runs = Arc::new(AtomicUsize::new(0)); + register_tool_execution_intercept( + "response_cache_outer_tool_execution_test", + 40, + Arc::new({ + let outer_runs = Arc::clone(&outer_runs); + move |_name, args, next| { + let outer_runs = Arc::clone(&outer_runs); + Box::pin(async move { + outer_runs.fetch_add(1, Ordering::SeqCst); + next(args).await.map(Into::into) + }) + } + }), + ) + .unwrap(); + activate_cache(cache_with_tools(one_cacheable_class(&["lookup"]))).await; + + let calls = Arc::new(AtomicUsize::new(0)); + let tool = counting_tool(Arc::clone(&calls), json!({"answer": "cached"})); + tool_call("lookup", &tool, json!({"q": "relay"})).await; + tool_call("lookup", &tool, json!({"q": "relay"})).await; + + assert_eq!(calls.load(Ordering::SeqCst), 1, "the second call must hit"); + assert_eq!( + outer_runs.load(Ordering::SeqCst), + 2, + "a lower-priority execution intercept wraps the cache and runs on hits" + ); + deregister_tool_execution_intercept("response_cache_outer_tool_execution_test").unwrap(); +} + +#[tokio::test] +async fn tool_hit_emits_a_surface_tool_mark_with_saved_invocations() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + activate_cache(cache_with_tools(one_cacheable_class(&["docs_lookup"]))).await; + + let captured = Arc::new(StdMutex::new(Vec::::new())); + let sink = Arc::clone(&captured); + register_subscriber( + "response_cache_tool_capture", + Arc::new(move |event: &Event| sink.lock().unwrap().push(event.clone())), + ) + .unwrap(); + + let calls = Arc::new(AtomicUsize::new(0)); + let tool = counting_tool(Arc::clone(&calls), json!({"doc": "x"})); + + tool_call("docs_lookup", &tool, json!({"q": "rust"})).await; // miss + tool_call("docs_lookup", &tool, json!({"q": "rust"})).await; // hit + flush_subscribers().unwrap(); + + let events = captured.lock().unwrap(); + 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 tool hit mark should be emitted"); + let metadata = hit_mark.metadata().expect("hit mark has metadata"); + assert_eq!( + metadata + .get("nemo_relay.response_cache.surface") + .and_then(Json::as_str), + Some("tool"), + "the tool hit mark must be tagged surface = tool" + ); + assert_eq!( + metadata + .get("nemo_relay.response_cache.saved_invocations") + .and_then(Json::as_u64), + Some(1), + "a tool hit reports one saved invocation" + ); + + drop(events); + deregister_subscriber("response_cache_tool_capture").unwrap(); +} + +#[tokio::test] +async fn invalid_tool_config_is_rejected_by_validation() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + register_adaptive_component().unwrap(); + + let mut classes = std::collections::BTreeMap::new(); + classes.insert( + "class_a".to_string(), + ToolClass { + cacheable: true, + members: vec!["dup".to_string()], + ..ToolClass::default() + }, + ); + classes.insert( + "class_b".to_string(), + ToolClass { + cacheable: true, + ttl_seconds: Some(0), + bypass_rate: Some(1.1), + members: vec!["dup".to_string()], + ..ToolClass::default() + }, + ); + let mut overrides = std::collections::BTreeMap::new(); + overrides.insert( + "docs_lookup".to_string(), + ToolOverride { + bypass_rate: Some(-0.1), + ..ToolOverride::default() + }, + ); + let adaptive = AdaptiveConfig { + response_cache: Some(cache_with_tools(ToolCacheConfig { + enabled: true, + default: ToolClass { + members: vec!["safe_lookup".to_string()], + ..ToolClass::default() + }, + classes, + overrides, + ..ToolCacheConfig::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.tool_multiple_classes"), + "a tool in multiple classes must be rejected: {:?}", + report.diagnostics + ); + assert!( + report + .diagnostics + .iter() + .any(|diagnostic| diagnostic.code == "response_cache.tool_invalid_ttl"), + "a zero class TTL must be rejected: {:?}", + report.diagnostics + ); + assert!( + report + .diagnostics + .iter() + .any(|diagnostic| diagnostic.code == "response_cache.tool_invalid_bypass_rate"), + "out-of-range class and override bypass rates must be rejected: {:?}", + report.diagnostics + ); + assert!( + report + .diagnostics + .iter() + .any(|diagnostic| diagnostic.code == "response_cache.tool_default_members"), + "members on the default bucket must be rejected: {:?}", + report.diagnostics + ); +} + +#[tokio::test] +async fn wildcard_member_validation_rules() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + register_adaptive_component().unwrap(); + + let validate = |classes: std::collections::BTreeMap| { + let adaptive = AdaptiveConfig { + response_cache: Some(cache_with_tools(ToolCacheConfig { + enabled: true, + classes, + ..ToolCacheConfig::default() + })), + ..AdaptiveConfig::default() + }; + validate_plugin_config(&PluginConfig { + components: vec![ComponentSpec::new(adaptive).into()], + ..PluginConfig::default() + }) + }; + let cacheable_class = |members: &[&str]| ToolClass { + cacheable: true, + members: members.iter().map(|member| member.to_string()).collect(), + ..ToolClass::default() + }; + + let mut classes = std::collections::BTreeMap::new(); + classes.insert("class_a".to_string(), cacheable_class(&["docs_*"])); + classes.insert("class_b".to_string(), cacheable_class(&["docs_*"])); + let report = validate(classes); + assert!( + report + .diagnostics + .iter() + .any(|diagnostic| diagnostic.code == "response_cache.tool_multiple_classes"), + "an identical pattern in two classes must be rejected: {:?}", + report.diagnostics + ); + + let mut classes = std::collections::BTreeMap::new(); + classes.insert( + "read_only".to_string(), + cacheable_class(&["docs_*", "docs_*"]), + ); + let report = validate(classes); + assert!( + !report + .diagnostics + .iter() + .any(|diagnostic| diagnostic.code == "response_cache.tool_multiple_classes"), + "a repeated member inside one class is inert, not a cross-class conflict: {:?}", + report.diagnostics + ); + + let mut classes = std::collections::BTreeMap::new(); + classes.insert("class_a".to_string(), cacheable_class(&["docs_*"])); + classes.insert("class_b".to_string(), cacheable_class(&["*_lookup"])); + let report = validate(classes); + assert!( + !report + .diagnostics + .iter() + .any(|diagnostic| diagnostic.code.starts_with("response_cache.tool")), + "distinct overlapping patterns must validate cleanly: {:?}", + report.diagnostics + ); + + let mut classes = std::collections::BTreeMap::new(); + classes.insert("safe".to_string(), cacheable_class(&["*_email"])); + classes.insert( + "effectful".to_string(), + ToolClass { + cacheable: false, + members: vec!["send_*".to_string()], + ..ToolClass::default() + }, + ); + let report = validate(classes); + assert!( + report + .diagnostics + .iter() + .any(|diagnostic| diagnostic.code == "response_cache.tool_conflicting_classes"), + "opposite cacheability on overlapping wildcard classes must be rejected: {:?}", + report.diagnostics + ); + + for catch_all in ["*", "**"] { + let mut classes = std::collections::BTreeMap::new(); + classes.insert("everything".to_string(), cacheable_class(&[catch_all])); + let report = validate(classes); + assert!( + report + .diagnostics + .iter() + .any(|diagnostic| diagnostic.code == "response_cache.tool_catch_all_member"), + "a cacheable '{catch_all}' member must warn: {:?}", + report.diagnostics + ); + } + + let mut overrides = std::collections::BTreeMap::new(); + overrides.insert( + "*".to_string(), + ToolOverride { + cacheable: Some(true), + ..ToolOverride::default() + }, + ); + let adaptive = AdaptiveConfig { + response_cache: Some(cache_with_tools(ToolCacheConfig { + enabled: true, + overrides, + ..ToolCacheConfig::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.tool_catch_all_override"), + "a cacheable '*' override must warn: {:?}", + report.diagnostics + ); + + let mut overrides = std::collections::BTreeMap::new(); + overrides.insert( + "*_email".to_string(), + ToolOverride { + cacheable: Some(true), + ..ToolOverride::default() + }, + ); + overrides.insert( + "send_*".to_string(), + ToolOverride { + cacheable: Some(false), + ..ToolOverride::default() + }, + ); + let adaptive = AdaptiveConfig { + response_cache: Some(cache_with_tools(ToolCacheConfig { + enabled: true, + overrides, + ..ToolCacheConfig::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.tool_conflicting_overrides"), + "opposite cacheability on overlapping wildcard overrides must be rejected: {:?}", + report.diagnostics + ); + + let mut overrides = std::collections::BTreeMap::new(); + overrides.insert( + "docs_*".to_string(), + ToolOverride { + cacheable: Some(false), + ..ToolOverride::default() + }, + ); + overrides.insert( + "*_private".to_string(), + ToolOverride { + ttl_seconds: Some(60), + ..ToolOverride::default() + }, + ); + let adaptive = AdaptiveConfig { + response_cache: Some(cache_with_tools(ToolCacheConfig { + enabled: true, + default: ToolClass { + cacheable: true, + ..ToolClass::default() + }, + overrides, + ..ToolCacheConfig::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.tool_conflicting_overrides"), + "a wildcard override that inherits cacheability must not outrank an explicit deny: {:?}", + report.diagnostics + ); +} + +#[tokio::test] +async fn unknown_tool_field_warns_but_valid_class_names_do_not() { + let _guard = TEST_MUTEX.lock().await; + reset_global(); + register_adaptive_component().unwrap(); + + let adaptive_json = json!({ + "response_cache": { + "tools": { + "enabled": true, + "cache_errors": false, + "classes": { + "read_only": { + "cacheable": true, + "members": ["docs_lookup"], + "not_a_field": 7 + } + } + } + } + }); + let component = nemo_relay::plugin::PluginComponentSpec { + kind: "adaptive".to_string(), + enabled: true, + config: adaptive_json.as_object().unwrap().clone(), + }; + let report = validate_plugin_config(&PluginConfig { + components: vec![component], + ..PluginConfig::default() + }); + + let unknown_field_diags: Vec<_> = report + .diagnostics + .iter() + .filter(|diagnostic| diagnostic.code == "adaptive.unknown_field") + .collect(); + assert_eq!( + unknown_field_diags.len(), + 1, + "exactly one unknown-field warning (the bogus field), not the class name: {:?}", + report.diagnostics + ); + assert_eq!( + unknown_field_diags[0].field.as_deref(), + Some("not_a_field"), + "the warning must point at the bogus field, never the class name" + ); +} diff --git a/crates/adaptive/tests/unit/cache_diagnostics_tests.rs b/crates/adaptive/tests/unit/cache_diagnostics_tests.rs index 433ab4718..75e17c19e 100644 --- a/crates/adaptive/tests/unit/cache_diagnostics_tests.rs +++ b/crates/adaptive/tests/unit/cache_diagnostics_tests.rs @@ -232,6 +232,34 @@ fn cache_request_facts_keeps_missing_facts_bounded_when_inputs_are_unavailable() assert_eq!(facts.stable_prefix_tokens, None); } +#[test] +fn cache_request_facts_rejects_a_truncated_stable_prefix() { + let hot_cache = make_hot_cache(Some(2)); + let mut tracker = CacheDiagnosticsTracker::default(); + let prompt_ir = make_prompt_ir(vec![("system-0", "You are a careful planner", Some(700))]); + + let facts = build_cache_request_facts_from_prompt_ir( + CacheFactsBuildInput { + agent_id: "agent-1", + provider: "openai", + model: Some("gpt-4o"), + prompt_ir: &prompt_ir, + hot_cache: &hot_cache, + profile_key: "test-profile", + now: sample_timestamp(), + }, + &mut tracker, + ); + + assert_eq!(facts.stable_prefix_length, 2); + assert_eq!(facts.stable_prefix_tokens, None); + assert!( + facts + .missing_facts + .contains(&"stable_prefix_tokens_unavailable".to_string()) + ); +} + #[test] fn cache_request_facts_populates_provider_thresholds_and_retention_defaults() { let hot_cache = make_hot_cache(Some(2)); diff --git a/crates/adaptive/tests/unit/config_tests.rs b/crates/adaptive/tests/unit/config_tests.rs index 0a5537518..e7b095d1d 100644 --- a/crates/adaptive/tests/unit/config_tests.rs +++ b/crates/adaptive/tests/unit/config_tests.rs @@ -7,6 +7,8 @@ use super::*; use nemo_relay::config_editor::{EditorConfig, EditorFieldKind}; use serde_json::json; +use crate::response_cache::config::ToolCacheConfig; + #[test] fn test_adaptive_config_defaults() { let config = AdaptiveConfig::default(); @@ -32,6 +34,22 @@ fn test_typed_section_helpers_default() { let response_cache = ResponseCacheConfig::default(); assert!(!response_cache.cache_nondeterministic); + + let tools = ToolCacheConfig::default(); + assert!(!tools.enabled); + assert!(!tools.cache_errors); + assert_eq!(tools.priority, 150); +} + +#[test] +fn test_tool_cache_deserializes_explicit_error_caching_opt_in() { + let tools: ToolCacheConfig = serde_json::from_value(json!({ + "enabled": true, + "cache_errors": true, + })) + .unwrap(); + assert!(tools.enabled); + assert!(tools.cache_errors); } #[test] @@ -39,6 +57,36 @@ fn test_backend_spec_in_memory_helper_uses_empty_config() { let backend = BackendSpec::in_memory(); assert_eq!(backend.kind, "in_memory"); assert!(backend.config.is_empty()); + + let default_backend = BackendSpec::default(); + assert_eq!(default_backend.kind, "in_memory"); + assert!(default_backend.config.is_empty()); +} + +#[cfg(not(feature = "redis-backend"))] +#[test] +fn test_response_cache_redis_backend_requires_the_redis_feature() { + let mut response_cache = ResponseCacheConfig { + namespace: "cache-tests".to_string(), + ..ResponseCacheConfig::default() + }; + response_cache.backend.kind = "redis".to_string(); + response_cache + .backend + .config + .insert("url".to_string(), json!("redis://127.0.0.1/")); + + let report = crate::runtime::features::AdaptiveRuntime::validate_config(&AdaptiveConfig { + response_cache: Some(response_cache), + ..AdaptiveConfig::default() + }); + + assert!( + report + .diagnostics + .iter() + .any(|diagnostic| diagnostic.code == "response_cache.backend_unavailable") + ); } #[cfg(feature = "redis-backend")] @@ -155,4 +203,21 @@ fn test_adaptive_editor_schema_covers_canonical_options() { response_cache_backend.field("kind").unwrap().enum_values, &["in_memory", "redis"] ); + + let tools = response_cache.field("tools").unwrap(); + assert_eq!(tools.kind, EditorFieldKind::Section); + assert!(tools.optional); + let tools = tools.schema().unwrap(); + assert_eq!( + tools.field("enabled").unwrap().kind, + EditorFieldKind::Boolean + ); + assert_eq!( + tools.field("priority").unwrap().kind, + EditorFieldKind::Integer + ); + assert_eq!( + tools.field("default").unwrap().kind, + EditorFieldKind::Section + ); } diff --git a/crates/adaptive/tests/unit/plugin_component_tests.rs b/crates/adaptive/tests/unit/plugin_component_tests.rs index e051bb301..4605e08f6 100644 --- a/crates/adaptive/tests/unit/plugin_component_tests.rs +++ b/crates/adaptive/tests/unit/plugin_component_tests.rs @@ -379,6 +379,42 @@ fn validate_adaptive_plugin_config_reports_component_specific_unknown_fields() { })); } +#[test] +fn response_cache_tool_policy_validation_checks_nested_classes_and_overrides() { + let config = json!({ + "version": 1, + "response_cache": { + "tools": { + "default": {"unexpected_default": true}, + "classes": { + "read_only": {"unexpected_class": true} + }, + "overrides": { + "docs_lookup": {"unexpected_override": true} + } + } + }, + "policy": {"unknown_field": "warn"} + }); + + let diagnostics = validate_adaptive_plugin_config(config.as_object().unwrap()); + for (component, field) in [ + ("response_cache.tools.default", "unexpected_default"), + ("response_cache.tools.classes.read_only", "unexpected_class"), + ( + "response_cache.tools.overrides.docs_lookup", + "unexpected_override", + ), + ] { + assert!(diagnostics.iter().any(|diagnostic| { + diagnostic.code == "adaptive.unknown_field" + && diagnostic.component.as_deref() == Some(component) + && diagnostic.field.as_deref() == Some(field) + && diagnostic.level == DiagnosticLevel::Warning + })); + } +} + #[tokio::test(flavor = "current_thread")] async fn adaptive_plugin_registers_runtime_and_rolls_back_registration() { let _guard = crate::TEST_GLOBAL_CONTEXT_MUTEX.lock().await; diff --git a/crates/adaptive/tests/unit/response_cache/intercept_tests.rs b/crates/adaptive/tests/unit/response_cache/intercept_tests.rs index 5469d4630..59e387147 100644 --- a/crates/adaptive/tests/unit/response_cache/intercept_tests.rs +++ b/crates/adaptive/tests/unit/response_cache/intercept_tests.rs @@ -5,6 +5,7 @@ use std::time::Duration; +use nemo_relay::api::runtime::LlmJsonStream; use serde_json::json; use tokio::sync::{oneshot, watch}; use tokio_stream::StreamExt; @@ -53,6 +54,66 @@ fn chat_stream_fidelity_gate_rejects_every_uncollected_non_null_shape() { } } +#[test] +fn malformed_stream_shapes_are_not_aggregated() { + for malformed in [ + json!(null), + json!({"choices": {}}), + json!({"choices": [null]}), + json!({"choices": [{"index": "first"}]}), + json!({"choices": [{"finish_reason": 1}]}), + json!({"choices": [{"unsupported": true}]}), + json!({"choices": [{"delta": "not-an-object"}]}), + json!({"choices": [{"delta": {"tool_calls": [null]}}]}), + json!({"choices": [{"delta": {"tool_calls": [{"id": 1}]}}]}), + json!({"choices": [{"delta": {"tool_calls": [{"unsupported": true}]}}]}), + json!({"choices": [{"delta": {"tool_calls": [{"function": "not-an-object"}]}}]}), + ] { + assert!( + chunk_has_uncollected_response_fields(&malformed), + "malformed stream chunk must not be cached: {malformed}" + ); + } + + assert!(!chunk_has_uncollected_response_fields(&json!({ + "type": "message_delta" + }))); + assert!(!chunk_has_uncollected_response_fields(&json!({ + "choices": null + }))); + assert!( + !chunk_has_uncollected_response_fields(&json!({ + "choices": [{"delta": {"tool_calls": [{"id": null}]}}] + })), + "null tool-call metadata is harmless when no uncollectable fields are present" + ); +} + +#[test] +fn replay_and_error_guards_reject_unfaithful_or_failed_responses() { + assert!(aggregate_replay_lossy(&json!({ + "choices": [{ + "message": {"role": "assistant", "content": null, "tool_calls": []} + }] + }))); + assert!(chunk_is_inband_error(&json!({"type": "response.failed"}))); + assert!(chunk_is_inband_error( + &json!({"error": {"message": "upstream failed"}}) + )); + assert!(!chunk_is_inband_error(&json!({"error": null}))); + assert!(!is_error_response(&json!("not-an-object"))); +} + +#[test] +fn sampled_bypass_uses_a_unit_interval_rng() { + assert_eq!(rng_seed() & 1, 1, "xorshift state must never be zero"); + + RNG_STATE.with(|state| state.set(1)); + let expected = next_unit_f64() < 0.5; + RNG_STATE.with(|state| state.set(1)); + assert_eq!(should_bypass(0.5), expected); +} + #[tokio::test] async fn write_behind_returns_eof_before_cache_commit_completes() { let (tx, rx) = tokio::sync::mpsc::channel(1); @@ -78,6 +139,10 @@ async fn write_behind_returns_eof_before_cache_commit_completes() { .expect("write-behind cache publication must not delay stream completion") .is_none() ); + assert!( + stream.next().await.is_none(), + "finished streams stay finished" + ); release .send(()) .expect("detached cache commit must still be waiting"); @@ -173,3 +238,28 @@ fn assert_error_response_and_bypass_detection() { let unit = next_unit_f64(); assert!((0.0..1.0).contains(&unit), "{unit}"); } + +#[tokio::test] +async fn stream_close_reports_when_the_cleanup_task_ends_early() { + let (cancel, _) = watch::channel(false); + let (closed_tx, closed) = watch::channel(None::>); + drop(closed_tx); + let (tx, rx) = tokio::sync::mpsc::channel(1); + drop(tx); + let mut stream = LlmJsonStream::from_closeable(ResponseCacheReceiver { + receiver: ReceiverStream::new(rx), + cancel, + closed, + finished: false, + }); + + let error = stream + .close() + .await + .expect_err("an unavailable cleanup result must be reported"); + assert!( + error + .to_string() + .contains("response-cache stream cleanup task ended early") + ); +} diff --git a/crates/adaptive/tests/unit/response_cache/key_tests.rs b/crates/adaptive/tests/unit/response_cache/key_tests.rs index 54912deb6..5e0a7f1ad 100644 --- a/crates/adaptive/tests/unit/response_cache/key_tests.rs +++ b/crates/adaptive/tests/unit/response_cache/key_tests.rs @@ -5,6 +5,8 @@ use super::*; use crate::acg::canonicalize::{canonicalize_value, sha256_hex}; +use sha2::{Digest, Sha256}; +use std::io::Write; #[test] fn fingerprint_matches_canonicalize_then_hash() { @@ -965,3 +967,310 @@ fn null_text_system_block_does_not_collide_with_no_system() { "a null-text system block must not key like an absent system" ); } + +fn tool_key( + namespace: &str, + tool: &str, + version: Option<&str>, + args: Json, + arg_skip: &[String], +) -> String { + tool_key_with_error_policy(namespace, tool, version, args, arg_skip, false) +} + +fn tool_key_with_error_policy( + namespace: &str, + tool: &str, + version: Option<&str>, + args: Json, + arg_skip: &[String], + cache_errors: bool, +) -> String { + match build_tool_cache_key(namespace, tool, version, &args, arg_skip, cache_errors) { + KeyOutcome::Key(key) => key, + other => panic!("expected a tool key, got {other:?}"), + } +} + +#[test] +fn same_tool_and_args_yield_the_same_key() { + let args = json!({"q": "weather", "units": "metric"}); + assert_eq!( + tool_key("", "get_weather", None, args.clone(), &[]), + tool_key( + "", + "get_weather", + None, + json!({"units": "metric", "q": "weather"}), + &[] + ) + ); +} + +#[test] +fn tool_name_args_namespace_and_version_each_separate_keys() { + let base = || json!({"q": "x"}); + let key = tool_key("", "t", None, base(), &[]); + assert_ne!(key, tool_key("", "t", None, json!({"q": "y"}), &[]), "args"); + assert_ne!(key, tool_key("", "other", None, base(), &[]), "tool name"); + assert_ne!(key, tool_key("ns", "t", None, base(), &[]), "namespace"); + assert_ne!(key, tool_key("", "t", Some("v1"), base(), &[]), "version"); +} + +#[test] +fn tool_keys_bypass_unrepresentable_integers() { + assert_eq!( + build_tool_cache_key( + "key-test", + "lookup", + None, + &json!({"id": 18014398509481985_i64}), + &[], + false, + ), + KeyOutcome::Bypass("unrepresentable_number") + ); +} + +#[test] +fn arg_skip_drops_only_the_listed_keys() { + let skip = vec!["request_id".to_string()]; + assert_eq!( + tool_key("", "t", None, json!({"q": "x", "request_id": "a"}), &skip), + tool_key("", "t", None, json!({"q": "x", "request_id": "b"}), &skip) + ); + assert_ne!( + tool_key("", "t", None, json!({"q": "x", "request_id": "a"}), &skip), + tool_key("", "t", None, json!({"q": "y", "request_id": "a"}), &skip) + ); +} + +#[test] +fn arg_skip_policy_partitions_keys_and_normalizes_order() { + let no_skip: Vec = Vec::new(); + let locale_only = vec!["locale".to_string()]; + assert_ne!( + tool_key("key-test", "lookup", None, json!({"q": "x"}), &no_skip), + tool_key("key-test", "lookup", None, json!({"q": "x"}), &locale_only), + "a policy change must not reuse an entry even when the newly skipped key is absent" + ); + + let reordered_and_duplicated = vec![ + "trace_id".to_string(), + "locale".to_string(), + "trace_id".to_string(), + ]; + let normalized = vec!["locale".to_string(), "trace_id".to_string()]; + assert_eq!( + tool_key( + "key-test", + "lookup", + None, + json!({"q": "x", "locale": "fr", "trace_id": "one"}), + &reordered_and_duplicated, + ), + tool_key( + "key-test", + "lookup", + None, + json!({"q": "x", "locale": "de", "trace_id": "two"}), + &normalized, + ), + "equivalent skip policies must keep their intended hit behavior" + ); +} + +#[test] +fn cache_error_policy_partitions_tool_keys() { + assert_ne!( + tool_key_with_error_policy("key-test", "lookup", None, json!({"q": "x"}), &[], false), + tool_key_with_error_policy("key-test", "lookup", None, json!({"q": "x"}), &[], true), + "an opt-in error-cache entry must not be replayed after the policy is disabled" + ); +} + +#[test] +fn header_allowlist_policy_partitions_keys_and_normalizes_case() { + let request = request(json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "temperature": 0.0, + })); + let unpartitioned = cache_all_config(); + let mut tenant_partitioned = cache_all_config(); + tenant_partitioned.header_allowlist = vec!["X-Tenant".to_string()]; + let mut duplicate_spelling = cache_all_config(); + duplicate_spelling.header_allowlist = vec![ + "x-tenant".to_string(), + "X-TENANT".to_string(), + "x-tenant".to_string(), + ]; + + assert_ne!( + key_of("openai", &request, &unpartitioned), + key_of("openai", &request, &tenant_partitioned), + "changing the header policy must partition keys even before a request supplies that header" + ); + assert_eq!( + key_of("openai", &request, &tenant_partitioned), + key_of("openai", &request, &duplicate_spelling), + "case-only and duplicate policy spellings are equivalent" + ); +} + +#[test] +fn tool_keys_are_disjoint_from_llm_keys() { + let llm = key_of( + "openai", + &request(json!({"model": "t", "messages": []})), + &cache_all_config(), + ); + let tool = tool_key("", "t", None, json!({"messages": []}), &[]); + assert_ne!(llm, tool); +} + +#[test] +fn non_object_request_bodies_stay_raw_and_cacheable() { + // Non-object requests have no stateful controls or normalized fields. They + // must still receive a deterministic raw-body key instead of being treated + // as an unparseable request. + let raw = request(json!(["opaque", {"request": "body"}])); + assert_eq!( + resolved_body("custom-provider", &raw), + (raw.content.clone(), None) + ); + assert!(matches!( + build_cache_key("custom-provider", &raw, &cache_all_config()), + KeyOutcome::Key(_) + )); +} + +#[test] +fn negative_integers_beyond_the_safe_json_range_bypass_tool_keys() { + // RFC 8785 canonicalization rounds integers through f64. Negative values + // need the same protection as the positive IDs covered above. + let too_large = -9_007_199_254_740_993_i64; + assert_eq!( + build_tool_cache_key("key-test", "lookup", None, &json!(too_large), &[], false), + KeyOutcome::Bypass("unrepresentable_number") + ); +} + +#[test] +fn hash_writer_flushes_after_streaming_canonical_bytes() { + let mut hasher = Sha256::new(); + { + let mut writer = HashWriter(&mut hasher); + writer.write_all(b"response-cache-key").unwrap(); + writer.flush().unwrap(); + } + + assert_eq!(hasher.finalize(), Sha256::digest(b"response-cache-key")); +} + +#[test] +fn key_headers_match_case_insensitively_and_exclude_unlisted_values() { + let mut headers = Map::new(); + headers.insert("X-Tenant".to_string(), json!("tenant-a")); + headers.insert("Authorization".to_string(), json!("secret")); + + let kept = allowlisted_headers(&headers, &["x-tenant".to_string()]); + assert_eq!(kept.len(), 1); + assert_eq!(kept.get("x-tenant"), Some(&json!("tenant-a"))); +} + +#[test] +fn tool_id_normalization_skips_nonobjects_and_nonstring_ids() { + let mut body = json!({ + "messages": [ + null, + {"role": "assistant", "tool_calls": [{"id": "call-raw"}, {"id": 7}]}, + {"role": "tool", "tool_call_id": "call-raw"}, + {"role": "tool", "tool_call_id": 42} + ] + }); + + normalize_tool_call_ids(body.as_object_mut().unwrap()); + assert_eq!( + body.pointer("/messages/1/tool_calls/0/id"), + Some(&json!("tcid_0")) + ); + assert_eq!(body.pointer("/messages/1/tool_calls/1/id"), Some(&json!(7))); + assert_eq!( + body.pointer("/messages/2/tool_call_id"), + Some(&json!("tcid_0")) + ); + assert_eq!(body.pointer("/messages/3/tool_call_id"), Some(&json!(42))); +} + +#[test] +fn lossy_shape_guards_handle_nonobjects_and_unmodeled_tool_choices() { + assert!( + !lossy_request_shape(ProviderSurface::OpenAIChat, &json!("opaque body")), + "a non-object has no normalized fields to lose" + ); + assert!( + lossy_system_block(&json!("not a system block")), + "a non-object system block cannot be faithfully normalized" + ); + + let request = request(json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "look it up"}], + "tool_choice": { + "type": "function", + "function": {"name": "lookup", "strict": true} + } + })); + assert_eq!( + resolved_body("openai", &request).1, + None, + "a lossy tool_choice must use the raw request body for its key" + ); +} + +#[test] +fn decode_round_trip_guards_fall_back_to_raw_tool_and_message_shapes() { + // Anthropic client-tool wire objects serialize differently from the shared + // normalized tool representation. Keeping their raw shape in the key is + // safer than silently treating a future schema variation as equivalent. + let anthropic_tool_request = request(json!({ + "model": "claude-test", + "max_tokens": 16, + "system": "Follow the tool contract.", + "messages": [{"role": "user", "content": "Look this up."}], + "tools": [{ + "name": "lookup", + "description": "Look up a document.", + "input_schema": {"type": "object", "properties": {}} + }] + })); + assert!( + decode_surface(ProviderSurface::AnthropicMessages, &anthropic_tool_request).is_none(), + "a non-round-tripping tool shape must use raw keying" + ); + assert_eq!( + resolved_body("anthropic", &anthropic_tool_request), + (anthropic_tool_request.content.clone(), None) + ); + + // Closed message types carry a provider-native value for legacy + // `function_call`; its normalized representation is intentionally not a + // wire-equivalent message, so it too must keep the raw key shape. + let legacy_message_request = request(json!({ + "model": "gpt-4o", + "messages": [{ + "role": "assistant", + "content": null, + "function_call": {"name": "lookup", "arguments": "{\"q\":\"docs\"}"} + }] + })); + assert!( + decode_surface(ProviderSurface::OpenAIChat, &legacy_message_request).is_none(), + "a non-round-tripping message shape must use raw keying" + ); + assert_eq!( + resolved_body("openai", &legacy_message_request), + (legacy_message_request.content.clone(), None) + ); +} diff --git a/crates/adaptive/tests/unit/response_cache/mark_tests.rs b/crates/adaptive/tests/unit/response_cache/mark_tests.rs index e65a005a1..849cb07bf 100644 --- a/crates/adaptive/tests/unit/response_cache/mark_tests.rs +++ b/crates/adaptive/tests/unit/response_cache/mark_tests.rs @@ -19,6 +19,12 @@ impl Drop for ResetPricingResolverGuard { } } +#[test] +fn cache_surfaces_have_stable_metadata_values() { + assert_eq!(CacheSurface::Llm.as_str(), "llm"); + assert_eq!(CacheSurface::Tool.as_str(), "tool"); +} + #[test] fn anthropic_shaped_bodies_price_through_the_catalog() { // A real Anthropic response body must yield a dollar figure when the @@ -89,3 +95,68 @@ fn savings_from_counts_anthropic_input_output_tokens() { "anthropic input+output tokens must be counted for savings" ); } + +#[test] +fn normalized_savings_uses_entry_model_and_derives_missing_total_tokens() { + // Providers can omit a model in the payload while the cache knows the + // request model. A Chat response with prompt/completion tokens but no + // total must still report its complete saved-token count. + let entry = CacheEntry::new( + json!({ + "id": "chatcmpl_1", + "object": "chat.completion", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 7, "completion_tokens": 5} + }), + Duration::from_secs(60), + "sha256:chat".to_string(), + Some("model-recorded-with-request".to_string()), + Some("openai".to_string()), + ); + + assert_eq!(savings_from(&entry).0, Some(12)); +} + +#[test] +fn normalized_empty_usage_falls_back_to_no_savings() { + // A recognized response with an empty usage object is not a zero-token + // hit: it is missing accounting, so diagnostics must leave savings unset. + let entry = CacheEntry::new( + json!({ + "id": "chatcmpl_2", + "object": "chat.completion", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop" + }], + "usage": {} + }), + Duration::from_secs(60), + "sha256:empty-usage".to_string(), + None, + None, + ); + + assert_eq!(normalized_savings(&entry), None); + assert_eq!(savings_from(&entry), (None, None)); +} + +#[test] +fn raw_usage_probe_derives_total_from_prompt_and_completion_tokens() { + // Unknown provider shapes still expose standard OpenAI-style usage fields; + // raw fallback must preserve their useful savings diagnostics. + let entry = CacheEntry::new( + json!({"usage": {"prompt_tokens": 11, "completion_tokens": 4}}), + Duration::from_secs(60), + "sha256:raw".to_string(), + None, + None, + ); + + assert_eq!(savings_from(&entry), (Some(15), None)); +} diff --git a/crates/adaptive/tests/unit/response_cache/replay_tests.rs b/crates/adaptive/tests/unit/response_cache/replay_tests.rs index 7d5fa65a8..3ec214d96 100644 --- a/crates/adaptive/tests/unit/response_cache/replay_tests.rs +++ b/crates/adaptive/tests/unit/response_cache/replay_tests.rs @@ -214,3 +214,61 @@ fn replay_of_an_unknown_shape_is_lossy_for_the_streaming_tier() { assert!(replay_is_lossy(&json!({"weird": true}))); assert!(replay_is_lossy(&json!("bare string"))); } + +#[test] +fn stripping_stream_metadata_leaves_nonobject_frames_unchanged() { + // The helper also runs against collector output. A malformed non-object + // frame must be a harmless no-op rather than preventing the lossiness + // check from completing. + let mut frame = json!("not an aggregate"); + strip_stream_metadata(&mut frame); + assert_eq!(frame, json!("not an aggregate")); +} + +#[test] +fn anthropic_replay_keeps_complete_unknown_blocks_and_stop_sequences() { + // Blocks without a delta representation (such as thinking/server blocks) + // must be sent intact at content-block start, while stop_sequence remains + // visible to strict Anthropic stream consumers. + let aggregate = json!({ + "id": "msg_2", + "type": "message", + "role": "assistant", + "model": "claude-test", + "content": [{"type": "thinking", "thinking": "reasoning"}], + "stop_reason": "end_turn", + "stop_sequence": "", + "usage": {"input_tokens": 3, "output_tokens": 2} + }); + + let chunks = synthesize_anthropic_chunks(&aggregate); + assert_eq!(chunks[1]["type"], json!("content_block_start")); + assert_eq!(chunks[1]["content_block"], aggregate["content"][0]); + assert_eq!(chunks[2]["type"], json!("content_block_stop")); + let message_delta = chunks + .iter() + .find(|chunk| chunk["type"] == "message_delta") + .expect("replay must finish with a message_delta"); + assert_eq!( + message_delta.pointer("/delta/stop_sequence"), + Some(&json!("")) + ); +} + +#[test] +fn responses_replay_omits_item_events_for_a_nonarray_output() { + // A partially formed stored Responses aggregate is still replayed with + // lifecycle framing, but only real output arrays produce item-done events. + let aggregate = json!({ + "id": "resp_2", + "object": "response", + "model": "gpt-test", + "output": {"unexpected": true} + }); + + let chunks = synthesize_responses_chunks(&aggregate); + assert_eq!(chunks.len(), 2); + assert_eq!(chunks[0]["type"], json!("response.created")); + assert_eq!(chunks[1]["type"], json!("response.completed")); + assert_eq!(chunks[1]["sequence_number"], json!(1)); +} diff --git a/crates/adaptive/tests/unit/response_cache/store_tests.rs b/crates/adaptive/tests/unit/response_cache/store_tests.rs index 50446d0b9..819043398 100644 --- a/crates/adaptive/tests/unit/response_cache/store_tests.rs +++ b/crates/adaptive/tests/unit/response_cache/store_tests.rs @@ -7,7 +7,11 @@ use super::*; use serde_json::json; #[cfg(feature = "redis-backend")] -use std::net::TcpListener; +use std::io::{Read, Write}; +#[cfg(feature = "redis-backend")] +use std::net::{TcpListener, TcpStream}; +#[cfg(feature = "redis-backend")] +use std::thread; fn entry(key: &str, created: u64, expires: u64) -> CacheEntry { CacheEntry { @@ -22,6 +26,70 @@ fn entry(key: &str, created: u64, expires: u64) -> CacheEntry { const BIG: usize = 1 << 20; // 1 MiB — never evicts in these tests +#[cfg(feature = "redis-backend")] +fn read_redis_command(stream: &mut TcpStream) -> Vec { + fn read_line(stream: &mut TcpStream, request: &mut Vec) -> String { + let start = request.len(); + loop { + let mut byte = [0_u8; 1]; + stream.read_exact(&mut byte).expect("read RESP command"); + request.push(byte[0]); + if request.ends_with(b"\r\n") { + return std::str::from_utf8(&request[start..request.len() - 2]) + .expect("RESP command must be UTF-8") + .to_string(); + } + } + } + + let mut request = Vec::new(); + let count = read_line(stream, &mut request) + .strip_prefix('*') + .expect("RESP command array") + .parse::() + .expect("RESP command count"); + for _ in 0..count { + let length = read_line(stream, &mut request) + .strip_prefix('$') + .expect("RESP bulk string") + .parse::() + .expect("RESP bulk string length"); + let mut argument = vec![0_u8; length + 2]; + stream + .read_exact(&mut argument) + .expect("RESP bulk string value"); + assert!(argument.ends_with(b"\r\n")); + request.extend(argument); + } + request +} + +#[cfg(feature = "redis-backend")] +fn start_redis_test_server(response: Vec) -> (String, thread::JoinHandle>) { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind test Redis peer"); + let url = format!( + "redis://{}/", + listener.local_addr().expect("test Redis address") + ); + let server = thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("accept Redis client"); + // redis-rs identifies itself with two `CLIENT SETINFO` commands before + // it allows normal commands on a new connection. + for _ in 0..2 { + let setup = read_redis_command(&mut stream); + assert!(setup.windows(6).any(|window| window == b"CLIENT")); + stream + .write_all(b"+OK\r\n") + .expect("acknowledge Redis client setup"); + } + let command = read_redis_command(&mut stream); + stream.write_all(&response).expect("write Redis response"); + stream.flush().expect("flush Redis response"); + command + }); + (url, server) +} + #[cfg(feature = "redis-backend")] #[tokio::test(start_paused = true)] async fn redis_initialization_times_out_for_a_silent_peer() { @@ -234,3 +302,102 @@ async fn repeated_replacement_compacts_stale_insertion_order_nodes() { assert_eq!(guard.order.len(), 4); assert_eq!(guard.next_generation, 70); } + +#[test] +fn evicting_an_empty_queue_is_a_noop() { + // An eviction loop can reach an empty queue after stale nodes have been + // skipped. It must report that nothing was removed rather than underflowing + // the byte accounting. + let mut inner = Inner::default(); + assert!(!evict_oldest(&mut inner)); + assert!(inner.map.is_empty()); + assert_eq!(inner.total_bytes, 0); +} + +#[tokio::test] +async fn an_unknown_backend_is_rejected_before_initialization() { + let mut config = ResponseCacheConfig::default(); + config.backend.kind = "not-a-cache".to_string(); + + let error = match build_store(&config).await { + Ok(_) => panic!("an unknown response-cache backend must be rejected"), + Err(error) => error, + }; + assert!(matches!( + error, + AdaptiveError::InvalidConfig(message) + if message == "response_cache: unknown backend kind 'not-a-cache'" + )); +} + +#[cfg(feature = "redis-backend")] +#[tokio::test] +async fn redis_backend_requires_a_url_before_connecting() { + // This validates configuration locally and never attempts a network + // connection, so it remains deterministic in the unit-test suite. + let mut config = ResponseCacheConfig::default(); + config.backend.kind = "redis".to_string(); + + let error = match build_store(&config).await { + Ok(_) => panic!("a Redis backend without a URL must be rejected"), + Err(error) => error, + }; + assert!(matches!( + error, + AdaptiveError::InvalidConfig(message) + if message == "response_cache: redis backend requires backend.config.url" + )); +} + +#[cfg(feature = "redis-backend")] +#[tokio::test] +async fn redis_get_treats_an_entry_past_its_own_expiry_as_a_miss() { + // Redis can retain a value briefly longer than the response-cache TTL. + // The entry stamp remains authoritative, so a stale serialized entry must + // not be served even when Redis returns it. + let expired = entry("expired", 0, 1); + let encoded = serde_json::to_vec(&expired).expect("serialize cache entry"); + let mut response = format!("${}\r\n", encoded.len()).into_bytes(); + response.extend(encoded); + response.extend(b"\r\n"); + let (url, server) = start_redis_test_server(response); + + let store = RedisCacheStore::new(&url, "response-cache:") + .await + .expect("connect test Redis peer"); + assert!( + store.get("expired").await.expect("Redis GET").is_none(), + "an entry whose embedded expiry elapsed must be a miss" + ); + + let command = server.join().expect("test Redis server"); + assert!(command.windows(3).any(|window| window == b"GET")); + assert!( + command + .windows(b"response-cache:expired".len()) + .any(|window| window == b"response-cache:expired") + ); +} + +#[cfg(feature = "redis-backend")] +#[tokio::test] +async fn configured_redis_backend_pings_and_reports_its_kind() { + // This minimal RESP peer validates the configured store's operational + // health path without relying on a host Redis service. + let (url, server) = start_redis_test_server(b"+PONG\r\n".to_vec()); + let mut config = ResponseCacheConfig::default(); + config.backend.kind = "redis".to_string(); + config + .backend + .config + .insert("url".to_string(), Json::String(url)); + + let store = build_store(&config) + .await + .expect("configured Redis backend builds"); + assert_eq!(store.backend_kind(), "redis"); + store.health().await.expect("Redis PING succeeds"); + + let command = server.join().expect("test Redis server"); + assert!(command.windows(4).any(|window| window == b"PING")); +} diff --git a/crates/adaptive/tests/unit/response_cache/tool_policy_tests.rs b/crates/adaptive/tests/unit/response_cache/tool_policy_tests.rs new file mode 100644 index 000000000..9f1900e05 --- /dev/null +++ b/crates/adaptive/tests/unit/response_cache/tool_policy_tests.rs @@ -0,0 +1,159 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Focused policy-resolution tests for the tool-result response cache. + +use super::*; +use crate::response_cache::config::ToolClass; +use std::collections::BTreeMap; + +fn response_cache() -> ResponseCacheConfig { + ResponseCacheConfig { + ttl_seconds: 3600, + bypass_rate: 0.0, + ..ResponseCacheConfig::default() + } +} + +fn class(cacheable: bool, members: &[&str]) -> ToolClass { + ToolClass { + cacheable, + members: members.iter().map(|member| member.to_string()).collect(), + ..ToolClass::default() + } +} + +#[test] +fn policy_resolution_inherits_class_values_and_honors_an_override() { + let mut classes = BTreeMap::new(); + classes.insert( + "read_only".to_string(), + ToolClass { + cacheable: true, + ttl_seconds: Some(300), + bypass_rate: Some(0.2), + arg_skip: vec!["request_id".to_string()], + members: vec!["docs_*".to_string()], + }, + ); + let mut overrides = BTreeMap::new(); + overrides.insert( + "docs_lookup".to_string(), + ToolOverride { + cacheable: Some(false), + arg_skip: Some(vec![]), + tool_version: Some("v1".to_string()), + ..ToolOverride::default() + }, + ); + let tools = ToolCacheConfig { + classes, + overrides, + ..ToolCacheConfig::default() + }; + + let unclassified = resolve_policy("send_email", &response_cache(), &tools); + assert!(!unclassified.cacheable); + assert_eq!(unclassified.ttl, Duration::from_secs(3600)); + + let class_only = resolve_policy("docs_search", &response_cache(), &tools); + assert!(class_only.cacheable); + assert_eq!(class_only.ttl, Duration::from_secs(300)); + assert_eq!(class_only.bypass_rate, 0.2); + assert_eq!(class_only.arg_skip, ["request_id"]); + + let overridden = resolve_policy("docs_lookup", &response_cache(), &tools); + assert!(!overridden.cacheable); + assert_eq!(overridden.ttl, Duration::from_secs(300)); + assert_eq!(overridden.bypass_rate, 0.2); + assert!(overridden.arg_skip.is_empty()); + assert_eq!(overridden.tool_version.as_deref(), Some("v1")); +} + +#[test] +fn exact_and_specific_pattern_rules_choose_one_policy() { + let mut classes = BTreeMap::new(); + classes.insert( + "catch_all".to_string(), + ToolClass { + ttl_seconds: Some(100), + members: vec!["*".to_string()], + ..class(true, &[]) + }, + ); + classes.insert( + "docs".to_string(), + ToolClass { + ttl_seconds: Some(60), + members: vec!["docs_*".to_string()], + ..class(true, &[]) + }, + ); + classes.insert( + "private".to_string(), + ToolClass { + ttl_seconds: Some(10), + members: vec!["docs_private".to_string()], + ..class(true, &[]) + }, + ); + let mut overrides = BTreeMap::new(); + overrides.insert( + "docs_*".to_string(), + ToolOverride { + ttl_seconds: Some(20), + ..ToolOverride::default() + }, + ); + overrides.insert( + "docs_private".to_string(), + ToolOverride { + ttl_seconds: Some(5), + ..ToolOverride::default() + }, + ); + let tools = ToolCacheConfig { + classes, + overrides, + ..ToolCacheConfig::default() + }; + + assert_eq!( + resolve_policy("docs_private", &response_cache(), &tools).ttl, + Duration::from_secs(5), + "exact class and override entries win" + ); + assert_eq!( + resolve_policy("docs_search", &response_cache(), &tools).ttl, + Duration::from_secs(20), + "the more-specific wildcard class and override win" + ); + assert_eq!( + resolve_policy("other", &response_cache(), &tools).ttl, + Duration::from_secs(100) + ); +} + +#[test] +fn wildcard_matching_and_overlap_cover_edge_cases() { + for (pattern, name, expected) in [ + ("*", "", true), + ("docs_*", "docs_lookup", true), + ("docs_*", "doc_lookup", false), + ("get_*_price", "get_stock_price", true), + ("get_*_price", "get_price", false), + ("a*a", "a", false), + ("a*a", "aba", true), + ("Docs_*", "docs_lookup", false), + ] { + assert_eq!( + wildcard_match(pattern, name), + expected, + "{pattern:?}, {name:?}" + ); + } + assert!(wildcard_patterns_overlap("*_email", "send_*")); + assert!(!wildcard_patterns_overlap("docs_*", "send_*")); + assert!(wildcard_patterns_overlap("é*", "*é")); + assert_eq!(wildcard_rank("*é*").0, 1); +} diff --git a/crates/adaptive/tests/unit/response_cache/tool_tests.rs b/crates/adaptive/tests/unit/response_cache/tool_tests.rs new file mode 100644 index 000000000..01850ff3f --- /dev/null +++ b/crates/adaptive/tests/unit/response_cache/tool_tests.rs @@ -0,0 +1,350 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Error-classification tests for the tool-result response cache. + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +use nemo_relay::api::runtime::ToolExecutionNextFn; + +use super::*; +use crate::config::ResponseCacheConfig; +use crate::response_cache::config::{ToolCacheConfig, ToolClass, ToolOverride}; +use crate::response_cache::store::{CacheEntry, CacheStore, InMemoryCacheStore}; + +#[test] +fn wildcard_matching_handles_literals_and_missing_middle_segments() { + assert!(wildcard_match("docs_lookup", "docs_lookup")); + assert!(!wildcard_match("docs_lookup", "docs_search")); + assert!(!wildcard_match("a*b*c", "axc")); +} + +#[derive(Default)] +struct FailingGetStore { + get_calls: AtomicUsize, + set_calls: AtomicUsize, +} + +impl CacheStore for FailingGetStore { + fn get<'a>( + &'a self, + _key: &'a str, + ) -> crate::response_cache::store::BoxCacheFuture<'a, Option>> { + self.get_calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async { + Err(crate::error::AdaptiveError::Storage( + "cache read unavailable".to_string(), + )) + }) + } + + fn set<'a>( + &'a self, + _key: &'a str, + _entry: CacheEntry, + _ttl: Duration, + ) -> crate::response_cache::store::BoxCacheFuture<'a, ()> { + self.set_calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async { Ok(()) }) + } + + fn health<'a>(&'a self) -> crate::response_cache::store::BoxCacheFuture<'a, ()> { + Box::pin(async { Ok(()) }) + } + + fn backend_kind(&self) -> &'static str { + "failing_test" + } +} + +fn cache_config() -> Arc { + Arc::new(ResponseCacheConfig { + namespace: "tool-cache-unit-tests".to_string(), + ttl_seconds: 60, + ..ResponseCacheConfig::default() + }) +} + +fn counting_next(calls: Arc, result: Json) -> ToolExecutionNextFn { + Arc::new(move |_args| { + let calls = Arc::clone(&calls); + let result = result.clone(); + Box::pin(async move { + calls.fetch_add(1, Ordering::SeqCst); + Ok(result) + }) + }) +} + +#[test] +fn conventional_tool_error_detection_is_deliberately_narrow() { + assert!(is_error_shaped_tool_result(&serde_json::json!({ + "error": "upstream unavailable" + }))); + assert!(is_error_shaped_tool_result(&serde_json::json!({ + "isError": true + }))); + assert!(is_error_shaped_tool_result(&serde_json::json!({ + "is_error": true + }))); + assert!(!is_error_shaped_tool_result(&serde_json::json!({ + "error": null + }))); + assert!(!is_error_shaped_tool_result(&serde_json::json!({ + "status": "failed" + }))); + assert!(!is_error_shaped_tool_result(&serde_json::json!("error"))); +} + +#[tokio::test] +async fn tool_cache_read_error_fails_open_without_writing() { + let store = Arc::new(FailingGetStore::default()); + let calls = Arc::new(AtomicUsize::new(0)); + let next = counting_next( + Arc::clone(&calls), + Json::String("live tool result".to_string()), + ); + let response_cache = cache_config(); + let tools = Arc::new(ToolCacheConfig { + enabled: true, + default: ToolClass { + cacheable: true, + ..ToolClass::default() + }, + ..ToolCacheConfig::default() + }); + + let outcome = run_tool_cache( + "docs_lookup".to_string(), + serde_json::json!({"query": "response cache"}), + next, + store.clone(), + response_cache, + tools, + ) + .await + .expect("a cache read failure must not fail the tool call"); + + assert_eq!(outcome.result, Json::String("live tool result".to_string())); + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(store.get_calls.load(Ordering::SeqCst), 1); + assert_eq!(store.set_calls.load(Ordering::SeqCst), 0); +} + +#[tokio::test] +async fn stale_error_entries_are_not_replayed_when_error_caching_is_disabled() { + let store: Arc = Arc::new(InMemoryCacheStore::new(1 << 20)); + let response_cache = Arc::new(ResponseCacheConfig { + namespace: "tool-cache-stale-error-test".to_string(), + ..ResponseCacheConfig::default() + }); + let tools = Arc::new(ToolCacheConfig { + enabled: true, + default: ToolClass { + cacheable: true, + ..ToolClass::default() + }, + ..ToolCacheConfig::default() + }); + let args = serde_json::json!({"query": "relay"}); + let key = match build_tool_cache_key( + &response_cache.namespace, + "docs_lookup", + None, + &args, + &[], + false, + ) { + KeyOutcome::Key(key) => key, + other => panic!("expected a cache key, got {other:?}"), + }; + let ttl = Duration::from_secs(60); + store + .set( + &key, + CacheEntry::new( + serde_json::json!({"is_error": true, "content": "stale"}), + ttl, + key.clone(), + None, + None, + ), + ttl, + ) + .await + .unwrap(); + + let calls = Arc::new(AtomicUsize::new(0)); + let next: ToolExecutionNextFn = Arc::new({ + let calls = Arc::clone(&calls); + move |_args| { + let calls = Arc::clone(&calls); + Box::pin(async move { + calls.fetch_add(1, Ordering::SeqCst); + Ok(serde_json::json!({"answer": "fresh"})) + }) + } + }); + let result = run_tool_cache( + "docs_lookup".to_string(), + args.clone(), + Arc::clone(&next), + Arc::clone(&store), + Arc::clone(&response_cache), + Arc::clone(&tools), + ) + .await + .unwrap(); + + assert_eq!(result.result, serde_json::json!({"answer": "fresh"})); + let hit = run_tool_cache( + "docs_lookup".to_string(), + args, + next, + store, + response_cache, + tools, + ) + .await + .unwrap(); + assert_eq!(hit.result, serde_json::json!({"answer": "fresh"})); + assert_eq!(calls.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn disabling_error_caching_does_not_replay_an_opt_in_error_entry() { + let store = Arc::new(InMemoryCacheStore::new(1 << 20)); + let response_cache = cache_config(); + let opt_in_tools = Arc::new(ToolCacheConfig { + enabled: true, + cache_errors: true, + default: ToolClass { + cacheable: true, + ..ToolClass::default() + }, + ..ToolCacheConfig::default() + }); + let default_tools = Arc::new(ToolCacheConfig { + enabled: true, + default: ToolClass { + cacheable: true, + ..ToolClass::default() + }, + ..ToolCacheConfig::default() + }); + let calls = Arc::new(AtomicUsize::new(0)); + let args = serde_json::json!({"query": "relay"}); + + let error = run_tool_cache( + "docs_lookup".to_string(), + args.clone(), + counting_next( + Arc::clone(&calls), + serde_json::json!({"error": "temporary outage"}), + ), + store.clone(), + Arc::clone(&response_cache), + opt_in_tools, + ) + .await + .unwrap(); + assert_eq!( + error.result, + serde_json::json!({"error": "temporary outage"}) + ); + + let success = run_tool_cache( + "docs_lookup".to_string(), + args.clone(), + counting_next(Arc::clone(&calls), serde_json::json!({"answer": "fresh"})), + store.clone(), + Arc::clone(&response_cache), + Arc::clone(&default_tools), + ) + .await + .unwrap(); + assert_eq!(success.result, serde_json::json!({"answer": "fresh"})); + + let hit = run_tool_cache( + "docs_lookup".to_string(), + args, + counting_next( + Arc::clone(&calls), + serde_json::json!({"answer": "unexpected"}), + ), + store, + response_cache, + default_tools, + ) + .await + .unwrap(); + assert_eq!(hit.result, serde_json::json!({"answer": "fresh"})); + assert_eq!(calls.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn per_tool_override_ttl_reaches_the_stored_entry() { + let store = Arc::new(InMemoryCacheStore::new(1 << 20)); + let response_cache = cache_config(); + let classes = std::collections::BTreeMap::from([( + "read_only".to_string(), + ToolClass { + cacheable: true, + ttl_seconds: Some(17), + members: vec!["docs_lookup".to_string()], + ..ToolClass::default() + }, + )]); + let overrides = std::collections::BTreeMap::from([( + "docs_lookup".to_string(), + ToolOverride { + ttl_seconds: Some(23), + ..ToolOverride::default() + }, + )]); + let tools = Arc::new(ToolCacheConfig { + enabled: true, + classes, + overrides, + ..ToolCacheConfig::default() + }); + let args = serde_json::json!({"query": "relay"}); + + run_tool_cache( + "docs_lookup".to_string(), + args.clone(), + counting_next( + Arc::new(AtomicUsize::new(0)), + serde_json::json!({"answer": "cached"}), + ), + store.clone(), + Arc::clone(&response_cache), + tools, + ) + .await + .unwrap(); + + let key = match build_tool_cache_key( + &response_cache.namespace, + "docs_lookup", + None, + &args, + &[], + false, + ) { + KeyOutcome::Key(key) => key, + other => panic!("expected tool key, got {other:?}"), + }; + let entry = store + .get(&key) + .await + .unwrap() + .expect("the successful result should be stored"); + assert_eq!( + entry.expires_unix_ms - entry.created_unix_ms, + Duration::from_secs(23).as_millis() as u64, + "the override TTL, not the class or parent TTL, controls the stored entry" + ); +} diff --git a/crates/adaptive/tests/unit/runtime_features_tests.rs b/crates/adaptive/tests/unit/runtime_features_tests.rs index eaf5b0265..fc0ef9256 100644 --- a/crates/adaptive/tests/unit/runtime_features_tests.rs +++ b/crates/adaptive/tests/unit/runtime_features_tests.rs @@ -5,13 +5,14 @@ use super::*; -use std::sync::Arc; +use std::sync::{Arc, Once}; use crate::acg::profile::{BlockStabilityScore, StabilityClass}; use crate::acg::prompt_ir::SpanId; use crate::acg::stability::StabilityAnalysisResult; use crate::config::{BackendSpec, StateConfig}; use crate::intercepts::AGENT_HINTS_HEADER_KEY; +use crate::response_cache::config::ToolCacheConfig; use crate::trie::accumulator::AccumulatorState; use crate::trie::serialization::TrieEnvelope; use crate::types::metadata::{AgentHints, MetadataEnvelope, ParallelHint}; @@ -26,13 +27,15 @@ use nemo_relay::api::registry::{ deregister_llm_stream_execution_intercept, deregister_tool_execution_intercept, register_llm_execution_intercept, register_llm_request_intercept, register_llm_stream_execution_intercept, register_tool_execution_intercept, + scope_deregister_llm_request_intercept, scope_register_llm_request_intercept, }; -use nemo_relay::api::runtime::LlmJsonStream; use nemo_relay::api::runtime::ToolExecutionNextFn; use nemo_relay::api::runtime::global_context; use nemo_relay::api::runtime::{ LlmExecutionNextFn, LlmStreamExecutionNextFn, NemoRelayContextState, }; +use nemo_relay::api::runtime::{LlmJsonStream, create_scope_stack, set_thread_scope_stack}; +use nemo_relay::api::scope::{PopScopeParams, PushScopeParams, ScopeType, pop_scope, push_scope}; use nemo_relay::api::subscriber::{deregister_subscriber, register_subscriber}; use nemo_relay::api::tool::tool_call_execute; use nemo_relay::error::FlowError; @@ -48,6 +51,28 @@ fn reset_global() { *state = NemoRelayContextState::new(); } +struct CoverageLogger; + +impl log::Log for CoverageLogger { + fn enabled(&self, metadata: &log::Metadata<'_>) -> bool { + metadata.level() <= log::Level::Warn + } + + fn log(&self, _record: &log::Record<'_>) {} + + fn flush(&self) {} +} + +static COVERAGE_LOGGER: CoverageLogger = CoverageLogger; +static COVERAGE_LOGGER_INIT: Once = Once::new(); + +fn enable_warning_logs() { + COVERAGE_LOGGER_INIT.call_once(|| { + let _ = log::set_logger(&COVERAGE_LOGGER); + }); + log::set_max_level(log::LevelFilter::Warn); +} + fn sample_plan(agent_id: &str) -> ExecutionPlan { ExecutionPlan { agent_id: agent_id.to_string(), @@ -283,6 +308,13 @@ impl StorageBackendDyn for SeedFailBackend { ) -> Pin>> + Send + 'a>> { Box::pin(async { Ok(None) }) } + + fn load_stability<'a>( + &'a self, + _agent_id: &'a str, + ) -> Pin>> + Send + 'a>> { + Box::pin(async { Err(AdaptiveError::Storage("ACG seed failed".into())) }) + } } struct PartiallyFailingFeature; @@ -493,10 +525,15 @@ async fn telemetry_feature_registers_subscriber_and_starts_drain_task() { rollback_registrations(&mut registrations); assert_subscriber_absent(&name); - - if let Some(handle) = runtime.drain_handle.take() { - handle.abort(); - } + let handle = runtime + .drain_handle + .take() + .expect("telemetry registration must start a drain task"); + drop(runtime); + tokio::time::timeout(Duration::from_secs(1), handle) + .await + .expect("drain task must stop after its subscriber is deregistered") + .expect("drain task must complete cleanly"); } #[tokio::test(flavor = "current_thread")] @@ -631,9 +668,14 @@ async fn tool_parallelism_feature_registers_execution_intercept() { async fn adaptive_runtime_register_survives_hot_cache_seed_failures() { let _lock = crate::TEST_GLOBAL_CONTEXT_MUTEX.lock().await; reset_global(); + enable_warning_logs(); let config = AdaptiveConfig { adaptive_hints: Some(AdaptiveHintsComponentConfig::default()), + acg: Some(AcgComponentConfig { + provider: "passthrough".to_string(), + ..AcgComponentConfig::default() + }), ..AdaptiveConfig::default() }; let report = validate_config(&config); @@ -934,6 +976,43 @@ async fn acg_feature_registers_execution_and_stream_intercepts() { assert_llm_stream_execution_intercept_absent(&stream_name); } +#[tokio::test(flavor = "current_thread")] +async fn acg_feature_reports_execution_registration_conflicts() { + let _lock = crate::TEST_GLOBAL_CONTEXT_MUTEX.lock().await; + reset_global(); + + let mut runtime = AdaptiveRuntime::new(AdaptiveConfig::default()) + .await + .unwrap(); + let mut feature = AcgFeature::new( + AcgComponentConfig { + provider: "passthrough".to_string(), + ..AcgComponentConfig::default() + }, + runtime.hot_cache.clone(), + runtime.bound_scopes.clone(), + "agent-acg-conflict".to_string(), + Uuid::now_v7(), + ); + let execution_name = feature.execution_name.clone(); + register_llm_execution_intercept( + &execution_name, + 1, + Arc::new(|_name, request, next| next(request)), + ) + .unwrap(); + + let error = { + let mut ctx = RegistrationContext::new(&mut runtime); + let error = feature.register(&mut ctx).await.unwrap_err(); + let mut registrations = ctx.finish(); + rollback_registrations(&mut registrations); + error + }; + assert!(error.to_string().contains(&execution_name)); + deregister_llm_execution_intercept(&execution_name).unwrap(); +} + #[tokio::test(flavor = "current_thread")] async fn adaptive_runtime_register_feature_rolls_back_partial_registrations_and_abort_handle() { let _lock = crate::TEST_GLOBAL_CONTEXT_MUTEX.lock().await; @@ -964,11 +1043,261 @@ async fn adaptive_runtime_register_feature_rolls_back_partial_registrations_and_ assert_subscriber_absent("partial_feature"); } +#[tokio::test(flavor = "current_thread")] +async fn response_cache_feature_registers_llm_stream_and_enabled_tool_intercepts() { + let _lock = crate::TEST_GLOBAL_CONTEXT_MUTEX.lock().await; + reset_global(); + + let mut runtime = AdaptiveRuntime::new(AdaptiveConfig::default()) + .await + .unwrap(); + let mut feature = ResponseCacheFeature::new( + ResponseCacheConfig { + namespace: "response-cache-feature-registration".into(), + priority: 17, + tools: Some(ToolCacheConfig { + enabled: true, + priority: 19, + ..ToolCacheConfig::default() + }), + ..ResponseCacheConfig::default() + }, + Uuid::now_v7(), + ); + let execution_name = feature.name.clone(); + let stream_name = feature.stream_name.clone(); + let tool_name = feature.tool_name.clone(); + + let mut ctx = RegistrationContext::new(&mut runtime); + feature.register(&mut ctx).await.unwrap(); + + assert_llm_execution_intercept_registered(&execution_name); + assert_llm_stream_execution_intercept_registered(&stream_name); + assert_tool_execution_intercept_registered(&tool_name); + + let mut registrations = ctx.finish(); + rollback_registrations(&mut registrations); + assert_llm_execution_intercept_absent(&execution_name); + assert_llm_stream_execution_intercept_absent(&stream_name); + assert_tool_execution_intercept_absent(&tool_name); +} + +#[tokio::test(flavor = "current_thread")] +async fn response_cache_feature_propagates_invalid_store_configuration() { + let _lock = crate::TEST_GLOBAL_CONTEXT_MUTEX.lock().await; + reset_global(); + + let mut config = ResponseCacheConfig { + namespace: "response-cache-invalid-store".into(), + ..ResponseCacheConfig::default() + }; + config.backend.kind = "unsupported-store".into(); + let mut feature = ResponseCacheFeature::new(config, Uuid::now_v7()); + let mut runtime = AdaptiveRuntime::new(AdaptiveConfig::default()) + .await + .unwrap(); + + let error = { + let mut ctx = RegistrationContext::new(&mut runtime); + feature.register(&mut ctx).await.unwrap_err() + }; + assert!(matches!( + error, + AdaptiveError::InvalidConfig(message) + if message.contains("unknown backend kind 'unsupported-store'") + )); +} + +#[tokio::test(flavor = "current_thread")] +async fn response_cache_feature_cleans_up_when_llm_registration_conflicts() { + let _lock = crate::TEST_GLOBAL_CONTEXT_MUTEX.lock().await; + reset_global(); + + let mut runtime = AdaptiveRuntime::new(AdaptiveConfig::default()) + .await + .unwrap(); + let mut feature = ResponseCacheFeature::new( + ResponseCacheConfig { + namespace: "response-cache-execution-conflict".into(), + ..ResponseCacheConfig::default() + }, + Uuid::now_v7(), + ); + let name = feature.name.clone(); + register_llm_execution_intercept(&name, 1, Arc::new(|_name, request, next| next(request))) + .unwrap(); + + let error = { + let mut ctx = RegistrationContext::new(&mut runtime); + let error = feature.register(&mut ctx).await.unwrap_err(); + let mut registrations = ctx.finish(); + rollback_registrations(&mut registrations); + error + }; + assert!(error.to_string().contains(&name)); + deregister_llm_execution_intercept(&name).unwrap(); +} + +#[tokio::test(flavor = "current_thread")] +async fn response_cache_feature_cleans_up_when_stream_registration_conflicts() { + let _lock = crate::TEST_GLOBAL_CONTEXT_MUTEX.lock().await; + reset_global(); + + let mut runtime = AdaptiveRuntime::new(AdaptiveConfig::default()) + .await + .unwrap(); + let mut feature = ResponseCacheFeature::new( + ResponseCacheConfig { + namespace: "response-cache-stream-conflict".into(), + ..ResponseCacheConfig::default() + }, + Uuid::now_v7(), + ); + let execution_name = feature.name.clone(); + let stream_name = feature.stream_name.clone(); + register_llm_stream_execution_intercept( + &stream_name, + 1, + Arc::new(|_name, request, next| next(request)), + ) + .unwrap(); + + let error = { + let mut ctx = RegistrationContext::new(&mut runtime); + let error = feature.register(&mut ctx).await.unwrap_err(); + let mut registrations = ctx.finish(); + rollback_registrations(&mut registrations); + error + }; + assert!(error.to_string().contains(&stream_name)); + assert_llm_execution_intercept_absent(&execution_name); + deregister_llm_stream_execution_intercept(&stream_name).unwrap(); +} + +#[tokio::test(flavor = "current_thread")] +async fn response_cache_feature_cleans_up_when_tool_registration_conflicts() { + let _lock = crate::TEST_GLOBAL_CONTEXT_MUTEX.lock().await; + reset_global(); + + let mut runtime = AdaptiveRuntime::new(AdaptiveConfig::default()) + .await + .unwrap(); + let mut feature = ResponseCacheFeature::new( + ResponseCacheConfig { + namespace: "response-cache-tool-conflict".into(), + tools: Some(ToolCacheConfig { + enabled: true, + ..ToolCacheConfig::default() + }), + ..ResponseCacheConfig::default() + }, + Uuid::now_v7(), + ); + let execution_name = feature.name.clone(); + let stream_name = feature.stream_name.clone(); + let tool_name = feature.tool_name.clone(); + register_tool_execution_intercept( + &tool_name, + 1, + Arc::new(|_name, args, next| Box::pin(async move { next(args).await.map(Into::into) })), + ) + .unwrap(); + + let error = { + let mut ctx = RegistrationContext::new(&mut runtime); + let error = feature.register(&mut ctx).await.unwrap_err(); + let mut registrations = ctx.finish(); + rollback_registrations(&mut registrations); + error + }; + assert!(error.to_string().contains(&tool_name)); + assert_llm_execution_intercept_absent(&execution_name); + assert_llm_stream_execution_intercept_absent(&stream_name); + deregister_tool_execution_intercept(&tool_name).unwrap(); +} + +#[tokio::test(flavor = "current_thread")] +async fn bind_scope_requires_an_agent_id_and_acg_configuration_after_registration() { + let _lock = crate::TEST_GLOBAL_CONTEXT_MUTEX.lock().await; + reset_global(); + + let mut runtime = AdaptiveRuntime::new(AdaptiveConfig::default()) + .await + .unwrap(); + runtime.registered = true; + let scope_uuid = Uuid::now_v7(); + + let error = runtime.bind_scope(scope_uuid).unwrap_err(); + assert!(matches!( + error, + AdaptiveError::Internal(message) if message.contains("missing registered agent id") + )); + + runtime.registered_agent_id = Some("agent-without-acg".to_string()); + let error = runtime.bind_scope(scope_uuid).unwrap_err(); + assert!(matches!( + error, + AdaptiveError::InvalidConfig(message) if message.contains("does not enable scope-bound ACG") + )); +} + +#[tokio::test(flavor = "current_thread")] +async fn bind_scope_reports_duplicate_scope_intercept_registration() { + let _lock = crate::TEST_GLOBAL_CONTEXT_MUTEX.lock().await; + reset_global(); + set_thread_scope_stack(create_scope_stack()); + + let mut runtime = AdaptiveRuntime::new(AdaptiveConfig { + agent_id: Some("scope-conflict-agent".into()), + state: Some(StateConfig { + backend: BackendSpec::in_memory(), + }), + acg: Some(AcgComponentConfig::default()), + ..AdaptiveConfig::default() + }) + .await + .unwrap(); + runtime.register().await.unwrap(); + let scope = push_scope( + PushScopeParams::builder() + .name("scope-conflict") + .scope_type(ScopeType::Agent) + .build(), + ) + .unwrap(); + let name = runtime.acg_scope_registration_name(scope.uuid); + scope_register_llm_request_intercept( + &scope.uuid, + &name, + 1, + false, + Arc::new(|_name, request, annotated| { + Box::pin(async move { + Ok(nemo_relay::api::llm::LlmRequestInterceptOutcome::new( + request, annotated, + )) + }) + }), + ) + .unwrap(); + + let error = runtime.bind_scope(scope.uuid).unwrap_err(); + assert!(matches!( + error, + AdaptiveError::RegistrationFailed(message) + if message.contains("scope-bound ACG llm request intercept") + )); + + assert!(scope_deregister_llm_request_intercept(&scope.uuid, &name).unwrap()); + pop_scope(PopScopeParams::builder().handle_uuid(&scope.uuid).build()).unwrap(); +} + #[cfg(feature = "redis-backend")] #[tokio::test(flavor = "current_thread")] async fn response_cache_store_initialization_failure_fails_open() { let _lock = crate::TEST_GLOBAL_CONTEXT_MUTEX.lock().await; reset_global(); + enable_warning_logs(); let mut response_cache = ResponseCacheConfig { namespace: "fail-open-test".into(), @@ -978,7 +1307,7 @@ async fn response_cache_store_initialization_failure_fails_open() { response_cache .backend .config - .insert("url".into(), json!("redis://127.0.0.1:0/")); + .insert("url".into(), json!("not-a-redis-url")); let mut runtime = AdaptiveRuntime::new(AdaptiveConfig { response_cache: Some(response_cache), diff --git a/crates/adaptive/tests/unit/runtime_tests.rs b/crates/adaptive/tests/unit/runtime_tests.rs index cfad601bd..266799688 100644 --- a/crates/adaptive/tests/unit/runtime_tests.rs +++ b/crates/adaptive/tests/unit/runtime_tests.rs @@ -627,6 +627,9 @@ async fn adaptive_runtime_bind_scope_requires_registration_and_passes_through_wi runtime .bind_scope(scope.uuid) .expect("registered runtime should bind acg to the active scope"); + runtime + .bind_scope(scope.uuid) + .expect("binding an already-bound scope should be idempotent"); let request = LlmRequest { headers: Map::new(), content: serde_json::json!({ diff --git a/crates/cli/src/diagnostics/mod.rs b/crates/cli/src/diagnostics/mod.rs index 943091456..9f1f53d8d 100644 --- a/crates/cli/src/diagnostics/mod.rs +++ b/crates/cli/src/diagnostics/mod.rs @@ -765,6 +765,35 @@ async fn collect_response_cache_component_checks( return; } checks.push(response_cache_backend_check(response_cache::check_backend_health(&config)).await); + if let Some(tools) = config.tools.as_ref() { + let details = if tools.enabled { + let cacheable_classes = tools + .classes + .values() + .filter(|class| class.cacheable) + .count(); + let cacheable_overrides = tools + .overrides + .values() + .filter(|override_| override_.cacheable == Some(true)) + .count(); + format!( + "on; {cacheable_classes} cacheable class(es); {cacheable_overrides} cacheable override(s); default {}", + if tools.default.cacheable { + "cacheable" + } else { + "uncached" + } + ) + } else { + "configured but disabled".to_string() + }; + checks.push(Check { + name: "Response cache (tools)", + status: Status::Info, + details, + }); + } } async fn response_cache_backend_check( diff --git a/crates/cli/tests/coverage/shared/doctor_tests.rs b/crates/cli/tests/coverage/shared/doctor_tests.rs index 4dd8fe117..ddc83e56f 100644 --- a/crates/cli/tests/coverage/shared/doctor_tests.rs +++ b/crates/cli/tests/coverage/shared/doctor_tests.rs @@ -1591,6 +1591,50 @@ async fn collect_observability_reports_response_cache_fail_when_config_invalid() ); } +#[tokio::test] +async fn collect_observability_reports_tool_cache_surface_for_cacheable_overrides() { + let gateway = GatewayConfig { + plugin_config: Some(serde_json::json!({ + "version": 1, + "components": [ + { + "kind": "adaptive", + "enabled": true, + "config": { + "response_cache": { + "ttl_seconds": 3600, + "namespace": "doctor-tool-cache-test", + "backend": { "kind": "in_memory" }, + "tools": { + "enabled": true, + "overrides": { + "docs_*": { "cacheable": true } + } + } + } + } + } + ] + })), + ..GatewayConfig::default() + }; + + let checks = collect_observability(&gateway).await; + + let tools = checks + .iter() + .find(|check| check.name == "Response cache (tools)") + .expect("a tool-surface check should be present when tools.enabled"); + assert_eq!(tools.status, Status::Info, "checks: {checks:?}"); + assert!( + tools.details.contains("on") + && tools.details.contains("0 cacheable class") + && tools.details.contains("1 cacheable override"), + "details: {}", + tools.details + ); +} + #[tokio::test] async fn collect_observability_registers_pii_redaction_before_validation() { let gateway = GatewayConfig { diff --git a/crates/core/tests/unit/codec/anthropic_tests.rs b/crates/core/tests/unit/codec/anthropic_tests.rs index 6eb85d273..2ca2024c6 100644 --- a/crates/core/tests/unit/codec/anthropic_tests.rs +++ b/crates/core/tests/unit/codec/anthropic_tests.rs @@ -1637,3 +1637,41 @@ fn anthropic_helpers_cover_invalid_and_provider_native_values() { assert!(AnthropicMessagesStreamingCodec::default().finalizer()().is_object()); } + +#[test] +fn anthropic_streaming_codec_ignores_incomplete_lifecycle_frames() { + // A disconnected SSE stream can leave any lifecycle event only partially + // populated. The collector must keep the valid usage snapshot while + // ignoring frames that cannot identify a message or content block. + let codec = AnthropicMessagesStreamingCodec::default(); + let mut collector = codec.collector(); + let finalizer = codec.finalizer(); + + for frame in [ + json!({"type": "message_start"}), + json!({"type": "content_block_start"}), + json!({"type": "content_block_start", "index": 0}), + json!({"type": "content_block_start", "index": 0, "content_block": []}), + json!({"type": "content_block_delta"}), + json!({"type": "content_block_delta", "index": 0}), + json!({ + "type": "content_block_delta", + "index": 5, + "delta": {"type": "text_delta", "text": "orphaned"} + }), + json!({ + "type": "message_delta", + "usage": {"input_tokens": 3, "output_tokens": 0} + }), + ] { + collector(frame).unwrap(); + } + + assert_eq!( + finalizer(), + json!({ + "content": [], + "usage": {"input_tokens": 3, "output_tokens": 0}, + }) + ); +} diff --git a/crates/core/tests/unit/codec/openai_chat_tests.rs b/crates/core/tests/unit/codec/openai_chat_tests.rs index cb3fdb3b8..0b9230076 100644 --- a/crates/core/tests/unit/codec/openai_chat_tests.rs +++ b/crates/core/tests/unit/codec/openai_chat_tests.rs @@ -1809,3 +1809,25 @@ fn openai_chat_helpers_cover_provider_edge_values() { } assert!(OpenAIChatStreamingCodec::default().finalizer()().is_object()); } + +#[test] +fn openai_chat_streaming_codec_keeps_sparse_choice_frames_replayable() { + // A provider can terminate or truncate a stream after declaring a choice + // index but before sending its delta. Response-cache must still be able to + // assemble a safe buffered body instead of panicking or inventing content. + let codec = OpenAIChatStreamingCodec::default(); + let mut collector = codec.collector(); + let finalizer = codec.finalizer(); + + collector(json!({"choices": [{"index": 2}]})).unwrap(); + + let assembled = finalizer(); + assert_eq!( + assembled["choices"], + json!([{ + "index": 2, + "message": {"role": "assistant", "content": null}, + "finish_reason": null, + }]) + ); +} diff --git a/crates/core/tests/unit/codec/openai_responses_tests.rs b/crates/core/tests/unit/codec/openai_responses_tests.rs index 0d978811e..cf264d7a3 100644 --- a/crates/core/tests/unit/codec/openai_responses_tests.rs +++ b/crates/core/tests/unit/codec/openai_responses_tests.rs @@ -1651,3 +1651,19 @@ fn responses_helpers_cover_invalid_and_provider_edge_values() { assert!(OpenAIResponsesStreamingCodec::default().finalizer()().is_object()); } + +#[test] +fn openai_responses_streaming_codec_ignores_incomplete_lifecycle_frames() { + // Truncated Responses streams can contain envelope events without their + // optional payload. Treat those frames as no-ops so cache aggregation stays + // fail-open and never creates a synthetic response or output item. + let codec = OpenAIResponsesStreamingCodec::default(); + let mut collector = codec.collector(); + let finalizer = codec.finalizer(); + + collector(json!({"type": "response.created"})).unwrap(); + collector(json!({"type": "response.output_item.done", "item": {"type": "message"}})).unwrap(); + collector(json!({"type": "response.output_item.done", "output_index": 0})).unwrap(); + + assert_eq!(finalizer(), json!({})); +} diff --git a/crates/core/tests/unit/llm_api_tests.rs b/crates/core/tests/unit/llm_api_tests.rs index 027960d1f..703b7277c 100644 --- a/crates/core/tests/unit/llm_api_tests.rs +++ b/crates/core/tests/unit/llm_api_tests.rs @@ -230,6 +230,24 @@ fn response_sanitizer_context_preserves_all_codec_identity_states() { ); } +#[test] +fn sanitizer_context_debug_includes_identity_without_codec_handles() { + let request = crate::api::runtime::LlmSanitizeRequestContext::for_request_codec(Some( + Arc::new(OpenAIChatCodec), + )); + let response = crate::api::runtime::LlmSanitizeResponseContext::for_response_codec(Some( + Arc::new(OpenAIChatCodec), + )); + + let request_debug = format!("{request:?}"); + assert!(request_debug.contains("BuiltIn(OpenAiChat)")); + assert!(!request_debug.contains("request_codec")); + + let response_debug = format!("{response:?}"); + assert!(response_debug.contains("BuiltIn(OpenAiChat)")); + assert!(!response_debug.contains("response_codec")); +} + impl LlmCodec for ProjectionFailingCodec { fn decode(&self, request: &LlmRequest) -> crate::error::Result { OpenAIChatCodec.decode(request) diff --git a/crates/node/adaptive.d.ts b/crates/node/adaptive.d.ts index 72f3a83c1..dd4018688 100644 --- a/crates/node/adaptive.d.ts +++ b/crates/node/adaptive.d.ts @@ -52,7 +52,7 @@ export interface AcgConfig { stability_thresholds?: AcgStabilityThresholds; } -/** Opt-in LLM response cache (exact-match) settings. */ +/** Opt-in exact-match LLM response and tool-result cache settings. */ export interface ResponseCacheConfig { ttlSeconds?: number; /** @@ -66,6 +66,8 @@ export interface ResponseCacheConfig { keyStrategy?: string; headerAllowlist?: string[]; backend?: BackendSpec; + /** Opt-in tool-result cache; omit to leave the tool surface off. */ + tools?: ToolCacheConfig; } interface ResponseCachePluginConfig { @@ -78,8 +80,62 @@ interface ResponseCachePluginConfig { key_strategy?: string; header_allowlist?: string[]; backend?: BackendSpec; + tools?: ToolCachePluginConfig; } +/** Shared policy; omitted TTL and bypass rate inherit response-cache defaults. */ +export interface ToolClass { + cacheable?: boolean; + ttlSeconds?: number; + bypassRate?: number; + argSkip?: string[]; + members?: string[]; +} + +/** Per-tool refinement; an explicit `argSkip: []` clears the class list. */ +export interface ToolOverride { + cacheable?: boolean; + ttlSeconds?: number; + bypassRate?: number; + toolVersion?: string; + argSkip?: string[]; +} + +/** Opt-in caching for tools that are read-only and stable for their TTL. */ +export interface ToolCacheConfig { + enabled?: boolean; + /** + * Tool execution-intercept priority; omit for Rust's default (150), which + * keeps standard priority-100 guardrails outside cache hits. + */ + priority?: number; + /** Whether error-shaped tool results may be cached; defaults to false. */ + cacheErrors?: boolean; + default?: ToolClass; + classes?: Record; + overrides?: Record; +} + +type ToolClassPluginConfig = Omit & { + ttl_seconds?: number; + bypass_rate?: number; + arg_skip?: string[]; +}; + +type ToolOverridePluginConfig = Omit & { + ttl_seconds?: number; + bypass_rate?: number; + tool_version?: string; + arg_skip?: string[]; +}; + +type ToolCachePluginConfig = Omit & { + cache_errors?: boolean; + default?: ToolClassPluginConfig; + classes?: Record; + overrides?: Record; +}; + /** Canonical config object for the top-level adaptive component. */ export interface Config { version?: number; @@ -280,8 +336,8 @@ 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 + * Merges caller-supplied overrides onto the opt-in LLM response and tool-result + * 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. diff --git a/crates/node/adaptive.js b/crates/node/adaptive.js index c96082012..0b6f947cd 100644 --- a/crates/node/adaptive.js +++ b/crates/node/adaptive.js @@ -150,8 +150,8 @@ 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 + * Merges caller-supplied overrides onto the opt-in LLM response and tool-result + * 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. @@ -185,16 +185,59 @@ const RESPONSE_CACHE_PLUGIN_FIELDS = { headerAllowlist: 'header_allowlist', }; +const TOOL_CLASS_PLUGIN_FIELDS = { + ttlSeconds: 'ttl_seconds', + bypassRate: 'bypass_rate', + argSkip: 'arg_skip', +}; + +const TOOL_OVERRIDE_PLUGIN_FIELDS = { + ...TOOL_CLASS_PLUGIN_FIELDS, + toolVersion: 'tool_version', +}; + +const TOOL_CACHE_PLUGIN_FIELDS = { + cacheErrors: 'cache_errors', +}; + +function mapPluginFields(config, fields) { + if (config === null || typeof config !== 'object' || Array.isArray(config)) return config; + return Object.fromEntries(Object.entries(config).map(([key, value]) => [fields[key] ?? key, value])); +} + +function mapPluginRecord(config, fields) { + if (config === null || typeof config !== 'object' || Array.isArray(config)) return config; + return Object.fromEntries(Object.entries(config).map(([key, value]) => [key, mapPluginFields(value, fields)])); +} + +function toToolCachePluginConfig(config) { + const serialized = mapPluginFields(config, TOOL_CACHE_PLUGIN_FIELDS); + if (serialized === config) return config; + if (serialized.default !== undefined) { + serialized.default = mapPluginFields(serialized.default, TOOL_CLASS_PLUGIN_FIELDS); + } + if (serialized.classes !== undefined) { + serialized.classes = mapPluginRecord(serialized.classes, TOOL_CLASS_PLUGIN_FIELDS); + } + if (serialized.overrides !== undefined) { + serialized.overrides = mapPluginRecord(serialized.overrides, TOOL_OVERRIDE_PLUGIN_FIELDS); + } + return serialized; +} + +function toResponseCachePluginConfig(config) { + const serialized = mapPluginFields(config, RESPONSE_CACHE_PLUGIN_FIELDS); + if (serialized === config) return config; + if (serialized.tools !== undefined) { + serialized.tools = toToolCachePluginConfig(serialized.tools); + } + return serialized; +} + function toPluginConfig(config) { const { responseCache, ...rest } = config; if (responseCache === undefined) return config; - const serialized = - responseCache !== null && typeof responseCache === 'object' && !Array.isArray(responseCache) - ? Object.fromEntries( - Object.entries(responseCache).map(([key, value]) => [RESPONSE_CACHE_PLUGIN_FIELDS[key] ?? key, value]), - ) - : responseCache; - return { ...rest, response_cache: serialized }; + return { ...rest, response_cache: toResponseCachePluginConfig(responseCache) }; } class AdaptiveRuntime extends lib.AdaptiveRuntime { diff --git a/crates/node/tests/adaptive_runtime_tests.mjs b/crates/node/tests/adaptive_runtime_tests.mjs index 57445671e..cedff8b4f 100644 --- a/crates/node/tests/adaptive_runtime_tests.mjs +++ b/crates/node/tests/adaptive_runtime_tests.mjs @@ -26,6 +26,27 @@ describe('adaptive runtime bridge', () => { assert.deepEqual(adaptive.validateConfig(adaptive.defaultConfig()).diagnostics, []); }); + it('rejects a tool listed in multiple classes', () => { + const config = { + version: 1, + responseCache: adaptive.responseCacheConfig({ + namespace: 'node-tool-cache-test', + tools: { + enabled: true, + classes: { + a: { cacheable: true, members: ['dup'] }, + b: { cacheable: true, members: ['dup'] }, + }, + }, + }), + }; + const codes = adaptive.validateConfig(config).diagnostics.map((diag) => diag.code); + assert.ok( + codes.includes('response_cache.tool_multiple_classes'), + `expected tool_multiple_classes, got ${JSON.stringify(codes)}`, + ); + }); + it('builds cache telemetry events from one options object', () => { const event = adaptive.buildCacheTelemetryEvent({ provider: 'openai', diff --git a/crates/node/tests/adaptive_tests.mjs b/crates/node/tests/adaptive_tests.mjs index 726fa69a1..e086d7dff 100644 --- a/crates/node/tests/adaptive_tests.mjs +++ b/crates/node/tests/adaptive_tests.mjs @@ -331,6 +331,28 @@ describe('adaptive helpers', () => { }); }); + it('serializes nested tool-cache config', () => { + const spec = adaptive.ComponentSpec({ + version: 1, + responseCache: { + tools: { + enabled: true, + cacheErrors: true, + default: { ttlSeconds: 30, bypassRate: 0.1, argSkip: ['trace'] }, + classes: { readOnly: { cacheable: true, members: ['search'] } }, + overrides: { search: { toolVersion: 'v1', argSkip: ['requestId'] } }, + }, + }, + }); + assert.deepEqual(spec.config.response_cache.tools, { + enabled: true, + cache_errors: true, + default: { ttl_seconds: 30, bypass_rate: 0.1, arg_skip: ['trace'] }, + classes: { readOnly: { cacheable: true, members: ['search'] } }, + overrides: { search: { tool_version: 'v1', arg_skip: ['requestId'] } }, + }); + }); + it('serializes response-cache config at both native boundaries', () => { const unscoped = adaptive.validateConfig({ version: 1, responseCache: {} }); assert.ok(unscoped.diagnostics.some(({ code }) => code === 'response_cache.missing_namespace')); diff --git a/go/nemo_relay/adaptive.go b/go/nemo_relay/adaptive.go index 2cb40707a..916c5b924 100644 --- a/go/nemo_relay/adaptive.go +++ b/go/nemo_relay/adaptive.go @@ -67,7 +67,7 @@ type AcgConfig struct { StabilityThresholds *AcgStabilityThresholds `json:"stability_thresholds,omitempty"` } -// ResponseCacheConfig configures the opt-in LLM response cache: a section +// ResponseCacheConfig configures the opt-in LLM response and tool-result 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. @@ -94,6 +94,40 @@ type ResponseCacheConfig struct { // 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"` + // Tools configures the optional tool-result cache. + Tools *ResponseCacheToolsConfig `json:"tools,omitempty"` +} + +// ResponseCacheToolsConfig configures caching for read-only, stable tools. +type ResponseCacheToolsConfig struct { + Enabled bool `json:"enabled,omitempty"` + // Priority is the execution-intercept priority. Nil delegates to Rust's + // default (150), which keeps standard priority-100 guardrails outside + // cache hits; a pointer to 0 selects outermost. + Priority *int32 `json:"priority,omitempty"` + // CacheErrors lets error-shaped tool results be cached (default false). + CacheErrors bool `json:"cache_errors"` + Default *ResponseCacheToolClass `json:"default,omitempty"` + Classes map[string]ResponseCacheToolClass `json:"classes,omitempty"` + Overrides map[string]ResponseCacheToolOverride `json:"overrides,omitempty"` +} + +// ResponseCacheToolClass defines a shared tool-cache policy. +type ResponseCacheToolClass struct { + Cacheable bool `json:"cacheable,omitempty"` + TTLSeconds *uint64 `json:"ttl_seconds,omitempty"` + BypassRate *float64 `json:"bypass_rate,omitempty"` + ArgSkip []string `json:"arg_skip,omitempty"` + Members []string `json:"members,omitempty"` +} + +// ResponseCacheToolOverride refines a resolved tool-cache policy. +type ResponseCacheToolOverride struct { + Cacheable *bool `json:"cacheable,omitempty"` + TTLSeconds *uint64 `json:"ttl_seconds,omitempty"` + BypassRate *float64 `json:"bypass_rate,omitempty"` + ToolVersion *string `json:"tool_version,omitempty"` + ArgSkip *[]string `json:"arg_skip,omitempty"` } // ResponseCacheBackendConfig selects the response-cache backend kind and options. @@ -212,6 +246,15 @@ func NewRedisResponseCacheBackend(url, keyPrefix string) ResponseCacheBackendCon } } +// NewResponseCacheToolsConfig returns a disabled tool-result cache config. +func NewResponseCacheToolsConfig() ResponseCacheToolsConfig { + priority := int32(150) + return ResponseCacheToolsConfig{ + CacheErrors: false, + Priority: &priority, + } +} + // 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 b6511c9b8..117eb01f3 100644 --- a/go/nemo_relay/adaptive/adaptive.go +++ b/go/nemo_relay/adaptive/adaptive.go @@ -52,12 +52,21 @@ type AcgStabilityThresholds = nemo_relay.AcgStabilityThresholds // AcgConfig configures the adaptive cache governor. type AcgConfig = nemo_relay.AcgConfig -// ResponseCacheConfig configures the opt-in LLM response cache. +// ResponseCacheConfig configures the opt-in LLM response and tool-result cache. type ResponseCacheConfig = nemo_relay.ResponseCacheConfig // ResponseCacheBackendConfig selects the response-cache backend kind and options. type ResponseCacheBackendConfig = nemo_relay.ResponseCacheBackendConfig +// ResponseCacheToolsConfig configures the opt-in tool-result cache surface. +type ResponseCacheToolsConfig = nemo_relay.ResponseCacheToolsConfig + +// ResponseCacheToolClass is one tool caching class (also the shape of default). +type ResponseCacheToolClass = nemo_relay.ResponseCacheToolClass + +// ResponseCacheToolOverride refines a single tool on top of its resolved class. +type ResponseCacheToolOverride = nemo_relay.ResponseCacheToolOverride + // CacheUsage is normalized LLM token usage for cache telemetry. type CacheUsage = nemo_relay.CacheUsage @@ -135,6 +144,11 @@ func NewRedisResponseCacheBackend(url, keyPrefix string) ResponseCacheBackendCon return nemo_relay.NewRedisResponseCacheBackend(url, keyPrefix) } +// NewResponseCacheToolsConfig returns a default (disabled) tool-result cache config. +func NewResponseCacheToolsConfig() ResponseCacheToolsConfig { + return nemo_relay.NewResponseCacheToolsConfig() +} + // 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 705779f60..c1ed0c067 100644 --- a/go/nemo_relay/adaptive_runtime_test.go +++ b/go/nemo_relay/adaptive_runtime_test.go @@ -253,6 +253,68 @@ func assertResponseCacheValidation(t *testing.T, responseCache ResponseCacheConf } } +func TestResponseCacheToolsConfigReachesTypedSurface(t *testing.T) { + rc := NewResponseCacheConfig() + rc.Namespace = "tool-cache-go-test" + tools := NewResponseCacheToolsConfig() + if tools.Priority == nil || *tools.Priority != 150 { + t.Fatalf("constructor tools priority default mismatch: %#v", tools.Priority) + } + if tools.CacheErrors { + t.Fatalf("constructor cache_errors default mismatch: %#v", tools.CacheErrors) + } + tools.Enabled = true + tools.CacheErrors = true + zero := int32(0) + tools.Priority = &zero + tools.Classes = map[string]ResponseCacheToolClass{ + "read_only": {Cacheable: true, Members: []string{"docs_lookup"}}, + } + rc.Tools = &tools + + config := NewAdaptiveConfig() + config.ResponseCache = &rc + + 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) + } + rcSection, ok := decoded["response_cache"].(map[string]any) + if !ok { + t.Fatalf("response_cache missing from marshaled config: %s", payload) + } + toolsSection, ok := rcSection["tools"].(map[string]any) + if !ok { + t.Fatalf("tools missing from marshaled response_cache: %#v", rcSection) + } + if enabled, _ := toolsSection["enabled"].(bool); !enabled { + t.Fatalf("tools.enabled not preserved: %#v", toolsSection) + } + if cacheErrors, ok := toolsSection["cache_errors"].(bool); !ok || !cacheErrors { + t.Fatalf("tools.cache_errors not preserved: %#v", toolsSection) + } + if priority, ok := toolsSection["priority"].(float64); !ok || priority != 0 { + t.Fatalf("explicit tools.priority = 0 must survive marshal: %#v", toolsSection) + } + classes, ok := toolsSection["classes"].(map[string]any) + if !ok || classes["read_only"] == nil { + t.Fatalf("tools.classes not preserved: %#v", toolsSection) + } + + 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) + } + +} + func TestResponseCacheConfigPreservesOmissionAndExplicitZero(t *testing.T) { t.Run("partial config delegates to Rust defaults", testPartialResponseCacheConfig) t.Run("missing namespace remains invalid", testMissingResponseCacheNamespace) diff --git a/python/nemo_relay/adaptive.py b/python/nemo_relay/adaptive.py index 5b58c9dc4..5cd397c7e 100644 --- a/python/nemo_relay/adaptive.py +++ b/python/nemo_relay/adaptive.py @@ -251,13 +251,127 @@ def to_dict(self) -> JsonObject: ) +@dataclass(slots=True) +class ToolClass: + """One tool caching class (also the shape of the ``default`` default bucket). + + Args: + cacheable: Whether tools in this class may be served from cache. Off by + default — a hit suppresses the real call, so caching must be opted in. + ttl_seconds: TTL for this class; inherits ``response_cache.ttl_seconds`` + when ``None``. + bypass_rate: Live-rerun probability for this class; inherits + ``response_cache.bypass_rate`` when ``None``. + arg_skip: Argument keys dropped before keying (default empty: key on all args). + members: Tool names in this class (unused for the ``default`` bucket). + Names may use ``*`` wildcards; an exact member wins over any + wildcard match, the most-specific pattern wins among wildcards, and + unmatched tools fall to ``default``. + """ + + cacheable: bool = False + ttl_seconds: int | None = None + bypass_rate: float | None = None + arg_skip: list[str] = field(default_factory=list) + members: list[str] = field(default_factory=list) + + def to_dict(self) -> JsonObject: + """Serialize this tool class to the canonical JSON object shape.""" + return _normalize_object( + { + "cacheable": self.cacheable, + "ttl_seconds": self.ttl_seconds, + "bypass_rate": self.bypass_rate, + "arg_skip": self.arg_skip, + "members": self.members, + } + ) + + +@dataclass(slots=True) +class ToolOverride: + """Per-tool refinement applied on top of the tool's resolved class. + + Args: + cacheable: Overrides the class ``cacheable`` for just this tool. + ttl_seconds: Overrides the class TTL for just this tool. + bypass_rate: Overrides the class bypass rate for just this tool. + tool_version: Version string folded into the key so a deployment can bust + stale entries before their TTL. + arg_skip: Replaces the class ``arg_skip`` when not ``None`` (``None`` + inherits the class list; ``[]`` clears it). + """ + + cacheable: bool | None = None + ttl_seconds: int | None = None + bypass_rate: float | None = None + tool_version: str | None = None + arg_skip: list[str] | None = None + + def to_dict(self) -> JsonObject: + """Serialize this tool override to the canonical JSON object shape.""" + return _normalize_object( + { + "cacheable": self.cacheable, + "ttl_seconds": self.ttl_seconds, + "bypass_rate": self.bypass_rate, + "tool_version": self.tool_version, + "arg_skip": self.arg_skip, + } + ) + + +@dataclass(slots=True) +class ToolCacheConfig: + """Opt-in tool-result cache settings. + + A separate surface under ``response_cache`` keyed on tool name + arguments and + gated by user-declared safety classes. Off until ``enabled`` is set; any tool + not listed in a class falls into ``default``, which defaults to not cached. + + Args: + enabled: Master switch for the tool surface. Off by default. + priority: Tool execution-intercept priority. Defaults to 150 so + standard priority-100 guardrails wrap cache hits; lower runs + first/outermost. + cache_errors: Whether error-shaped tool results may be cached. Off by + default. + default: Policy for tools not listed in any class (defaults to not cached). + classes: Named tool classes, each with its own policy and member list. + overrides: Per-tool refinements applied on top of the resolved class. + Keys may be exact tool names or ``*`` patterns; an exact key wins + outright, then the most-specific matching pattern applies. + """ + + enabled: bool = False + priority: int = 150 + cache_errors: bool = False + default: ToolClass = field(default_factory=ToolClass) + classes: dict[str, ToolClass] = field(default_factory=dict) + overrides: dict[str, ToolOverride] = field(default_factory=dict) + + def to_dict(self) -> JsonObject: + """Serialize this tool-cache config to the canonical JSON object shape.""" + return _normalize_object( + { + "enabled": self.enabled, + "priority": self.priority, + "cache_errors": self.cache_errors, + "default": _normalize(self.default), + "classes": {name: _normalize(cls) for name, cls in self.classes.items()}, + "overrides": {name: _normalize(ov) for name, ov in self.overrides.items()}, + } + ) + + @dataclass(slots=True) class ResponseCacheConfig: - """Opt-in LLM response cache (exact-match) settings. + """Opt-in exact-match LLM response and tool-result cache 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. + intercepts that reuse earlier LLM answers and, when ``tools.enabled`` is + set, explicitly classified tool results. Args: ttl_seconds: How long a stored answer stays reusable, in seconds. @@ -271,6 +385,7 @@ class ResponseCacheConfig: key_strategy: Key strategy. Only ``"exact_request"`` is supported. header_allowlist: Request headers folded into the key; never auth headers. backend: Cache storage backend (``in_memory`` or ``redis``). + tools: Opt-in tool-result cache; ``None`` leaves it off. """ ttl_seconds: int = 3600 @@ -281,6 +396,7 @@ class ResponseCacheConfig: key_strategy: str = "exact_request" header_allowlist: list[str] = field(default_factory=list) backend: BackendSpec = field(default_factory=BackendSpec.in_memory) + tools: ToolCacheConfig | None = None def to_dict(self) -> JsonObject: """Serialize this response-cache config to the canonical JSON object shape.""" @@ -294,6 +410,7 @@ def to_dict(self) -> JsonObject: "key_strategy": self.key_strategy, "header_allowlist": self.header_allowlist, "backend": _normalize(self.backend), + "tools": _normalize(self.tools), } ) @@ -311,7 +428,7 @@ class AdaptiveConfig: tool_parallelism: Built-in tool scheduling settings. acg: Adaptive Cache Governor settings. policy: Unsupported-config policy applied within the adaptive config. - response_cache: Opt-in LLM response cache settings. + response_cache: Opt-in LLM response and tool-result cache settings. Behavior: This document configures only the adaptive component. Plugins are @@ -438,6 +555,9 @@ def set_latency_sensitivity(level: int) -> None: "ResponseCacheConfig", "StateConfig", "TelemetryConfig", + "ToolCacheConfig", + "ToolClass", + "ToolOverride", "ToolParallelismConfig", "set_latency_sensitivity", "UnsupportedBehavior", diff --git a/python/nemo_relay/adaptive.pyi b/python/nemo_relay/adaptive.pyi index 415c4df2c..6319cb8e2 100644 --- a/python/nemo_relay/adaptive.pyi +++ b/python/nemo_relay/adaptive.pyi @@ -181,9 +181,52 @@ class AcgConfig: """Serialize this ACG config to the canonical JSON object shape.""" ... +@dataclass(slots=True) +class ToolClass: + """One tool caching class (also the shape of the ``default`` default bucket).""" + + cacheable: bool = ... + ttl_seconds: int | None = ... + bypass_rate: float | None = ... + arg_skip: list[str] = ... + members: list[str] = ... + + def to_dict(self) -> JsonObject: + """Serialize this tool class to the canonical JSON object shape.""" + ... + +@dataclass(slots=True) +class ToolOverride: + """Per-tool refinement applied on top of the tool's resolved class.""" + + cacheable: bool | None = ... + ttl_seconds: int | None = ... + bypass_rate: float | None = ... + tool_version: str | None = ... + arg_skip: list[str] | None = ... + + def to_dict(self) -> JsonObject: + """Serialize this tool override to the canonical JSON object shape.""" + ... + +@dataclass(slots=True) +class ToolCacheConfig: + """Opt-in tool-result cache settings.""" + + enabled: bool = ... + priority: int = ... + cache_errors: bool = ... + default: ToolClass = ... + classes: dict[str, ToolClass] = ... + overrides: dict[str, ToolOverride] = ... + + def to_dict(self) -> JsonObject: + """Serialize this tool-cache config to the canonical JSON object shape.""" + ... + @dataclass(slots=True) class ResponseCacheConfig: - """Opt-in LLM response cache (exact-match) settings. + """Opt-in exact-match LLM response and tool-result cache settings. A section of the adaptive component, not a standalone plugin kind. @@ -199,6 +242,7 @@ class ResponseCacheConfig: key_strategy: Key strategy. Only ``"exact_request"`` is supported. header_allowlist: Request headers folded into the key. backend: Cache storage backend (``in_memory`` or ``redis``). + tools: Opt-in tool-result cache; ``None`` leaves the tool surface off. """ ttl_seconds: int = ... @@ -209,6 +253,7 @@ class ResponseCacheConfig: key_strategy: str = ... header_allowlist: list[str] = ... backend: BackendSpec = ... + tools: ToolCacheConfig | None = ... def to_dict(self) -> JsonObject: """Serialize this response-cache config to the canonical JSON object shape.""" @@ -227,7 +272,7 @@ class AdaptiveConfig: tool_parallelism: Built-in adaptive tool-scheduling configuration. acg: Adaptive Cache Governor configuration. policy: Policy for unsupported adaptive configuration. - response_cache: Opt-in LLM response cache configuration. + response_cache: Opt-in LLM response and tool-result cache configuration. """ version: int = ... diff --git a/python/tests/test_adaptive_config.py b/python/tests/test_adaptive_config.py index 414ed517b..37e4fb201 100644 --- a/python/tests/test_adaptive_config.py +++ b/python/tests/test_adaptive_config.py @@ -18,6 +18,9 @@ ResponseCacheConfig, StateConfig, TelemetryConfig, + ToolCacheConfig, + ToolClass, + ToolOverride, ToolParallelismConfig, ) @@ -218,6 +221,23 @@ def test_invalid_response_cache_section_is_rejected(self): assert "response_cache.invalid_ttl" in codes assert "response_cache.invalid_bypass_rate" in codes + def test_tool_cache_config_serializes_and_omits_unset_optionals(self): + tools = ToolCacheConfig( + enabled=True, + cache_errors=True, + classes={"read_only": ToolClass(cacheable=True, members=["docs_lookup"])}, + overrides={"docs_lookup": ToolOverride(tool_version="v1")}, + ) + serialized = ResponseCacheConfig(tools=tools).to_dict()["tools"] + assert serialized == { + "enabled": True, + "cache_errors": True, + "priority": 150, + "default": {"cacheable": False, "arg_skip": [], "members": []}, + "classes": {"read_only": {"cacheable": True, "arg_skip": [], "members": ["docs_lookup"]}}, + "overrides": {"docs_lookup": {"tool_version": "v1"}}, + } + def test_canonical_cache_telemetry_helper_supports_openai_provider(self): event = adaptive_module.build_cache_telemetry_event( provider="openai",