From cba26f6157a4e423dde9991cdabc668f04b4c199 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Thu, 21 May 2026 09:09:51 -0400 Subject: [PATCH] refactor: narrow core runtime registry surface Signed-off-by: Will Killian --- crates/adaptive/src/acg_component.rs | 2 +- .../adaptive/src/adaptive_hints_intercept.rs | 2 +- .../integration/runtime_integration_tests.rs | 4 +- .../tests/unit/runtime_features_tests.rs | 210 +++++++++------- crates/cli/tests/coverage/server_tests.rs | 2 +- crates/core/src/api/llm.rs | 6 +- crates/core/src/api/registry.rs | 109 ++++---- crates/core/src/api/runtime.rs | 3 +- crates/core/src/api/runtime/callbacks.rs | 14 +- crates/core/src/api/runtime/scope_stack.rs | 27 +- crates/core/src/api/runtime/state.rs | 199 +++++---------- crates/core/src/api/tool.rs | 4 +- crates/core/src/context/registries.rs | 90 +++---- crates/core/src/lib.rs | 4 +- .../src/observability/plugin_component.rs | 1 + crates/core/src/registry.rs | 116 +++------ .../tests/integration/api_surface_tests.rs | 44 ++-- .../integration/context_isolation_tests.rs | 6 - .../tests/integration/middleware_tests.rs | 101 ++++---- .../core/tests/integration/pipeline_tests.rs | 24 +- .../tests/integration/scope_local_tests.rs | 24 +- crates/core/tests/unit/context_tests.rs | 153 +++++++----- crates/core/tests/unit/plugin_tests.rs | 50 ++-- crates/core/tests/unit/registry_tests.rs | 236 ++++-------------- crates/core/tests/unit/shared_tests.rs | 4 +- crates/ffi/src/callable.rs | 28 ++- crates/node/src/callable.rs | 19 +- crates/python/src/py_callable.rs | 23 +- crates/wasm/src/callable.rs | 45 ++-- 29 files changed, 672 insertions(+), 878 deletions(-) diff --git a/crates/adaptive/src/acg_component.rs b/crates/adaptive/src/acg_component.rs index 666c80e80..40183779f 100644 --- a/crates/adaptive/src/acg_component.rs +++ b/crates/adaptive/src/acg_component.rs @@ -594,7 +594,7 @@ pub(crate) fn create_acg_llm_request_intercept( provider: String, plugin: Arc, ) -> LlmRequestInterceptFn { - Box::new(move |_name: &str, request: LlmRequest, annotated| { + Arc::new(move |_name: &str, request: LlmRequest, annotated| { let translated = translate_request(&request, &agent_id, &provider, plugin.as_ref(), &hot_cache) .unwrap_or(request); diff --git a/crates/adaptive/src/adaptive_hints_intercept.rs b/crates/adaptive/src/adaptive_hints_intercept.rs index 341ebafcc..33875f4fb 100644 --- a/crates/adaptive/src/adaptive_hints_intercept.rs +++ b/crates/adaptive/src/adaptive_hints_intercept.rs @@ -171,7 +171,7 @@ impl AdaptiveHintsIntercept { /// transformed request. pub fn into_request_fn(self) -> LlmRequestInterceptFn { let this = Arc::new(self); - Box::new( + Arc::new( move |_name: &str, mut request: LlmRequest, annotated: Option| { let scope_path = extract_scope_path(); let manual_ls = read_manual_latency_sensitivity(); diff --git a/crates/adaptive/tests/integration/runtime_integration_tests.rs b/crates/adaptive/tests/integration/runtime_integration_tests.rs index 22e593488..5302a8328 100644 --- a/crates/adaptive/tests/integration/runtime_integration_tests.rs +++ b/crates/adaptive/tests/integration/runtime_integration_tests.rs @@ -718,7 +718,7 @@ impl Plugin for HeaderPlugin { "header_plugin", priority, false, - Box::new(|_name, mut request, annotated| { + Arc::new(|_name, mut request, annotated| { request.headers.insert("x-plugin".into(), json!("set")); Ok((request, annotated)) }), @@ -727,7 +727,7 @@ impl Plugin for HeaderPlugin { "tool_request_plugin", priority, false, - Box::new(|_name, mut args| { + Arc::new(|_name, mut args| { if let Json::Object(ref mut map) = args { map.insert("x-tool-plugin".into(), json!(true)); } diff --git a/crates/adaptive/tests/unit/runtime_features_tests.rs b/crates/adaptive/tests/unit/runtime_features_tests.rs index afe266bc6..1c441381e 100644 --- a/crates/adaptive/tests/unit/runtime_features_tests.rs +++ b/crates/adaptive/tests/unit/runtime_features_tests.rs @@ -9,10 +9,18 @@ use std::sync::Arc; use nemo_flow::api::llm::LlmRequest; use nemo_flow::api::llm::llm_request_intercepts; +use nemo_flow::api::registry::{ + deregister_llm_execution_intercept, deregister_llm_request_intercept, + deregister_llm_stream_execution_intercept, deregister_tool_execution_intercept, + register_llm_execution_intercept, register_llm_request_intercept, + register_llm_stream_execution_intercept, register_tool_execution_intercept, +}; use nemo_flow::api::runtime::NemoFlowContextState; use nemo_flow::api::runtime::ToolExecutionNextFn; use nemo_flow::api::runtime::global_context; +use nemo_flow::api::subscriber::{deregister_subscriber, register_subscriber}; use nemo_flow::api::tool::tool_call_execute; +use nemo_flow::error::FlowError; use nemo_flow::plugin::{ConfigPolicy, UnsupportedBehavior}; use nemo_flow::plugin::{clear_plugin_configuration, rollback_registrations}; use serde_json::json; @@ -55,6 +63,100 @@ fn sample_plan(agent_id: &str) -> ExecutionPlan { } } +fn assert_already_registered(result: nemo_flow::error::Result<()>, name: &str) { + match result { + Err(FlowError::AlreadyExists(message)) => assert!(message.contains(name)), + other => panic!("expected {name} to be registered, got {other:?}"), + } +} + +fn assert_subscriber_registered(name: &str) { + assert_already_registered(register_subscriber(name, Arc::new(|_event| {})), name); +} + +fn assert_subscriber_absent(name: &str) { + register_subscriber(name, Arc::new(|_event| {})).unwrap(); + deregister_subscriber(name).unwrap(); +} + +fn assert_llm_request_intercept_registered(name: &str) { + assert_already_registered( + register_llm_request_intercept( + name, + i32::MAX, + false, + Arc::new(|_name, request, annotated| Ok((request, annotated))), + ), + name, + ); +} + +fn assert_llm_request_intercept_absent(name: &str) { + register_llm_request_intercept( + name, + i32::MAX, + false, + Arc::new(|_name, request, annotated| Ok((request, annotated))), + ) + .unwrap(); + deregister_llm_request_intercept(name).unwrap(); +} + +fn assert_llm_execution_intercept_registered(name: &str) { + assert_already_registered( + register_llm_execution_intercept( + name, + i32::MAX, + Arc::new(|_name, request, next| next(request)), + ), + name, + ); +} + +fn assert_llm_execution_intercept_absent(name: &str) { + register_llm_execution_intercept( + name, + i32::MAX, + Arc::new(|_name, request, next| next(request)), + ) + .unwrap(); + deregister_llm_execution_intercept(name).unwrap(); +} + +fn assert_llm_stream_execution_intercept_registered(name: &str) { + assert_already_registered( + register_llm_stream_execution_intercept( + name, + i32::MAX, + Arc::new(|_name, request, next| next(request)), + ), + name, + ); +} + +fn assert_llm_stream_execution_intercept_absent(name: &str) { + register_llm_stream_execution_intercept( + name, + i32::MAX, + Arc::new(|_name, request, next| next(request)), + ) + .unwrap(); + deregister_llm_stream_execution_intercept(name).unwrap(); +} + +fn assert_tool_execution_intercept_registered(name: &str) { + assert_already_registered( + register_tool_execution_intercept(name, i32::MAX, Arc::new(|_name, args, next| next(args))), + name, + ); +} + +fn assert_tool_execution_intercept_absent(name: &str) { + register_tool_execution_intercept(name, i32::MAX, Arc::new(|_name, args, next| next(args))) + .unwrap(); + deregister_tool_execution_intercept(name).unwrap(); +} + struct SeedFailBackend; impl StorageBackendDyn for SeedFailBackend { @@ -201,22 +303,10 @@ async fn telemetry_feature_registers_subscriber_and_starts_drain_task() { ctx.finish() }; assert!(runtime.drain_handle.is_some()); - assert!( - global_context() - .read() - .unwrap() - .event_subscribers - .contains_key(&name) - ); + assert_subscriber_registered(&name); rollback_registrations(&mut registrations); - assert!( - !global_context() - .read() - .unwrap() - .event_subscribers - .contains_key(&name) - ); + assert_subscriber_absent(&name); if let Some(handle) = runtime.drain_handle.take() { handle.abort(); @@ -284,13 +374,7 @@ async fn adaptive_hints_feature_registers_request_intercept() { let mut ctx = RegistrationContext::new(&mut runtime); feature.register(&mut ctx).await.unwrap(); - assert!( - global_context() - .read() - .unwrap() - .llm_request_intercepts - .contains(&name) - ); + assert_llm_request_intercept_registered(&name); let request = llm_request_intercepts( "model", @@ -304,13 +388,7 @@ async fn adaptive_hints_feature_registers_request_intercept() { let mut registrations = ctx.finish(); rollback_registrations(&mut registrations); - assert!( - !global_context() - .read() - .unwrap() - .llm_request_intercepts - .contains(&name) - ); + assert_llm_request_intercept_absent(&name); } #[tokio::test(flavor = "current_thread")] @@ -343,13 +421,7 @@ async fn tool_parallelism_feature_registers_execution_intercept() { let mut ctx = RegistrationContext::new(&mut runtime); feature.register(&mut ctx).await.unwrap(); - assert!( - global_context() - .read() - .unwrap() - .tool_execution_intercepts - .contains(&name) - ); + assert_tool_execution_intercept_registered(&name); let next: ToolExecutionNextFn = Arc::new(|args| Box::pin(async move { Ok(args) })); let result = tool_call_execute( @@ -365,13 +437,7 @@ async fn tool_parallelism_feature_registers_execution_intercept() { let mut registrations = ctx.finish(); rollback_registrations(&mut registrations); - assert!( - !global_context() - .read() - .unwrap() - .tool_execution_intercepts - .contains(&name) - ); + assert_tool_execution_intercept_absent(&name); } #[tokio::test(flavor = "current_thread")] @@ -481,7 +547,7 @@ async fn registration_context_registers_all_supported_callback_types() { "adaptive_test_request", 5, false, - Box::new(|_name, request, annotated| Ok((request, annotated))), + Arc::new(|_name, request, annotated| Ok((request, annotated))), ) .unwrap(); ctx.register_llm_execution_intercept( @@ -515,34 +581,11 @@ async fn registration_context_registers_all_supported_callback_types() { .unwrap(); let mut registrations = ctx.finish(); - let global = global_context(); - let state = global.read().unwrap(); - assert!( - state - .event_subscribers - .contains_key("adaptive_test_subscriber") - ); - assert!( - state - .llm_request_intercepts - .contains("adaptive_test_request") - ); - assert!( - state - .llm_execution_intercepts - .contains("adaptive_test_execution") - ); - assert!( - state - .llm_stream_execution_intercepts - .contains("adaptive_test_stream") - ); - assert!( - state - .tool_execution_intercepts - .contains("adaptive_test_tool") - ); - drop(state); + assert_subscriber_registered("adaptive_test_subscriber"); + assert_llm_request_intercept_registered("adaptive_test_request"); + assert_llm_execution_intercept_registered("adaptive_test_execution"); + assert_llm_stream_execution_intercept_registered("adaptive_test_stream"); + assert_tool_execution_intercept_registered("adaptive_test_tool"); rollback_registrations(&mut registrations); } @@ -622,14 +665,13 @@ async fn acg_feature_registers_execution_and_stream_intercepts() { let mut ctx = RegistrationContext::new(&mut runtime); feature.register(&mut ctx).await.unwrap(); - let global = global_context(); - let state = global.read().unwrap(); - assert!(state.llm_execution_intercepts.contains(&execution_name)); - assert!(state.llm_stream_execution_intercepts.contains(&stream_name)); - drop(state); + assert_llm_execution_intercept_registered(&execution_name); + assert_llm_stream_execution_intercept_registered(&stream_name); let mut registrations = ctx.finish(); rollback_registrations(&mut registrations); + assert_llm_execution_intercept_absent(&execution_name); + assert_llm_stream_execution_intercept_absent(&stream_name); } #[tokio::test(flavor = "current_thread")] @@ -658,20 +700,8 @@ async fn adaptive_runtime_register_feature_rolls_back_partial_registrations_and_ assert!(!runtime.registered); assert!(runtime.drain_handle.is_none()); assert!(runtime.registrations.is_empty()); - assert!( - !global_context() - .read() - .unwrap() - .event_subscribers - .contains_key("existing_feature") - ); - assert!( - !global_context() - .read() - .unwrap() - .event_subscribers - .contains_key("partial_feature") - ); + assert_subscriber_absent("existing_feature"); + assert_subscriber_absent("partial_feature"); } #[tokio::test(flavor = "current_thread")] diff --git a/crates/cli/tests/coverage/server_tests.rs b/crates/cli/tests/coverage/server_tests.rs index 792c0ab37..39368d079 100644 --- a/crates/cli/tests/coverage/server_tests.rs +++ b/crates/cli/tests/coverage/server_tests.rs @@ -531,7 +531,7 @@ async fn pre_tool_hook_rejects_when_conditional_guardrail_blocks() { register_tool_conditional_execution_guardrail( "cli-pre-tool-blocker", 1, - Box::new(|name, _args| { + Arc::new(|name, _args| { Ok((name == BLOCKED_TEST_TOOL).then(|| "blocked by policy".to_string())) }), ) diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index 2cc3a2363..c7b894e2a 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -532,7 +532,7 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { }; if let Some(error) = NemoFlowContextState::llm_conditional_execution_snapshot_chain( &request, - entries, + &entries, &subscribers, parent_uuid, guardrail_metadata, @@ -680,7 +680,7 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu }; if let Some(error) = NemoFlowContextState::llm_conditional_execution_snapshot_chain( &request, - entries, + &entries, &subscribers, parent_uuid, guardrail_metadata, @@ -822,7 +822,7 @@ pub fn llm_conditional_execution(request: &LlmRequest) -> Result<()> { }; if let Some(error) = NemoFlowContextState::llm_conditional_execution_snapshot_chain( request, - entries, + &entries, &subscribers, parent_uuid, None, diff --git a/crates/core/src/api/registry.rs b/crates/core/src/api/registry.rs index e497c78ca..527e216c8 100644 --- a/crates/core/src/api/registry.rs +++ b/crates/core/src/api/registry.rs @@ -12,33 +12,74 @@ use crate::api::runtime::{ use crate::api::runtime::{current_scope_stack, global_context}; use crate::api::shared::ensure_runtime_owner; use crate::error::{FlowError, Result}; +use crate::registry::RegistryEntry; -/// A priority-ordered request intercept registration entry. -pub struct Intercept { +/// A priority-ordered registration record. +/// +/// Registry records carry the common entry metadata used by all middleware +/// registries plus the caller-provided payload. +#[derive(Clone)] +pub(crate) struct RegistryRecord { + /// Unique middleware name within its registry. + pub(crate) name: String, /// Lower values run earlier in the chain. - pub priority: i32, - /// Whether this intercept stops later request intercepts after it returns. - pub break_chain: bool, - /// The caller-provided intercept callback. - pub callable: F, + pub(crate) priority: i32, + /// The caller-provided registry payload. + pub(crate) payload: T, } -/// A priority-ordered execution intercept registration entry. -pub struct ExecutionIntercept { - /// Lower values run earlier in the chain. - pub priority: i32, - /// The caller-provided execution intercept callback. - pub callable: F, +impl RegistryRecord { + /// Create a new priority-ordered registry record. + pub(crate) fn new(name: impl Into, priority: i32, payload: T) -> Self { + Self { + name: name.into(), + priority, + payload, + } + } } -/// A priority-ordered guardrail registration entry. -pub struct GuardrailEntry { - /// Lower values run earlier in the chain. - pub priority: i32, - /// The caller-provided guardrail callback. - pub guardrail: F, +impl RegistryEntry for RegistryRecord { + fn name(&self) -> &str { + &self.name + } + + fn priority(&self) -> i32 { + self.priority + } +} + +/// Request-intercept-specific registration payload. +/// +/// Request intercepts carry one extra chain-control flag that does not apply +/// to guardrails or execution intercepts. +#[derive(Clone)] +pub(crate) struct RequestIntercept { + /// Whether this intercept stops later request intercepts after it returns. + pub(crate) break_chain: bool, + /// The caller-provided request intercept callback. + pub(crate) callable: F, +} + +impl RequestIntercept { + /// Create a new request intercept payload. + pub(crate) fn new(break_chain: bool, callable: F) -> Self { + Self { + break_chain, + callable, + } + } } +/// A priority-ordered guardrail registration record. +pub(crate) type Guardrail = RegistryRecord; + +/// A priority-ordered request intercept registration record. +pub(crate) type Intercept = RegistryRecord>; + +/// A priority-ordered execution intercept registration record. +pub(crate) type ExecutionIntercept = RegistryRecord; + macro_rules! global_guardrail_registry_api { ( $(#[$register_meta:meta])* @@ -70,11 +111,7 @@ macro_rules! global_guardrail_registry_api { state .$field .register( - name.to_string(), - GuardrailEntry { - priority, - guardrail: guardrail.into(), - }, + Guardrail::new(name, priority, guardrail), ) .map_err(FlowError::AlreadyExists) } @@ -139,12 +176,7 @@ macro_rules! global_intercept_registry_api { state .$field .register( - name.to_string(), - Intercept { - priority, - break_chain, - callable, - }, + Intercept::new(name, priority, RequestIntercept::new(break_chain, callable)), ) .map_err(FlowError::AlreadyExists) } @@ -201,7 +233,7 @@ macro_rules! global_execution_registry_api { .map_err(|error| FlowError::Internal(error.to_string()))?; state .$field - .register(name.to_string(), ExecutionIntercept { priority, callable }) + .register(ExecutionIntercept::new(name, priority, callable)) .map_err(FlowError::AlreadyExists) } @@ -266,11 +298,7 @@ macro_rules! scope_guardrail_registry_api { registries .$field .register( - name.to_string(), - GuardrailEntry { - priority, - guardrail: guardrail.into(), - }, + Guardrail::new(name, priority, guardrail), ) .map_err(FlowError::AlreadyExists) } @@ -342,12 +370,7 @@ macro_rules! scope_intercept_registry_api { registries .$field .register( - name.to_string(), - Intercept { - priority, - break_chain, - callable, - }, + Intercept::new(name, priority, RequestIntercept::new(break_chain, callable)), ) .map_err(FlowError::AlreadyExists) } @@ -415,7 +438,7 @@ macro_rules! scope_execution_registry_api { .ok_or_else(|| FlowError::NotFound(format!("scope {scope_uuid} not found")))?; registries .$field - .register(name.to_string(), ExecutionIntercept { priority, callable }) + .register(ExecutionIntercept::new(name, priority, callable)) .map_err(FlowError::AlreadyExists) } diff --git a/crates/core/src/api/runtime.rs b/crates/core/src/api/runtime.rs index 9583046ed..28b83db63 100644 --- a/crates/core/src/api/runtime.rs +++ b/crates/core/src/api/runtime.rs @@ -11,8 +11,7 @@ pub mod state; pub use callbacks::{ EventSubscriberFn, LlmCollectorFn, LlmConditionalFn, LlmExecutionFn, LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestFn, - LlmSanitizeResponseFn, LlmStreamExecutionFn, LlmStreamExecutionNextFn, - LlmStreamExecutionRegistryRef, LlmStreamExecutionRegistryRefs, ToolConditionalFn, + LlmSanitizeResponseFn, LlmStreamExecutionFn, LlmStreamExecutionNextFn, ToolConditionalFn, ToolExecutionFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; pub use global::global_context; diff --git a/crates/core/src/api/runtime/callbacks.rs b/crates/core/src/api/runtime/callbacks.rs index 43be6109b..980c47d74 100644 --- a/crates/core/src/api/runtime/callbacks.rs +++ b/crates/core/src/api/runtime/callbacks.rs @@ -32,7 +32,7 @@ use crate::json::Json; /// /// # Returns /// Sanitized JSON payload for the emitted event. -pub type ToolSanitizeFn = Box Json + Send + Sync>; +pub type ToolSanitizeFn = Arc Json + Send + Sync>; /// Decide whether a tool call is allowed to continue. /// /// The callback receives the tool name and the current argument payload. It can @@ -70,7 +70,7 @@ pub type ToolConditionalFn = Arc Result> + /// # Errors /// The callback can return any [`FlowError`](crate::error::FlowError) to abort /// the request-intercept chain. -pub type ToolInterceptFn = Box Result + Send + Sync>; +pub type ToolInterceptFn = Arc Result + Send + Sync>; /// Continuation type invoked by tool execution intercepts. /// /// Execution intercepts receive this callable as their `next` continuation and @@ -120,7 +120,7 @@ pub type ToolExecutionFn = Arc< /// /// # Returns /// Sanitized [`LlmRequest`] for the emitted event. -pub type LlmSanitizeRequestFn = Box LlmRequest + Send + Sync>; +pub type LlmSanitizeRequestFn = Arc LlmRequest + Send + Sync>; /// Sanitize an LLM response before the runtime records it. /// /// These callbacks rewrite the JSON response payload captured on LLM-end @@ -131,7 +131,7 @@ pub type LlmSanitizeRequestFn = Box LlmRequest + Send + Sy /// /// # Returns /// Sanitized JSON response payload for the emitted event. -pub type LlmSanitizeResponseFn = Box Json + Send + Sync>; +pub type LlmSanitizeResponseFn = Arc Json + Send + Sync>; /// Decide whether an LLM call is allowed to continue. /// /// The callback receives the current [`LlmRequest`] and can allow execution, @@ -168,7 +168,7 @@ pub type LlmConditionalFn = Arc Result> + /// # Errors /// The callback can return any [`FlowError`](crate::error::FlowError) to abort /// the request-intercept chain. -pub type LlmRequestInterceptFn = Box< +pub type LlmRequestInterceptFn = Arc< dyn Fn( &str, LlmRequest, @@ -243,14 +243,14 @@ pub type LlmFinalizerFn = Box Json + Send>; /// /// # Returns /// A shared reference to a scope-local streaming execution registry. -pub type LlmStreamExecutionRegistryRef<'a> = &'a crate::registry::SortedRegistry< +pub(crate) type LlmStreamExecutionRegistryRef<'a> = &'a crate::registry::SortedRegistry< crate::api::registry::ExecutionIntercept, >; /// Slice of scope-local streaming execution registries. /// /// # Returns /// A borrowed slice of scope-local streaming execution registry references. -pub type LlmStreamExecutionRegistryRefs<'a> = &'a [LlmStreamExecutionRegistryRef<'a>]; +pub(crate) type LlmStreamExecutionRegistryRefs<'a> = &'a [LlmStreamExecutionRegistryRef<'a>]; /// Continuation type invoked by streaming LLM execution intercepts. /// diff --git a/crates/core/src/api/runtime/scope_stack.rs b/crates/core/src/api/runtime/scope_stack.rs index a678325cb..413407455 100644 --- a/crates/core/src/api/runtime/scope_stack.rs +++ b/crates/core/src/api/runtime/scope_stack.rs @@ -5,8 +5,8 @@ //! //! The runtime tracks the current scope hierarchy through a shared //! [`ScopeStack`] stored in task-local or thread-local state. Advanced callers -//! can use this module to inspect the active scope chain, attach scope-local -//! middleware, or propagate scope context into worker threads. +//! can use this module to inspect the active scope chain or propagate scope +//! context into worker threads. use std::cell::RefCell; use std::sync::{Arc, RwLock}; @@ -17,7 +17,7 @@ use crate::api::runtime::callbacks::EventSubscriberFn; use crate::api::scope::{ScopeHandle, ScopeType}; use crate::context::registries::ScopeLocalRegistries; use crate::error::{FlowError, Result}; -use crate::registry::SortedRegistry; +use crate::registry::{RegistryEntry, SortedRegistry}; /// Mutable stack of active scopes plus their scope-local registries. /// @@ -161,7 +161,10 @@ impl ScopeStack { /// # Notes /// When the scope is active but has no registries yet, this function /// creates an empty scope-local registry set first. - pub fn local_registries_mut(&mut self, uuid: &Uuid) -> Option<&mut ScopeLocalRegistries> { + pub(crate) fn local_registries_mut( + &mut self, + uuid: &Uuid, + ) -> Option<&mut ScopeLocalRegistries> { if !self.stack.iter().any(|handle| handle.uuid == *uuid) { return None; } @@ -177,7 +180,7 @@ impl ScopeStack { /// # Returns /// A vector of registry references ordered from root toward the current /// top-most scope. - pub fn collect_scope_local_registries<'a, T>( + pub(crate) fn collect_scope_local_registries<'a, T: RegistryEntry>( &'a self, field: impl Fn(&'a ScopeLocalRegistries) -> &'a SortedRegistry, ) -> Vec<&'a SortedRegistry> { @@ -193,25 +196,13 @@ impl ScopeStack { /// # Returns /// A vector of subscribers collected from each active scope that owns /// scope-local registries. - pub fn collect_scope_local_subscribers(&self) -> Vec { + pub(crate) fn collect_scope_local_subscribers(&self) -> Vec { self.stack .iter() .filter_map(|handle| self.scope_registries.get(&handle.uuid)) .flat_map(|registries| registries.event_subscribers.values().cloned()) .collect() } - - /// Return the scope-local registries for `uuid` without creating them. - /// - /// # Parameters - /// - `uuid`: UUID of the scope whose registries should be borrowed. - /// - /// # Returns - /// `Some(&ScopeLocalRegistries)` when registries already exist for that - /// scope and `None` otherwise. - pub fn scope_registries_get(&self, uuid: &Uuid) -> Option<&ScopeLocalRegistries> { - self.scope_registries.get(uuid) - } } impl std::fmt::Debug for ScopeStack { diff --git a/crates/core/src/api/runtime/state.rs b/crates/core/src/api/runtime/state.rs index 9c65e6682..943ab198e 100644 --- a/crates/core/src/api/runtime/state.rs +++ b/crates/core/src/api/runtime/state.rs @@ -18,7 +18,7 @@ use crate::api::event::{ }; use crate::api::llm::{CreateLlmHandleParams, EndLlmHandleParams}; use crate::api::llm::{LlmHandle, LlmRequest}; -use crate::api::registry::{ExecutionIntercept, GuardrailEntry, Intercept}; +use crate::api::registry::{ExecutionIntercept, Guardrail, Intercept}; use crate::api::runtime::callbacks::{ EventSubscriberFn, LlmConditionalFn, LlmExecutionFn, LlmExecutionNextFn, LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionFn, LlmStreamExecutionNextFn, @@ -32,7 +32,6 @@ use crate::codec::request::AnnotatedLlmRequest; use crate::codec::response::AnnotatedLlmResponse; use crate::context::registries::{ merge_execution_intercept_callables, merge_guardrail_entries, merge_intercept_entries, - merge_named_guardrail_entries, }; use crate::json::{Json, merge_json}; use crate::registry::SortedRegistry; @@ -40,11 +39,6 @@ use chrono::{Duration, Utc}; use serde_json::json; use uuid::Uuid; -pub(crate) struct ConditionalGuardrailSnapshot { - name: String, - guardrail: F, -} - /// Process-global runtime state backing middleware and event emission. /// /// The public API layer stores one shared instance of this type for the @@ -52,31 +46,32 @@ pub(crate) struct ConditionalGuardrailSnapshot { /// and arbitrary extension slots used by bindings or integrations. pub struct NemoFlowContextState { /// Global tool request sanitizers applied to emitted tool-start payloads. - pub tool_sanitize_request_guardrails: SortedRegistry>, + pub(crate) tool_sanitize_request_guardrails: SortedRegistry>, /// Global tool response sanitizers applied to emitted tool-end payloads. - pub tool_sanitize_response_guardrails: SortedRegistry>, + pub(crate) tool_sanitize_response_guardrails: SortedRegistry>, /// Global tool guardrails that can reject execution before the callback runs. - pub tool_conditional_execution_guardrails: SortedRegistry>, + pub(crate) tool_conditional_execution_guardrails: SortedRegistry>, /// Global tool request intercepts that can rewrite arguments before execution. - pub tool_request_intercepts: SortedRegistry>, + pub(crate) tool_request_intercepts: SortedRegistry>, /// Global tool execution intercepts that wrap or replace callback execution. - pub tool_execution_intercepts: SortedRegistry>, + pub(crate) tool_execution_intercepts: SortedRegistry>, /// Global LLM request sanitizers applied to emitted LLM-start payloads. - pub llm_sanitize_request_guardrails: SortedRegistry>, + pub(crate) llm_sanitize_request_guardrails: SortedRegistry>, /// Global LLM response sanitizers applied to emitted LLM-end payloads. - pub llm_sanitize_response_guardrails: SortedRegistry>, + pub(crate) llm_sanitize_response_guardrails: SortedRegistry>, /// Global LLM guardrails that can reject execution before the provider callback runs. - pub llm_conditional_execution_guardrails: SortedRegistry>, + pub(crate) llm_conditional_execution_guardrails: SortedRegistry>, /// Global LLM request intercepts that can rewrite or annotate requests. - pub llm_request_intercepts: SortedRegistry>, + pub(crate) llm_request_intercepts: SortedRegistry>, /// Global non-streaming LLM execution intercepts that wrap callback execution. - pub llm_execution_intercepts: SortedRegistry>, + pub(crate) llm_execution_intercepts: SortedRegistry>, /// Global streaming LLM execution intercepts that wrap stream-producing callbacks. - pub llm_stream_execution_intercepts: SortedRegistry>, + pub(crate) llm_stream_execution_intercepts: + SortedRegistry>, /// Global lifecycle subscribers notified after runtime events are emitted. - pub event_subscribers: HashMap, + pub(crate) event_subscribers: HashMap, /// Arbitrary binding- or integration-specific runtime extensions. - pub extensions: HashMap>, + pub(crate) extensions: HashMap>, } impl NemoFlowContextState { @@ -87,17 +82,17 @@ impl NemoFlowContextState { /// extensions. pub fn new() -> Self { Self { - tool_sanitize_request_guardrails: SortedRegistry::new(|entry| entry.priority), - tool_sanitize_response_guardrails: SortedRegistry::new(|entry| entry.priority), - tool_conditional_execution_guardrails: SortedRegistry::new(|entry| entry.priority), - tool_request_intercepts: SortedRegistry::new(|entry| entry.priority), - tool_execution_intercepts: SortedRegistry::new(|entry| entry.priority), - llm_sanitize_request_guardrails: SortedRegistry::new(|entry| entry.priority), - llm_sanitize_response_guardrails: SortedRegistry::new(|entry| entry.priority), - llm_conditional_execution_guardrails: SortedRegistry::new(|entry| entry.priority), - llm_request_intercepts: SortedRegistry::new(|entry| entry.priority), - llm_execution_intercepts: SortedRegistry::new(|entry| entry.priority), - llm_stream_execution_intercepts: SortedRegistry::new(|entry| entry.priority), + tool_sanitize_request_guardrails: SortedRegistry::new(), + tool_sanitize_response_guardrails: SortedRegistry::new(), + tool_conditional_execution_guardrails: SortedRegistry::new(), + tool_request_intercepts: SortedRegistry::new(), + tool_execution_intercepts: SortedRegistry::new(), + llm_sanitize_request_guardrails: SortedRegistry::new(), + llm_sanitize_response_guardrails: SortedRegistry::new(), + llm_conditional_execution_guardrails: SortedRegistry::new(), + llm_request_intercepts: SortedRegistry::new(), + llm_execution_intercepts: SortedRegistry::new(), + llm_stream_execution_intercepts: SortedRegistry::new(), event_subscribers: HashMap::new(), extensions: HashMap::new(), } @@ -164,7 +159,7 @@ impl NemoFlowContextState { /// # Returns /// A vector containing all global subscribers followed by the provided /// scope-local subscribers. - pub fn collect_event_subscribers( + pub(crate) fn collect_event_subscribers( &self, scope_local_subscribers: &[EventSubscriberFn], ) -> Vec { @@ -180,7 +175,7 @@ impl NemoFlowContextState { /// # Parameters /// - `event`: Fully constructed lifecycle event to deliver. /// - `subscribers`: Subscribers that should observe the event. - pub fn emit_event(event: &Event, subscribers: &[EventSubscriberFn]) { + pub(crate) fn emit_event(event: &Event, subscribers: &[EventSubscriberFn]) { for subscriber in subscribers { subscriber(event); } @@ -598,16 +593,16 @@ impl NemoFlowContextState { /// /// # Returns /// The sanitized JSON payload after every matching guardrail has run. - pub fn tool_sanitize_request_chain( + pub(crate) fn tool_sanitize_request_chain( &self, name: &str, args: Json, - scope_locals: &[&SortedRegistry>], + scope_locals: &[&SortedRegistry>], ) -> Json { let entries = merge_guardrail_entries(&self.tool_sanitize_request_guardrails, scope_locals); let mut value = args; for entry in entries { - value = (entry.guardrail)(name, value); + value = (entry.payload)(name, value); } value } @@ -622,17 +617,17 @@ impl NemoFlowContextState { /// /// # Returns /// The sanitized JSON payload after every matching guardrail has run. - pub fn tool_sanitize_response_chain( + pub(crate) fn tool_sanitize_response_chain( &self, name: &str, result: Json, - scope_locals: &[&SortedRegistry>], + scope_locals: &[&SortedRegistry>], ) -> Json { let entries = merge_guardrail_entries(&self.tool_sanitize_response_guardrails, scope_locals); let mut value = result; for entry in entries { - value = (entry.guardrail)(name, value); + value = (entry.payload)(name, value); } value } @@ -644,51 +639,18 @@ impl NemoFlowContextState { /// from the active scope stack. /// /// # Returns - /// Owned guardrail snapshots that can be evaluated after registry locks + /// Named guardrail snapshots that can be evaluated after registry locks /// are released. pub(crate) fn tool_conditional_execution_entries( &self, - scope_locals: &[&SortedRegistry>], - ) -> Vec> { - merge_named_guardrail_entries(&self.tool_conditional_execution_guardrails, scope_locals) + scope_locals: &[&SortedRegistry>], + ) -> Vec> { + merge_guardrail_entries(&self.tool_conditional_execution_guardrails, scope_locals) .into_iter() - .map(|(name, entry)| ConditionalGuardrailSnapshot { - name: name.to_string(), - guardrail: entry.guardrail.clone(), - }) + .cloned() .collect() } - /// Evaluate tool conditional-execution guardrails in priority order. - /// - /// # Parameters - /// - `name`: Tool name associated with the request. - /// - `args`: Tool arguments to validate. - /// - `scope_locals`: Scope-local conditional guardrail registries collected - /// from the active scope stack. - /// - /// # Returns - /// A [`Result`] containing `Ok(None)` when execution is allowed or - /// `Ok(Some(reason))` when a guardrail rejects the call. - /// - /// # Errors - /// Propagates any error returned by a guardrail callback. - pub fn tool_conditional_execution_chain( - &self, - name: &str, - args: &Json, - scope_locals: &[&SortedRegistry>], - ) -> crate::error::Result> { - let entries = - merge_guardrail_entries(&self.tool_conditional_execution_guardrails, scope_locals); - for entry in entries { - if let Some(error) = (entry.guardrail)(name, args)? { - return Ok(Some(error)); - } - } - Ok(None) - } - /// Evaluate a snapshot of tool conditional-execution guardrails in priority order. /// /// This function emits guardrail scope start/end events while evaluating @@ -701,7 +663,7 @@ impl NemoFlowContextState { /// # Parameters /// - `name`: Tool name associated with the request. /// - `args`: Tool arguments to validate. - /// - `entries`: Owned conditional guardrail snapshots to evaluate. + /// - `entries`: Borrowed conditional guardrail snapshots to evaluate. /// - `subscribers`: Event subscribers that should observe guardrail scope /// start/end events. /// - `parent_uuid`: Optional parent scope UUID for emitted guardrail @@ -718,7 +680,7 @@ impl NemoFlowContextState { pub(crate) fn tool_conditional_execution_snapshot_chain( name: &str, args: &Json, - entries: Vec>, + entries: &[Guardrail], subscribers: &[EventSubscriberFn], parent_uuid: Option, metadata: Option, @@ -734,7 +696,7 @@ impl NemoFlowContextState { }), subscribers, ); - let result = (entry.guardrail)(name, args); + let result = (entry.payload)(name, args); let output = match &result { Ok(Some(reason)) => json!({ "allowed": false, @@ -775,7 +737,7 @@ impl NemoFlowContextState { /// # Notes /// If an intercept entry has `break_chain` enabled, later intercepts are /// skipped after that entry runs. - pub fn tool_request_intercepts_chain( + pub(crate) fn tool_request_intercepts_chain( &self, name: &str, args: Json, @@ -784,8 +746,8 @@ impl NemoFlowContextState { let entries = merge_intercept_entries(&self.tool_request_intercepts, scope_locals); let mut value = args; for entry in entries { - value = (entry.callable)(name, value)?; - if entry.break_chain { + value = (entry.payload.callable)(name, value)?; + if entry.payload.break_chain { break; } } @@ -803,7 +765,7 @@ impl NemoFlowContextState { /// # Returns /// A composed [`ToolExecutionNextFn`] that wraps `default_fn` in every /// matching execution intercept. - pub fn tool_build_execution_chain( + pub(crate) fn tool_build_execution_chain( &self, name: &str, default_fn: ToolExecutionNextFn, @@ -830,15 +792,15 @@ impl NemoFlowContextState { /// /// # Returns /// The sanitized [`LlmRequest`] after every matching guardrail has run. - pub fn llm_sanitize_request_chain( + pub(crate) fn llm_sanitize_request_chain( &self, request: LlmRequest, - scope_locals: &[&SortedRegistry>], + scope_locals: &[&SortedRegistry>], ) -> LlmRequest { let entries = merge_guardrail_entries(&self.llm_sanitize_request_guardrails, scope_locals); let mut value = request; for entry in entries { - value = (entry.guardrail)(value); + value = (entry.payload)(value); } value } @@ -852,15 +814,15 @@ impl NemoFlowContextState { /// /// # Returns /// The sanitized response payload after every matching guardrail has run. - pub fn llm_sanitize_response_chain( + pub(crate) fn llm_sanitize_response_chain( &self, response: Json, - scope_locals: &[&SortedRegistry>], + scope_locals: &[&SortedRegistry>], ) -> Json { let entries = merge_guardrail_entries(&self.llm_sanitize_response_guardrails, scope_locals); let mut value = response; for entry in entries { - value = (entry.guardrail)(value); + value = (entry.payload)(value); } value } @@ -872,49 +834,18 @@ impl NemoFlowContextState { /// from the active scope stack. /// /// # Returns - /// Owned guardrail snapshots that can be evaluated after registry locks + /// Named guardrail snapshots that can be evaluated after registry locks /// are released. pub(crate) fn llm_conditional_execution_entries( &self, - scope_locals: &[&SortedRegistry>], - ) -> Vec> { - merge_named_guardrail_entries(&self.llm_conditional_execution_guardrails, scope_locals) + scope_locals: &[&SortedRegistry>], + ) -> Vec> { + merge_guardrail_entries(&self.llm_conditional_execution_guardrails, scope_locals) .into_iter() - .map(|(name, entry)| ConditionalGuardrailSnapshot { - name: name.to_string(), - guardrail: entry.guardrail.clone(), - }) + .cloned() .collect() } - /// Evaluate LLM conditional-execution guardrails in priority order. - /// - /// # Parameters - /// - `request`: LLM request to validate. - /// - `scope_locals`: Scope-local conditional guardrail registries collected - /// from the active scope stack. - /// - /// # Returns - /// A [`Result`] containing `Ok(None)` when execution is allowed or - /// `Ok(Some(reason))` when a guardrail rejects the call. - /// - /// # Errors - /// Propagates any error returned by a guardrail callback. - pub fn llm_conditional_execution_chain( - &self, - request: &LlmRequest, - scope_locals: &[&SortedRegistry>], - ) -> crate::error::Result> { - let entries = - merge_guardrail_entries(&self.llm_conditional_execution_guardrails, scope_locals); - for entry in entries { - if let Some(error) = (entry.guardrail)(request)? { - return Ok(Some(error)); - } - } - Ok(None) - } - /// Evaluate a snapshot of LLM conditional-execution guardrails in priority order. /// /// This function emits guardrail scope start/end events while evaluating @@ -926,7 +857,7 @@ impl NemoFlowContextState { /// /// # Parameters /// - `request`: LLM request to validate. - /// - `entries`: Owned conditional guardrail snapshots to evaluate. + /// - `entries`: Borrowed conditional guardrail snapshots to evaluate. /// - `subscribers`: Event subscribers that should observe guardrail scope /// start/end events. /// - `parent_uuid`: Optional parent scope UUID for emitted guardrail @@ -942,7 +873,7 @@ impl NemoFlowContextState { /// corresponding guardrail scope end event. pub(crate) fn llm_conditional_execution_snapshot_chain( request: &LlmRequest, - entries: Vec>, + entries: &[Guardrail], subscribers: &[EventSubscriberFn], parent_uuid: Option, metadata: Option, @@ -957,7 +888,7 @@ impl NemoFlowContextState { }), subscribers, ); - let result = (entry.guardrail)(request); + let result = (entry.payload)(request); let output = match &result { Ok(Some(reason)) => json!({ "allowed": false, @@ -1000,7 +931,7 @@ impl NemoFlowContextState { /// # Notes /// If an intercept entry has `break_chain` enabled, later intercepts are /// skipped after that entry runs. - pub fn llm_request_intercepts_chain( + pub(crate) fn llm_request_intercepts_chain( &self, name: &str, request: LlmRequest, @@ -1012,10 +943,10 @@ impl NemoFlowContextState { let mut annotated_value = annotated; for entry in entries { let (new_request, new_annotated) = - (entry.callable)(name, request_value, annotated_value)?; + (entry.payload.callable)(name, request_value, annotated_value)?; request_value = new_request; annotated_value = new_annotated; - if entry.break_chain { + if entry.payload.break_chain { break; } } @@ -1035,7 +966,7 @@ impl NemoFlowContextState { /// # Returns /// A composed [`LlmExecutionNextFn`] that wraps `default_fn` in every /// matching execution intercept. - pub fn llm_build_execution_chain( + pub(crate) fn llm_build_execution_chain( &self, name: &str, default_fn: LlmExecutionNextFn, @@ -1066,7 +997,7 @@ impl NemoFlowContextState { /// # Returns /// A composed [`LlmStreamExecutionNextFn`] that wraps `default_fn` in every /// matching execution intercept. - pub fn llm_stream_build_execution_chain( + pub(crate) fn llm_stream_build_execution_chain( &self, name: &str, default_fn: LlmStreamExecutionNextFn, diff --git a/crates/core/src/api/tool.rs b/crates/core/src/api/tool.rs index 5056bd974..fc6495bde 100644 --- a/crates/core/src/api/tool.rs +++ b/crates/core/src/api/tool.rs @@ -386,7 +386,7 @@ pub async fn tool_call_execute(params: ToolCallExecuteParams) -> Result { if let Some(error) = NemoFlowContextState::tool_conditional_execution_snapshot_chain( &name, &args, - entries, + &entries, &subscribers, parent_uuid, guardrail_metadata, @@ -533,7 +533,7 @@ pub fn tool_conditional_execution(name: &str, args: &Json) -> Result<()> { if let Some(error) = NemoFlowContextState::tool_conditional_execution_snapshot_chain( name, args, - entries, + &entries, &subscribers, parent_uuid, None, diff --git a/crates/core/src/context/registries.rs b/crates/core/src/context/registries.rs index 308ee6abe..2a0d2fde9 100644 --- a/crates/core/src/context/registries.rs +++ b/crates/core/src/context/registries.rs @@ -9,7 +9,7 @@ use std::collections::HashMap; -use crate::api::registry::{ExecutionIntercept, GuardrailEntry, Intercept}; +use crate::api::registry::{ExecutionIntercept, Guardrail, Intercept}; use crate::api::runtime::{ EventSubscriberFn, LlmConditionalFn, LlmExecutionFn, LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionFn, ToolConditionalFn, @@ -23,31 +23,32 @@ use crate::registry::SortedRegistry; /// subscribers. These registrations are merged with the global runtime /// registries when the runtime resolves the effective middleware chain for a /// tool or LLM call executed inside that scope. -pub struct ScopeLocalRegistries { +pub(crate) struct ScopeLocalRegistries { /// Tool request sanitizers applied to emitted tool-start payloads. - pub tool_sanitize_request_guardrails: SortedRegistry>, + pub(crate) tool_sanitize_request_guardrails: SortedRegistry>, /// Tool response sanitizers applied to emitted tool-end payloads. - pub tool_sanitize_response_guardrails: SortedRegistry>, + pub(crate) tool_sanitize_response_guardrails: SortedRegistry>, /// Tool guardrails that can reject execution before the callback runs. - pub tool_conditional_execution_guardrails: SortedRegistry>, + pub(crate) tool_conditional_execution_guardrails: SortedRegistry>, /// Tool request intercepts that can rewrite arguments before execution. - pub tool_request_intercepts: SortedRegistry>, + pub(crate) tool_request_intercepts: SortedRegistry>, /// Tool execution intercepts that wrap or replace callback execution. - pub tool_execution_intercepts: SortedRegistry>, + pub(crate) tool_execution_intercepts: SortedRegistry>, /// LLM request sanitizers applied to emitted LLM-start payloads. - pub llm_sanitize_request_guardrails: SortedRegistry>, + pub(crate) llm_sanitize_request_guardrails: SortedRegistry>, /// LLM response sanitizers applied to emitted LLM-end payloads. - pub llm_sanitize_response_guardrails: SortedRegistry>, + pub(crate) llm_sanitize_response_guardrails: SortedRegistry>, /// LLM guardrails that can reject execution before the provider callback runs. - pub llm_conditional_execution_guardrails: SortedRegistry>, + pub(crate) llm_conditional_execution_guardrails: SortedRegistry>, /// LLM request intercepts that can rewrite or annotate requests. - pub llm_request_intercepts: SortedRegistry>, + pub(crate) llm_request_intercepts: SortedRegistry>, /// Non-streaming LLM execution intercepts that wrap callback execution. - pub llm_execution_intercepts: SortedRegistry>, + pub(crate) llm_execution_intercepts: SortedRegistry>, /// Streaming LLM execution intercepts that wrap stream-producing callbacks. - pub llm_stream_execution_intercepts: SortedRegistry>, + pub(crate) llm_stream_execution_intercepts: + SortedRegistry>, /// Scope-local lifecycle subscribers visible while the owning scope is active. - pub event_subscribers: HashMap, + pub(crate) event_subscribers: HashMap, } impl ScopeLocalRegistries { @@ -56,19 +57,19 @@ impl ScopeLocalRegistries { /// # Returns /// A [`ScopeLocalRegistries`] value with no registered guardrails, /// intercepts, or subscribers. - pub fn new() -> Self { + pub(crate) fn new() -> Self { Self { - tool_sanitize_request_guardrails: SortedRegistry::new(|entry| entry.priority), - tool_sanitize_response_guardrails: SortedRegistry::new(|entry| entry.priority), - tool_conditional_execution_guardrails: SortedRegistry::new(|entry| entry.priority), - tool_request_intercepts: SortedRegistry::new(|entry| entry.priority), - tool_execution_intercepts: SortedRegistry::new(|entry| entry.priority), - llm_sanitize_request_guardrails: SortedRegistry::new(|entry| entry.priority), - llm_sanitize_response_guardrails: SortedRegistry::new(|entry| entry.priority), - llm_conditional_execution_guardrails: SortedRegistry::new(|entry| entry.priority), - llm_request_intercepts: SortedRegistry::new(|entry| entry.priority), - llm_execution_intercepts: SortedRegistry::new(|entry| entry.priority), - llm_stream_execution_intercepts: SortedRegistry::new(|entry| entry.priority), + tool_sanitize_request_guardrails: SortedRegistry::new(), + tool_sanitize_response_guardrails: SortedRegistry::new(), + tool_conditional_execution_guardrails: SortedRegistry::new(), + tool_request_intercepts: SortedRegistry::new(), + tool_execution_intercepts: SortedRegistry::new(), + llm_sanitize_request_guardrails: SortedRegistry::new(), + llm_sanitize_response_guardrails: SortedRegistry::new(), + llm_conditional_execution_guardrails: SortedRegistry::new(), + llm_request_intercepts: SortedRegistry::new(), + llm_execution_intercepts: SortedRegistry::new(), + llm_stream_execution_intercepts: SortedRegistry::new(), event_subscribers: HashMap::new(), } } @@ -92,10 +93,10 @@ impl Default for ScopeLocalRegistries { /// /// # Returns /// A vector of guardrail entries sorted by ascending priority. -pub fn merge_guardrail_entries<'a, F>( - global: &'a SortedRegistry>, - scope_locals: &'a [&'a SortedRegistry>], -) -> Vec<&'a GuardrailEntry> { +pub(crate) fn merge_guardrail_entries<'a, F>( + global: &'a SortedRegistry>, + scope_locals: &'a [&'a SortedRegistry>], +) -> Vec<&'a Guardrail> { let mut all = Vec::new(); all.extend(global.sorted_values()); for registry in scope_locals { @@ -105,27 +106,6 @@ pub fn merge_guardrail_entries<'a, F>( all } -/// Merge named global and scope-local guardrail entries in priority order. -/// -/// # Parameters -/// - `global`: Process-global guardrail registry. -/// - `scope_locals`: Scope-local registries collected from active scopes. -/// -/// # Returns -/// A vector of `(name, guardrail entry)` pairs sorted by ascending priority. -pub(crate) fn merge_named_guardrail_entries<'a, F>( - global: &'a SortedRegistry>, - scope_locals: &'a [&'a SortedRegistry>], -) -> Vec<(&'a str, &'a GuardrailEntry)> { - let mut all = Vec::new(); - all.extend(global.sorted_entries()); - for registry in scope_locals { - all.extend(registry.sorted_entries()); - } - all.sort_by_key(|(_, entry)| entry.priority); - all -} - /// Merge global and scope-local intercept entries into one priority-sorted list. /// /// # Parameters @@ -134,7 +114,7 @@ pub(crate) fn merge_named_guardrail_entries<'a, F>( /// /// # Returns /// A vector of intercept entries sorted by ascending priority. -pub fn merge_intercept_entries<'a, F>( +pub(crate) fn merge_intercept_entries<'a, F>( global: &'a SortedRegistry>, scope_locals: &'a [&'a SortedRegistry>], ) -> Vec<&'a Intercept> { @@ -158,17 +138,17 @@ pub fn merge_intercept_entries<'a, F>( /// /// # Returns /// A vector of `(callable, priority)` pairs sorted by ascending priority. -pub fn merge_execution_intercept_callables( +pub(crate) fn merge_execution_intercept_callables( global: &SortedRegistry>, scope_locals: &[&SortedRegistry>], ) -> Vec<(F, i32)> { let mut all = Vec::new(); for entry in global.sorted_values() { - all.push((entry.callable.clone(), entry.priority)); + all.push((entry.payload.clone(), entry.priority)); } for registry in scope_locals { for entry in registry.sorted_values() { - all.push((entry.callable.clone(), entry.priority)); + all.push((entry.payload.clone(), entry.priority)); } } all.sort_by_key(|(_, priority)| *priority); diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index 5b3fb5fa7..13aed13c8 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -38,8 +38,6 @@ //! [`atif::AtifExporter`](observability::atif::AtifExporter), //! [`otel::OpenTelemetrySubscriber`](observability::otel::OpenTelemetrySubscriber), //! and [`openinference::OpenInferenceSubscriber`](observability::openinference::OpenInferenceSubscriber). -//! - [`registry`] — [`SortedRegistry`](registry::SortedRegistry) — a priority-sorted, named collection used for -//! all guardrail and intercept registries. //! - [`stream`] — [`stream::LlmStreamWrapper`] — a stream adapter that applies per-chunk //! intercepts and aggregates streaming LLM responses. //! @@ -64,7 +62,7 @@ pub mod json; pub mod observability; pub mod plugin; pub mod plugins; -pub mod registry; +mod registry; #[doc(hidden)] pub mod shared_runtime; pub mod stream; diff --git a/crates/core/src/observability/plugin_component.rs b/crates/core/src/observability/plugin_component.rs index 4322948c8..09e9fd43e 100644 --- a/crates/core/src/observability/plugin_component.rs +++ b/crates/core/src/observability/plugin_component.rs @@ -21,6 +21,7 @@ use std::future::Future; use std::path::PathBuf; use std::pin::Pin; use std::sync::{Arc, Mutex}; +#[cfg(any(feature = "otel", feature = "openinference"))] use std::time::Duration; use serde::{Deserialize, Serialize}; diff --git a/crates/core/src/registry.rs b/crates/core/src/registry.rs index 3798b5eeb..3712aaf06 100644 --- a/crates/core/src/registry.rs +++ b/crates/core/src/registry.rs @@ -3,17 +3,30 @@ //! Priority-sorted named registry. //! -//! [`SortedRegistry`] is the backbone data structure for all guardrail and intercept -//! registries in the NeMo Flow runtime. It stores entries by unique name and provides -//! iteration in ascending priority order, with eager re-sorting on every mutation. +//! [`SortedRegistry`] is the backbone data structure for all guardrail and +//! intercept registries in the NeMo Flow runtime. It stores self-describing +//! entries by unique name and provides iteration in ascending priority order, +//! with eager re-sorting on every mutation. use std::collections::HashMap; +/// A named, priority-ordered registry entry. +/// +/// Registry entries carry their own identity so snapshots can be cloned or +/// copied without pairing a separately stored map key with the entry value. +pub(crate) trait RegistryEntry { + /// Unique entry name within one registry. + fn name(&self) -> &str; + + /// Entry priority. Lower values run earlier. + fn priority(&self) -> i32; +} + /// A named registry that maintains a sorted order by priority. /// -/// Items are stored by unique string name and sorted by an integer priority -/// extracted via a caller-provided function. The sort is performed eagerly: -/// every [`register`](SortedRegistry::register) or +/// Items are stored by their embedded [`RegistryEntry::name`] value and sorted +/// by [`RegistryEntry::priority`]. The sort is performed eagerly: every +/// [`register`](SortedRegistry::register) or /// [`deregister`](SortedRegistry::deregister) call re-sorts immediately, so /// [`sorted_values`](SortedRegistry::sorted_values) is a read-only lookup. /// @@ -27,45 +40,41 @@ use std::collections::HashMap; /// Names must be unique within a registry. Attempting to [`register`](SortedRegistry::register) /// a duplicate name returns an error. Use [`deregister`](SortedRegistry::deregister) first /// to remove an existing entry before re-registering. -pub struct SortedRegistry { +pub(crate) struct SortedRegistry { entries: HashMap, sorted_keys: Vec, - priority_fn: fn(&T) -> i32, } -impl SortedRegistry { - /// Create a new empty registry with the given priority extraction function. - /// - /// The runtime calls `priority_fn` on each stored entry to determine its - /// sort key. Lower values are ordered first. - /// - /// # Parameters - /// - `priority_fn`: Function used to extract the integer priority from a - /// stored entry. +impl SortedRegistry { + /// Create a new empty registry. /// /// # Returns /// A new empty [`SortedRegistry`] with no entries. - pub fn new(priority_fn: fn(&T) -> i32) -> Self { + pub(crate) fn new() -> Self { Self { entries: HashMap::new(), sorted_keys: Vec::new(), - priority_fn, } } /// Re-sorts the cached key order by priority. Called eagerly on every mutation. fn resort(&mut self) { - let pf = self.priority_fn; let entries = &self.entries; let mut keys: Vec = entries.keys().cloned().collect(); - keys.sort_by_key(|k| pf(entries.get(k).unwrap())); + keys.sort_by(|left, right| { + let left_entry = entries.get(left).unwrap(); + let right_entry = entries.get(right).unwrap(); + left_entry + .priority() + .cmp(&right_entry.priority()) + .then_with(|| left.cmp(right)) + }); self.sorted_keys = keys; } - /// Register a new entry under a unique name. + /// Register a new entry under its embedded name. /// /// # Parameters - /// - `name`: Unique name used to address the entry later. /// - `entry`: Value to store in the registry. /// /// # Returns @@ -76,7 +85,8 @@ impl SortedRegistry { /// /// # Notes /// Successful registration eagerly re-sorts the cached priority order. - pub fn register(&mut self, name: String, entry: T) -> Result<(), String> { + pub(crate) fn register(&mut self, entry: T) -> Result<(), String> { + let name = entry.name().to_string(); if self.entries.contains_key(&name) { return Err(format!("{name} already exists")); } @@ -96,7 +106,7 @@ impl SortedRegistry { /// /// # Notes /// Successful removal eagerly re-sorts the cached priority order. - pub fn deregister(&mut self, name: &str) -> bool { + pub(crate) fn deregister(&mut self, name: &str) -> bool { if self.entries.remove(name).is_some() { self.resort(); true @@ -113,63 +123,17 @@ impl SortedRegistry { /// # Returns /// A newly allocated [`Vec`] of shared references ordered from lowest /// priority to highest priority. - pub fn sorted_values(&self) -> Vec<&T> { + pub(crate) fn sorted_values(&self) -> Vec<&T> { self.sorted_keys .iter() .filter_map(|k| self.entries.get(k)) .collect() } +} - /// Return named entries sorted by priority (ascending). - /// - /// # Returns - /// A newly allocated [`Vec`] of `(name, entry)` pairs ordered from lowest - /// priority to highest priority. - pub(crate) fn sorted_entries(&self) -> Vec<(&str, &T)> { - self.sorted_keys - .iter() - .filter_map(|key| self.entries.get(key).map(|entry| (key.as_str(), entry))) - .collect() - } - - /// Return a shared reference to an entry by name. - /// - /// # Parameters - /// - `name`: Name of the entry to resolve. - /// - /// # Returns - /// `Some(&T)` when an entry exists under `name`, otherwise `None`. - pub fn get(&self, name: &str) -> Option<&T> { - self.entries.get(name) - } - - /// Remove and return an entry by name. - /// - /// # Parameters - /// - `name`: Name of the entry to remove. - /// - /// # Returns - /// `Some(T)` when an entry was removed, otherwise `None`. - /// - /// # Notes - /// Successful removal eagerly re-sorts the cached priority order. - pub fn remove(&mut self, name: &str) -> Option { - let removed = self.entries.remove(name); - if removed.is_some() { - self.resort(); - } - removed - } - - /// Report whether an entry with the given name exists. - /// - /// # Parameters - /// - `name`: Name to test for membership. - /// - /// # Returns - /// `true` when the registry contains `name`, otherwise `false`. - pub fn contains(&self, name: &str) -> bool { - self.entries.contains_key(name) +impl Default for SortedRegistry { + fn default() -> Self { + Self::new() } } diff --git a/crates/core/tests/integration/api_surface_tests.rs b/crates/core/tests/integration/api_surface_tests.rs index ee6a31f04..b65339114 100644 --- a/crates/core/tests/integration/api_surface_tests.rs +++ b/crates/core/tests/integration/api_surface_tests.rs @@ -345,14 +345,14 @@ fn test_global_registry_and_subscriber_wrappers_cover_success_and_duplicates() { register_tool_sanitize_request_guardrail( "tool-sanitize-request", 1, - Box::new(|_name, args| args), + Arc::new(|_name, args| args), ) .unwrap(); expect_already_exists( register_tool_sanitize_request_guardrail( "tool-sanitize-request", 1, - Box::new(|_name, args| args), + Arc::new(|_name, args| args), ) .unwrap_err(), "tool-sanitize-request", @@ -363,7 +363,7 @@ fn test_global_registry_and_subscriber_wrappers_cover_success_and_duplicates() { register_tool_sanitize_response_guardrail( "tool-sanitize-response", 1, - Box::new(|_name, args| args), + Arc::new(|_name, args| args), ) .unwrap(); assert!(deregister_tool_sanitize_response_guardrail("tool-sanitize-response").unwrap()); @@ -376,7 +376,7 @@ fn test_global_registry_and_subscriber_wrappers_cover_success_and_duplicates() { .unwrap(); assert!(deregister_tool_conditional_execution_guardrail("tool-conditional").unwrap()); - register_tool_request_intercept("tool-request", 1, false, Box::new(|_name, args| Ok(args))) + register_tool_request_intercept("tool-request", 1, false, Arc::new(|_name, args| Ok(args))) .unwrap(); assert!(deregister_tool_request_intercept("tool-request").unwrap()); @@ -388,14 +388,14 @@ fn test_global_registry_and_subscriber_wrappers_cover_success_and_duplicates() { .unwrap(); assert!(deregister_tool_execution_intercept("tool-execution").unwrap()); - register_llm_sanitize_request_guardrail("llm-sanitize-request", 1, Box::new(|request| request)) + register_llm_sanitize_request_guardrail("llm-sanitize-request", 1, Arc::new(|request| request)) .unwrap(); assert!(deregister_llm_sanitize_request_guardrail("llm-sanitize-request").unwrap()); register_llm_sanitize_response_guardrail( "llm-sanitize-response", 1, - Box::new(|response| response), + Arc::new(|response| response), ) .unwrap(); assert!(deregister_llm_sanitize_response_guardrail("llm-sanitize-response").unwrap()); @@ -412,7 +412,7 @@ fn test_global_registry_and_subscriber_wrappers_cover_success_and_duplicates() { "llm-request", 1, false, - Box::new(|_name, request, annotated| Ok((request, annotated))), + Arc::new(|_name, request, annotated| Ok((request, annotated))), ) .unwrap(); assert!(deregister_llm_request_intercept("llm-request").unwrap()); @@ -465,7 +465,7 @@ fn test_scope_registry_and_subscriber_wrappers_cover_success_duplicates_and_miss &scope.uuid, "tool-sanitize-request", 1, - Box::new(|_name, args| args), + Arc::new(|_name, args| args), ) .unwrap(); expect_already_exists( @@ -473,7 +473,7 @@ fn test_scope_registry_and_subscriber_wrappers_cover_success_duplicates_and_miss &scope.uuid, "tool-sanitize-request", 1, - Box::new(|_name, args| args), + Arc::new(|_name, args| args), ) .unwrap_err(), "tool-sanitize-request", @@ -487,7 +487,7 @@ fn test_scope_registry_and_subscriber_wrappers_cover_success_duplicates_and_miss &scope.uuid, "tool-sanitize-response", 1, - Box::new(|_name, args| args), + Arc::new(|_name, args| args), ) .unwrap(); assert!( @@ -512,7 +512,7 @@ fn test_scope_registry_and_subscriber_wrappers_cover_success_duplicates_and_miss "tool-request", 1, false, - Box::new(|_name, args| Ok(args)), + Arc::new(|_name, args| Ok(args)), ) .unwrap(); assert!(scope_deregister_tool_request_intercept(&scope.uuid, "tool-request").unwrap()); @@ -530,7 +530,7 @@ fn test_scope_registry_and_subscriber_wrappers_cover_success_duplicates_and_miss &scope.uuid, "llm-sanitize-request", 1, - Box::new(|request| request), + Arc::new(|request| request), ) .unwrap(); assert!( @@ -542,7 +542,7 @@ fn test_scope_registry_and_subscriber_wrappers_cover_success_duplicates_and_miss &scope.uuid, "llm-sanitize-response", 1, - Box::new(|response| response), + Arc::new(|response| response), ) .unwrap(); assert!( @@ -567,7 +567,7 @@ fn test_scope_registry_and_subscriber_wrappers_cover_success_duplicates_and_miss "llm-request", 1, false, - Box::new(|_name, request, annotated| Ok((request, annotated))), + Arc::new(|_name, request, annotated| Ok((request, annotated))), ) .unwrap(); assert!(scope_deregister_llm_request_intercept(&scope.uuid, "llm-request").unwrap()); @@ -616,7 +616,7 @@ fn test_scope_registry_and_subscriber_wrappers_cover_success_duplicates_and_miss &scope.uuid, "missing-tool-sanitize", 1, - Box::new(|_name, args| args), + Arc::new(|_name, args| args), ) .unwrap_err(), "scope", @@ -627,7 +627,7 @@ fn test_scope_registry_and_subscriber_wrappers_cover_success_duplicates_and_miss "missing-tool-request", 1, false, - Box::new(|_name, args| Ok(args)), + Arc::new(|_name, args| Ok(args)), ) .unwrap_err(), "scope", @@ -660,7 +660,7 @@ async fn test_tool_api_emits_sanitized_events_and_covers_error_paths() { register_tool_sanitize_request_guardrail( "tool-sanitize-request", 1, - Box::new(|_name, mut args| { + Arc::new(|_name, mut args| { args.as_object_mut() .unwrap() .insert("sanitized_request".into(), json!(true)); @@ -671,7 +671,7 @@ async fn test_tool_api_emits_sanitized_events_and_covers_error_paths() { register_tool_sanitize_response_guardrail( "tool-sanitize-response", 1, - Box::new(|_name, mut result| { + Arc::new(|_name, mut result| { result .as_object_mut() .unwrap() @@ -728,7 +728,7 @@ async fn test_tool_api_emits_sanitized_events_and_covers_error_paths() { "tool-request", 1, false, - Box::new(|_name, mut args| { + Arc::new(|_name, mut args| { args.as_object_mut() .unwrap() .insert("intercepted".into(), json!(true)); @@ -824,7 +824,7 @@ async fn test_llm_api_emits_sanitized_events_and_covers_error_paths() { register_llm_sanitize_request_guardrail( "llm-sanitize-request", 1, - Box::new(|mut request| { + Arc::new(|mut request| { request.headers.insert("x-sanitized".into(), json!(true)); request }), @@ -833,7 +833,7 @@ async fn test_llm_api_emits_sanitized_events_and_covers_error_paths() { register_llm_sanitize_response_guardrail( "llm-sanitize-response", 1, - Box::new(|mut response| { + Arc::new(|mut response| { response .as_object_mut() .unwrap() @@ -891,7 +891,7 @@ async fn test_llm_api_emits_sanitized_events_and_covers_error_paths() { "llm-request", 1, false, - Box::new(|_name, mut request, annotated| { + Arc::new(|_name, mut request, annotated| { request.headers.insert("x-intercepted".into(), json!(true)); Ok((request, annotated)) }), diff --git a/crates/core/tests/integration/context_isolation_tests.rs b/crates/core/tests/integration/context_isolation_tests.rs index 051b31e45..6657a2ff3 100644 --- a/crates/core/tests/integration/context_isolation_tests.rs +++ b/crates/core/tests/integration/context_isolation_tests.rs @@ -312,11 +312,6 @@ fn test_scope_stack_helpers_cover_lookup_mutation_and_remove_paths() { stack.top_mut().name = "root-renamed".into(); assert_eq!(stack.top().name, "root-renamed"); - assert!(stack.local_registries_mut(&Uuid::now_v7()).is_none()); - assert!(stack.scope_registries_get(&root_uuid).is_none()); - assert!(stack.local_registries_mut(&root_uuid).is_some()); - assert!(stack.scope_registries_get(&root_uuid).is_some()); - let child = ScopeHandle::builder() .name("child") .scope_type(ScopeType::Function) @@ -337,7 +332,6 @@ fn test_scope_stack_helpers_cover_lookup_mutation_and_remove_paths() { let removed = stack.remove(&child_uuid).unwrap(); assert_eq!(removed.name, "child"); assert!(stack.find(&child_uuid).is_none()); - assert!(stack.scope_registries_get(&child_uuid).is_none()); match stack.remove(&root_uuid) { Err(FlowError::InvalidArgument(message)) => { diff --git a/crates/core/tests/integration/middleware_tests.rs b/crates/core/tests/integration/middleware_tests.rs index 01b7e4f06..f8b47342a 100644 --- a/crates/core/tests/integration/middleware_tests.rs +++ b/crates/core/tests/integration/middleware_tests.rs @@ -99,7 +99,7 @@ fn test_sanitize_guardrail_priority_ordering() { register_tool_sanitize_request_guardrail( "g_p1", 1, - Box::new(move |_name, args| { + Arc::new(move |_name, args| { o1.lock().unwrap().push(1); args }), @@ -111,7 +111,7 @@ fn test_sanitize_guardrail_priority_ordering() { register_tool_sanitize_request_guardrail( "g_p3", 3, - Box::new(move |_name, args| { + Arc::new(move |_name, args| { o3.lock().unwrap().push(3); args }), @@ -123,7 +123,7 @@ fn test_sanitize_guardrail_priority_ordering() { register_tool_sanitize_request_guardrail( "g_p2", 2, - Box::new(move |_name, args| { + Arc::new(move |_name, args| { o2.lock().unwrap().push(2); args }), @@ -167,7 +167,7 @@ fn test_request_intercept_priority_ordering() { "i_p1", 1, false, - Box::new(move |_name, args| { + Arc::new(move |_name, args| { o1.lock().unwrap().push(1); Ok(args) }), @@ -179,7 +179,7 @@ fn test_request_intercept_priority_ordering() { "i_p3", 3, false, - Box::new(move |_name, args| { + Arc::new(move |_name, args| { o3.lock().unwrap().push(3); Ok(args) }), @@ -191,7 +191,7 @@ fn test_request_intercept_priority_ordering() { "i_p2", 2, false, - Box::new(move |_name, args| { + Arc::new(move |_name, args| { o2.lock().unwrap().push(2); Ok(args) }), @@ -228,7 +228,7 @@ fn test_re_registration_at_different_priority_re_sorts() { "intercept_a", 10, false, - Box::new(move |_name, args| { + Arc::new(move |_name, args| { o_a.lock().unwrap().push("a_p10".into()); Ok(args) }), @@ -240,7 +240,7 @@ fn test_re_registration_at_different_priority_re_sorts() { "intercept_b", 20, false, - Box::new(move |_name, args| { + Arc::new(move |_name, args| { o_b.lock().unwrap().push("b_p20".into()); Ok(args) }), @@ -261,7 +261,7 @@ fn test_re_registration_at_different_priority_re_sorts() { "intercept_a", 30, false, - Box::new(move |_name, args| { + Arc::new(move |_name, args| { o_a2.lock().unwrap().push("a_p30".into()); Ok(args) }), @@ -303,7 +303,7 @@ fn test_break_chain_stops_subsequent_intercepts() { "breaker", 1, true, // break_chain = true - Box::new(|_name, mut args| { + Arc::new(|_name, mut args| { args.as_object_mut() .unwrap() .insert("breaker_ran".into(), json!(true)); @@ -317,7 +317,7 @@ fn test_break_chain_stops_subsequent_intercepts() { "after_breaker", 2, false, - Box::new(move |_name, mut args| { + Arc::new(move |_name, mut args| { sc.store(true, Ordering::SeqCst); args.as_object_mut() .unwrap() @@ -360,7 +360,7 @@ fn test_no_break_chain_runs_all_intercepts() { "first", 1, false, - Box::new(move |_name, args| { + Arc::new(move |_name, args| { c1.fetch_add(1, Ordering::SeqCst); Ok(args) }), @@ -372,7 +372,7 @@ fn test_no_break_chain_runs_all_intercepts() { "second", 2, false, - Box::new(move |_name, args| { + Arc::new(move |_name, args| { c2.fetch_add(1, Ordering::SeqCst); Ok(args) }), @@ -709,7 +709,6 @@ async fn test_tool_conditional_guardrail_emits_guardrail_scope() { Arc::new(|_, _| Ok(Some("blocked by tool guardrail".to_string()))), ) .unwrap(); - assert!(global_context().read().unwrap().extensions.is_empty()); let func: ToolExecutionNextFn = Arc::new(|args| Box::pin(async move { Ok(args) })); let allowed = tool_call_execute( @@ -889,7 +888,7 @@ fn test_scope_local_guardrail_lifecycle() { &handle.uuid, "scoped_guardrail", 1, - Box::new(move |_name, args| { + Arc::new(move |_name, args| { cc.fetch_add(1, Ordering::SeqCst); args }), @@ -1013,7 +1012,7 @@ fn test_scope_local_and_global_guardrail_merge_priority() { register_tool_sanitize_request_guardrail( "global_g", 5, - Box::new(move |_name, mut args| { + Arc::new(move |_name, mut args| { og.lock().unwrap().push("global".into()); args.as_object_mut() .unwrap() @@ -1029,7 +1028,7 @@ fn test_scope_local_and_global_guardrail_merge_priority() { &handle.uuid, "local_g", 3, - Box::new(move |_name, mut args| { + Arc::new(move |_name, mut args| { ol.lock().unwrap().push("local".into()); args.as_object_mut() .unwrap() @@ -1197,7 +1196,7 @@ async fn test_conditional_rejection_prevents_intercepts() { "should_not_run", 1, false, - Box::new(move |_name, args| { + Arc::new(move |_name, args| { ic.store(true, Ordering::SeqCst); Ok(args) }), @@ -1299,7 +1298,7 @@ fn test_sanitize_guardrails_pipe_data() { register_tool_sanitize_request_guardrail( "add_a", 1, - Box::new(|_name, mut args| { + Arc::new(|_name, mut args| { args.as_object_mut() .unwrap() .insert("field_a".into(), json!(true)); @@ -1312,7 +1311,7 @@ fn test_sanitize_guardrails_pipe_data() { register_tool_sanitize_request_guardrail( "add_b", 2, - Box::new(|_name, mut args| { + Arc::new(|_name, mut args| { // Verify field_a was added by the previous guardrail let has_a = args.get("field_a").is_some(); args.as_object_mut() @@ -1370,7 +1369,7 @@ fn test_response_sanitize_guardrails_pipe() { register_tool_sanitize_response_guardrail( "resp_g1", 1, - Box::new(|_name, mut result| { + Arc::new(|_name, mut result| { result .as_object_mut() .unwrap() @@ -1445,7 +1444,7 @@ fn test_concurrent_register_deregister() { let res = register_tool_sanitize_request_guardrail( &name, i, - Box::new(|_name, args| args), + Arc::new(|_name, args| args), ); assert!(res.is_ok(), "Registration should succeed for {name}"); @@ -1463,16 +1462,13 @@ fn test_concurrent_register_deregister() { h.join().expect("Thread should not panic"); } - // Verify the context is still in a consistent state - let ctx = global_context(); - let state = ctx.read().unwrap(); - assert!( - state - .tool_sanitize_request_guardrails - .sorted_values() - .is_empty(), - "All guardrails should be deregistered" - ); + for i in 0..10i32 { + let name = format!("concurrent_guardrail_{i}"); + assert!( + !deregister_tool_sanitize_request_guardrail(&name).unwrap(), + "{name} should already be deregistered" + ); + } } /// Concurrent register/deregister of intercepts across multiple threads. @@ -1494,7 +1490,7 @@ fn test_concurrent_intercept_mutations() { &name, i, false, - Box::new(|_name, args| Ok(args)), + Arc::new(|_name, args| Ok(args)), ); assert!(res.is_ok()); @@ -1510,12 +1506,13 @@ fn test_concurrent_intercept_mutations() { h.join().expect("Thread should not panic"); } - let ctx = global_context(); - let state = ctx.read().unwrap(); - assert!( - state.tool_request_intercepts.sorted_values().is_empty(), - "All intercepts should be deregistered" - ); + for i in 0..10i32 { + let name = format!("concurrent_intercept_{i}"); + assert!( + !deregister_tool_request_intercept(&name).unwrap(), + "{name} should already be deregistered" + ); + } } /// Interleaved register and tool call execution from multiple threads. @@ -1529,7 +1526,7 @@ fn test_concurrent_register_and_read() { register_tool_sanitize_request_guardrail( &format!("stable_{i}"), i, - Box::new(|_name, args| args), + Arc::new(|_name, args| args), ) .unwrap(); } @@ -1548,7 +1545,7 @@ fn test_concurrent_register_and_read() { let _ = register_tool_sanitize_request_guardrail( &name, 100 + i, - Box::new(|_name, args| args), + Arc::new(|_name, args| args), ); std::thread::yield_now(); let _ = deregister_tool_sanitize_request_guardrail(&name); @@ -1599,7 +1596,7 @@ async fn test_full_pipeline_integration() { "req_intercept", 1, false, - Box::new(move |_name, mut args| { + Arc::new(move |_name, mut args| { o1.lock().unwrap().push("request_intercept".into()); args.as_object_mut() .unwrap() @@ -1614,7 +1611,7 @@ async fn test_full_pipeline_integration() { register_tool_sanitize_request_guardrail( "sanitize_req", 1, - Box::new(move |_name, args| { + Arc::new(move |_name, args| { o2.lock().unwrap().push("sanitize_request".into()); args }), @@ -1653,7 +1650,7 @@ async fn test_full_pipeline_integration() { register_tool_sanitize_response_guardrail( "sanitize_resp", 1, - Box::new(move |_name, result| { + Arc::new(move |_name, result| { o5.lock().unwrap().push("sanitize_response".into()); result }), @@ -1718,10 +1715,10 @@ fn test_duplicate_guardrail_registration_returns_error() { let _lock = TEST_MUTEX.lock().unwrap(); reset_global(); - register_tool_sanitize_request_guardrail("duplicate", 1, Box::new(|_name, args| args)).unwrap(); + register_tool_sanitize_request_guardrail("duplicate", 1, Arc::new(|_name, args| args)).unwrap(); let err = - register_tool_sanitize_request_guardrail("duplicate", 2, Box::new(|_name, args| args)); + register_tool_sanitize_request_guardrail("duplicate", 2, Arc::new(|_name, args| args)); assert!(err.is_err()); match err.unwrap_err() { @@ -1741,14 +1738,14 @@ fn test_duplicate_intercept_registration_returns_error() { let _lock = TEST_MUTEX.lock().unwrap(); reset_global(); - register_tool_request_intercept("dup_intercept", 1, false, Box::new(|_name, args| Ok(args))) + register_tool_request_intercept("dup_intercept", 1, false, Arc::new(|_name, args| Ok(args))) .unwrap(); let err = register_tool_request_intercept( "dup_intercept", 2, false, - Box::new(|_name, args| Ok(args)), + Arc::new(|_name, args| Ok(args)), ); assert!(err.is_err()); @@ -1793,7 +1790,7 @@ fn test_deregister_removes_from_chain() { register_tool_sanitize_request_guardrail( "removable", 1, - Box::new(move |_name, args| { + Arc::new(move |_name, args| { cc.fetch_add(1, Ordering::SeqCst); args }), @@ -1988,7 +1985,7 @@ async fn test_llm_request_intercept_transforms() { "llm_req_i", 1, false, - Box::new(|_name: &str, mut req: LlmRequest, annotated| { + Arc::new(|_name: &str, mut req: LlmRequest, annotated| { req.headers.insert("x-intercepted".into(), json!(true)); Ok((req, annotated)) }), @@ -2086,7 +2083,7 @@ async fn test_llm_start_emits_before_short_circuit_execution_intercept() { "llm_short_circuit_request", 1, false, - Box::new(|_name, mut req, annotated| { + Arc::new(|_name, mut req, annotated| { req.content .as_object_mut() .unwrap() @@ -2192,7 +2189,7 @@ async fn test_llm_stream_start_emits_before_short_circuit_execution_intercept() "llm_stream_short_circuit_request", 1, false, - Box::new(|_name, mut req, annotated| { + Arc::new(|_name, mut req, annotated| { req.content .as_object_mut() .unwrap() diff --git a/crates/core/tests/integration/pipeline_tests.rs b/crates/core/tests/integration/pipeline_tests.rs index 71d2d4000..f8e1a48ea 100644 --- a/crates/core/tests/integration/pipeline_tests.rs +++ b/crates/core/tests/integration/pipeline_tests.rs @@ -225,7 +225,7 @@ async fn test_decode_runs_before_intercepts() { "ann_i", 1, false, - Box::new(move |_name, req, annotated| { + Arc::new(move |_name, req, annotated| { *cap.lock().unwrap() = Some(annotated.clone()); Ok((req, annotated)) }), @@ -276,7 +276,7 @@ async fn test_encode_runs_after_intercepts() { "modify_model", 1, false, - Box::new(|_name, req, annotated| { + Arc::new(|_name, req, annotated| { let mut ann = annotated.unwrap(); ann.model = Some("modified".into()); Ok((req, Some(ann))) @@ -339,7 +339,7 @@ async fn test_annotated_intercept_receives_both() { "capture_both", 1, false, - Box::new(move |_name, req, annotated| { + Arc::new(move |_name, req, annotated| { *cp.lock().unwrap() = Some((req.clone(), annotated.clone())); Ok((req, annotated)) }), @@ -394,7 +394,7 @@ async fn test_legacy_intercept_backward_compat() { "legacy_1", 1, false, - Box::new(move |_name, mut req, annotated| { + Arc::new(move |_name, mut req, annotated| { *lc1.lock().unwrap() = true; req.headers.insert("x-legacy".into(), json!("was-here")); Ok((req, annotated)) @@ -438,7 +438,7 @@ async fn test_legacy_intercept_backward_compat() { "legacy_2", 1, false, - Box::new(move |_name, mut req, annotated| { + Arc::new(move |_name, mut req, annotated| { *lc2.lock().unwrap() = true; req.headers.insert("x-legacy-2".into(), json!("also-here")); Ok((req, annotated)) @@ -495,7 +495,7 @@ async fn test_stream_path_also_decodes() { "stream_ann", 1, false, - Box::new(move |_name, req, annotated| { + Arc::new(move |_name, req, annotated| { *ca.lock().unwrap() = Some(annotated.clone()); Ok((req, annotated)) }), @@ -559,7 +559,7 @@ async fn test_shared_helper_both_paths() { "shared_ann", 1, false, - Box::new(move |_name, req, annotated| { + Arc::new(move |_name, req, annotated| { *acc.lock().unwrap() += 1; Ok((req, annotated)) }), @@ -634,7 +634,7 @@ async fn test_explicit_codec_param_overrides() { "check_model", 1, false, - Box::new(move |_name, req, annotated| { + Arc::new(move |_name, req, annotated| { if let Some(ref ann) = annotated { *cm.lock().unwrap() = ann.model.clone(); } @@ -682,7 +682,7 @@ async fn test_encode_merge_not_replace() { "merge_mod", 1, false, - Box::new(|_name, req, annotated| { + Arc::new(|_name, req, annotated| { let mut ann = annotated.unwrap(); ann.model = Some("new_model".into()); Ok((req, Some(ann))) @@ -751,7 +751,7 @@ async fn test_unified_chain_priority_order() { "legacy_p10", 10, false, - Box::new(move |_name, req, annotated| { + Arc::new(move |_name, req, annotated| { cl1.lock().unwrap().push("legacy_p10".into()); Ok((req, annotated)) }), @@ -764,7 +764,7 @@ async fn test_unified_chain_priority_order() { "annotated_p5", 5, false, - Box::new(move |_name, req, annotated| { + Arc::new(move |_name, req, annotated| { cl2.lock().unwrap().push("annotated_p5".into()); Ok((req, annotated)) }), @@ -812,7 +812,7 @@ async fn test_no_codec_annotated_intercept_receives_none() { "no_codec_ann", 1, false, - Box::new(move |_name, req, annotated| { + Arc::new(move |_name, req, annotated| { *ca.lock().unwrap() = Some(annotated.clone()); Ok((req, annotated)) }), diff --git a/crates/core/tests/integration/scope_local_tests.rs b/crates/core/tests/integration/scope_local_tests.rs index 2e5fee807..34ee4db8f 100644 --- a/crates/core/tests/integration/scope_local_tests.rs +++ b/crates/core/tests/integration/scope_local_tests.rs @@ -75,7 +75,7 @@ fn test_scope_local_guardrail_registration_and_execution() { &handle.uuid, "local_sanitizer", 10, - Box::new(|_name, mut args| { + Arc::new(|_name, mut args| { args.as_object_mut() .unwrap() .insert("scope_sanitized".into(), json!(true)); @@ -159,7 +159,7 @@ async fn test_auto_cleanup_on_scope_pop() { "ephemeral_intercept", 1, false, - Box::new(|_name, mut args| { + Arc::new(|_name, mut args| { args.as_object_mut() .unwrap() .insert("ephemeral".into(), json!(true)); @@ -226,7 +226,7 @@ async fn test_priority_merge_global_and_scope_local() { "global_p10", 10, false, - Box::new(move |_name, mut args| { + Arc::new(move |_name, mut args| { o1.lock().unwrap().push(10); args.as_object_mut() .unwrap() @@ -242,7 +242,7 @@ async fn test_priority_merge_global_and_scope_local() { "global_p30", 30, false, - Box::new(move |_name, mut args| { + Arc::new(move |_name, mut args| { o3.lock().unwrap().push(30); args.as_object_mut() .unwrap() @@ -259,7 +259,7 @@ async fn test_priority_merge_global_and_scope_local() { "local_p20", 20, false, - Box::new(move |_name, mut args| { + Arc::new(move |_name, mut args| { o2.lock().unwrap().push(20); args.as_object_mut() .unwrap() @@ -321,7 +321,7 @@ fn test_name_coexistence_global_and_scope_local() { register_tool_sanitize_request_guardrail( "shared_name", 1, - Box::new(move |_name, args| { + Arc::new(move |_name, args| { c1.fetch_add(1, Ordering::SeqCst); args }), @@ -334,7 +334,7 @@ fn test_name_coexistence_global_and_scope_local() { &handle.uuid, "shared_name", 2, - Box::new(move |_name, args| { + Arc::new(move |_name, args| { c2.fetch_add(1, Ordering::SeqCst); args }), @@ -393,7 +393,7 @@ async fn test_scope_isolation_between_stacks() { "a_intercept", 1, false, - Box::new(|_name, mut args| { + Arc::new(|_name, mut args| { args.as_object_mut() .unwrap() .insert("agent".into(), json!("a")); @@ -419,7 +419,7 @@ async fn test_scope_isolation_between_stacks() { "b_intercept", 1, false, - Box::new(|_name, mut args| { + Arc::new(|_name, mut args| { args.as_object_mut() .unwrap() .insert("agent".into(), json!("b")); @@ -498,7 +498,7 @@ async fn test_nested_scope_inheritance() { "global_intercept", 1, false, - Box::new(move |_name, mut args| { + Arc::new(move |_name, mut args| { og.lock().unwrap().push("global".into()); args.as_object_mut() .unwrap() @@ -522,7 +522,7 @@ async fn test_nested_scope_inheritance() { "a_intercept", 5, false, - Box::new(move |_name, mut args| { + Arc::new(move |_name, mut args| { oa.lock().unwrap().push("scope_a".into()); args.as_object_mut() .unwrap() @@ -547,7 +547,7 @@ async fn test_nested_scope_inheritance() { "b_intercept", 10, false, - Box::new(move |_name, mut args| { + Arc::new(move |_name, mut args| { ob.lock().unwrap().push("scope_b".into()); args.as_object_mut() .unwrap() diff --git a/crates/core/tests/unit/context_tests.rs b/crates/core/tests/unit/context_tests.rs index c1bd18cdb..084a7a2c3 100644 --- a/crates/core/tests/unit/context_tests.rs +++ b/crates/core/tests/unit/context_tests.rs @@ -10,7 +10,7 @@ use uuid::{Uuid, Version}; use crate::api::event::Event; use crate::api::llm::LlmRequest; -use crate::api::registry::{ExecutionIntercept, GuardrailEntry, Intercept}; +use crate::api::registry::{ExecutionIntercept, Guardrail, Intercept, RequestIntercept}; use crate::api::runtime::EventSubscriberFn; use crate::api::runtime::NemoFlowContextState; use crate::api::runtime::ScopeStack; @@ -41,14 +41,14 @@ fn scope_stack_tracks_scope_local_registries_and_subscribers() { .insert("sub".to_string(), subscriber.clone()); registries .tool_request_intercepts - .register( - "tool".to_string(), - Intercept { - priority: 10, + .register(Intercept { + name: "tool".to_string(), + priority: 10, + payload: RequestIntercept { break_chain: false, - callable: Box::new(|_, value| Ok(value)), + callable: Arc::new(|_, value| Ok(value)), }, - ) + }) .unwrap(); assert_eq!(stack.collect_scope_local_subscribers().len(), 1); @@ -58,11 +58,8 @@ fn scope_stack_tracks_scope_local_registries_and_subscribers() { .len(), 1 ); - assert!(stack.scope_registries_get(&child_uuid).is_some()); - let removed = stack.remove(&child_uuid).unwrap(); assert_eq!(removed.uuid, child_uuid); - assert!(stack.scope_registries_get(&child_uuid).is_none()); } #[test] @@ -100,28 +97,22 @@ fn scope_stack_rejects_removing_non_top_or_root_scopes() { #[test] fn merge_helpers_preserve_global_and_scope_local_priority_order() { - let mut global_guardrails = - SortedRegistry::new(|entry: &GuardrailEntry<&'static str>| entry.priority); + let mut global_guardrails = SortedRegistry::new(); global_guardrails - .register( - "global".to_string(), - GuardrailEntry { - priority: 20, - guardrail: "global", - }, - ) + .register(Guardrail { + name: "global".to_string(), + priority: 20, + payload: "global", + }) .unwrap(); - let mut local_guardrails = - SortedRegistry::new(|entry: &GuardrailEntry<&'static str>| entry.priority); + let mut local_guardrails = SortedRegistry::new(); local_guardrails - .register( - "local".to_string(), - GuardrailEntry { - priority: 5, - guardrail: "local", - }, - ) + .register(Guardrail { + name: "local".to_string(), + priority: 5, + payload: "local", + }) .unwrap(); let local_guardrail_refs = [&local_guardrails]; @@ -129,35 +120,33 @@ fn merge_helpers_preserve_global_and_scope_local_priority_order() { assert_eq!( merged_guardrails .iter() - .map(|entry| entry.guardrail) + .map(|entry| entry.payload) .collect::>(), vec!["local", "global"] ); - let mut global_intercepts = - SortedRegistry::new(|entry: &Intercept<&'static str>| entry.priority); + let mut global_intercepts = SortedRegistry::new(); global_intercepts - .register( - "global".to_string(), - Intercept { - priority: 40, + .register(Intercept { + name: "global".to_string(), + priority: 40, + payload: RequestIntercept { break_chain: false, callable: "global", }, - ) + }) .unwrap(); - let mut local_intercepts = - SortedRegistry::new(|entry: &Intercept<&'static str>| entry.priority); + let mut local_intercepts = SortedRegistry::new(); local_intercepts - .register( - "local".to_string(), - Intercept { - priority: 10, + .register(Intercept { + name: "local".to_string(), + priority: 10, + payload: RequestIntercept { break_chain: false, callable: "local", }, - ) + }) .unwrap(); let local_intercept_refs = [&local_intercepts]; @@ -165,39 +154,79 @@ fn merge_helpers_preserve_global_and_scope_local_priority_order() { assert_eq!( merged_intercepts .iter() - .map(|entry| entry.callable) + .map(|entry| entry.payload.callable) .collect::>(), vec!["local", "global"] ); - let mut global_exec = - SortedRegistry::new(|entry: &ExecutionIntercept<&'static str>| entry.priority); + let mut global_exec = SortedRegistry::new(); global_exec - .register( - "global".to_string(), - ExecutionIntercept { - priority: 15, - callable: "global", - }, - ) + .register(ExecutionIntercept { + name: "global".to_string(), + priority: 15, + payload: "global", + }) .unwrap(); - let mut local_exec = - SortedRegistry::new(|entry: &ExecutionIntercept<&'static str>| entry.priority); + let mut local_exec = SortedRegistry::new(); local_exec - .register( - "local".to_string(), - ExecutionIntercept { - priority: 1, - callable: "local", - }, - ) + .register(ExecutionIntercept { + name: "local".to_string(), + priority: 1, + payload: "local", + }) .unwrap(); let merged_exec = merge_execution_intercept_callables(&global_exec, &[&local_exec]); assert_eq!(merged_exec, vec![("local", 1), ("global", 15)]); } +#[test] +fn conditional_guardrail_snapshots_keep_names_and_callbacks_after_deregister() { + let mut state = NemoFlowContextState::new(); + state + .tool_conditional_execution_guardrails + .register(Guardrail { + name: "snapshot_guardrail".to_string(), + priority: 1, + payload: Arc::new(|name, _args| Ok(Some(format!("{name} blocked")))), + }) + .unwrap(); + + let entries = state.tool_conditional_execution_entries(&[]); + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].name, "snapshot_guardrail"); + assert!( + state + .tool_conditional_execution_guardrails + .deregister("snapshot_guardrail") + ); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + let subscriber: EventSubscriberFn = Arc::new(move |event| { + captured.lock().unwrap().push(event.clone()); + }); + let subscribers = [subscriber]; + + let rejection = NemoFlowContextState::tool_conditional_execution_snapshot_chain( + "snapshot_target", + &json!({}), + &entries, + &subscribers, + None, + None, + ) + .unwrap(); + + assert_eq!(rejection.as_deref(), Some("snapshot_target blocked")); + let events = events.lock().unwrap(); + assert_eq!( + events.iter().map(Event::name).collect::>(), + vec!["snapshot_guardrail", "snapshot_guardrail"] + ); +} + #[test] fn context_state_supports_extensions_events_and_builders() { let mut state = NemoFlowContextState::new(); diff --git a/crates/core/tests/unit/plugin_tests.rs b/crates/core/tests/unit/plugin_tests.rs index 0fa8fee4d..1335cf748 100644 --- a/crates/core/tests/unit/plugin_tests.rs +++ b/crates/core/tests/unit/plugin_tests.rs @@ -95,7 +95,7 @@ impl Plugin for TestPlugin { "intercept", 1, false, - Box::new(|_name, mut request, annotated| { + Arc::new(|_name, mut request, annotated| { request.headers.insert("x-plugin".into(), json!(true)); Ok((request, annotated)) }), @@ -647,7 +647,7 @@ fn test_plugin_registration_context_covers_all_registration_helpers() { let mut ctx = PluginRegistrationContext::with_namespace("demo::"); ctx.register_subscriber("subscriber", Arc::new(|_event| {})) .unwrap(); - ctx.register_tool_request_intercept("tool-request", 1, false, Box::new(|_name, args| Ok(args))) + ctx.register_tool_request_intercept("tool-request", 1, false, Arc::new(|_name, args| Ok(args))) .unwrap(); ctx.register_tool_execution_intercept( "tool-exec", @@ -659,7 +659,7 @@ fn test_plugin_registration_context_covers_all_registration_helpers() { "llm-request", 1, false, - Box::new(|_name, request, annotated| Ok((request, annotated))), + Arc::new(|_name, request, annotated| Ok((request, annotated))), ) .unwrap(); ctx.register_llm_execution_intercept( @@ -884,13 +884,13 @@ fn test_plugin_registration_context_supports_guardrail_helpers() { ctx.register_tool_sanitize_request_guardrail( "tool_sanitize_request", 1, - Box::new(|_, args| args), + Arc::new(|_, args| args), ) .unwrap(); ctx.register_tool_sanitize_response_guardrail( "tool_sanitize_response", 1, - Box::new(|_, response| response), + Arc::new(|_, response| response), ) .unwrap(); ctx.register_tool_conditional_execution_guardrail( @@ -902,13 +902,13 @@ fn test_plugin_registration_context_supports_guardrail_helpers() { ctx.register_llm_sanitize_request_guardrail( "llm_sanitize_request", 1, - Box::new(|request| request), + Arc::new(|request| request), ) .unwrap(); ctx.register_llm_sanitize_response_guardrail( "llm_sanitize_response", 1, - Box::new(|response| response), + Arc::new(|response| response), ) .unwrap(); ctx.register_llm_conditional_execution_guardrail( @@ -959,7 +959,7 @@ fn test_plugin_registration_context_maps_duplicate_registration_errors() { "llm-request", 1, false, - Box::new(|_name, request, annotated| Ok((request, annotated))), + Arc::new(|_name, request, annotated| Ok((request, annotated))), ) .unwrap(); expect_registration_failed( @@ -967,7 +967,7 @@ fn test_plugin_registration_context_maps_duplicate_registration_errors() { "llm-request", 1, false, - Box::new(|_name, request, annotated| Ok((request, annotated))), + Arc::new(|_name, request, annotated| Ok((request, annotated))), ), "llm request intercept:", ); @@ -975,14 +975,14 @@ fn test_plugin_registration_context_maps_duplicate_registration_errors() { ctx.register_tool_sanitize_request_guardrail( "tool-sanitize-request", 1, - Box::new(|_, args| args), + Arc::new(|_, args| args), ) .unwrap(); expect_registration_failed( ctx.register_tool_sanitize_request_guardrail( "tool-sanitize-request", 1, - Box::new(|_, args| args), + Arc::new(|_, args| args), ), "tool sanitize request guardrail:", ); @@ -990,14 +990,14 @@ fn test_plugin_registration_context_maps_duplicate_registration_errors() { ctx.register_tool_sanitize_response_guardrail( "tool-sanitize-response", 1, - Box::new(|_, response| response), + Arc::new(|_, response| response), ) .unwrap(); expect_registration_failed( ctx.register_tool_sanitize_response_guardrail( "tool-sanitize-response", 1, - Box::new(|_, response| response), + Arc::new(|_, response| response), ), "tool sanitize response guardrail:", ); @@ -1020,14 +1020,14 @@ fn test_plugin_registration_context_maps_duplicate_registration_errors() { ctx.register_llm_sanitize_request_guardrail( "llm-sanitize-request", 1, - Box::new(|request| request), + Arc::new(|request| request), ) .unwrap(); expect_registration_failed( ctx.register_llm_sanitize_request_guardrail( "llm-sanitize-request", 1, - Box::new(|request| request), + Arc::new(|request| request), ), "llm sanitize request guardrail:", ); @@ -1035,14 +1035,14 @@ fn test_plugin_registration_context_maps_duplicate_registration_errors() { ctx.register_llm_sanitize_response_guardrail( "llm-sanitize-response", 1, - Box::new(|response| response), + Arc::new(|response| response), ) .unwrap(); expect_registration_failed( ctx.register_llm_sanitize_response_guardrail( "llm-sanitize-response", 1, - Box::new(|response| response), + Arc::new(|response| response), ), "llm sanitize response guardrail:", ); @@ -1102,14 +1102,14 @@ fn test_plugin_registration_context_maps_duplicate_registration_errors() { "llm stream execution intercept:", ); - ctx.register_tool_request_intercept("tool-request", 1, false, Box::new(|_name, args| Ok(args))) + ctx.register_tool_request_intercept("tool-request", 1, false, Arc::new(|_name, args| Ok(args))) .unwrap(); expect_registration_failed( ctx.register_tool_request_intercept( "tool-request", 1, false, - Box::new(|_name, args| Ok(args)), + Arc::new(|_name, args| Ok(args)), ), "tool request intercept:", ); @@ -1146,19 +1146,19 @@ fn test_plugin_registration_context_maps_deregistration_errors() { "llm-request", 1, false, - Box::new(|_name, request, annotated| Ok((request, annotated))), + Arc::new(|_name, request, annotated| Ok((request, annotated))), ) .unwrap(); ctx.register_tool_sanitize_request_guardrail( "tool-sanitize-request", 1, - Box::new(|_, args| args), + Arc::new(|_, args| args), ) .unwrap(); ctx.register_tool_sanitize_response_guardrail( "tool-sanitize-response", 1, - Box::new(|_, response| response), + Arc::new(|_, response| response), ) .unwrap(); ctx.register_tool_conditional_execution_guardrail( @@ -1170,13 +1170,13 @@ fn test_plugin_registration_context_maps_deregistration_errors() { ctx.register_llm_sanitize_request_guardrail( "llm-sanitize-request", 1, - Box::new(|request| request), + Arc::new(|request| request), ) .unwrap(); ctx.register_llm_sanitize_response_guardrail( "llm-sanitize-response", 1, - Box::new(|response| response), + Arc::new(|response| response), ) .unwrap(); ctx.register_llm_conditional_execution_guardrail("llm-conditional", 1, Arc::new(|_| Ok(None))) @@ -1200,7 +1200,7 @@ fn test_plugin_registration_context_maps_deregistration_errors() { }), ) .unwrap(); - ctx.register_tool_request_intercept("tool-request", 1, false, Box::new(|_name, args| Ok(args))) + ctx.register_tool_request_intercept("tool-request", 1, false, Arc::new(|_name, args| Ok(args))) .unwrap(); ctx.register_tool_execution_intercept( "tool-exec", diff --git a/crates/core/tests/unit/registry_tests.rs b/crates/core/tests/unit/registry_tests.rs index 6e8cc4dbe..e2eb64248 100644 --- a/crates/core/tests/unit/registry_tests.rs +++ b/crates/core/tests/unit/registry_tests.rs @@ -6,40 +6,38 @@ use super::*; struct PriorityItem { + name: String, priority: i32, value: String, } +impl RegistryEntry for PriorityItem { + fn name(&self) -> &str { + &self.name + } + + fn priority(&self) -> i32 { + self.priority + } +} + +fn item(name: &str, priority: i32, value: &str) -> PriorityItem { + PriorityItem { + name: name.into(), + priority, + value: value.into(), + } +} + #[test] fn test_sorted_registry() { - let mut reg = SortedRegistry::new(|item: &PriorityItem| item.priority); + let mut reg = SortedRegistry::new(); - reg.register( - "b".into(), - PriorityItem { - priority: 20, - value: "B".into(), - }, - ) - .unwrap(); + reg.register(item("b", 20, "B")).unwrap(); - reg.register( - "a".into(), - PriorityItem { - priority: 10, - value: "A".into(), - }, - ) - .unwrap(); + reg.register(item("a", 10, "A")).unwrap(); - reg.register( - "c".into(), - PriorityItem { - priority: 30, - value: "C".into(), - }, - ) - .unwrap(); + reg.register(item("c", 30, "C")).unwrap(); let sorted: Vec<&str> = reg .sorted_values() @@ -49,16 +47,7 @@ fn test_sorted_registry() { assert_eq!(sorted, vec!["A", "B", "C"]); // duplicate - assert!( - reg.register( - "a".into(), - PriorityItem { - priority: 5, - value: "A2".into(), - }, - ) - .is_err() - ); + assert!(reg.register(item("a", 5, "A2")).is_err()); // deregister assert!(reg.deregister("b")); @@ -74,56 +63,17 @@ fn test_sorted_registry() { #[test] fn test_empty_registry() { - let reg = SortedRegistry::new(|item: &PriorityItem| item.priority); + let reg = SortedRegistry::::new(); let sorted = reg.sorted_values(); assert!(sorted.is_empty()); } -#[test] -fn test_contains() { - let mut reg = SortedRegistry::new(|item: &PriorityItem| item.priority); - assert!(!reg.contains("x")); - reg.register( - "x".into(), - PriorityItem { - priority: 1, - value: "X".into(), - }, - ) - .unwrap(); - assert!(reg.contains("x")); - assert!(!reg.contains("y")); - reg.deregister("x"); - assert!(!reg.contains("x")); -} - #[test] fn test_negative_priorities() { - let mut reg = SortedRegistry::new(|item: &PriorityItem| item.priority); - reg.register( - "pos".into(), - PriorityItem { - priority: 10, - value: "P".into(), - }, - ) - .unwrap(); - reg.register( - "neg".into(), - PriorityItem { - priority: -5, - value: "N".into(), - }, - ) - .unwrap(); - reg.register( - "zero".into(), - PriorityItem { - priority: 0, - value: "Z".into(), - }, - ) - .unwrap(); + let mut reg = SortedRegistry::new(); + reg.register(item("pos", 10, "P")).unwrap(); + reg.register(item("neg", -5, "N")).unwrap(); + reg.register(item("zero", 0, "Z")).unwrap(); let sorted: Vec<&str> = reg .sorted_values() @@ -135,24 +85,10 @@ fn test_negative_priorities() { #[test] fn test_re_register_after_deregister() { - let mut reg = SortedRegistry::new(|item: &PriorityItem| item.priority); - reg.register( - "a".into(), - PriorityItem { - priority: 10, - value: "A1".into(), - }, - ) - .unwrap(); + let mut reg = SortedRegistry::new(); + reg.register(item("a", 10, "A1")).unwrap(); reg.deregister("a"); - reg.register( - "a".into(), - PriorityItem { - priority: 5, - value: "A2".into(), - }, - ) - .unwrap(); + reg.register(item("a", 5, "A2")).unwrap(); let sorted: Vec<&str> = reg .sorted_values() .iter() @@ -163,45 +99,23 @@ fn test_re_register_after_deregister() { #[test] fn test_deregister_nonexistent() { - let mut reg = SortedRegistry::::new(|item| item.priority); + let mut reg = SortedRegistry::::new(); assert!(!reg.deregister("nope")); } #[test] fn test_duplicate_error_message() { - let mut reg = SortedRegistry::new(|item: &PriorityItem| item.priority); - reg.register( - "dup".into(), - PriorityItem { - priority: 1, - value: "D".into(), - }, - ) - .unwrap(); - let err = reg - .register( - "dup".into(), - PriorityItem { - priority: 2, - value: "D2".into(), - }, - ) - .unwrap_err(); + let mut reg = SortedRegistry::new(); + reg.register(item("dup", 1, "D")).unwrap(); + let err = reg.register(item("dup", 2, "D2")).unwrap_err(); assert!(err.contains("dup")); assert!(err.contains("already exists")); } #[test] fn test_sorted_values_caching() { - let mut reg = SortedRegistry::new(|item: &PriorityItem| item.priority); - reg.register( - "a".into(), - PriorityItem { - priority: 1, - value: "A".into(), - }, - ) - .unwrap(); + let mut reg = SortedRegistry::new(); + reg.register(item("a", 1, "A")).unwrap(); // Sort order is already maintained eagerly by register() let s1: Vec<&str> = reg .sorted_values() @@ -220,15 +134,13 @@ fn test_sorted_values_caching() { #[test] fn test_many_entries_ordering() { - let mut reg = SortedRegistry::new(|item: &PriorityItem| item.priority); + let mut reg = SortedRegistry::new(); for i in (0..20).rev() { - reg.register( - format!("item_{i}"), - PriorityItem { - priority: i, - value: format!("V{i}"), - }, - ) + reg.register(PriorityItem { + name: format!("item_{i}"), + priority: i, + value: format!("V{i}"), + }) .unwrap(); } let sorted: Vec = reg.sorted_values().iter().map(|i| i.priority).collect(); @@ -238,65 +150,13 @@ fn test_many_entries_ordering() { #[test] fn test_same_priority_stable() { - let mut reg = SortedRegistry::new(|item: &PriorityItem| item.priority); - reg.register( - "x".into(), - PriorityItem { - priority: 1, - value: "X".into(), - }, - ) - .unwrap(); - reg.register( - "y".into(), - PriorityItem { - priority: 1, - value: "Y".into(), - }, - ) - .unwrap(); - // Both should be present + let mut reg = SortedRegistry::new(); + reg.register(item("x", 1, "X")).unwrap(); + reg.register(item("y", 1, "Y")).unwrap(); let sorted: Vec<&str> = reg .sorted_values() .iter() .map(|i| i.value.as_str()) .collect(); - assert_eq!(sorted.len(), 2); - assert!(sorted.contains(&"X")); - assert!(sorted.contains(&"Y")); -} - -#[test] -fn test_get_and_remove_cover_lookup_and_resort_paths() { - let mut reg = SortedRegistry::new(|item: &PriorityItem| item.priority); - reg.register( - "alpha".into(), - PriorityItem { - priority: 2, - value: "A".into(), - }, - ) - .unwrap(); - reg.register( - "beta".into(), - PriorityItem { - priority: 1, - value: "B".into(), - }, - ) - .unwrap(); - - assert_eq!(reg.get("beta").map(|item| item.value.as_str()), Some("B")); - assert!(reg.get("missing").is_none()); - - let removed = reg.remove("beta").unwrap(); - assert_eq!(removed.value, "B"); - assert!(reg.remove("beta").is_none()); - - let sorted: Vec<&str> = reg - .sorted_values() - .iter() - .map(|item| item.value.as_str()) - .collect(); - assert_eq!(sorted, vec!["A"]); + assert_eq!(sorted, vec!["X", "Y"]); } diff --git a/crates/core/tests/unit/shared_tests.rs b/crates/core/tests/unit/shared_tests.rs index 81690e3f7..d423e06e1 100644 --- a/crates/core/tests/unit/shared_tests.rs +++ b/crates/core/tests/unit/shared_tests.rs @@ -122,7 +122,7 @@ fn test_run_request_intercepts_with_codec_none_and_codec_paths() { "shared-none", 1, false, - Box::new(|_name, mut request, annotated| { + Arc::new(|_name, mut request, annotated| { assert!(annotated.is_none()); request.headers.insert("x-no-codec".into(), json!(true)); Ok((request, None)) @@ -150,7 +150,7 @@ fn test_run_request_intercepts_with_codec_none_and_codec_paths() { "shared-codec", 1, false, - Box::new(|_name, mut request, annotated| { + Arc::new(|_name, mut request, annotated| { let mut annotated = annotated.expect("codec should provide annotated request"); annotated.model = Some("intercepted-model".into()); request.headers.insert("x-codec".into(), json!(true)); diff --git a/crates/ffi/src/callable.rs b/crates/ffi/src/callable.rs index 90d3425ce..c55fad3e6 100644 --- a/crates/ffi/src/callable.rs +++ b/crates/ffi/src/callable.rs @@ -10,10 +10,11 @@ //! generated `nemo_flow.h` header. //! //! The `wrap_*` functions convert C callbacks (with opaque `user_data` pointers) -//! into Rust closures (`Box`) that the core runtime can invoke. -//! Each wrapper captures the user data and its optional free function in an -//! `Arc` so the closure is `Send + Sync` and the free function is -//! called exactly once when all references are dropped. +//! into Rust closures that the core runtime can invoke. Registry-stored +//! callbacks return `Arc`-backed closures, while one-shot or mutable callback +//! shapes remain boxed. Each wrapper captures the user data and its optional +//! free function in an `Arc` so the closure is `Send + Sync` and the +//! free function is called exactly once when all references are dropped. use std::ffi::{CStr, CString}; use std::future::Future; @@ -23,7 +24,8 @@ use std::sync::Arc; use libc::c_char; use nemo_flow::api::runtime::{ EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, LlmRequestInterceptFn, - LlmStreamExecutionNextFn, ToolConditionalFn, ToolExecutionNextFn, ToolInterceptFn, + LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, ToolConditionalFn, + ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; use serde_json::Value as Json; use tokio_stream::{Stream, StreamExt}; @@ -248,9 +250,9 @@ pub fn wrap_tool_sanitize_fn( cb: NemoFlowToolSanitizeCb, user_data: *mut libc::c_void, free_fn: NemoFlowFreeFn, -) -> Box Json + Send + Sync> { +) -> ToolSanitizeFn { let ud = make_user_data(user_data, free_fn); - Box::new(move |name: &str, args: Json| { + Arc::new(move |name: &str, args: Json| { let c_name = CString::new(name).unwrap_or_default(); let c_args = json_to_c_string(&args); let result_ptr = unsafe { cb(ud.ptr, c_name.as_ptr(), c_args) }; @@ -294,7 +296,7 @@ pub fn wrap_tool_request_intercept_fn( free_fn: NemoFlowFreeFn, ) -> ToolInterceptFn { let ud = make_user_data(user_data, free_fn); - Box::new(move |name: &str, args: Json| { + Arc::new(move |name: &str, args: Json| { clear_last_error(); let c_name = CString::new(name).unwrap_or_default(); let c_args = json_to_c_string(&args); @@ -565,7 +567,7 @@ pub fn wrap_llm_request_intercept_fn( free_fn: NemoFlowFreeFn, ) -> LlmRequestInterceptFn { let ud = make_user_data(user_data, free_fn); - Box::new( + Arc::new( move |name: &str, request: LlmRequest, annotated: Option| { clear_last_error(); let c_name = CString::new(name).unwrap_or_default(); @@ -641,9 +643,9 @@ pub fn wrap_llm_response_fn( cb: NemoFlowJsonCb, user_data: *mut libc::c_void, free_fn: NemoFlowFreeFn, -) -> Box Json + Send + Sync> { +) -> LlmSanitizeResponseFn { let ud = make_user_data(user_data, free_fn); - Box::new(move |response: Json| { + Arc::new(move |response: Json| { let c_json = json_to_c_string(&response); let result_ptr = unsafe { cb(ud.ptr, c_json) }; unsafe { nemo_flow_string_free_internal(c_json) }; @@ -658,9 +660,9 @@ pub fn wrap_llm_sanitize_request_fn( cb: NemoFlowLlmRequestCb, user_data: *mut libc::c_void, free_fn: NemoFlowFreeFn, -) -> Box LlmRequest + Send + Sync> { +) -> LlmSanitizeRequestFn { let ud = make_user_data(user_data, free_fn); - Box::new(move |request: LlmRequest| { + Arc::new(move |request: LlmRequest| { let ffi_req = Box::into_raw(Box::new(FfiLLMRequest(request))); let result_ptr = unsafe { cb(ud.ptr, ffi_req) }; // Free the input request diff --git a/crates/node/src/callable.rs b/crates/node/src/callable.rs index 19cba2a8a..8a079bbdd 100644 --- a/crates/node/src/callable.rs +++ b/crates/node/src/callable.rs @@ -16,7 +16,8 @@ use std::sync::Arc; use napi::threadsafe_function::{ErrorStrategy, ThreadsafeFunction, ThreadsafeFunctionCallMode}; use nemo_flow::api::runtime::{ EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, LlmRequestInterceptFn, - LlmStreamExecutionNextFn, ToolConditionalFn, ToolExecutionNextFn, ToolInterceptFn, + LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, ToolConditionalFn, + ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; use serde_json::Value as Json; use tokio_stream::StreamExt; @@ -97,9 +98,9 @@ fn recv_llm_request_result( /// Wrap a JS function `(name: string, args: object) => object` for tool sanitize/intercept. pub fn wrap_js_tool_fn( func: ThreadsafeFunction<(String, Json), ErrorStrategy::Fatal>, -) -> Box Json + Send + Sync> { +) -> ToolSanitizeFn { let func = Arc::new(func); - Box::new(move |name: &str, args: Json| { + Arc::new(move |name: &str, args: Json| { let func = func.clone(); let name = name.to_string(); let (tx, rx) = std::sync::mpsc::channel(); @@ -155,7 +156,7 @@ pub fn wrap_js_tool_request_intercept_fn( func: ThreadsafeFunction<(String, Json), ErrorStrategy::Fatal>, ) -> ToolInterceptFn { let func = Arc::new(func); - Box::new(move |name: &str, args: Json| { + Arc::new(move |name: &str, args: Json| { let func = func.clone(); let name = name.to_string(); let (tx, rx) = std::sync::mpsc::channel(); @@ -212,7 +213,7 @@ pub fn wrap_js_llm_request_intercept_fn( func: ThreadsafeFunction, ) -> LlmRequestInterceptFn { let func = Arc::new(func); - Box::new( + Arc::new( move |name: &str, request: LlmRequest, annotated: Option| @@ -278,9 +279,9 @@ pub fn wrap_js_llm_request_intercept_fn( /// Since ThreadsafeFunction requires serde-serializable args, we serialize the request as JSON. pub fn wrap_js_llm_sanitize_request_fn( func: ThreadsafeFunction, -) -> Box LlmRequest + Send + Sync> { +) -> LlmSanitizeRequestFn { let func = Arc::new(func); - Box::new(move |request: LlmRequest| { + Arc::new(move |request: LlmRequest| { let func = func.clone(); let req_json = serde_json::to_value(&request).unwrap_or(Json::Null); let (tx, rx) = std::sync::mpsc::channel(); @@ -311,9 +312,9 @@ pub fn wrap_js_llm_sanitize_request_fn( /// Wrap a JS function for LLM sanitize response: `(response: Json) => Json`. pub fn wrap_js_llm_response_fn( func: ThreadsafeFunction, -) -> Box Json + Send + Sync> { +) -> LlmSanitizeResponseFn { let func = Arc::new(func); - Box::new(move |response: Json| { + Arc::new(move |response: Json| { let func = func.clone(); let (tx, rx) = std::sync::mpsc::channel(); let status = func.call_with_return_value( diff --git a/crates/python/src/py_callable.rs b/crates/python/src/py_callable.rs index 99f63e721..79f6f09a9 100644 --- a/crates/python/src/py_callable.rs +++ b/crates/python/src/py_callable.rs @@ -26,7 +26,8 @@ use std::sync::Arc; use nemo_flow::api::runtime::{ EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, LlmRequestInterceptFn, - LlmStreamExecutionNextFn, ToolConditionalFn, ToolExecutionNextFn, ToolInterceptFn, + LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, ToolConditionalFn, + ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; use nemo_flow::error::{FlowError, Result as FlowResult}; use pyo3::prelude::*; @@ -182,8 +183,8 @@ fn stream_from_async_iter( } /// Wrap a Python callable `(str, Json) -> Json` for tool sanitize/intercept fns. -pub fn wrap_py_tool_fn(py_fn: Py) -> Box Json + Send + Sync> { - Box::new(move |name: &str, args: Json| { +pub fn wrap_py_tool_fn(py_fn: Py) -> ToolSanitizeFn { + Arc::new(move |name: &str, args: Json| { Python::attach(|py| { let py_args = match json_to_py(py, &args) { Ok(v) => v, @@ -237,7 +238,7 @@ pub fn wrap_py_tool_conditional_fn(py_fn: Py) -> ToolConditionalFn { /// Wrap a Python callable `(str, Json) -> Json` for tool request intercepts. pub fn wrap_py_tool_request_intercept_fn(py_fn: Py) -> ToolInterceptFn { - Box::new(move |name: &str, args: Json| { + Arc::new(move |name: &str, args: Json| { Python::attach(|py| { let py_args = json_to_py(py, &args).map_err(|e| { FlowError::Internal(format!("tool callback json_to_py failed for '{name}': {e}")) @@ -587,10 +588,8 @@ pub fn wrap_py_llm_stream_exec_intercept_fn( } /// Wrap a Python callable `(LlmRequest) -> LlmRequest` for LLM sanitize request guardrails. -pub fn wrap_py_llm_sanitize_request_fn( - py_fn: Py, -) -> Box LlmRequest + Send + Sync> { - Box::new(move |request: LlmRequest| { +pub fn wrap_py_llm_sanitize_request_fn(py_fn: Py) -> LlmSanitizeRequestFn { + Arc::new(move |request: LlmRequest| { Python::attach(|py| { let py_req = PyLLMRequest { inner: request.clone(), @@ -646,7 +645,7 @@ pub fn wrap_py_llm_conditional_fn(py_fn: Py) -> LlmConditionalFn { /// The Python function receives ``(name: str, request: LlmRequest, annotated: AnnotatedLLMRequest | None)`` /// and must return ``(LlmRequest, AnnotatedLLMRequest | None)``. pub fn wrap_py_llm_request_intercept_fn(py_fn: Py) -> LlmRequestInterceptFn { - Box::new( + Arc::new( move |name: &str, request: LlmRequest, annotated: Option| @@ -819,10 +818,8 @@ pub fn wrap_py_finalizer_fn(py_fn: Py) -> Box Json + Send } /// Wrap a Python callable `(dict) -> dict` for LLM sanitize response guardrails. -pub fn wrap_py_llm_sanitize_response_fn( - py_fn: Py, -) -> Box Json + Send + Sync> { - Box::new(move |response: Json| { +pub fn wrap_py_llm_sanitize_response_fn(py_fn: Py) -> LlmSanitizeResponseFn { + Arc::new(move |response: Json| { Python::attach(|py| { let py_resp = match json_to_py(py, &response) { Ok(v) => v, diff --git a/crates/wasm/src/callable.rs b/crates/wasm/src/callable.rs index dba7ab2c3..f8e0a8643 100644 --- a/crates/wasm/src/callable.rs +++ b/crates/wasm/src/callable.rs @@ -5,9 +5,9 @@ //! Wrappers that adapt JavaScript callback functions into Rust closures. //! //! Each wrapper takes a `js_sys::Function`, wraps it with `SendWrapper` (since -//! JS functions are not `Send`), and returns a boxed closure matching the -//! signature expected by the core runtime for guardrails, intercepts, -//! execution functions, and event subscribers. +//! JS functions are not `Send`), and returns the closure shape expected by the +//! core runtime. Registry-stored callbacks return `Arc`-backed closures, while +//! one-shot or mutable callback shapes remain boxed. use std::future::Future; use std::pin::Pin; @@ -32,7 +32,8 @@ use nemo_flow::api::event::Event; use nemo_flow::api::llm::LlmRequest; use nemo_flow::api::runtime::{ EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, LlmRequestInterceptFn, - LlmStreamExecutionNextFn, ToolConditionalFn, ToolExecutionNextFn, ToolInterceptFn, + LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, ToolConditionalFn, + ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; use nemo_flow::codec::request::AnnotatedLlmRequest; #[cfg(target_arch = "wasm32")] @@ -153,14 +154,14 @@ fn wasm_only_error() -> FlowError { /// Wrap a JS function `(name, args) => result` for tool sanitize/intercept. #[cfg(not(target_arch = "wasm32"))] -pub fn wrap_js_tool_fn(_func: Function) -> Box Json + Send + Sync> { - Box::new(move |_name: &str, _args: Json| Json::Null) +pub fn wrap_js_tool_fn(_func: Function) -> ToolSanitizeFn { + Arc::new(move |_name: &str, _args: Json| Json::Null) } #[cfg(target_arch = "wasm32")] -pub fn wrap_js_tool_fn(func: Function) -> Box Json + Send + Sync> { +pub fn wrap_js_tool_fn(func: Function) -> ToolSanitizeFn { let func = SendWrapper::new(func); - Box::new(move |name: &str, args: Json| { + Arc::new(move |name: &str, args: Json| { let js_name = JsValue::from_str(name); let js_args = json_to_js(&args); // TODO: This closure returns Json (not Result), so we cannot propagate @@ -206,13 +207,13 @@ pub fn wrap_js_tool_conditional_fn(func: Function) -> ToolConditionalFn { /// Wrap a JS function `(name, args) => result` for fallible tool request intercepts. #[cfg(not(target_arch = "wasm32"))] pub fn wrap_js_tool_request_intercept_fn(_func: Function) -> ToolInterceptFn { - Box::new(move |_name: &str, args: Json| Ok(args)) + Arc::new(move |_name: &str, args: Json| Ok(args)) } #[cfg(target_arch = "wasm32")] pub fn wrap_js_tool_request_intercept_fn(func: Function) -> ToolInterceptFn { let func = SendWrapper::new(func); - Box::new(move |name: &str, args: Json| { + Arc::new(move |name: &str, args: Json| { let js_name = JsValue::from_str(name); let js_args = json_to_js(&args); let result = func @@ -263,7 +264,7 @@ pub fn wrap_js_tool_exec_fn( /// `({ name, request, annotated }) => { request, annotated }`. #[cfg(not(target_arch = "wasm32"))] pub fn wrap_js_llm_request_intercept_fn(_func: Function) -> LlmRequestInterceptFn { - Box::new( + Arc::new( move |_name: &str, request: LlmRequest, annotated: Option| { Ok((request, annotated)) }, @@ -273,7 +274,7 @@ pub fn wrap_js_llm_request_intercept_fn(_func: Function) -> LlmRequestInterceptF #[cfg(target_arch = "wasm32")] pub fn wrap_js_llm_request_intercept_fn(func: Function) -> LlmRequestInterceptFn { let func = SendWrapper::new(func); - Box::new( + Arc::new( move |name: &str, request: LlmRequest, annotated: Option| @@ -346,18 +347,14 @@ pub fn wrap_js_llm_request_intercept_fn(func: Function) -> LlmRequestInterceptFn /// Wrap a JS function for LLM sanitize request: `(request) => request`. #[cfg(not(target_arch = "wasm32"))] -pub fn wrap_js_llm_sanitize_request_fn( - _func: Function, -) -> Box LlmRequest + Send + Sync> { - Box::new(move |request: LlmRequest| request) +pub fn wrap_js_llm_sanitize_request_fn(_func: Function) -> LlmSanitizeRequestFn { + Arc::new(move |request: LlmRequest| request) } #[cfg(target_arch = "wasm32")] -pub fn wrap_js_llm_sanitize_request_fn( - func: Function, -) -> Box LlmRequest + Send + Sync> { +pub fn wrap_js_llm_sanitize_request_fn(func: Function) -> LlmSanitizeRequestFn { let func = SendWrapper::new(func); - Box::new(move |request: LlmRequest| { + Arc::new(move |request: LlmRequest| { let req_json = serde_json::to_value(&request).unwrap_or(Json::Null); let js_req = json_to_js(&req_json); // TODO: This closure returns LlmRequest (not Result), so we cannot propagate @@ -894,15 +891,15 @@ pub fn wrap_js_response_codec(decode_response_fn: Function) -> Arc Box Json + Send + Sync> { +pub fn wrap_js_llm_response_fn(func: Function) -> LlmSanitizeResponseFn { let _ = func; - Box::new(move |response: Json| response) + Arc::new(move |response: Json| response) } #[cfg(target_arch = "wasm32")] -pub fn wrap_js_llm_response_fn(func: Function) -> Box Json + Send + Sync> { +pub fn wrap_js_llm_response_fn(func: Function) -> LlmSanitizeResponseFn { let func = SendWrapper::new(func); - Box::new(move |response: Json| { + Arc::new(move |response: Json| { let js_resp = json_to_js(&response); // TODO: This closure returns Json (not Result), so we cannot propagate // errors through the type system. Log errors and fall back to original response.