diff --git a/crates/adaptive/src/config.rs b/crates/adaptive/src/config.rs index d5db2445a..27e8172bb 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,7 @@ nemo_relay::editor_config! { nested: BackendConfig, default: BackendConfig, }, + tools => { label: "tools", kind: Json, optional: true }, } } 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..51cd7c973 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,76 @@ 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(), + } + } +} + +/// 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>, +} diff --git a/crates/adaptive/src/response_cache/key.rs b/crates/adaptive/src/response_cache/key.rs index 36eeca3d1..9770b9a87 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 @@ -400,6 +445,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..286e5a4a7 100644 --- a/crates/adaptive/src/response_cache/mark.rs +++ b/crates/adaptive/src/response_cache/mark.rs @@ -103,12 +103,14 @@ fn probed_savings(entry: &CacheEntry) -> (Option, Option) { pub(crate) struct CacheMark<'a> { status: &'a str, reason: Option<&'a str>, + surface: &'a str, 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 +118,14 @@ impl<'a> CacheMark<'a> { Self { status, reason: None, + surface: "llm", backend, key_hash: None, age_ms: None, ttl_ms: None, saved_tokens: None, saved_cost_usd: None, + saved_invocations: None, } } @@ -130,6 +134,12 @@ impl<'a> CacheMark<'a> { self } + /// Overrides the surface label (defaults to `"llm"`; tool marks set `"tool"`). + pub(crate) fn surface(mut self, surface: &'a str) -> 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 +160,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 +174,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), ); metadata.insert( "nemo_relay.response_cache.backend".to_string(), @@ -200,6 +216,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..e2122942a --- /dev/null +++ b/crates/adaptive/src/response_cache/tool.rs @@ -0,0 +1,359 @@ +// 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, emit_cache_mark}; +use crate::response_cache::store::{CacheEntry, CacheStore, now_unix_ms}; + +const TOOL_SURFACE: &str = "tool"; + +#[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(TOOL_SURFACE) + .reason(reason), + ); + return next(args).await.map(Into::into); + } + }; + + if should_bypass(policy.bypass_rate) { + emit_cache_mark( + CacheMark::new("bypass", backend) + .surface(TOOL_SURFACE) + .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(TOOL_SURFACE) + .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(TOOL_SURFACE) + .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(TOOL_SURFACE) + .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(TOOL_SURFACE) + .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)] +#[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..0f5bcaf20 100644 --- a/crates/adaptive/tests/integration/response_cache_tests.rs +++ b/crates/adaptive/tests/integration/response_cache_tests.rs @@ -16,12 +16,16 @@ 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, @@ -29,6 +33,7 @@ use nemo_relay::plugin::{ 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; @@ -1777,6 +1782,7 @@ async fn cache_coexists_with_acg_execution_intercept() { async fn redis_backend_shares_entries_across_store_instances() { use std::time::Duration; + use nemo_relay_adaptive::response_cache::check_backend_health; use nemo_relay_adaptive::response_cache::store::{CacheEntry, CacheStore, RedisCacheStore}; let enabled = std::env::var("NEMO_RELAY_RUN_REDIS_TESTS") @@ -1801,6 +1807,23 @@ async fn redis_backend_shares_entries_across_store_instances() { .await .expect("connect redis (reader)"); + let mut config = ResponseCacheConfig::default(); + config.backend.kind = "redis".to_string(); + config + .backend + .config + .insert("url".to_string(), Json::String(url.to_string())); + config + .backend + .config + .insert("key_prefix".to_string(), Json::String(prefix.to_string())); + assert_eq!( + check_backend_health(&config) + .await + .expect("configured Redis health check"), + "redis" + ); + let key = "sha256:shared-cache-test-key"; let _ = writer.delete(key).await; @@ -1824,6 +1847,578 @@ async fn redis_backend_shares_entries_across_store_instances() { ); assert_eq!(got.unwrap().response["answer"], json!("shared")); + let stale_key = "sha256:shared-cache-stale-key"; + let _ = writer.delete(stale_key).await; + writer + .set( + stale_key, + CacheEntry { + response: json!({"answer": "stale"}), + created_unix_ms: 0, + expires_unix_ms: 1, + key_hash: stale_key.to_string(), + model_name: None, + provider_name: None, + }, + Duration::from_secs(60), + ) + .await + .expect("set stale entry"); + assert!( + reader + .get(stale_key) + .await + .expect("get stale entry") + .is_none(), + "the response-cache entry expiry must remain authoritative over Redis TTL" + ); + writer.delete(stale_key).await.expect("delete stale entry"); + 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_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), + members: vec!["dup".to_string()], + ..ToolClass::default() + }, + ); + let adaptive = AdaptiveConfig { + response_cache: Some(cache_with_tools(ToolCacheConfig { + enabled: true, + default: ToolClass { + members: vec!["safe_lookup".to_string()], + ..ToolClass::default() + }, + classes, + ..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_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( + "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/config_tests.rs b/crates/adaptive/tests/unit/config_tests.rs index 0a5537518..c6a6699a8 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,21 @@ fn test_typed_section_helpers_default() { let response_cache = ResponseCacheConfig::default(); assert!(!response_cache.cache_nondeterministic); + + let tools = ToolCacheConfig::default(); + 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] diff --git a/crates/adaptive/tests/unit/response_cache/key_tests.rs b/crates/adaptive/tests/unit/response_cache/key_tests.rs index aa03d30f5..214d5cb8d 100644 --- a/crates/adaptive/tests/unit/response_cache/key_tests.rs +++ b/crates/adaptive/tests/unit/response_cache/key_tests.rs @@ -816,3 +816,164 @@ 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); +} 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..4255bd5f7 --- /dev/null +++ b/crates/adaptive/tests/unit/response_cache/tool_tests.rs @@ -0,0 +1,114 @@ +// 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}; +use crate::response_cache::store::{CacheEntry, CacheStore, InMemoryCacheStore}; + +#[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" + }))); +} + +#[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); +} diff --git a/crates/cli/src/diagnostics/mod.rs b/crates/cli/src/diagnostics/mod.rs index 6df3026ae..98c3b4fee 100644 --- a/crates/cli/src/diagnostics/mod.rs +++ b/crates/cli/src/diagnostics/mod.rs @@ -699,6 +699,35 @@ async fn collect_response_cache_component_checks( } }; 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 10d3bc147..3706ede6b 100644 --- a/crates/cli/tests/coverage/shared/doctor_tests.rs +++ b/crates/cli/tests/coverage/shared/doctor_tests.rs @@ -1303,6 +1303,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/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",