diff --git a/crates/core/src/api/registry.rs b/crates/core/src/api/registry.rs index 984ac43d8..afda43950 100644 --- a/crates/core/src/api/registry.rs +++ b/crates/core/src/api/registry.rs @@ -7,7 +7,7 @@ use crate::api::runtime::{ EventSanitizeFn, LlmConditionalFn, LlmExecutionFn, LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionFn, ToolConditionalFn, ToolExecutionFn, - ToolInterceptFn, ToolSanitizeFn, + ToolExecutionFrameFn, ToolInterceptFn, ToolSanitizeFn, }; use crate::api::runtime::{current_scope_stack, global_context}; use crate::api::shared::ensure_runtime_owner; @@ -203,12 +203,10 @@ macro_rules! global_intercept_registry_api { }; } -macro_rules! global_execution_registry_api { +macro_rules! global_execution_registry_register_api { ( $(#[$register_meta:meta])* $register_name:ident, - $(#[$deregister_meta:meta])* - $deregister_name:ident, $field:ident, $fn_type:ty ) => { @@ -233,10 +231,18 @@ macro_rules! global_execution_registry_api { .map_err(|error| FlowError::Internal(error.to_string()))?; state .$field - .register(ExecutionIntercept::new(name, priority, callable)) + .register(ExecutionIntercept::new(name, priority, callable.into())) .map_err(FlowError::AlreadyExists) } + }; +} +macro_rules! global_execution_registry_deregister_api { + ( + $(#[$deregister_meta:meta])* + $deregister_name:ident, + $field:ident + ) => { $(#[$deregister_meta])* /// /// # Parameters @@ -259,6 +265,29 @@ macro_rules! global_execution_registry_api { }; } +macro_rules! global_execution_registry_api { + ( + $(#[$register_meta:meta])* + $register_name:ident, + $(#[$deregister_meta:meta])* + $deregister_name:ident, + $field:ident, + $fn_type:ty + ) => { + global_execution_registry_register_api!( + $(#[$register_meta])* + $register_name, + $field, + $fn_type + ); + global_execution_registry_deregister_api!( + $(#[$deregister_meta])* + $deregister_name, + $field + ); + }; +} + macro_rules! scope_guardrail_registry_api { ( $(#[$register_meta:meta])* @@ -400,12 +429,10 @@ macro_rules! scope_intercept_registry_api { }; } -macro_rules! scope_execution_registry_api { +macro_rules! scope_execution_registry_register_api { ( $(#[$register_meta:meta])* $register_name:ident, - $(#[$deregister_meta:meta])* - $deregister_name:ident, $field:ident, $fn_type:ty ) => { @@ -438,10 +465,18 @@ macro_rules! scope_execution_registry_api { .ok_or_else(|| FlowError::NotFound(format!("scope {scope_uuid} not found")))?; registries .$field - .register(ExecutionIntercept::new(name, priority, callable)) + .register(ExecutionIntercept::new(name, priority, callable.into())) .map_err(FlowError::AlreadyExists) } + }; +} +macro_rules! scope_execution_registry_deregister_api { + ( + $(#[$deregister_meta:meta])* + $deregister_name:ident, + $field:ident + ) => { $(#[$deregister_meta])* /// /// # Parameters @@ -467,6 +502,29 @@ macro_rules! scope_execution_registry_api { }; } +macro_rules! scope_execution_registry_api { + ( + $(#[$register_meta:meta])* + $register_name:ident, + $(#[$deregister_meta:meta])* + $deregister_name:ident, + $field:ident, + $fn_type:ty + ) => { + scope_execution_registry_register_api!( + $(#[$register_meta])* + $register_name, + $field, + $fn_type + ); + scope_execution_registry_deregister_api!( + $(#[$deregister_meta])* + $deregister_name, + $field + ); + }; +} + global_guardrail_registry_api!( /// Register a global mark event sanitizer. register_mark_sanitize_guardrail, @@ -535,7 +593,7 @@ global_intercept_registry_api!( global_execution_registry_api!( /// Register a global tool execution intercept. /// Execution intercepts can wrap or replace the tool callback. Each - /// callback returns a canonical tool execution outcome, while its + /// callback returns Relay's tool execution outcome wrapper, while its /// continuation resolves to the raw downstream result JSON. register_tool_execution_intercept, /// Deregister a global tool execution intercept. @@ -543,6 +601,14 @@ global_execution_registry_api!( tool_execution_intercepts, ToolExecutionFn ); +global_execution_registry_register_api!( + /// Register a global annotation-aware tool execution intercept. + /// Frame intercepts share the existing tool execution registry, namespace, + /// and priority order with raw-JSON intercepts. + register_tool_execution_frame_intercept, + tool_execution_intercepts, + ToolExecutionFrameFn +); global_guardrail_registry_api!( /// Register a global LLM sanitize-request guardrail. @@ -669,7 +735,7 @@ scope_intercept_registry_api!( scope_execution_registry_api!( /// Register a scope-local tool execution intercept. /// Execution intercepts can wrap or replace the tool callback inside the - /// owning scope. Each callback returns a canonical tool execution outcome, + /// owning scope. Each callback returns Relay's tool execution outcome wrapper, /// while its continuation resolves to the raw downstream result JSON. scope_register_tool_execution_intercept, /// Deregister a scope-local tool execution intercept. @@ -677,6 +743,14 @@ scope_execution_registry_api!( tool_execution_intercepts, ToolExecutionFn ); +scope_execution_registry_register_api!( + /// Register a scope-local annotation-aware tool execution intercept. + /// Frame intercepts share the existing tool execution registry, namespace, + /// and priority order with raw-JSON intercepts. + scope_register_tool_execution_frame_intercept, + tool_execution_intercepts, + ToolExecutionFrameFn +); scope_guardrail_registry_api!( /// Register a scope-local LLM sanitize-request guardrail. diff --git a/crates/core/src/api/runtime.rs b/crates/core/src/api/runtime.rs index 12d6612c3..1619dfda3 100644 --- a/crates/core/src/api/runtime.rs +++ b/crates/core/src/api/runtime.rs @@ -16,7 +16,8 @@ pub use callbacks::{ LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn, LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionFn, LlmStreamExecutionNextFn, LlmStreamInner, ToolConditionalFn, ToolExecutionFn, - ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, + ToolExecutionFrameFn, ToolExecutionFrameNextFn, ToolExecutionNextFn, ToolInterceptFn, + ToolSanitizeFn, }; #[doc(hidden)] pub use continuation_context::MiddlewareContinuationContext; diff --git a/crates/core/src/api/runtime/callbacks.rs b/crates/core/src/api/runtime/callbacks.rs index f1d55a6cf..521a52e72 100644 --- a/crates/core/src/api/runtime/callbacks.rs +++ b/crates/core/src/api/runtime/callbacks.rs @@ -17,7 +17,9 @@ use tokio_stream::Stream; use crate::api::event::{Event, EventSanitizeFields}; use crate::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; -use crate::api::tool::ToolExecutionInterceptOutcome; +use crate::api::tool::{ + ToolExecutionFrame, ToolExecutionFrameOutcome, ToolExecutionInterceptOutcome, +}; use crate::codec::request::AnnotatedLlmRequest; use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::error::Result; @@ -128,8 +130,9 @@ pub type ToolExecutionNextFn = /// - Third argument: Continuation for the remaining execution chain. /// /// # Returns -/// A future resolving to the canonical tool execution outcome, containing the -/// tool result and any pending lifecycle marks produced by this intercept. +/// A future resolving to Relay's execution outcome wrapper, containing the +/// harness-owned tool result and any pending lifecycle marks produced by this +/// intercept. /// /// # Errors /// The future resolves to an error when the intercept or remaining execution @@ -144,13 +147,61 @@ pub type ToolExecutionFn = Arc< + Sync, >; -/// Internal continuation carrying both a tool result and accumulated marks. -pub(crate) type ToolExecutionOutcomeNextFn = Arc< - dyn Fn(Json) -> Pin> + Send>> +/// Annotation-aware continuation invoked by tool execution frame intercepts. +/// +/// The continuation exposes the raw downstream result together with its +/// optional opaque annotation. Relay retains downstream pending marks +/// internally, matching [`ToolExecutionNextFn`]. +pub type ToolExecutionFrameNextFn = Arc< + dyn Fn(Json) -> Pin> + Send>> + Send + Sync, +>; + +/// Annotation-aware tool execution intercept. +/// +/// This callback participates in the same priority-ordered chain as +/// [`ToolExecutionFn`], but its continuation and outcome carry a +/// [`ToolExecutionFrame`]. +pub type ToolExecutionFrameFn = Arc< + dyn Fn( + &str, + Json, + ToolExecutionFrameNextFn, + ) -> Pin> + Send>> + + Send + + Sync, +>; + +/// Internal continuation carrying a tool frame and accumulated marks. +pub(crate) type ToolExecutionFrameOutcomeNextFn = Arc< + dyn Fn(Json) -> Pin> + Send>> + Send + Sync, >; +/// One registry payload for legacy and annotation-aware tool intercepts. +/// +/// Keeping both callback forms in this enum preserves a single namespace and +/// priority order rather than creating a second middleware chain. +#[derive(Clone)] +pub(crate) enum ToolExecutionCallable { + /// Existing raw-JSON execution intercept. + Legacy(ToolExecutionFn), + /// Annotation-aware frame execution intercept. + Frame(ToolExecutionFrameFn), +} + +impl From for ToolExecutionCallable { + fn from(value: ToolExecutionFn) -> Self { + Self::Legacy(value) + } +} + +impl From for ToolExecutionCallable { + fn from(value: ToolExecutionFrameFn) -> Self { + Self::Frame(value) + } +} + /// Relay's built-in LLM codec identities. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum BuiltinLlmCodec { diff --git a/crates/core/src/api/runtime/state.rs b/crates/core/src/api/runtime/state.rs index 0dfbcfe61..24b72db5a 100644 --- a/crates/core/src/api/runtime/state.rs +++ b/crates/core/src/api/runtime/state.rs @@ -32,8 +32,8 @@ use crate::api::runtime::callbacks::{ LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn, LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionFn, LlmStreamExecutionNextFn, LlmStreamExecutionRegistryRefs, LlmStreamInner, ToolConditionalFn, - ToolExecutionFn, ToolExecutionNextFn, ToolExecutionOutcomeNextFn, ToolInterceptFn, - ToolSanitizeFn, + ToolExecutionCallable, ToolExecutionFrameNextFn, ToolExecutionFrameOutcomeNextFn, + ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; use crate::api::runtime::continuation_context::{ MiddlewareContinuationContext, MiddlewareContinuationGuard, MiddlewareContinuationLease, @@ -43,7 +43,7 @@ use crate::api::scope::{CreateScopeHandleParams, EndScopeHandleParams, ScopeHand use crate::api::shared::snapshot_event_sanitizers; use crate::api::tool::ToolHandle; use crate::api::tool::{ - CreateToolHandleParams, EndToolHandleParams, ToolExecutionInterceptOutcome, + CreateToolHandleParams, EndToolHandleParams, ToolExecutionFrame, ToolExecutionFrameOutcome, }; use crate::codec::request::AnnotatedLlmRequest; use crate::codec::response::AnnotatedLlmResponse; @@ -220,7 +220,7 @@ pub struct NemoRelayContextState { /// Global tool request intercepts that can rewrite arguments before execution. pub(crate) tool_request_intercepts: SortedRegistry>, /// Global tool execution intercepts that wrap or replace callback execution. - pub(crate) tool_execution_intercepts: SortedRegistry>, + pub(crate) tool_execution_intercepts: SortedRegistry>, /// Global LLM request sanitizers applied to emitted LLM-start payloads. pub(crate) llm_sanitize_request_guardrails: SortedRegistry>, /// Global LLM response sanitizers applied to emitted LLM-end payloads. @@ -1123,73 +1123,147 @@ impl NemoRelayContextState { /// from the active scope stack. /// /// # Returns - /// A composed [`ToolExecutionOutcomeNextFn`] that wraps `default_fn` in + /// A composed [`ToolExecutionFrameOutcomeNextFn`] that wraps `default_fn` in /// every matching execution intercept. pub(crate) fn tool_build_execution_chain( &self, name: &str, - default_fn: ToolExecutionNextFn, - scope_locals: &[&SortedRegistry>], - ) -> ToolExecutionOutcomeNextFn { + default_fn: ToolExecutionFrameNextFn, + scope_locals: &[&SortedRegistry>], + ) -> ToolExecutionFrameOutcomeNextFn { let matching = merge_execution_intercept_callables(&self.tool_execution_intercepts, scope_locals); - let mut next: ToolExecutionOutcomeNextFn = Arc::new(move |args| { + let mut next: ToolExecutionFrameOutcomeNextFn = Arc::new(move |args| { let default_fn = default_fn.clone(); - Box::pin(async move { - default_fn(args) - .await - .map(ToolExecutionInterceptOutcome::new) - }) + Box::pin(async move { default_fn(args).await.map(ToolExecutionFrameOutcome::new) }) }); let name = name.to_string(); for (callable, _) in matching.into_iter().rev() { let current_next = next.clone(); let current_name = name.clone(); - next = Arc::new(move |args| { - let callable = callable.clone(); - let current_name = current_name.clone(); - let (continuation, continuation_guard) = MiddlewareContinuationLease::capture(); - let next_sequence = Arc::new(AtomicUsize::new(0)); - let downstream_marks = Arc::new(Mutex::new(Vec::new())); - let raw_next: ToolExecutionNextFn = { - let current_next = current_next.clone(); - let continuation = continuation.clone(); - let next_sequence = next_sequence.clone(); - let downstream_marks = downstream_marks.clone(); - Arc::new(move |args| { - let sequence = next_sequence.fetch_add(1, Ordering::Relaxed); + next = match callable { + ToolExecutionCallable::Legacy(callable) => Arc::new(move |args| { + let callable = callable.clone(); + let current_name = current_name.clone(); + let (continuation, continuation_guard) = MiddlewareContinuationLease::capture(); + let next_sequence = Arc::new(AtomicUsize::new(0)); + let downstream_outcomes = Arc::new(Mutex::new(Vec::new())); + let raw_next: ToolExecutionNextFn = { let current_next = current_next.clone(); - let invocation = continuation.begin(); - let downstream_marks = downstream_marks.clone(); - Box::pin(async move { - let outcome = invocation?.invoke(move || current_next(args)).await?; - downstream_marks + let continuation = continuation.clone(); + let next_sequence = next_sequence.clone(); + let downstream_outcomes = downstream_outcomes.clone(); + Arc::new(move |args| { + let sequence = next_sequence.fetch_add(1, Ordering::Relaxed); + let current_next = current_next.clone(); + let invocation = continuation.begin(); + let downstream_outcomes = downstream_outcomes.clone(); + Box::pin(async move { + let outcome = + invocation?.invoke(move || current_next(args)).await?; + let ToolExecutionFrameOutcome { + frame, + pending_marks, + } = outcome; + let ToolExecutionFrame { result, annotation } = frame; + // Existing raw-only chains should not pay to + // clone potentially large tool results. Retain + // a comparison copy only when there is an + // annotation that might be preserved. + let downstream_result = annotation.as_ref().map(|_| result.clone()); + downstream_outcomes + .lock() + .expect("tool downstream outcome accumulator lock poisoned") + .push((sequence, downstream_result, annotation, pending_marks)); + Ok(result) + }) + }) + }; + Box::pin(async move { + let callback_result = callable(¤t_name, args, raw_next).await; + drop(continuation_guard); + let mut legacy_outcome = callback_result?; + let mut downstream = std::mem::take( + &mut *downstream_outcomes .lock() - .expect("tool pending mark accumulator lock poisoned") - .push((sequence, outcome.pending_marks)); - Ok(outcome.result) + .expect("tool downstream outcome accumulator lock poisoned"), + ); + downstream.sort_by_key(|(sequence, _, _, _)| *sequence); + + // A legacy callback cannot identify which downstream + // annotation belongs to its output after multiple + // continuation calls. Preserve only the unambiguous + // single-call, unchanged-result case. + let annotation = if next_sequence.load(Ordering::Relaxed) == 1 + && downstream.len() == 1 + && downstream[0].1.as_ref() == Some(&legacy_outcome.result) + { + downstream[0].2.clone() + } else { + None + }; + let mut marks = downstream + .into_iter() + .flat_map(|(_, _, _, pending_marks)| pending_marks) + .collect::>(); + marks.append(&mut legacy_outcome.pending_marks); + Ok(ToolExecutionFrameOutcome { + frame: ToolExecutionFrame { + result: legacy_outcome.result, + annotation, + }, + pending_marks: marks, }) }) - }; - Box::pin(async move { - let outcome = callable(¤t_name, args, raw_next).await; - drop(continuation_guard); - let mut outcome = outcome?; - let mut downstream_batches = std::mem::take( - &mut *downstream_marks - .lock() - .expect("tool pending mark accumulator lock poisoned"), - ); - downstream_batches.sort_by_key(|(sequence, _)| *sequence); - let mut marks = downstream_batches - .into_iter() - .flat_map(|(_, marks)| marks) - .collect::>(); - marks.append(&mut outcome.pending_marks); - outcome.pending_marks = marks; - Ok(outcome) - }) - }); + }), + ToolExecutionCallable::Frame(callable) => Arc::new(move |args| { + let callable = callable.clone(); + let current_name = current_name.clone(); + let (continuation, continuation_guard) = MiddlewareContinuationLease::capture(); + let next_sequence = Arc::new(AtomicUsize::new(0)); + let downstream_marks = Arc::new(Mutex::new(Vec::new())); + let frame_next: ToolExecutionFrameNextFn = { + let current_next = current_next.clone(); + let continuation = continuation.clone(); + let next_sequence = next_sequence.clone(); + let downstream_marks = downstream_marks.clone(); + Arc::new(move |args| { + let sequence = next_sequence.fetch_add(1, Ordering::Relaxed); + let current_next = current_next.clone(); + let invocation = continuation.begin(); + let downstream_marks = downstream_marks.clone(); + Box::pin(async move { + let outcome = + invocation?.invoke(move || current_next(args)).await?; + downstream_marks + .lock() + .expect("tool pending mark accumulator lock poisoned") + .push((sequence, outcome.pending_marks)); + Ok(outcome.frame) + }) + }) + }; + Box::pin(async move { + let callback_result = callable(¤t_name, args, frame_next).await; + drop(continuation_guard); + let mut outcome = callback_result?; + outcome.frame = outcome.frame.normalized(); + let mut downstream_batches = std::mem::take( + &mut *downstream_marks + .lock() + .expect("tool pending mark accumulator lock poisoned"), + ); + downstream_batches.sort_by_key(|(sequence, _)| *sequence); + let mut marks = downstream_batches + .into_iter() + .flat_map(|(_, marks)| marks) + .collect::>(); + marks.append(&mut outcome.pending_marks); + outcome.pending_marks = marks; + Ok(outcome) + }) + }), + }; } next } diff --git a/crates/core/src/api/runtime/subscriber_dispatcher.rs b/crates/core/src/api/runtime/subscriber_dispatcher.rs index 872d17aea..c6e591c86 100644 --- a/crates/core/src/api/runtime/subscriber_dispatcher.rs +++ b/crates/core/src/api/runtime/subscriber_dispatcher.rs @@ -470,7 +470,7 @@ mod native { subscribers: &[EventSubscriberFn], scope_stack: ScopeStackHandle, ) -> bool { - if subscribers.is_empty() { + if subscribers.is_empty() && sanitizers.is_empty() { return true; } let Some(scope_stack) = immutable_scope_stack(&scope_stack) else { @@ -494,7 +494,7 @@ mod native { subscribers: &[EventSubscriberFn], scope_stack: ScopeStackHandle, ) -> bool { - if subscribers.is_empty() { + if subscribers.is_empty() && sanitizers.is_empty() { return true; } let Some(scope_stack) = immutable_scope_stack(&scope_stack) else { diff --git a/crates/core/src/api/tool.rs b/crates/core/src/api/tool.rs index dd67a0b53..4f0070a4b 100644 --- a/crates/core/src/api/tool.rs +++ b/crates/core/src/api/tool.rs @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use serde_json::json; +use std::sync::Arc; use crate::api::event::{BaseEvent, Event, MarkEvent, PendingMarkSpec}; use crate::api::runtime::NemoRelayContextState; @@ -11,7 +12,8 @@ use crate::api::runtime::subscriber_dispatcher::{ dispatch_sanitized_event, dispatch_transformed_event, }; use crate::api::runtime::{ - EventSubscriberFn, ScopeStackHandle, ToolExecutionNextFn, with_active_event_uuid, + EventSubscriberFn, ScopeStackHandle, ToolExecutionFrameNextFn, ToolExecutionNextFn, + with_active_event_uuid, }; use crate::api::scope::event; use crate::api::scope::{EmitMarkEventParams, ScopeHandle}; @@ -27,7 +29,11 @@ use serde::{Deserialize, Serialize}; use typed_builder::TypedBuilder; use uuid::Uuid; -pub use nemo_relay_types::api::tool::{ToolAttributes, ToolExecutionInterceptOutcome}; +pub use nemo_relay_types::api::tool::{ + TOOL_EXECUTION_FRAME_OUTCOME_SCHEMA, TOOL_EXECUTION_FRAME_SCHEMA, + TOOL_RESULT_ANNOTATION_PROFILE_KEY, ToolAttributes, ToolExecutionFrame, + ToolExecutionFrameOutcome, ToolExecutionInterceptOutcome, +}; fn queue_sanitized_event(event: Event, subscribers: &[EventSubscriberFn]) -> bool { let scope_stack = current_scope_stack(); @@ -194,6 +200,31 @@ pub struct ToolCallExecuteParams { pub metadata: Option, } +/// Builder parameters for [`tool_call_execute_frame`]. +#[derive(TypedBuilder)] +#[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))] +pub struct ToolCallExecuteFrameParams { + /// Tool name recorded on emitted lifecycle events. + #[builder(setter(into))] + pub name: String, + /// Raw tool arguments passed into the managed pipeline. + pub args: Json, + /// Annotation-aware tool callback or execution continuation. + pub func: ToolExecutionFrameNextFn, + /// Optional explicit parent scope for the emitted tool span. + #[builder(default)] + pub parent: Option, + /// Tool attribute bitflags applied to the managed span. + #[builder(default = ToolAttributes::empty())] + pub attributes: ToolAttributes, + /// Optional application payload stored on the managed tool handle. + #[builder(default)] + pub data: Option, + /// Optional JSON metadata recorded on emitted events. + #[builder(default)] + pub metadata: Option, +} + /// Builder parameters for [`tool_call_end`]. #[derive(TypedBuilder)] #[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))] @@ -216,6 +247,25 @@ pub struct ToolCallEndParams<'a> { pub timestamp: Option>, } +/// Builder parameters for [`tool_call_end_frame`]. +#[derive(TypedBuilder)] +#[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))] +pub struct ToolCallEndFrameParams<'a> { + /// Tool handle to close. + pub handle: &'a ToolHandle, + /// Raw tool result and optional opaque annotation. + pub frame: ToolExecutionFrame, + /// Optional application payload retained for compatibility. + #[builder(default)] + pub data: Option, + /// Optional JSON metadata recorded on the end event. + #[builder(default)] + pub metadata: Option, + /// Optional timestamp recorded on the emitted end event. + #[builder(default)] + pub timestamp: Option>, +} + /// Start a manual tool lifecycle span. /// /// This submits a tool-start event for queued sanitize-request guardrails and @@ -433,6 +483,35 @@ async fn tool_call_with_subscriber_snapshot( /// Sanitize-response guardrails affect only the emitted end-event payload, not /// the caller-owned `result` value. pub fn tool_call_end(params: ToolCallEndParams<'_>) -> Result<()> { + queue_tool_call_end(params, None) +} + +/// Finish a manual tool lifecycle span with an opaque result annotation. +/// +/// The raw frame result is sanitized and emitted exactly like +/// [`tool_call_end`]. The optional annotation is attached to the tool +/// category profile and remains subject to the existing event sanitizer chain. +pub fn tool_call_end_frame(params: ToolCallEndFrameParams<'_>) -> Result<()> { + let ToolCallEndFrameParams { + handle, + frame, + data, + metadata, + timestamp, + } = params; + queue_tool_call_end( + ToolCallEndParams::builder() + .handle(handle) + .result(frame.result) + .data_opt(data) + .metadata_opt(metadata) + .timestamp_opt(timestamp) + .build(), + frame.annotation, + ) +} + +fn queue_tool_call_end(params: ToolCallEndParams<'_>, annotation: Option) -> Result<()> { ensure_runtime_owner()?; let scope_stack = current_scope_stack(); let (entries, subscribers) = { @@ -470,6 +549,7 @@ pub fn tool_call_end(params: ToolCallEndParams<'_>) -> Result<()> { ) }; let tool_name = params.handle.name.clone(); + let annotation = annotation.filter(|value| !value.is_null()); let event_sanitizers = snapshot_event_sanitizers(&event, &scope_stack).unwrap_or_default(); dispatch_transformed_event( event, @@ -486,6 +566,13 @@ pub fn tool_call_end(params: ToolCallEndParams<'_>) -> Result<()> { Some(sanitized) }; event.apply_sanitize_fields(fields); + if let Some(annotation) = annotation + && let Some(profile) = event.category_profile_mut() + { + profile + .extra + .insert(TOOL_RESULT_ANNOTATION_PROFILE_KEY.into(), annotation); + } event }) }), @@ -498,6 +585,7 @@ pub fn tool_call_end(params: ToolCallEndParams<'_>) -> Result<()> { async fn tool_call_end_with_pending_marks( params: ToolCallEndParams<'_>, + annotation: Option, pending_marks: Vec, lifecycle_subscribers: Option<&[EventSubscriberFn]>, ) -> Result<()> { @@ -532,7 +620,7 @@ async fn tool_call_end_with_pending_marks( } else { Some(sanitized_result) }; - let event = { + let mut event = { let context = global_context(); let state = context .read() @@ -546,6 +634,13 @@ async fn tool_call_end_with_pending_marks( .build(), ) }; + if let Some(annotation) = annotation.filter(|value| !value.is_null()) + && let Some(profile) = event.category_profile_mut() + { + profile + .extra + .insert(TOOL_RESULT_ANNOTATION_PROFILE_KEY.into(), annotation); + } let marks = pending_marks .into_iter() .enumerate() @@ -674,6 +769,42 @@ pub async fn tool_call_execute(params: ToolCallExecuteParams) -> Result { data, metadata, } = params; + let frame_func: ToolExecutionFrameNextFn = Arc::new(move |args| { + let func = func.clone(); + Box::pin(async move { func(args).await.map(ToolExecutionFrame::new) }) + }); + let frame = tool_call_execute_frame( + ToolCallExecuteFrameParams::builder() + .name(name) + .args(args) + .func(frame_func) + .parent_opt(parent) + .attributes(attributes) + .data_opt(data) + .metadata_opt(metadata) + .build(), + ) + .await?; + Ok(frame.result) +} + +/// Execute a tool call while carrying an opaque result annotation. +/// +/// This uses the same guardrail and intercept registries as +/// [`tool_call_execute`]. Annotation-aware and raw-JSON execution intercepts +/// are resolved in one priority order. +pub async fn tool_call_execute_frame( + params: ToolCallExecuteFrameParams, +) -> Result { + let ToolCallExecuteFrameParams { + name, + args, + func, + parent, + attributes, + data, + metadata, + } = params; ensure_runtime_owner()?; { let (entries, subscribers, parent_uuid, guardrail_metadata) = { @@ -778,24 +909,25 @@ pub async fn tool_call_execute(params: ToolCallExecuteParams) -> Result { .await; match execution { Ok(outcome) => { - let ToolExecutionInterceptOutcome { - result, + let ToolExecutionFrameOutcome { + frame, pending_marks, } = outcome; let end_metadata = metadata_with_otel_status(metadata, "OK", None); tool_call_end_with_pending_marks( ToolCallEndParams::builder() .handle(&handle) - .result(result.clone()) + .result(frame.result.clone()) .data_opt(data) .metadata_opt(end_metadata) .build(), + frame.annotation.clone(), pending_marks, Some(&lifecycle_subscribers), ) .await?; completion.disarm(); - Ok(result) + Ok(frame) } Err(error) => { let end_metadata = diff --git a/crates/core/src/context/registries.rs b/crates/core/src/context/registries.rs index 098d818c8..9466516f7 100644 --- a/crates/core/src/context/registries.rs +++ b/crates/core/src/context/registries.rs @@ -10,10 +10,11 @@ use std::collections::HashMap; use crate::api::registry::{ExecutionIntercept, Guardrail, Intercept}; +use crate::api::runtime::callbacks::ToolExecutionCallable; use crate::api::runtime::{ EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionFn, LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionFn, ToolConditionalFn, - ToolExecutionFn, ToolInterceptFn, ToolSanitizeFn, + ToolInterceptFn, ToolSanitizeFn, }; use crate::registry::SortedRegistry; @@ -40,7 +41,7 @@ pub(crate) struct ScopeLocalRegistries { /// Tool request intercepts that can rewrite arguments before execution. pub(crate) tool_request_intercepts: SortedRegistry>, /// Tool execution intercepts that wrap or replace callback execution. - pub(crate) tool_execution_intercepts: SortedRegistry>, + pub(crate) tool_execution_intercepts: SortedRegistry>, /// LLM request sanitizers applied to emitted LLM-start payloads. pub(crate) llm_sanitize_request_guardrails: SortedRegistry>, /// LLM response sanitizers applied to emitted LLM-end payloads. diff --git a/crates/core/src/observability/atif.rs b/crates/core/src/observability/atif.rs index 9d5402e52..0e84cf02f 100644 --- a/crates/core/src/observability/atif.rs +++ b/crates/core/src/observability/atif.rs @@ -560,10 +560,13 @@ fn observation_content_value(value: &Json) -> Option { fn observation_extra(event: &Event, output: &Json) -> Json { let mut extra = event_extra(event); - if let Some(tool_result) = observation_tool_result_extra(output) - && let Json::Object(extra_object) = &mut extra - { - extra_object.insert("tool_result".to_string(), tool_result); + if let Json::Object(extra_object) = &mut extra { + if let Some(tool_result) = observation_tool_result_extra(output) { + extra_object.insert("tool_result".to_string(), tool_result); + } + if let Some(annotation) = super::tool_result_annotation(event) { + extra_object.insert("tool_result_annotation".to_string(), annotation.clone()); + } } extra } diff --git a/crates/core/src/observability/mod.rs b/crates/core/src/observability/mod.rs index 55f29400c..f2a97b049 100644 --- a/crates/core/src/observability/mod.rs +++ b/crates/core/src/observability/mod.rs @@ -284,6 +284,37 @@ pub(crate) fn push_top_level_json_attributes( } } +/// Projects an optional opaque tool-result annotation as one JSON-valued +/// attribute. +/// +/// Unlike [`push_top_level_json_attributes`], this helper deliberately does +/// not inspect or flatten the annotation. Serializing the complete value keeps +/// its schema application-owned and makes scalar, array, and object annotations +/// distinguishable to downstream consumers. +pub(crate) fn push_tool_result_annotation_attribute( + attributes: &mut Vec, + event: &crate::api::event::Event, +) { + use opentelemetry::KeyValue; + + let Some(annotation) = tool_result_annotation(event) else { + return; + }; + if let Ok(value) = serde_json::to_string(annotation) { + attributes.push(KeyValue::new("nemo_relay.tool.result.annotation", value)); + } +} + +/// Returns the opaque result annotation carried on a tool lifecycle event. +pub(crate) fn tool_result_annotation( + event: &crate::api::event::Event, +) -> Option<&crate::json::Json> { + event + .category_profile()? + .extra + .get(crate::api::tool::TOOL_RESULT_ANNOTATION_PROFILE_KEY) +} + /// Adds canonical session-correlation attributes from event metadata and the /// active scope-stack instance. pub(crate) fn push_session_identity_attributes( diff --git a/crates/core/src/observability/openinference.rs b/crates/core/src/observability/openinference.rs index de0b32b8b..15d2235ab 100644 --- a/crates/core/src/observability/openinference.rs +++ b/crates/core/src/observability/openinference.rs @@ -8,7 +8,7 @@ use super::{ estimate_cost_for_response_or_model, estimate_cost_for_response_or_requested_model, manual, merge_usage, model_name_for_llm_event, push_serialized_top_level_attributes, - push_top_level_json_attributes, + push_tool_result_annotation_attribute, push_top_level_json_attributes, }; use crate::api::event::{Event, EventNormalizationExt}; use crate::api::scope::ScopeType; @@ -126,6 +126,7 @@ pub(super) fn end_attributes(event: &Event) -> Vec { } push_top_level_json_attributes(&mut attributes, "openinference.metadata", event.metadata()); push_top_level_json_attributes(&mut attributes, "nemo_relay.end.output", event.output()); + push_tool_result_annotation_attribute(&mut attributes, event); if let Some((output, mime_type)) = openinference_output_value(event) { attributes.push(KeyValue::new("output.value", output)); attributes.push(KeyValue::new("output.mime_type", mime_type)); diff --git a/crates/core/src/observability/otel.rs b/crates/core/src/observability/otel.rs index ccd428a39..b9ed28427 100644 --- a/crates/core/src/observability/otel.rs +++ b/crates/core/src/observability/otel.rs @@ -27,7 +27,8 @@ use super::{ effective_mark_projection, estimate_cost_for_response_or_model, estimate_cost_for_response_or_requested_model, manual, model_name_for_llm_event, push_serialized_top_level_attributes, push_session_identity_attributes, - push_top_level_json_attributes, relay_span_id, relay_trace_id, validate_attribute_mappings, + push_tool_result_annotation_attribute, push_top_level_json_attributes, relay_span_id, + relay_trace_id, validate_attribute_mappings, }; use crate::api::event::{Event, EventNormalizationExt, ScopeCategory}; use crate::api::runtime::{EventSubscriberFn, current_scope_stack}; @@ -1235,6 +1236,7 @@ fn end_attributes(event: &Event) -> Vec { push_top_level_json_attributes(&mut attributes, "nemo_relay.end.data", event.data()); push_top_level_json_attributes(&mut attributes, "nemo_relay.end.metadata", event.metadata()); push_top_level_json_attributes(&mut attributes, "nemo_relay.end.output", event.output()); + push_tool_result_annotation_attribute(&mut attributes, event); if event .category() .is_some_and(|category| category.as_str() == "llm") diff --git a/crates/core/src/plugin.rs b/crates/core/src/plugin.rs index 8d298e396..3e7c500ec 100644 --- a/crates/core/src/plugin.rs +++ b/crates/core/src/plugin.rs @@ -35,13 +35,14 @@ use crate::api::registry::{ register_llm_sanitize_response_guardrail, register_llm_stream_execution_intercept, register_mark_sanitize_guardrail, register_scope_sanitize_end_guardrail, register_scope_sanitize_start_guardrail, register_tool_conditional_execution_guardrail, - register_tool_execution_intercept, register_tool_request_intercept, - register_tool_sanitize_request_guardrail, register_tool_sanitize_response_guardrail, + register_tool_execution_frame_intercept, register_tool_execution_intercept, + register_tool_request_intercept, register_tool_sanitize_request_guardrail, + register_tool_sanitize_response_guardrail, }; use crate::api::runtime::{ EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionFn, LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionFn, ToolConditionalFn, - ToolExecutionFn, ToolInterceptFn, ToolSanitizeFn, + ToolExecutionFn, ToolExecutionFrameFn, ToolInterceptFn, ToolSanitizeFn, }; use crate::api::subscriber::{deregister_subscriber, register_subscriber}; pub use nemo_relay_types::plugin::{ConfigDiagnostic, DiagnosticLevel}; @@ -810,6 +811,35 @@ impl PluginRegistrationContext { Ok(()) } + /// Registers an annotation-aware tool execution intercept and records its rollback closure. + pub fn register_tool_execution_frame_intercept( + &mut self, + name: &str, + priority: i32, + callback: ToolExecutionFrameFn, + ) -> Result<()> { + let qualified_name = self.qualify_name(name); + register_tool_execution_frame_intercept(&qualified_name, priority, callback).map_err( + |err| PluginError::RegistrationFailed(format!("tool execution frame intercept: {err}")), + )?; + + let name_owned = qualified_name; + self.registrations.push(PluginRegistration::new( + "plugin", + name_owned.clone(), + Box::new(move || { + deregister_tool_execution_intercept(&name_owned) + .map(|_| ()) + .map_err(|err| { + PluginError::RegistrationFailed(format!( + "tool execution frame intercept deregistration failed: {err}" + )) + }) + }), + )); + Ok(()) + } + /// Adds a prebuilt registration to the context. pub fn add_registration(&mut self, registration: PluginRegistration) { self.registrations.push(registration); diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index d6fde6414..dbd2dc25e 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -28,7 +28,8 @@ use crate::api::runtime::{ LlmExecutionNextFn, LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn, LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionFn, LlmStreamExecutionNextFn, MiddlewareContinuationContext, ToolConditionalFn, ToolExecutionFn, - ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, + ToolExecutionFrameFn, ToolExecutionFrameNextFn, ToolExecutionNextFn, ToolInterceptFn, + ToolSanitizeFn, }; use crate::api::runtime::{ ScopeStackHandle, ThreadScopeStackBinding, capture_thread_scope_stack, create_scope_stack, @@ -38,7 +39,7 @@ use crate::api::scope::{ EmitMarkEventParams, PopScopeParams, PushScopeParams, ScopeAttributes, ScopeHandle, ScopeType, }; use crate::api::scope::{event as emit_scope_mark, get_handle, pop_scope, push_scope}; -use crate::api::tool::ToolExecutionInterceptOutcome; +use crate::api::tool::{ToolExecutionFrameOutcome, ToolExecutionInterceptOutcome}; use crate::codec::request::AnnotatedLlmRequest; use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::error::{FlowError, Result as FlowResult}; @@ -55,8 +56,8 @@ use nemo_relay_plugin::{ NemoRelayNativeAsyncNextResultCb, NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, NemoRelayNativeAsyncStreamMiddlewareCb, NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, - NemoRelayNativeHostApiV3, NemoRelayNativeLlmCodecKind, NemoRelayNativeLlmConditionalCb, - NemoRelayNativeLlmExecutionCb, NemoRelayNativeLlmRequestCodec, + NemoRelayNativeHostApiV3, NemoRelayNativeHostApiV3ToolFrames, NemoRelayNativeLlmCodecKind, + NemoRelayNativeLlmConditionalCb, NemoRelayNativeLlmExecutionCb, NemoRelayNativeLlmRequestCodec, NemoRelayNativeLlmRequestInterceptCb, NemoRelayNativeLlmResponseCodec, NemoRelayNativeLlmSanitizeRequestCb, NemoRelayNativeLlmSanitizeRequestContext, NemoRelayNativeLlmSanitizeResponseCb, NemoRelayNativeLlmSanitizeResponseContext, @@ -64,7 +65,8 @@ use nemo_relay_plugin::{ NemoRelayNativePluginEntry, NemoRelayNativePluginV1, NemoRelayNativeScopeHandle, NemoRelayNativeScopeStack, NemoRelayNativeScopeStackBinding, NemoRelayNativeScopeType, NemoRelayNativeString, NemoRelayNativeToolConditionalCb, NemoRelayNativeToolExecutionCb, - NemoRelayNativeToolJsonCb, NemoRelayNativeWithScopeStackCb, NemoRelayStatus, + NemoRelayNativeToolExecutionFrameCb, NemoRelayNativeToolJsonCb, + NemoRelayNativeWithScopeStackCb, NemoRelayStatus, }; use semver::{Version, VersionReq}; use serde_json::{Map, Value as Json}; @@ -797,8 +799,8 @@ unsafe extern "C" fn native_llm_response_codec_decode( } fn native_host_api() -> *const NemoRelayNativeHostApiV1 { - static HOST_API: OnceLock = OnceLock::new(); - &HOST_API.get_or_init(build_native_host_api_v3).v1 as *const NemoRelayNativeHostApiV1 + static HOST_API: OnceLock = OnceLock::new(); + &HOST_API.get_or_init(build_native_host_api_v3).v3.v1 as *const NemoRelayNativeHostApiV1 } fn native_host_api_legacy() -> *const NemoRelayNativeHostApiV1 { @@ -866,28 +868,33 @@ fn build_native_host_api_legacy() -> NemoRelayNativeHostApiV1 { } } -fn build_native_host_api_v3() -> NemoRelayNativeHostApiV3 { +fn build_native_host_api_v3() -> NemoRelayNativeHostApiV3ToolFrames { let mut v1 = build_native_host_api_legacy(); v1.abi_version = NEMO_RELAY_NATIVE_ABI_VERSION; - v1.struct_size = std::mem::size_of::(); - NemoRelayNativeHostApiV3 { - v1, - async_completion_resolve_json: native_async_completion_resolve_json, - async_completion_reject: native_async_completion_reject, - async_completion_is_cancelled: native_async_completion_is_cancelled, - async_completion_release: native_async_completion_release, - async_next_invoke: native_async_next_invoke, - async_next_release: native_async_next_release, - plugin_context_register_async_middleware: native_plugin_context_register_async_middleware, - async_stream_push_json: native_async_stream_push_json, - async_stream_finish: native_async_stream_finish, - async_stream_reject: native_async_stream_reject, - async_stream_is_cancelled: native_async_stream_is_cancelled, - async_stream_release: native_async_stream_release, - async_next_invoke_stream: native_async_next_invoke_stream, - plugin_context_register_async_stream_middleware: - native_plugin_context_register_async_stream_middleware, - async_next_invoke_result: native_async_next_invoke_result, + v1.struct_size = std::mem::size_of::(); + NemoRelayNativeHostApiV3ToolFrames { + v3: NemoRelayNativeHostApiV3 { + v1, + async_completion_resolve_json: native_async_completion_resolve_json, + async_completion_reject: native_async_completion_reject, + async_completion_is_cancelled: native_async_completion_is_cancelled, + async_completion_release: native_async_completion_release, + async_next_invoke: native_async_next_invoke, + async_next_release: native_async_next_release, + plugin_context_register_async_middleware: + native_plugin_context_register_async_middleware, + async_stream_push_json: native_async_stream_push_json, + async_stream_finish: native_async_stream_finish, + async_stream_reject: native_async_stream_reject, + async_stream_is_cancelled: native_async_stream_is_cancelled, + async_stream_release: native_async_stream_release, + async_next_invoke_stream: native_async_next_invoke_stream, + plugin_context_register_async_stream_middleware: + native_plugin_context_register_async_stream_middleware, + async_next_invoke_result: native_async_next_invoke_result, + }, + plugin_context_register_tool_execution_frame_intercept: + native_plugin_context_register_tool_execution_frame_intercept, } } @@ -2995,6 +3002,35 @@ unsafe extern "C" fn native_plugin_context_register_tool_execution_intercept( } } +unsafe extern "C" fn native_plugin_context_register_tool_execution_frame_intercept( + ctx: *mut NemoRelayNativePluginContext, + name: *const NemoRelayNativeString, + priority: i32, + cb: NemoRelayNativeToolExecutionFrameCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> NemoRelayStatus { + clear_native_last_error(); + let host_ctx = match host_ctx_mut(ctx) { + Ok(ctx) => ctx, + Err(status) => return status, + }; + let instance = host_ctx.instance.clone(); + let ctx = unsafe { &mut *host_ctx.ctx }; + let name = match read_name(name) { + Ok(name) => name, + Err(status) => return status, + }; + match ctx.register_tool_execution_frame_intercept( + &name, + priority, + wrap_tool_execution_frame_fn(instance, cb, user_data, free_fn), + ) { + Ok(()) => NemoRelayStatus::Ok, + Err(err) => status_from_plugin_error(err), + } +} + unsafe extern "C" fn native_plugin_context_register_llm_sanitize_request_guardrail( ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, @@ -3472,6 +3508,93 @@ unsafe extern "C" fn native_tool_next( } } +fn wrap_tool_execution_frame_fn( + instance: Arc, + cb: NemoRelayNativeToolExecutionFrameCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> ToolExecutionFrameFn { + let user_data = make_user_data(instance, user_data, free_fn); + Arc::new(move |name, args, next| { + let name = name.to_owned(); + let user_data = user_data.clone(); + Box::pin(async move { + clear_native_last_error(); + let name_string = native_string_from_str(&name) + .ok_or_else(|| FlowError::Internal("failed to allocate native name".into()))?; + let args_string = native_string_from_json(&args) + .ok_or_else(|| FlowError::Internal("failed to allocate native args".into()))?; + let next_ctx = Box::into_raw(Box::new(next)) as *mut c_void; + let mut out_outcome = ptr::null_mut(); + let status = unsafe { + cb( + user_data.ptr, + name_string, + args_string, + native_tool_frame_next, + next_ctx, + &mut out_outcome, + ) + }; + unsafe { + drop(Box::from_raw(next_ctx as *mut ToolExecutionFrameNextFn)); + native_string_free(name_string); + native_string_free(args_string); + } + if status != NemoRelayStatus::Ok { + if !out_outcome.is_null() { + unsafe { native_string_free(out_outcome) }; + } + return Err(flow_error_from_status( + status, + "native tool execution frame failed", + )); + } + let outcome_json = take_json_from_native_string( + out_outcome, + "native tool execution frame returned null outcome", + )?; + serde_json::from_value::(outcome_json).map_err(|err| { + FlowError::Internal(format!( + "invalid native tool execution frame outcome JSON: {err}" + )) + }) + }) + }) +} + +unsafe extern "C" fn native_tool_frame_next( + args_json: *const NemoRelayNativeString, + next_ctx: *mut c_void, + out_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + if next_ctx.is_null() || out_json.is_null() { + set_native_last_error("native tool frame next received null pointer"); + return NemoRelayStatus::NullPointer; + } + let args = match parse_json_arg(args_json, "native tool frame next args") { + Ok(args) => args, + Err(status) => return status, + }; + let next = unsafe { (*(next_ctx as *const ToolExecutionFrameNextFn)).clone() }; + let context = MiddlewareContinuationContext::capture(); + let result = spawn_with_continuation_context(context, move || next(args)).join(); + match result { + Ok(Ok(frame)) => match serde_json::to_value(frame) { + Ok(frame) => write_native_json(&frame, out_json), + Err(err) => { + set_native_last_error(format!("failed to serialize native tool frame: {err}")); + NemoRelayStatus::Internal + } + }, + Ok(Err(err)) => status_from_flow_error(err), + Err(_) => { + set_native_last_error("native tool frame next panicked"); + NemoRelayStatus::Internal + } + } +} + fn wrap_llm_sanitize_request_fn( instance: Arc, cb: NemoRelayNativeLlmSanitizeRequestCb, diff --git a/crates/core/src/plugin/dynamic/worker.rs b/crates/core/src/plugin/dynamic/worker.rs index b99e861b9..341db0f72 100644 --- a/crates/core/src/plugin/dynamic/worker.rs +++ b/crates/core/src/plugin/dynamic/worker.rs @@ -27,8 +27,8 @@ use nemo_relay_worker_proto::v1::{ LlmSanitizeRequestContext as ProtoLlmSanitizeRequestContext, LlmSanitizeResponseContext as ProtoLlmSanitizeResponseContext, LlmStreamNextRequest, PopScopeRequest, PushScopeRequest, PushScopeResponse, RegisterRequest, RegisterResponse, - Registration, RegistrationSurface, ScopeContext, ShutdownRequest, StreamChunk, ToolInvocation, - ToolNextRequest, ValidateRequest, WorkerError, + Registration, RegistrationSurface, ScopeContext, ShutdownRequest, StreamChunk, + ToolFrameNextRequest, ToolInvocation, ToolNextRequest, ValidateRequest, WorkerError, }; use nemo_relay_worker_proto::{WORKER_PROTOCOL_GRPC_V1, decode_json_envelope, json_envelope}; use semver::{Version, VersionReq}; @@ -66,13 +66,17 @@ use crate::api::runtime::subscriber_dispatcher::{ use crate::api::runtime::{ EventSanitizeFn, LlmCodecIdentity, LlmExecutionNextFn, LlmJsonStream, LlmSanitizeRequestContext, LlmSanitizeResponseContext, LlmStreamExecutionNextFn, - MiddlewareContinuationContext, ToolExecutionNextFn, current_scope_stack, with_scope_stack, + MiddlewareContinuationContext, ToolExecutionFrameNextFn, ToolExecutionNextFn, + current_scope_stack, with_scope_stack, }; use crate::api::scope::{ EmitMarkEventParams, PopScopeParams, PushScopeParams, ScopeAttributes, ScopeHandle, ScopeType, event as emit_scope_mark, pop_scope, push_scope, }; -use crate::api::tool::ToolExecutionInterceptOutcome; +use crate::api::tool::{ + TOOL_EXECUTION_FRAME_OUTCOME_SCHEMA, TOOL_EXECUTION_FRAME_SCHEMA, ToolExecutionFrameOutcome, + ToolExecutionInterceptOutcome, +}; use crate::codec::request::{ANNOTATED_LLM_REQUEST_SCHEMA, AnnotatedLlmRequest}; use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::error::{FlowError, Result as FlowResult}; @@ -1076,7 +1080,8 @@ impl WorkerPluginInstance { | RegistrationSurface::ToolSanitizeResponseGuardrail | RegistrationSurface::ToolConditionalExecutionGuardrail | RegistrationSurface::ToolRequestIntercept - | RegistrationSurface::ToolExecutionIntercept => { + | RegistrationSurface::ToolExecutionIntercept + | RegistrationSurface::ToolExecutionFrameIntercept => { self.install_tool_registration(ctx, registration, surface)? } RegistrationSurface::LlmSanitizeRequestGuardrail @@ -1125,9 +1130,11 @@ impl WorkerPluginInstance { let instance = Arc::new(self.clone_for_callback()); let callback_name = name.to_owned(); let callback: EventSanitizeFn = - Arc::new(move |event: Arc, _fields: EventSanitizeFields| { + Arc::new(move |event: Arc, fields: EventSanitizeFields| { let instance = instance.clone(); let callback_name = callback_name.clone(); + let mut event = (*event).clone(); + event.apply_sanitize_fields(fields); Box::pin(async move { instance .invoke_event_sanitize(&callback_name, surface, &event) @@ -1229,6 +1236,26 @@ impl WorkerPluginInstance { }) }), ), + RegistrationSurface::ToolExecutionFrameIntercept => ctx + .register_tool_execution_frame_intercept( + name, + priority, + Arc::new(move |tool_name, value, next| { + let instance = instance.clone(); + let callback_name = callback_name.clone(); + let tool_name = tool_name.to_owned(); + Box::pin(async move { + instance + .invoke_tool_execution_frame( + &callback_name, + &tool_name, + value, + next, + ) + .await + }) + }), + ), _ => Err(PluginError::RegistrationFailed(format!( "worker plugin '{}' cannot install registration surface {} as a tool callback", self.plugin_kind, @@ -1589,6 +1616,46 @@ impl WorkerPluginCallback { } } + async fn invoke_tool_execution_frame( + &self, + registration_name: &str, + tool_name: &str, + value: Json, + next: ToolExecutionFrameNextFn, + ) -> FlowResult { + let continuation_id = self + .host_state + .insert_continuation(Continuation::tool_frame(next))?; + let request = self.base_request( + registration_name, + RegistrationSurface::ToolExecutionFrameIntercept, + Some(continuation_id), + Some(invoke_request_payload_tool(tool_name, value)), + ); + let response = self.invoke_async(request).await?; + match response.result { + Some(invoke_response_result::Result::ToolExecutionFrame(result)) => { + let outcome = + required_envelope(result.outcome, "tool execution frame intercept outcome")?; + if outcome.schema != TOOL_EXECUTION_FRAME_OUTCOME_SCHEMA { + return Err(FlowError::Internal(format!( + "worker returned unsupported tool execution frame intercept outcome schema: {}", + outcome.schema + ))); + } + decode_json_envelope(&outcome).map_err(|err| { + FlowError::Internal(format!( + "worker returned invalid tool execution frame intercept outcome: {err}" + )) + }) + } + Some(invoke_response_result::Result::Error(error)) => Err(worker_error_to_flow(error)), + _ => Err(FlowError::Internal( + "worker tool execution frame intercept returned unexpected result".into(), + )), + } + } + async fn invoke_llm_sanitize_request( &self, registration_name: &str, @@ -2437,6 +2504,10 @@ enum Continuation { next: ToolExecutionNextFn, context: MiddlewareContinuationContext, }, + ToolFrame { + next: ToolExecutionFrameNextFn, + context: MiddlewareContinuationContext, + }, Llm { next: LlmExecutionNextFn, context: MiddlewareContinuationContext, @@ -2455,6 +2526,13 @@ impl Continuation { } } + fn tool_frame(next: ToolExecutionFrameNextFn) -> Self { + Self::ToolFrame { + next, + context: MiddlewareContinuationContext::capture(), + } + } + fn llm(next: LlmExecutionNextFn) -> Self { Self::Llm { next, @@ -2656,6 +2734,40 @@ impl RelayHostRuntime for WorkerHostRuntimeService { Ok(Response::new(json_result(result))) } + async fn tool_frame_next( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + self.state + .authorize(&request.activation_id, &request.auth_token)?; + let continuation = self.state.continuation(&request.continuation_id)?; + let Continuation::ToolFrame { next, context } = continuation else { + return Err(Status::invalid_argument( + "continuation is not a tool frame continuation", + )); + }; + let context = self.isolated_continuation_context(&context, request.scope.as_ref())?; + let value = + required_envelope(request.value, "tool frame next value").map_err(status_from_flow)?; + let value = decode_json_envelope::(&value).map_err(|err| { + Status::invalid_argument(format!("invalid tool frame next JSON: {err}")) + })?; + let result = AssertUnwindSafe(context.invoke(move || next(value))) + .catch_unwind() + .await + .unwrap_or_else(|payload| { + Err(FlowError::Internal(format!( + "worker tool frame continuation panicked: {}", + panic_payload_message(payload.as_ref()) + ))) + }); + Ok(Response::new(typed_json_result( + TOOL_EXECUTION_FRAME_SCHEMA, + result, + ))) + } + async fn llm_next( &self, request: Request, diff --git a/crates/core/tests/fixtures/native_plugin/src/lib.rs b/crates/core/tests/fixtures/native_plugin/src/lib.rs index a8c47a1bc..4013d3c2a 100644 --- a/crates/core/tests/fixtures/native_plugin/src/lib.rs +++ b/crates/core/tests/fixtures/native_plugin/src/lib.rs @@ -13,7 +13,7 @@ use nemo_relay_plugin::{ NemoRelayNativeHostApiV1, NemoRelayNativeHostApiV3, NemoRelayNativePluginContext, NemoRelayNativePluginV1, NemoRelayNativeString, NemoRelayNativeToolNextFn, NemoRelayStatus, PendingMarkSpec, PluginContext, PluginRuntime, ScopeCategory, ScopeType, - ToolExecutionInterceptOutcome, + ToolExecutionFrameOutcome, ToolExecutionInterceptOutcome, }; use serde_json::{Map, json}; @@ -94,7 +94,8 @@ impl NativePlugin for FixtureNativePlugin { })?; ctx.register_tool_execution_intercept("fixture_tool_execution", 0, { let runtime = runtime.clone(); - move |_name, args, next| { + move |name, args, next| { + let preserve_result = name == "native-fixture-tool-frame"; let args = mark_json(args, "native_plugin_tool_execution_request"); let result = if args .get("use_isolated_next") @@ -118,7 +119,11 @@ impl NativePlugin for FixtureNativePlugin { } else { next.call(args)? }; - let result = mark_json(result, "native_plugin_tool_execution"); + let result = if preserve_result { + result + } else { + mark_json(result, "native_plugin_tool_execution") + }; Ok( ToolExecutionInterceptOutcome::new(result).with_pending_mark( PendingMarkSpec::builder() @@ -135,6 +140,28 @@ impl NativePlugin for FixtureNativePlugin { ) } })?; + ctx.register_tool_execution_frame_intercept( + "fixture_tool_execution_frame", + -1, + |name, args, next| { + let mark_name = if name == "native-fixture-tool-frame" { + "fixture.native.tool_execution_frame.mark" + } else { + "fixture.native.tool_execution_frame.raw.mark" + }; + let mut frame = next.call(mark_json( + args, + "native_plugin_tool_execution_frame_request", + ))?; + frame.result = mark_json(frame.result, "native_plugin_tool_execution_frame"); + let mut annotation = frame.annotation.take().unwrap_or_else(|| json!({})); + annotation["native_plugin_tool_execution_frame"] = json!(true); + frame.annotation = Some(annotation); + Ok(ToolExecutionFrameOutcome::new(frame).with_pending_mark( + PendingMarkSpec::builder().name(mark_name).build(), + )) + }, + )?; ctx.register_llm_sanitize_request_guardrail( "fixture_llm_sanitize_request", diff --git a/crates/core/tests/fixtures/worker_plugin/src/main.rs b/crates/core/tests/fixtures/worker_plugin/src/main.rs index a2276e123..9f9741d9a 100644 --- a/crates/core/tests/fixtures/worker_plugin/src/main.rs +++ b/crates/core/tests/fixtures/worker_plugin/src/main.rs @@ -6,7 +6,7 @@ use nemo_relay_worker::{ }; use nemo_relay_worker::{ JsonStream, LlmNext, LlmStreamNext, PluginContext, ScopeType, ToolExecutionInterceptOutcome, - ToolNext, WorkerPlugin, WorkerSdkError, serve_plugin, + ToolExecutionFrameOutcome, ToolFrameNext, ToolNext, WorkerPlugin, WorkerSdkError, serve_plugin, }; use serde_json::json; @@ -195,19 +195,50 @@ fn register_fixture_tool_hooks( ctx.register_tool_execution_intercept( "fixture_tool_execution", 0, - |_name, args, next: ToolNext| async move { + |name, args, next: ToolNext| { + let preserve_result = name == "worker-fixture-tool-frame"; + async move { let result = next .call(mark_json(args, "worker_plugin_tool_execution_request")) .await?; - Ok(ToolExecutionInterceptOutcome::new(mark_json( - result, - "worker_plugin_tool_execution", - )) + let result = if preserve_result { + result + } else { + mark_json(result, "worker_plugin_tool_execution") + }; + Ok(ToolExecutionInterceptOutcome::new(result) .with_pending_mark( PendingMarkSpec::builder() .name("fixture.worker.tool_execution.mark") .build(), )) + } + }, + ); + ctx.register_tool_execution_frame_intercept( + "fixture_tool_execution_frame", + -1, + |name, args, next: ToolFrameNext| { + let mark_name = if name == "worker-fixture-tool-frame" { + "fixture.worker.tool_execution_frame.mark" + } else { + "fixture.worker.tool_execution_frame.raw.mark" + }; + async move { + let mut frame = next + .call(mark_json( + args, + "worker_plugin_tool_execution_frame_request", + )) + .await?; + frame.result = mark_json(frame.result, "worker_plugin_tool_execution_frame"); + let mut annotation = frame.annotation.take().unwrap_or_else(|| json!({})); + annotation["worker_plugin_tool_execution_frame"] = json!(true); + frame.annotation = Some(annotation); + Ok(ToolExecutionFrameOutcome::new(frame).with_pending_mark( + PendingMarkSpec::builder().name(mark_name).build(), + )) + } }, ); } diff --git a/crates/core/tests/integration/middleware_tests.rs b/crates/core/tests/integration/middleware_tests.rs index 8fd5810ae..d71804e0d 100644 --- a/crates/core/tests/integration/middleware_tests.rs +++ b/crates/core/tests/integration/middleware_tests.rs @@ -42,12 +42,13 @@ use nemo_relay::api::registry::{ register_llm_sanitize_response_guardrail, register_llm_stream_execution_intercept, register_mark_sanitize_guardrail, register_scope_sanitize_end_guardrail, register_scope_sanitize_start_guardrail, register_tool_conditional_execution_guardrail, - register_tool_execution_intercept, register_tool_request_intercept, - register_tool_sanitize_request_guardrail, register_tool_sanitize_response_guardrail, - scope_register_llm_conditional_execution_guardrail, scope_register_llm_execution_intercept, - scope_register_llm_request_intercept, scope_register_llm_sanitize_request_guardrail, - scope_register_llm_sanitize_response_guardrail, scope_register_llm_stream_execution_intercept, - scope_register_mark_sanitize_guardrail, scope_register_scope_sanitize_end_guardrail, + register_tool_execution_frame_intercept, register_tool_execution_intercept, + register_tool_request_intercept, register_tool_sanitize_request_guardrail, + register_tool_sanitize_response_guardrail, scope_register_llm_conditional_execution_guardrail, + scope_register_llm_execution_intercept, scope_register_llm_request_intercept, + scope_register_llm_sanitize_request_guardrail, scope_register_llm_sanitize_response_guardrail, + scope_register_llm_stream_execution_intercept, scope_register_mark_sanitize_guardrail, + scope_register_scope_sanitize_end_guardrail, scope_register_tool_conditional_execution_guardrail, scope_register_tool_execution_intercept, scope_register_tool_request_intercept, scope_register_tool_sanitize_request_guardrail, scope_register_tool_sanitize_response_guardrail, @@ -56,15 +57,16 @@ use nemo_relay::api::runtime::NemoRelayContextState; use nemo_relay::api::runtime::global_context; use nemo_relay::api::runtime::{ LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, LlmStreamInner, TASK_SCOPE_STACK, - ToolExecutionNextFn, capture_propagation_context, task_scope_top, + ToolExecutionFrameNextFn, ToolExecutionNextFn, capture_propagation_context, task_scope_top, }; use nemo_relay::api::runtime::{create_scope_stack, current_scope_stack, set_thread_scope_stack}; use nemo_relay::api::scope::{EmitMarkEventParams, ScopeHandle, ScopeType, event}; use nemo_relay::api::scope::{pop_scope, push_scope}; use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; use nemo_relay::api::tool::{ + TOOL_RESULT_ANNOTATION_PROFILE_KEY, ToolExecutionFrame, ToolExecutionFrameOutcome, ToolExecutionInterceptOutcome, tool_call, tool_call_end, tool_call_execute, - tool_conditional_execution, tool_request_intercepts, + tool_call_execute_frame, tool_conditional_execution, tool_request_intercepts, }; use nemo_relay::codec::optimization::{ LlmOptimizationContribution, LlmOptimizationEvidenceQuality, LlmOptimizationTokenImpact, @@ -708,6 +710,567 @@ async fn test_execution_intercept_modifies_args() { deregister_tool_execution_intercept("arg_modifier").unwrap(); } +#[tokio::test] +async fn test_frame_and_legacy_intercepts_share_one_chain_and_preserve_annotation() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + let order = Arc::new(Mutex::new(Vec::<&'static str>::new())); + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "tool_frame_observer", + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + + let outer_order = order.clone(); + register_tool_execution_frame_intercept( + "frame_outer", + 1, + Arc::new(move |_name, args, next| { + let order = outer_order.clone(); + Box::pin(async move { + order.lock().unwrap().push("frame_before"); + let mut frame = next(args).await?; + assert_eq!( + frame.annotation.as_ref().unwrap()["producer_status"], + "failed" + ); + frame.result["frame_seen"] = json!(true); + frame.annotation.as_mut().unwrap()["observed_by"] = json!("frame_outer"); + order.lock().unwrap().push("frame_after"); + Ok(ToolExecutionFrameOutcome::new(frame)) + }) + }), + ) + .unwrap(); + + let legacy_order = order.clone(); + register_tool_execution_intercept( + "legacy_middle", + 2, + Arc::new(move |_name, args, next| { + let order = legacy_order.clone(); + Box::pin(async move { + order.lock().unwrap().push("legacy_before"); + let result = next(args).await?; + order.lock().unwrap().push("legacy_after"); + Ok(ToolExecutionInterceptOutcome::new(result)) + }) + }), + ) + .unwrap(); + + let producer_order = order.clone(); + let producer: ToolExecutionFrameNextFn = Arc::new(move |_args| { + let order = producer_order.clone(); + Box::pin(async move { + order.lock().unwrap().push("producer"); + Ok(ToolExecutionFrame::annotated( + json!({"raw": true}), + json!({ + "producer_status": "failed", + "representation": { + "media_type": "application/json", + "data_schema": "example.failure@1" + }, + "tool_call_id": "producer-call" + }), + )) + }) + }); + + let frame = tool_call_execute_frame( + nemo_relay::api::tool::ToolCallExecuteFrameParams::builder() + .name("annotated-tool") + .args(json!({})) + .func(producer) + .build(), + ) + .await + .unwrap(); + + assert_eq!( + *order.lock().unwrap(), + vec![ + "frame_before", + "legacy_before", + "producer", + "legacy_after", + "frame_after", + ] + ); + assert_eq!(frame.result, json!({"raw": true, "frame_seen": true})); + let annotation = frame.annotation.as_ref().unwrap(); + assert_eq!(annotation["observed_by"], "frame_outer"); + + let captured = captured_events_snapshot(&events); + let end = captured + .iter() + .find(|event| { + event.name() == "annotated-tool" && event.scope_category() == Some(ScopeCategory::End) + }) + .unwrap(); + assert_eq!( + end.category_profile().unwrap().extra[TOOL_RESULT_ANNOTATION_PROFILE_KEY]["producer_status"], + "failed" + ); + assert_eq!( + end.category_profile().unwrap().extra[TOOL_RESULT_ANNOTATION_PROFILE_KEY]["representation"] + ["data_schema"], + "example.failure@1" + ); + + deregister_tool_execution_intercept("frame_outer").unwrap(); + deregister_tool_execution_intercept("legacy_middle").unwrap(); + deregister_subscriber("tool_frame_observer").unwrap(); +} + +#[tokio::test] +async fn test_frame_intercept_explicitly_preserves_replaces_and_removes_annotation() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + let producer = || -> ToolExecutionFrameNextFn { + Arc::new(|_args| { + Box::pin(async move { + Ok(ToolExecutionFrame::annotated( + json!({"raw": true}), + json!({"source": "producer"}), + )) + }) + }) + }; + + register_tool_execution_frame_intercept( + "frame_preserve", + 1, + Arc::new(|_name, args, next| { + Box::pin(async move { next(args).await.map(ToolExecutionFrameOutcome::new) }) + }), + ) + .unwrap(); + let preserved = tool_call_execute_frame( + nemo_relay::api::tool::ToolCallExecuteFrameParams::builder() + .name("preserved-tool") + .args(json!({})) + .func(producer()) + .build(), + ) + .await + .unwrap(); + assert_eq!(preserved.annotation, Some(json!({"source": "producer"}))); + deregister_tool_execution_intercept("frame_preserve").unwrap(); + + register_tool_execution_frame_intercept( + "frame_replace", + 1, + Arc::new(|_name, args, next| { + Box::pin(async move { + let frame = next(args) + .await? + .with_annotation(json!({"source": "middleware"})); + Ok(ToolExecutionFrameOutcome::new(frame)) + }) + }), + ) + .unwrap(); + let replaced = tool_call_execute_frame( + nemo_relay::api::tool::ToolCallExecuteFrameParams::builder() + .name("replaced-tool") + .args(json!({})) + .func(producer()) + .build(), + ) + .await + .unwrap(); + assert_eq!(replaced.annotation, Some(json!({"source": "middleware"}))); + deregister_tool_execution_intercept("frame_replace").unwrap(); + + register_tool_execution_frame_intercept( + "frame_remove", + 1, + Arc::new(|_name, args, next| { + Box::pin(async move { + let frame = next(args).await?.without_annotation(); + Ok(ToolExecutionFrameOutcome::new(frame)) + }) + }), + ) + .unwrap(); + let removed = tool_call_execute_frame( + nemo_relay::api::tool::ToolCallExecuteFrameParams::builder() + .name("removed-tool") + .args(json!({})) + .func(producer()) + .build(), + ) + .await + .unwrap(); + assert!(removed.annotation.is_none()); + deregister_tool_execution_intercept("frame_remove").unwrap(); +} + +#[tokio::test] +async fn test_tool_result_annotation_uses_existing_event_sanitizer_chain() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "tool_annotation_sanitizer_observer", + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + register_scope_sanitize_end_guardrail( + "tool_annotation_sanitizer", + 1, + Arc::new(|event, mut fields| { + Box::pin(async move { + if event.name() == "sanitized-annotation-tool" + && let Some(profile) = fields.category_profile.as_mut() + && let Some(annotation) = profile + .extra + .get_mut(TOOL_RESULT_ANNOTATION_PROFILE_KEY) + .and_then(Json::as_object_mut) + { + annotation.insert("secret".into(), json!("[redacted]")); + annotation.remove("future_secret"); + } + Ok(fields) + }) + }), + ) + .unwrap(); + + let frame = tool_call_execute_frame( + nemo_relay::api::tool::ToolCallExecuteFrameParams::builder() + .name("sanitized-annotation-tool") + .args(json!({})) + .func(Arc::new(|_args| { + Box::pin(async move { + Ok(ToolExecutionFrame::annotated( + json!({"raw": true}), + json!({ + "secret": "classified", + "future_secret": "classified" + }), + )) + }) + })) + .build(), + ) + .await + .unwrap(); + + assert_eq!(frame.annotation.as_ref().unwrap()["secret"], "classified"); + assert_eq!( + frame.annotation.as_ref().unwrap()["future_secret"], + "classified" + ); + let captured = captured_events_snapshot(&events); + let end = captured + .iter() + .find(|event| { + event.name() == "sanitized-annotation-tool" + && event.scope_category() == Some(ScopeCategory::End) + }) + .unwrap(); + let annotation = &end.category_profile().unwrap().extra[TOOL_RESULT_ANNOTATION_PROFILE_KEY]; + assert_eq!(annotation["secret"], "[redacted]"); + assert!(annotation.get("future_secret").is_none()); + + deregister_scope_sanitize_end_guardrail("tool_annotation_sanitizer").unwrap(); + deregister_subscriber("tool_annotation_sanitizer_observer").unwrap(); +} + +#[tokio::test] +async fn test_legacy_result_mutation_invalidates_downstream_annotation() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + register_tool_execution_intercept( + "legacy_mutator", + 1, + Arc::new(|_name, args, next| { + Box::pin(async move { + let mut result = next(args).await?; + result["mutated"] = json!(true); + Ok(ToolExecutionInterceptOutcome::new(result)) + }) + }), + ) + .unwrap(); + let producer: ToolExecutionFrameNextFn = Arc::new(|_args| { + Box::pin(async move { + Ok(ToolExecutionFrame::annotated( + json!({"raw": true}), + json!({"producer_status": "succeeded"}), + )) + }) + }); + + let frame = tool_call_execute_frame( + nemo_relay::api::tool::ToolCallExecuteFrameParams::builder() + .name("mutated-tool") + .args(json!({})) + .func(producer) + .build(), + ) + .await + .unwrap(); + + assert_eq!(frame.result, json!({"raw": true, "mutated": true})); + assert!( + frame.annotation.is_none(), + "legacy mutation must not leave stale producer semantics attached" + ); +} + +#[tokio::test] +async fn test_legacy_multiple_next_calls_conservatively_invalidate_annotation() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + register_tool_execution_intercept( + "legacy_retry", + 1, + Arc::new(|_name, args, next| { + Box::pin(async move { + let first = next(args.clone()).await?; + let _second = next(args).await?; + Ok(ToolExecutionInterceptOutcome::new(first)) + }) + }), + ) + .unwrap(); + let producer: ToolExecutionFrameNextFn = Arc::new(|args| { + Box::pin(async move { + Ok(ToolExecutionFrame::annotated( + args, + json!({"producer_status": "succeeded"}), + )) + }) + }); + + let frame = tool_call_execute_frame( + nemo_relay::api::tool::ToolCallExecuteFrameParams::builder() + .name("retried-tool") + .args(json!({"attempt": 1})) + .func(producer) + .build(), + ) + .await + .unwrap(); + + assert_eq!(frame.result, json!({"attempt": 1})); + assert!( + frame.annotation.is_none(), + "legacy middleware cannot unambiguously select annotation after multiple next calls" + ); +} + +#[tokio::test] +async fn test_legacy_failed_second_next_call_still_invalidates_annotation() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + register_tool_execution_intercept( + "legacy_retry_after_error", + 1, + Arc::new(|_name, args, next| { + Box::pin(async move { + let first = next(args).await?; + let second = next(json!({"fail": true})).await; + assert!(second.is_err()); + Ok(ToolExecutionInterceptOutcome::new(first)) + }) + }), + ) + .unwrap(); + let producer: ToolExecutionFrameNextFn = Arc::new(|args| { + Box::pin(async move { + if args.get("fail").and_then(Json::as_bool) == Some(true) { + return Err(FlowError::Internal("producer rejected retry".into())); + } + Ok(ToolExecutionFrame::annotated( + args, + json!({"producer_status": "succeeded"}), + )) + }) + }); + + let frame = tool_call_execute_frame( + nemo_relay::api::tool::ToolCallExecuteFrameParams::builder() + .name("retry-after-error-tool") + .args(json!({"attempt": 1})) + .func(producer) + .build(), + ) + .await + .unwrap(); + + assert_eq!(frame.result, json!({"attempt": 1})); + assert!( + frame.annotation.is_none(), + "every continuation invocation must count, including failed calls" + ); +} + +#[tokio::test] +async fn test_frame_continuation_expires_when_callback_returns() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + let retained = Arc::new(Mutex::new(None::)); + let retained_from_callback = retained.clone(); + register_tool_execution_frame_intercept( + "frame_retained_next", + 1, + Arc::new(move |_name, _args, next| { + *retained_from_callback.lock().unwrap() = Some(next); + Box::pin(async move { + Ok(ToolExecutionFrameOutcome::new(ToolExecutionFrame::new( + json!({"short_circuit": true}), + ))) + }) + }), + ) + .unwrap(); + + let producer_called = Arc::new(AtomicBool::new(false)); + let producer_called_from_callback = producer_called.clone(); + let producer: ToolExecutionFrameNextFn = Arc::new(move |args| { + let producer_called = producer_called_from_callback.clone(); + Box::pin(async move { + producer_called.store(true, Ordering::SeqCst); + Ok(ToolExecutionFrame::new(args)) + }) + }); + + let frame = tool_call_execute_frame( + nemo_relay::api::tool::ToolCallExecuteFrameParams::builder() + .name("retained-next-tool") + .args(json!({})) + .func(producer) + .build(), + ) + .await + .unwrap(); + assert_eq!(frame.result, json!({"short_circuit": true})); + assert!(!producer_called.load(Ordering::SeqCst)); + + let retained = retained.lock().unwrap().take().unwrap(); + let error = retained(json!({"late": true})).await.unwrap_err(); + assert!(error.to_string().contains("no longer active")); + assert!( + !producer_called.load(Ordering::SeqCst), + "a retained continuation must not run the downstream tool" + ); +} + +#[tokio::test] +async fn test_legacy_continuation_is_revoked_after_callback_settles() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + let retained = Arc::new(Mutex::new(None::)); + let retained_from_callback = retained.clone(); + register_tool_execution_intercept( + "legacy_retained_next", + 1, + Arc::new(move |_name, _args, next| { + *retained_from_callback.lock().unwrap() = Some(next); + Box::pin(async move { + Ok(ToolExecutionInterceptOutcome::new( + json!({"short_circuit": true}), + )) + }) + }), + ) + .unwrap(); + + let producer_called = Arc::new(AtomicBool::new(false)); + let producer_called_from_callback = producer_called.clone(); + let result = tool_call_execute( + nemo_relay::api::tool::ToolCallExecuteParams::builder() + .name("legacy-retained-next-tool") + .args(json!({})) + .func(Arc::new(move |args| { + let producer_called = producer_called_from_callback.clone(); + Box::pin(async move { + producer_called.store(true, Ordering::SeqCst); + Ok(args) + }) + })) + .build(), + ) + .await + .unwrap(); + assert_eq!(result, json!({"short_circuit": true})); + assert!(!producer_called.load(Ordering::SeqCst)); + + let retained = retained.lock().unwrap().take().unwrap(); + let error = retained(json!({"late": true})).await.unwrap_err(); + assert!(error.to_string().contains("no longer active")); + assert!( + !producer_called.load(Ordering::SeqCst), + "a retained continuation must not run the downstream tool" + ); + deregister_tool_execution_intercept("legacy_retained_next").unwrap(); +} + +#[tokio::test] +async fn test_legacy_short_circuit_has_no_implicit_annotation() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + setup_isolated_thread(); + + register_tool_execution_intercept( + "legacy_short_circuit", + 1, + Arc::new(|_name, _args, _next| { + Box::pin(async move { + Ok(ToolExecutionInterceptOutcome::new( + json!({"short_circuit": true}), + )) + }) + }), + ) + .unwrap(); + let producer: ToolExecutionFrameNextFn = Arc::new(|_args| { + Box::pin(async move { + Ok(ToolExecutionFrame::annotated( + json!({"producer": true}), + json!({"producer_status": "succeeded"}), + )) + }) + }); + + let frame = tool_call_execute_frame( + nemo_relay::api::tool::ToolCallExecuteFrameParams::builder() + .name("short-circuited-tool") + .args(json!({})) + .func(producer) + .build(), + ) + .await + .unwrap(); + + assert_eq!(frame.result, json!({"short_circuit": true})); + assert!(frame.annotation.is_none()); +} + #[tokio::test] async fn test_tool_execution_outcome_marks_follow_end_with_tool_parentage() { let _lock = TEST_MUTEX.lock().unwrap(); @@ -3978,6 +4541,43 @@ async fn test_duplicate_intercept_registration_returns_error() { deregister_tool_request_intercept("dup_intercept").unwrap(); } +/// Raw-result and frame-aware callbacks share one name namespace and +/// deregistration operation. +#[test] +fn test_tool_execution_callback_forms_share_namespace_and_deregistration() { + let _lock = TEST_MUTEX.lock().unwrap(); + reset_global(); + + register_tool_execution_intercept( + "shared_exec_name", + 1, + Arc::new(|_name, args, next| { + Box::pin(async move { next(args).await.map(ToolExecutionInterceptOutcome::new) }) + }), + ) + .unwrap(); + + let duplicate = register_tool_execution_frame_intercept( + "shared_exec_name", + 2, + Arc::new(|_name, args, next| { + Box::pin(async move { next(args).await.map(ToolExecutionFrameOutcome::new) }) + }), + ); + assert!(matches!(duplicate, Err(FlowError::AlreadyExists(_)))); + assert!(deregister_tool_execution_intercept("shared_exec_name").unwrap()); + + register_tool_execution_frame_intercept( + "shared_exec_name", + 2, + Arc::new(|_name, args, next| { + Box::pin(async move { next(args).await.map(ToolExecutionFrameOutcome::new) }) + }), + ) + .unwrap(); + assert!(deregister_tool_execution_intercept("shared_exec_name").unwrap()); +} + // ========================================================================= // Deregistration Tests // ========================================================================= diff --git a/crates/core/tests/integration/native_plugin_tests.rs b/crates/core/tests/integration/native_plugin_tests.rs index 82fc2168e..3f4ad30d6 100644 --- a/crates/core/tests/integration/native_plugin_tests.rs +++ b/crates/core/tests/integration/native_plugin_tests.rs @@ -22,7 +22,10 @@ use nemo_relay::api::scope::{ pop_scope, push_scope, }; use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; -use nemo_relay::api::tool::{ToolCallExecuteParams, tool_call_execute, tool_request_intercepts}; +use nemo_relay::api::tool::{ + TOOL_RESULT_ANNOTATION_PROFILE_KEY, ToolCallExecuteFrameParams, ToolCallExecuteParams, + ToolExecutionFrame, tool_call_execute, tool_call_execute_frame, tool_request_intercepts, +}; use nemo_relay::codec::response::AnnotatedLlmResponse; use nemo_relay::plugin::dynamic::{ DynamicPluginActivationSpec, DynamicPluginKind, NativePluginLoadSpec, PluginHostActivation, @@ -228,7 +231,7 @@ async fn sdk_cdylib_registers_tool_request_intercept() { cleanup.mark_subscriber_registered("native_plugin_fixture_events"); let stack = create_scope_stack(); - let (outer_uuid, rewritten, tool_result) = TASK_SCOPE_STACK + let (outer_uuid, rewritten, tool_result, tool_frame) = TASK_SCOPE_STACK .scope(stack, async { let outer = push_scope( PushScopeParams::builder() @@ -252,20 +255,54 @@ async fn sdk_cdylib_registers_tool_request_intercept() { ) .await .expect("native tool middleware should run"); + let tool_frame = tool_call_execute_frame( + ToolCallExecuteFrameParams::builder() + .name("native-fixture-tool-frame") + .args(json!({ "input": "execute-frame" })) + .func(Arc::new(|args| { + Box::pin(async move { + Ok(ToolExecutionFrame::annotated( + json!({ "tool_frame_callback": true, "args": args }), + json!({ "producer_annotation": true }), + )) + }) + })) + .build(), + ) + .await + .expect("native frame-aware tool middleware should run"); pop_scope(PopScopeParams::builder().handle_uuid(&outer.uuid).build()) .expect("outer scope should pop"); - (outer_uuid, rewritten, tool_result) + (outer_uuid, rewritten, tool_result, tool_frame) }) .await; assert_eq!(rewritten["input"], "value"); assert_eq!(rewritten["native_plugin"], true); assert_eq!(tool_result["tool_callback"], true); assert_eq!(tool_result["native_plugin_tool_execution"], true); + assert_eq!(tool_result["native_plugin_tool_execution_frame"], true); assert_eq!( tool_result["args"]["native_plugin_tool_execution_request"], true ); assert!(tool_result.get("pending_marks").is_none()); + assert_eq!(tool_frame.result["tool_frame_callback"], true); + assert_eq!( + tool_frame.result["native_plugin_tool_execution_frame"], + true + ); + assert_eq!( + tool_frame.result["args"]["native_plugin_tool_execution_frame_request"], + true + ); + assert_eq!( + tool_frame.annotation.as_ref().unwrap()["producer_annotation"], + true + ); + assert_eq!( + tool_frame.annotation.as_ref().unwrap()["native_plugin_tool_execution_frame"], + true + ); flush_subscribers().expect("native fixture events should flush"); let first_events = events.lock().unwrap().clone(); @@ -370,6 +407,26 @@ async fn sdk_cdylib_registers_tool_request_intercept() { .position(|event| event.name() == "fixture.native.tool_execution.mark") .unwrap(); assert!(tool_end_index < tool_mark_index); + let tool_frame_start = find_event( + &first_events, + "native-fixture-tool-frame", + Some(ScopeCategory::Start), + ); + let tool_frame_end = find_event( + &first_events, + "native-fixture-tool-frame", + Some(ScopeCategory::End), + ); + assert_eq!( + tool_frame_end.category_profile().unwrap().extra[TOOL_RESULT_ANNOTATION_PROFILE_KEY]["producer_annotation"], + true + ); + let tool_frame_mark = find_event( + &first_events, + "fixture.native.tool_execution_frame.mark", + None, + ); + assert_eq!(tool_frame_mark.parent_uuid(), Some(tool_frame_start.uuid())); events.lock().unwrap().clear(); let isolated_next_stack = create_scope_stack(); diff --git a/crates/core/tests/integration/worker_plugin_tests.rs b/crates/core/tests/integration/worker_plugin_tests.rs index 358447ae9..151ea9d3d 100644 --- a/crates/core/tests/integration/worker_plugin_tests.rs +++ b/crates/core/tests/integration/worker_plugin_tests.rs @@ -18,7 +18,10 @@ use nemo_relay::api::scope::{ EmitMarkEventParams, PopScopeParams, PushScopeParams, ScopeType, event, pop_scope, push_scope, }; use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; -use nemo_relay::api::tool::{ToolCallExecuteParams, tool_call_execute, tool_request_intercepts}; +use nemo_relay::api::tool::{ + ToolCallExecuteFrameParams, ToolCallExecuteParams, ToolExecutionFrame, tool_call_execute, + tool_call_execute_frame, tool_request_intercepts, +}; use nemo_relay::codec::request::AnnotatedLlmRequest; use nemo_relay::codec::traits::LlmCodec; use nemo_relay::error::Result as FlowResult; @@ -151,7 +154,7 @@ async fn rust_worker_registers_and_invokes_all_current_surfaces() { .expect("test subscriber should register"); let stack = create_scope_stack(); - let (outer_uuid, rewritten, tool_result) = TASK_SCOPE_STACK + let (outer_uuid, rewritten, tool_result, tool_frame) = TASK_SCOPE_STACK .scope(stack, async { let outer = push_scope( PushScopeParams::builder() @@ -175,19 +178,53 @@ async fn rust_worker_registers_and_invokes_all_current_surfaces() { ) .await .expect("worker tool middleware should run"); + let tool_frame = tool_call_execute_frame( + ToolCallExecuteFrameParams::builder() + .name("worker-fixture-tool-frame") + .args(json!({ "input": "execute-frame" })) + .func(Arc::new(|args| { + Box::pin(async move { + Ok(ToolExecutionFrame::annotated( + json!({ "tool_frame_callback": true, "args": args }), + json!({ "producer_annotation": true }), + )) + }) + })) + .build(), + ) + .await + .expect("worker frame-aware tool middleware should run"); pop_scope(PopScopeParams::builder().handle_uuid(&outer.uuid).build()) .expect("outer scope should pop"); - (outer_uuid, rewritten, tool_result) + (outer_uuid, rewritten, tool_result, tool_frame) }) .await; assert_eq!(rewritten["worker_plugin"], true); assert_eq!(tool_result["tool_callback"], true); assert_eq!(tool_result["worker_plugin_tool_execution"], true); + assert_eq!(tool_result["worker_plugin_tool_execution_frame"], true); assert_eq!( tool_result["args"]["worker_plugin_tool_execution_request"], true ); + assert_eq!(tool_frame.result["tool_frame_callback"], true); + assert_eq!( + tool_frame.result["worker_plugin_tool_execution_frame"], + true + ); + assert_eq!( + tool_frame.result["args"]["worker_plugin_tool_execution_frame_request"], + true + ); + assert_eq!( + tool_frame.annotation.as_ref().unwrap()["producer_annotation"], + true + ); + assert_eq!( + tool_frame.annotation.as_ref().unwrap()["worker_plugin_tool_execution_frame"], + true + ); flush_subscribers().expect("worker fixture events should flush"); let captured_events = events.lock().unwrap().clone(); @@ -277,6 +314,17 @@ async fn rust_worker_registers_and_invokes_all_current_surfaces() { assert_eq!(tool_mark.parent_uuid(), Some(tool_start.uuid())); assert!(tool_mark.timestamp() > tool_end.timestamp()); assert_eq!(tool_mark.metadata().unwrap()["worker_plugin_mark"], true); + let tool_frame_mark = find_event( + &captured_events, + "fixture.worker.tool_execution_frame.mark", + None, + ); + let tool_frame_start = find_event( + &captured_events, + "worker-fixture-tool-frame", + Some(ScopeCategory::Start), + ); + assert_eq!(tool_frame_mark.parent_uuid(), Some(tool_frame_start.uuid())); let llm_execute_response = llm_call_execute( LlmCallExecuteParams::builder() diff --git a/crates/core/tests/unit/atif_tests.rs b/crates/core/tests/unit/atif_tests.rs index 050134f3a..85909f5b6 100644 --- a/crates/core/tests/unit/atif_tests.rs +++ b/crates/core/tests/unit/atif_tests.rs @@ -10,7 +10,7 @@ use crate::api::event::{ }; use crate::api::llm::{LlmAttributes, LlmRequest}; use crate::api::scope::{HandleAttributes, ScopeAttributes, ScopeType}; -use crate::api::tool::ToolAttributes; +use crate::api::tool::{TOOL_RESULT_ANNOTATION_PROFILE_KEY, ToolAttributes}; use crate::codec::anthropic::AnthropicMessagesCodec; use crate::codec::model_pricing::pricing_test_mutex; use crate::codec::openai_chat::OpenAIChatCodec; @@ -687,6 +687,39 @@ fn test_exporter_moves_structured_tool_close_result_to_observation_extra() { assert!(content.is_none()); } +#[test] +fn test_exporter_preserves_opaque_tool_result_annotation_in_observation_extra() { + let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); + let tool_uuid = Uuid::now_v7(); + let annotation = json!({ + "status": "failed", + "nested": {"code": 17}, + "items": [true, null, "opaque"], + }); + let mut end = event_builder(tool_uuid, EventType::End) + .name("terminal") + .scope_type(ScopeType::Tool) + .output(json!({"raw": "result"})) + .tool_call_id("call_123") + .build(); + end.category_profile_mut().unwrap().extra.insert( + TOOL_RESULT_ANNOTATION_PROFILE_KEY.to_string(), + annotation.clone(), + ); + + { + let mut state = exporter.state.lock().unwrap(); + state.events.push(end); + } + + let trajectory = exporter.export().unwrap(); + let result = &trajectory.steps[0].observation.as_ref().unwrap().results[0]; + assert_eq!( + result.extra.as_ref().unwrap()["tool_result_annotation"], + annotation + ); +} + #[test] fn test_exporter_preserves_tool_content_part_array_observation_content() { let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); diff --git a/crates/core/tests/unit/observability/openinference_tests.rs b/crates/core/tests/unit/observability/openinference_tests.rs index 652311287..64cf3baea 100644 --- a/crates/core/tests/unit/observability/openinference_tests.rs +++ b/crates/core/tests/unit/observability/openinference_tests.rs @@ -12,7 +12,7 @@ use crate::api::runtime::NemoRelayContextState; use crate::api::runtime::global_context; use crate::api::scope::ScopeType; use crate::api::scope::{event, pop_scope, push_scope}; -use crate::api::tool::ToolAttributes; +use crate::api::tool::{TOOL_RESULT_ANNOTATION_PROFILE_KEY, ToolAttributes}; use crate::codec::model_pricing::pricing_test_mutex; use crate::codec::request::{ AnnotatedLlmRequest, ContentPart, FunctionDefinition, GenerationParams, Message, @@ -55,6 +55,43 @@ fn reset_global() { *context.write().unwrap() = NemoRelayContextState::new(); } +#[test] +fn tool_result_annotation_is_exported_as_one_opaque_json_attribute() { + let annotation = json!({ + "status": "failed", + "nested": {"code": 17}, + "items": [true, null, "opaque"], + }); + let mut profile = CategoryProfile::builder() + .tool_call_id("call-annotation") + .build(); + profile.extra.insert( + TOOL_RESULT_ANNOTATION_PROFILE_KEY.to_string(), + annotation.clone(), + ); + let event = Event::Scope(ScopeEvent::new( + BaseEvent::builder() + .name("annotated-tool") + .data(json!({"raw": "result"})) + .build(), + ScopeCategory::End, + Vec::new(), + EventCategory::tool(), + Some(profile), + )); + + let attributes = attr_map(&end_attributes(&event)); + assert_eq!( + attributes.get("nemo_relay.tool.result.annotation"), + Some(&serde_json::to_string(&annotation).unwrap()) + ); + assert!( + attributes + .keys() + .all(|key| !key.starts_with("nemo_relay.tool.result.annotation.")) + ); +} + fn make_provider() -> ( SdkTracerProvider, opentelemetry_sdk::trace::InMemorySpanExporter, diff --git a/crates/core/tests/unit/observability/otel_tests.rs b/crates/core/tests/unit/observability/otel_tests.rs index 9123373eb..544b764da 100644 --- a/crates/core/tests/unit/observability/otel_tests.rs +++ b/crates/core/tests/unit/observability/otel_tests.rs @@ -15,7 +15,7 @@ use crate::api::runtime::{ }; use crate::api::scope::ScopeType; use crate::api::scope::{event, pop_scope, push_scope}; -use crate::api::tool::ToolAttributes; +use crate::api::tool::{TOOL_RESULT_ANNOTATION_PROFILE_KEY, ToolAttributes}; use crate::codec::model_pricing::pricing_test_mutex; use crate::codec::response::{ AnnotatedLlmResponse, CostEstimate, CostSource, FinishReason, PricingCatalog, PricingResolver, @@ -114,6 +114,43 @@ fn optimization_summary_emits_namespaced_otel_attributes() { ); } +#[test] +fn tool_result_annotation_is_exported_as_one_opaque_json_attribute() { + let annotation = json!({ + "status": "failed", + "nested": {"code": 17}, + "items": [true, null, "opaque"], + }); + let mut profile = CategoryProfile::builder() + .tool_call_id("call-annotation") + .build(); + profile.extra.insert( + TOOL_RESULT_ANNOTATION_PROFILE_KEY.to_string(), + annotation.clone(), + ); + let event = Event::Scope(ScopeEvent::new( + BaseEvent::builder() + .name("annotated-tool") + .data(json!({"raw": "result"})) + .build(), + ScopeCategory::End, + Vec::new(), + EventCategory::tool(), + Some(profile), + )); + + let attributes = attr_map(&end_attributes(&event)); + assert_eq!( + attributes.get("nemo_relay.tool.result.annotation"), + Some(&serde_json::to_string(&annotation).unwrap()) + ); + assert!( + attributes + .keys() + .all(|key| !key.starts_with("nemo_relay.tool.result.annotation.")) + ); +} + #[test] fn optimization_cost_attributes_preserve_independent_currency_and_provenance() { let summary: crate::codec::optimization::LlmOptimizationSummary = diff --git a/crates/ffi/nemo_relay.h b/crates/ffi/nemo_relay.h index 9a17fe485..a9d9b6872 100644 --- a/crates/ffi/nemo_relay.h +++ b/crates/ffi/nemo_relay.h @@ -438,6 +438,24 @@ typedef char *(*NemoRelayToolExecInterceptCb)(void *user_data, NemoRelayToolExecNextFn next_fn, void *next_ctx); +/** + * Runtime-provided annotation-aware continuation. + * + * The returned JSON serializes a `ToolExecutionFrame`. + */ +typedef char *(*NemoRelayToolExecFrameNextFn)(const char *args_json, void *next_ctx); + +/** + * Annotation-aware tool execution intercept callback. + * + * `next_fn` returns a serialized `ToolExecutionFrame`; this callback must + * return a serialized `ToolExecutionFrameOutcome`. + */ +typedef char *(*NemoRelayToolExecFrameInterceptCb)(void *user_data, + const char *args_json, + NemoRelayToolExecFrameNextFn next_fn, + void *next_ctx); + /** * Callback for tool execution (default callable). Receives arguments as JSON, * returns result as JSON. The returned string must be allocated with `malloc` @@ -445,6 +463,13 @@ typedef char *(*NemoRelayToolExecInterceptCb)(void *user_data, */ typedef char *(*NemoRelayToolExecCb)(void *user_data, const char *args_json); +/** + * Annotation-aware tool execution callback. + * + * The returned JSON must serialize a `ToolExecutionFrame`. + */ +typedef char *(*NemoRelayToolExecFrameCb)(void *user_data, const char *args_json); + /** * Run the registered tool request intercept chain on the given arguments. * @@ -1869,6 +1894,21 @@ NemoRelayStatus nemo_relay_plugin_context_register_tool_execution_intercept(stru void *user_data, NemoRelayFreeFn free_fn); +/** + * Register an annotation-aware tool execution intercept into the plugin + * registration context. + * + * # Safety + * `ctx` and `name` must be valid pointers and the callback must remain valid + * for the duration of the plugin registration lifetime. + */ +NemoRelayStatus nemo_relay_plugin_context_register_tool_execution_frame_intercept(struct FfiPluginContext *ctx, + const char *name, + int32_t priority, + NemoRelayToolExecFrameInterceptCb cb, + void *user_data, + NemoRelayFreeFn free_fn); + /** * Retrieve the current scope handle from the thread-local scope stack. * @@ -2042,6 +2082,20 @@ NemoRelayStatus nemo_relay_scope_register_tool_execution_intercept(const char *s void *exec_user_data, NemoRelayFreeFn exec_free); +/** + * Register a scope-local annotation-aware tool execution intercept in the + * existing tool execution chain. + * + * # Safety + * `scope_uuid` and `name` must be valid C strings. Callback pointers must be valid. + */ +NemoRelayStatus nemo_relay_scope_register_tool_execution_frame_intercept(const char *scope_uuid, + const char *name, + int32_t priority, + NemoRelayToolExecFrameInterceptCb exec_cb, + void *exec_user_data, + NemoRelayFreeFn exec_free); + /** * Deregister a scope-local tool execution intercept by name. * @@ -2462,6 +2516,23 @@ NemoRelayStatus nemo_relay_tool_call_end(const struct FfiToolHandle *handle, const char *metadata_json, const int64_t *timestamp_unix_micros); +/** + * End a manual tool call with a serialized `ToolExecutionFrame`. + * + * The frame's raw result follows the existing response-sanitization path. Its + * optional producer annotation is carried in the emitted tool category + * profile. + * + * # Safety + * `handle` and `frame_json` must be valid, non-null pointers. Optional pointer + * arguments follow [`nemo_relay_tool_call_end`]. + */ +NemoRelayStatus nemo_relay_tool_call_end_frame(const struct FfiToolHandle *handle, + const char *frame_json, + const char *data_json, + const char *metadata_json, + const int64_t *timestamp_unix_micros); + /** * Execute a tool call end-to-end: run conditional-execution guardrails (on raw * args), then request intercepts, sanitize-request guardrails, execution @@ -2497,6 +2568,27 @@ NemoRelayStatus nemo_relay_tool_call_execute(const char *name, const char *metadata_json, char **out); +/** + * Execute a tool call with an annotation-aware producer callback. + * + * The callback and output use serialized `ToolExecutionFrame` values. This + * function participates in the same execution-intercept chain as + * [`nemo_relay_tool_call_execute`]. + * + * # Safety + * `name`, `args_json`, and `out` must be valid, non-null pointers. + */ +NemoRelayStatus nemo_relay_tool_call_execute_frame(const char *name, + const char *args_json, + NemoRelayToolExecFrameCb func, + void *func_user_data, + NemoRelayFreeFn func_free, + const struct FfiScopeHandle *parent, + uint32_t attributes, + const char *data_json, + const char *metadata_json, + char **out); + /** * Register a tool conditional execution guardrail. The callback decides whether * a tool call should proceed. Returns an error message to reject, or null to allow. @@ -2551,6 +2643,22 @@ NemoRelayStatus nemo_relay_register_tool_execution_intercept(const char *name, void *exec_user_data, NemoRelayFreeFn exec_free); +/** + * Register an annotation-aware tool execution intercept in the existing chain. + * + * The callback receives a continuation that returns serialized + * `ToolExecutionFrame` JSON and must return serialized + * `ToolExecutionFrameOutcome` JSON. + * + * # Safety + * `name` must be a valid C string. Callback pointers must be valid. + */ +NemoRelayStatus nemo_relay_register_tool_execution_frame_intercept(const char *name, + int32_t priority, + NemoRelayToolExecFrameInterceptCb exec_cb, + void *exec_user_data, + NemoRelayFreeFn exec_free); + /** * Deregister a tool execution intercept by name. * diff --git a/crates/ffi/src/api/mod.rs b/crates/ffi/src/api/mod.rs index cd08d8d03..581f8ee5d 100644 --- a/crates/ffi/src/api/mod.rs +++ b/crates/ffi/src/api/mod.rs @@ -19,12 +19,14 @@ use crate::callable::{ NemoRelayLlmExecCb, NemoRelayLlmExecInterceptCb, NemoRelayLlmRequestInterceptCb, NemoRelayLlmSanitizeRequestCb, NemoRelayLlmSanitizeResponseCb, NemoRelayPluginRegisterCb, NemoRelayPluginValidateCb, NemoRelayToolConditionalCb, NemoRelayToolExecCb, - NemoRelayToolExecInterceptCb, NemoRelayToolSanitizeCb, wrap_codec_fn, wrap_collector_fn, - wrap_event_sanitize_fn, wrap_event_subscriber, wrap_finalizer_fn, wrap_llm_conditional_fn, - wrap_llm_exec_fn, wrap_llm_exec_intercept_fn, wrap_llm_request_intercept_fn, - wrap_llm_sanitize_request_fn, wrap_llm_sanitize_response_fn, wrap_llm_stream_exec_fn, - wrap_llm_stream_exec_intercept_fn, wrap_tool_conditional_fn, wrap_tool_exec_fn, - wrap_tool_exec_intercept_fn, wrap_tool_request_intercept_fn, wrap_tool_sanitize_fn, + NemoRelayToolExecFrameCb, NemoRelayToolExecFrameInterceptCb, NemoRelayToolExecInterceptCb, + NemoRelayToolSanitizeCb, wrap_codec_fn, wrap_collector_fn, wrap_event_sanitize_fn, + wrap_event_subscriber, wrap_finalizer_fn, wrap_llm_conditional_fn, wrap_llm_exec_fn, + wrap_llm_exec_intercept_fn, wrap_llm_request_intercept_fn, wrap_llm_sanitize_request_fn, + wrap_llm_sanitize_response_fn, wrap_llm_stream_exec_fn, wrap_llm_stream_exec_intercept_fn, + wrap_tool_conditional_fn, wrap_tool_exec_fn, wrap_tool_exec_frame_fn, + wrap_tool_exec_frame_intercept_fn, wrap_tool_exec_intercept_fn, wrap_tool_request_intercept_fn, + wrap_tool_sanitize_fn, }; use crate::convert::{ c_str_to_json, c_str_to_opt_json, c_str_to_string, json_to_c_string, nemo_relay_string_free, @@ -45,7 +47,9 @@ use libc::c_char; use nemo_relay::api::llm as core_llm_api; use nemo_relay::api::llm::{LlmAttributes, LlmRequest, LlmRequestInterceptOutcome}; use nemo_relay::api::registry as core_registry_api; -use nemo_relay::api::runtime::{LlmExecutionNextFn, LlmStreamExecutionNextFn, ToolExecutionNextFn}; +use nemo_relay::api::runtime::{ + LlmExecutionNextFn, LlmStreamExecutionNextFn, ToolExecutionFrameNextFn, ToolExecutionNextFn, +}; use nemo_relay::api::runtime::{ TASK_SCOPE_STACK, capture_thread_scope_stack, create_scope_stack, current_scope_stack, restore_thread_scope_stack, scope_stack_active, set_thread_scope_stack, with_scope_stack, diff --git a/crates/ffi/src/api/plugin.rs b/crates/ffi/src/api/plugin.rs index 09cd151f8..ebbb358bc 100644 --- a/crates/ffi/src/api/plugin.rs +++ b/crates/ffi/src/api/plugin.rs @@ -7,16 +7,17 @@ use super::{ NemoRelayFreeFn, NemoRelayLlmConditionalCb, NemoRelayLlmExecInterceptCb, NemoRelayLlmRequestInterceptCb, NemoRelayLlmSanitizeRequestCb, NemoRelayLlmSanitizeResponseCb, NemoRelayPluginRegisterCb, NemoRelayPluginValidateCb, NemoRelayStatus, - NemoRelayToolConditionalCb, NemoRelayToolExecInterceptCb, NemoRelayToolSanitizeCb, Pin, Plugin, - PluginConfig, PluginError, PluginHostActivation, PluginRegistrationContext, - active_plugin_report, c_char, c_str_to_json, c_str_to_string, clear_last_error, - clear_plugin_configuration, deregister_plugin, initialize_plugins, json_to_c_string, - last_error_message, list_plugin_kinds, nemo_relay_string_free, register_adaptive_component, - register_plugin, set_last_error, status_from_plugin_error, tokio_runtime, - validate_plugin_config, wrap_event_sanitize_fn, wrap_event_subscriber, wrap_llm_conditional_fn, - wrap_llm_exec_intercept_fn, wrap_llm_request_intercept_fn, wrap_llm_sanitize_request_fn, - wrap_llm_sanitize_response_fn, wrap_llm_stream_exec_intercept_fn, wrap_tool_conditional_fn, - wrap_tool_exec_intercept_fn, wrap_tool_request_intercept_fn, wrap_tool_sanitize_fn, + NemoRelayToolConditionalCb, NemoRelayToolExecFrameInterceptCb, NemoRelayToolExecInterceptCb, + NemoRelayToolSanitizeCb, Pin, Plugin, PluginConfig, PluginError, PluginHostActivation, + PluginRegistrationContext, active_plugin_report, c_char, c_str_to_json, c_str_to_string, + clear_last_error, clear_plugin_configuration, deregister_plugin, initialize_plugins, + json_to_c_string, last_error_message, list_plugin_kinds, nemo_relay_string_free, + register_adaptive_component, register_plugin, set_last_error, status_from_plugin_error, + tokio_runtime, validate_plugin_config, wrap_event_sanitize_fn, wrap_event_subscriber, + wrap_llm_conditional_fn, wrap_llm_exec_intercept_fn, wrap_llm_request_intercept_fn, + wrap_llm_sanitize_request_fn, wrap_llm_sanitize_response_fn, wrap_llm_stream_exec_intercept_fn, + wrap_tool_conditional_fn, wrap_tool_exec_frame_intercept_fn, wrap_tool_exec_intercept_fn, + wrap_tool_request_intercept_fn, wrap_tool_sanitize_fn, }; use crate::api::event_registry::Surface; use nemo_relay_pii_redaction::component::register_pii_redaction_component; @@ -980,3 +981,36 @@ pub unsafe extern "C" fn nemo_relay_plugin_context_register_tool_execution_inter Err(err) => status_from_plugin_error(&err), } } + +/// Register an annotation-aware tool execution intercept into the plugin +/// registration context. +/// +/// # Safety +/// `ctx` and `name` must be valid pointers and the callback must remain valid +/// for the duration of the plugin registration lifetime. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_plugin_context_register_tool_execution_frame_intercept( + ctx: *mut FfiPluginContext, + name: *const c_char, + priority: i32, + cb: NemoRelayToolExecFrameInterceptCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, +) -> NemoRelayStatus { + clear_last_error(); + if ctx.is_null() { + set_last_error("plugin context is null"); + return NemoRelayStatus::NullPointer; + } + let name = match c_str_to_string(name) { + Ok(value) => value, + Err(status) => return status, + }; + let wrapped = wrap_tool_exec_frame_intercept_fn(cb, user_data, free_fn); + match unsafe { &mut *((*ctx).0) } + .register_tool_execution_frame_intercept(&name, priority, wrapped) + { + Ok(()) => NemoRelayStatus::Ok, + Err(error) => status_from_plugin_error(&error), + } +} diff --git a/crates/ffi/src/api/scope_registry.rs b/crates/ffi/src/api/scope_registry.rs index 50efd644b..3a6efab4b 100644 --- a/crates/ffi/src/api/scope_registry.rs +++ b/crates/ffi/src/api/scope_registry.rs @@ -5,12 +5,13 @@ use super::{ NemoRelayEventSubscriberCb, NemoRelayFreeFn, NemoRelayLlmConditionalCb, NemoRelayLlmExecInterceptCb, NemoRelayLlmRequestInterceptCb, NemoRelayLlmSanitizeRequestCb, NemoRelayLlmSanitizeResponseCb, NemoRelayStatus, NemoRelayToolConditionalCb, - NemoRelayToolExecInterceptCb, NemoRelayToolSanitizeCb, c_char, c_str_to_string, - clear_last_error, core_registry_api, core_subscriber_api, set_last_error, status_from_error, - wrap_event_subscriber, wrap_llm_conditional_fn, wrap_llm_exec_intercept_fn, - wrap_llm_request_intercept_fn, wrap_llm_sanitize_request_fn, wrap_llm_sanitize_response_fn, - wrap_llm_stream_exec_intercept_fn, wrap_tool_conditional_fn, wrap_tool_exec_intercept_fn, - wrap_tool_request_intercept_fn, wrap_tool_sanitize_fn, + NemoRelayToolExecFrameInterceptCb, NemoRelayToolExecInterceptCb, NemoRelayToolSanitizeCb, + c_char, c_str_to_string, clear_last_error, core_registry_api, core_subscriber_api, + set_last_error, status_from_error, wrap_event_subscriber, wrap_llm_conditional_fn, + wrap_llm_exec_intercept_fn, wrap_llm_request_intercept_fn, wrap_llm_sanitize_request_fn, + wrap_llm_sanitize_response_fn, wrap_llm_stream_exec_intercept_fn, wrap_tool_conditional_fn, + wrap_tool_exec_frame_intercept_fn, wrap_tool_exec_intercept_fn, wrap_tool_request_intercept_fn, + wrap_tool_sanitize_fn, }; // --------------------------------------------------------------------------- @@ -318,6 +319,38 @@ pub unsafe extern "C" fn nemo_relay_scope_register_tool_execution_intercept( } } +/// Register a scope-local annotation-aware tool execution intercept in the +/// existing tool execution chain. +/// +/// # Safety +/// `scope_uuid` and `name` must be valid C strings. Callback pointers must be valid. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_scope_register_tool_execution_frame_intercept( + scope_uuid: *const c_char, + name: *const c_char, + priority: i32, + exec_cb: NemoRelayToolExecFrameInterceptCb, + exec_user_data: *mut libc::c_void, + exec_free: NemoRelayFreeFn, +) -> NemoRelayStatus { + clear_last_error(); + let uuid = match parse_scope_uuid(scope_uuid) { + Ok(value) => value, + Err(status) => return status, + }; + let name = match c_str_to_string(name) { + Ok(value) => value, + Err(status) => return status, + }; + let exec = wrap_tool_exec_frame_intercept_fn(exec_cb, exec_user_data, exec_free); + match core_registry_api::scope_register_tool_execution_frame_intercept( + &uuid, &name, priority, exec, + ) { + Ok(()) => NemoRelayStatus::Ok, + Err(error) => status_from_error(&error), + } +} + /// Deregister a scope-local tool execution intercept by name. /// /// # Safety diff --git a/crates/ffi/src/api/tool_lifecycle.rs b/crates/ffi/src/api/tool_lifecycle.rs index c8a0fd644..dbe699ff2 100644 --- a/crates/ffi/src/api/tool_lifecycle.rs +++ b/crates/ffi/src/api/tool_lifecycle.rs @@ -3,10 +3,11 @@ use super::{ Arc, FfiScopeHandle, FfiToolHandle, NemoRelayFreeFn, NemoRelayStatus, NemoRelayToolExecCb, - TASK_SCOPE_STACK, ToolAttributes, ToolExecutionNextFn, c_char, c_str_to_json, - c_str_to_opt_json, c_str_to_string, clear_last_error, core_tool_api, current_scope_stack, - json_to_c_string, set_last_error, status_from_error, tokio_runtime, - unix_micros_to_opt_timestamp, wrap_tool_exec_fn, + NemoRelayToolExecFrameCb, TASK_SCOPE_STACK, ToolAttributes, ToolExecutionFrameNextFn, + ToolExecutionNextFn, c_char, c_str_to_json, c_str_to_opt_json, c_str_to_string, + clear_last_error, core_tool_api, current_scope_stack, json_to_c_string, set_last_error, + status_from_error, tokio_runtime, unix_micros_to_opt_timestamp, wrap_tool_exec_fn, + wrap_tool_exec_frame_fn, }; // --------------------------------------------------------------------------- @@ -187,6 +188,66 @@ pub unsafe extern "C" fn nemo_relay_tool_call_end( } } +/// End a manual tool call with a serialized `ToolExecutionFrame`. +/// +/// The frame's raw result follows the existing response-sanitization path. Its +/// optional producer annotation is carried in the emitted tool category +/// profile. +/// +/// # Safety +/// `handle` and `frame_json` must be valid, non-null pointers. Optional pointer +/// arguments follow [`nemo_relay_tool_call_end`]. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_tool_call_end_frame( + handle: *const FfiToolHandle, + frame_json: *const c_char, + data_json: *const c_char, + metadata_json: *const c_char, + timestamp_unix_micros: *const i64, +) -> NemoRelayStatus { + clear_last_error(); + if handle.is_null() { + set_last_error("handle is null"); + return NemoRelayStatus::NullPointer; + } + let frame_json = match c_str_to_json(frame_json) { + Some(value) => value, + None => return NemoRelayStatus::InvalidJson, + }; + let frame = match serde_json::from_value::(frame_json) { + Ok(frame) => frame, + Err(error) => { + set_last_error(&format!("invalid tool execution frame JSON: {error}")); + return NemoRelayStatus::InvalidJson; + } + }; + let data = match c_str_to_opt_json(data_json) { + Some(value) => value, + None => return NemoRelayStatus::InvalidJson, + }; + let metadata = match c_str_to_opt_json(metadata_json) { + Some(value) => value, + None => return NemoRelayStatus::InvalidJson, + }; + let timestamp = match unix_micros_to_opt_timestamp(timestamp_unix_micros) { + Some(value) => value, + None => return NemoRelayStatus::InvalidArg, + }; + + match core_tool_api::tool_call_end_frame( + core_tool_api::ToolCallEndFrameParams::builder() + .handle(&unsafe { &*handle }.0) + .frame(frame) + .data_opt(data) + .metadata_opt(metadata) + .timestamp_opt(timestamp) + .build(), + ) { + Ok(()) => NemoRelayStatus::Ok, + Err(error) => status_from_error(&error), + } +} + /// Execute a tool call end-to-end: run conditional-execution guardrails (on raw /// args), then request intercepts, sanitize-request guardrails, execution /// intercepts, the callback, and sanitize-response @@ -277,3 +338,85 @@ pub unsafe extern "C" fn nemo_relay_tool_call_execute( Err(e) => status_from_error(&e), } } + +/// Execute a tool call with an annotation-aware producer callback. +/// +/// The callback and output use serialized `ToolExecutionFrame` values. This +/// function participates in the same execution-intercept chain as +/// [`nemo_relay_tool_call_execute`]. +/// +/// # Safety +/// `name`, `args_json`, and `out` must be valid, non-null pointers. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_tool_call_execute_frame( + name: *const c_char, + args_json: *const c_char, + func: NemoRelayToolExecFrameCb, + func_user_data: *mut libc::c_void, + func_free: NemoRelayFreeFn, + parent: *const FfiScopeHandle, + attributes: u32, + data_json: *const c_char, + metadata_json: *const c_char, + out: *mut *mut c_char, +) -> NemoRelayStatus { + clear_last_error(); + if out.is_null() { + set_last_error("out pointer is null"); + return NemoRelayStatus::NullPointer; + } + let name = match c_str_to_string(name) { + Ok(value) => value, + Err(status) => return status, + }; + let args = match c_str_to_json(args_json) { + Some(value) => value, + None => return NemoRelayStatus::InvalidJson, + }; + let parent_handle = if parent.is_null() { + None + } else { + Some(unsafe { &*parent }.0.clone()) + }; + let attributes = ToolAttributes::from_bits_truncate(attributes); + let data = match c_str_to_opt_json(data_json) { + Some(value) => value, + None => return NemoRelayStatus::InvalidJson, + }; + let metadata = match c_str_to_opt_json(metadata_json) { + Some(value) => value, + None => return NemoRelayStatus::InvalidJson, + }; + + let exec_fn = wrap_tool_exec_frame_fn(func, func_user_data, func_free); + let default_fn: ToolExecutionFrameNextFn = Arc::new(move |args| exec_fn(args)); + let scope_stack = current_scope_stack(); + let result = tokio_runtime().block_on(TASK_SCOPE_STACK.scope(scope_stack, async { + core_tool_api::tool_call_execute_frame( + core_tool_api::ToolCallExecuteFrameParams::builder() + .name(name) + .args(args) + .func(default_fn) + .parent_opt(parent_handle) + .attributes(attributes) + .data_opt(data) + .metadata_opt(metadata) + .build(), + ) + .await + })); + + match result { + Ok(frame) => match serde_json::to_value(frame) { + Ok(frame) => { + unsafe { *out = json_to_c_string(&frame) }; + NemoRelayStatus::Ok + } + Err(error) => { + set_last_error(&error.to_string()); + NemoRelayStatus::Internal + } + }, + Err(error) => status_from_error(&error), + } +} diff --git a/crates/ffi/src/api/tool_registry.rs b/crates/ffi/src/api/tool_registry.rs index 5d5cefb40..aa5417985 100644 --- a/crates/ffi/src/api/tool_registry.rs +++ b/crates/ffi/src/api/tool_registry.rs @@ -2,9 +2,10 @@ // SPDX-License-Identifier: Apache-2.0 use super::{ - NemoRelayFreeFn, NemoRelayStatus, NemoRelayToolConditionalCb, NemoRelayToolExecInterceptCb, - NemoRelayToolSanitizeCb, c_char, c_str_to_string, clear_last_error, core_registry_api, - status_from_error, wrap_tool_conditional_fn, wrap_tool_exec_intercept_fn, + NemoRelayFreeFn, NemoRelayStatus, NemoRelayToolConditionalCb, + NemoRelayToolExecFrameInterceptCb, NemoRelayToolExecInterceptCb, NemoRelayToolSanitizeCb, + c_char, c_str_to_string, clear_last_error, core_registry_api, status_from_error, + wrap_tool_conditional_fn, wrap_tool_exec_frame_intercept_fn, wrap_tool_exec_intercept_fn, wrap_tool_request_intercept_fn, wrap_tool_sanitize_fn, }; @@ -270,6 +271,34 @@ pub unsafe extern "C" fn nemo_relay_register_tool_execution_intercept( } } +/// Register an annotation-aware tool execution intercept in the existing chain. +/// +/// The callback receives a continuation that returns serialized +/// `ToolExecutionFrame` JSON and must return serialized +/// `ToolExecutionFrameOutcome` JSON. +/// +/// # Safety +/// `name` must be a valid C string. Callback pointers must be valid. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_register_tool_execution_frame_intercept( + name: *const c_char, + priority: i32, + exec_cb: NemoRelayToolExecFrameInterceptCb, + exec_user_data: *mut libc::c_void, + exec_free: NemoRelayFreeFn, +) -> NemoRelayStatus { + clear_last_error(); + let name = match c_str_to_string(name) { + Ok(value) => value, + Err(status) => return status, + }; + let exec = wrap_tool_exec_frame_intercept_fn(exec_cb, exec_user_data, exec_free); + match core_registry_api::register_tool_execution_frame_intercept(&name, priority, exec) { + Ok(()) => NemoRelayStatus::Ok, + Err(error) => status_from_error(&error), + } +} + /// Deregister a tool execution intercept by name. /// /// # Safety diff --git a/crates/ffi/src/callable.rs b/crates/ffi/src/callable.rs index 054c098bd..7072c08bd 100644 --- a/crates/ffi/src/callable.rs +++ b/crates/ffi/src/callable.rs @@ -26,14 +26,17 @@ use nemo_relay::api::runtime::{ EventSanitizeFn, EventSubscriberFn, LlmCodecIdentity, LlmConditionalFn, LlmExecutionNextFn, LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn, LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, ToolConditionalFn, - ToolExecutionFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, + ToolExecutionFn, ToolExecutionFrameFn, ToolExecutionFrameNextFn, ToolExecutionNextFn, + ToolInterceptFn, ToolSanitizeFn, }; use serde_json::Value as Json; use tokio_stream::StreamExt; use nemo_relay::api::event::{Event, EventSanitizeFields}; use nemo_relay::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; -use nemo_relay::api::tool::ToolExecutionInterceptOutcome; +use nemo_relay::api::tool::{ + ToolExecutionFrame, ToolExecutionFrameOutcome, ToolExecutionInterceptOutcome, +}; use nemo_relay::codec::request::AnnotatedLlmRequest as AnnotatedLLMRequest; use nemo_relay::codec::traits::LlmCodec; use nemo_relay::error::{FlowError, Result}; @@ -78,6 +81,12 @@ pub type NemoRelayToolConditionalCb = unsafe extern "C" fn( pub type NemoRelayToolExecCb = unsafe extern "C" fn(user_data: *mut libc::c_void, args_json: *const c_char) -> *mut c_char; +/// Annotation-aware tool execution callback. +/// +/// The returned JSON must serialize a `ToolExecutionFrame`. +pub type NemoRelayToolExecFrameCb = + unsafe extern "C" fn(user_data: *mut libc::c_void, args_json: *const c_char) -> *mut c_char; + /// Runtime-provided "next" callback for tool execution middleware chain. /// Call this from an intercept to invoke the next layer (or original function). /// `next_ctx` is borrowed and valid only until the intercept callback returns; @@ -87,6 +96,12 @@ pub type NemoRelayToolExecCb = pub type NemoRelayToolExecNextFn = unsafe extern "C" fn(args_json: *const c_char, next_ctx: *mut libc::c_void) -> *mut c_char; +/// Runtime-provided annotation-aware continuation. +/// +/// The returned JSON serializes a `ToolExecutionFrame`. +pub type NemoRelayToolExecFrameNextFn = + unsafe extern "C" fn(args_json: *const c_char, next_ctx: *mut libc::c_void) -> *mut c_char; + /// Callback for tool execution intercepts. Receives arguments as JSON plus /// a `next` callback and its context. Call `next_fn(args, next_ctx)` to invoke /// the next layer in the middleware chain, or return directly to short-circuit. @@ -105,6 +120,17 @@ pub type NemoRelayToolExecInterceptCb = unsafe extern "C" fn( next_ctx: *mut libc::c_void, ) -> *mut c_char; +/// Annotation-aware tool execution intercept callback. +/// +/// `next_fn` returns a serialized `ToolExecutionFrame`; this callback must +/// return a serialized `ToolExecutionFrameOutcome`. +pub type NemoRelayToolExecFrameInterceptCb = unsafe extern "C" fn( + user_data: *mut libc::c_void, + args_json: *const c_char, + next_fn: NemoRelayToolExecFrameNextFn, + next_ctx: *mut libc::c_void, +) -> *mut c_char; + /// Codec identity kind supplied to an LLM sanitizer. #[repr(u32)] #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -420,6 +446,31 @@ pub fn wrap_tool_exec_fn( }) } +/// Wrap a C annotation-aware tool execution callback into an async Rust closure. +pub fn wrap_tool_exec_frame_fn( + cb: NemoRelayToolExecFrameCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, +) -> Box< + dyn Fn(Json) -> Pin> + Send>> + Send + Sync, +> { + let ud = make_user_data(user_data, free_fn); + Box::new(move |args: Json| { + let ud = ud.clone(); + Box::pin(async move { + let c_args = json_to_c_string(&args); + let result_ptr = unsafe { cb(ud.ptr, c_args) }; + unsafe { nemo_relay_string_free_internal(c_args) }; + let frame_json = + json_result_from_ptr(result_ptr, "tool execution frame callback failed")?; + unsafe { nemo_relay_string_free_internal(result_ptr) }; + serde_json::from_value::(frame_json).map_err(|error| { + FlowError::Internal(format!("invalid tool execution frame JSON: {error}")) + }) + }) + }) +} + /// Wrap a C tool execution intercept callback into a [`ToolExecutionFn`]. /// /// The wrapper packages the Rust `ToolExecutionNextFn` into a C-callable @@ -483,6 +534,69 @@ pub fn wrap_tool_exec_intercept_fn( }) } +/// Wrap a C annotation-aware tool execution intercept into a [`ToolExecutionFrameFn`]. +pub fn wrap_tool_exec_frame_intercept_fn( + cb: NemoRelayToolExecFrameInterceptCb, + user_data: *mut libc::c_void, + free_fn: NemoRelayFreeFn, +) -> ToolExecutionFrameFn { + let ud = make_user_data(user_data, free_fn); + Arc::new( + move |_name: &str, args: Json, next: ToolExecutionFrameNextFn| { + let ud = ud.clone(); + Box::pin(async move { + let next_box = Box::new(next); + let next_ctx = Box::into_raw(next_box) as *mut libc::c_void; + + unsafe extern "C" fn tool_frame_next_trampoline( + args_json: *const c_char, + next_ctx: *mut libc::c_void, + ) -> *mut c_char { + let next_arc = unsafe { &*(next_ctx as *const ToolExecutionFrameNextFn) }; + let next = next_arc.clone(); + let args = if args_json.is_null() { + Json::Null + } else { + let value = unsafe { CStr::from_ptr(args_json) }.to_string_lossy(); + serde_json::from_str(&value).unwrap_or(Json::Null) + }; + let handle = tokio::runtime::Handle::current(); + let result = tokio::task::block_in_place(|| handle.block_on(next(args))); + match result { + Ok(frame) => match serde_json::to_value(frame) { + Ok(value) => json_to_c_string(&value), + Err(error) => { + set_last_error(&error.to_string()); + std::ptr::null_mut() + } + }, + Err(error) => { + set_last_error(&error.to_string()); + std::ptr::null_mut() + } + } + } + + let c_args = json_to_c_string(&args); + let result_ptr = + unsafe { cb(ud.ptr, c_args, tool_frame_next_trampoline, next_ctx) }; + unsafe { drop(Box::from_raw(next_ctx as *mut ToolExecutionFrameNextFn)) }; + unsafe { nemo_relay_string_free_internal(c_args) }; + let outcome_json = json_result_from_ptr( + result_ptr, + "tool execution frame intercept callback failed", + )?; + unsafe { nemo_relay_string_free_internal(result_ptr) }; + serde_json::from_value::(outcome_json).map_err(|error| { + FlowError::Internal(format!( + "invalid tool execution frame outcome JSON: {error}" + )) + }) + }) + }, + ) +} + /// Wrap a C LLM execution intercept callback into an `Arc ...>`. pub fn wrap_llm_exec_intercept_fn( cb: NemoRelayLlmExecInterceptCb, diff --git a/crates/ffi/tests/unit/api/execution_tests.rs b/crates/ffi/tests/unit/api/execution_tests.rs index 0d158544d..bc84eb6e7 100644 --- a/crates/ffi/tests/unit/api/execution_tests.rs +++ b/crates/ffi/tests/unit/api/execution_tests.rs @@ -5,6 +5,50 @@ use super::*; +unsafe extern "C" fn tool_exec_frame_cb( + _user_data: *mut libc::c_void, + args_json: *const c_char, +) -> *mut c_char { + let args: Json = serde_json::from_str( + unsafe { CStr::from_ptr(args_json) } + .to_str() + .unwrap_or("null"), + ) + .unwrap(); + CString::new( + json!({ + "result": args, + "annotation": { + "producer": "ffi", + "status": "failed" + } + }) + .to_string(), + ) + .unwrap() + .into_raw() +} + +unsafe extern "C" fn tool_exec_frame_intercept_cb( + _user_data: *mut libc::c_void, + args_json: *const c_char, + next_fn: NemoRelayToolExecFrameNextFn, + next_ctx: *mut libc::c_void, +) -> *mut c_char { + let frame_ptr = unsafe { next_fn(args_json, next_ctx) }; + if frame_ptr.is_null() { + return ptr::null_mut(); + } + let mut frame: Json = + serde_json::from_str(unsafe { CStr::from_ptr(frame_ptr) }.to_str().unwrap()).unwrap(); + unsafe { nemo_relay_string_free(frame_ptr) }; + frame["result"]["frame_seen"] = json!(true); + frame["annotation"]["observed_by"] = json!("ffi_frame_outer"); + CString::new(json!({ "frame": frame, "pending_marks": [] }).to_string()) + .unwrap() + .into_raw() +} + #[test] fn test_ffi_tool_execute_parent_data_and_error_paths() { let _lock = TEST_MUTEX.lock().unwrap_or_else(|e| e.into_inner()); @@ -82,6 +126,90 @@ fn test_ffi_tool_execute_parent_data_and_error_paths() { } } +#[test] +fn test_ffi_tool_execute_frame_uses_shared_chain_and_deregistration() { + let _lock = TEST_MUTEX.lock().unwrap_or_else(|e| e.into_inner()); + reset_globals(); + + unsafe { + let stack = fresh_scope_stack(); + let frame_name = cstring(&unique_name("ffi_frame_outer")); + let legacy_name = cstring(&unique_name("ffi_frame_legacy")); + + assert_eq!( + nemo_relay_register_tool_execution_frame_intercept( + frame_name.as_ptr(), + 1, + tool_exec_frame_intercept_cb, + ptr::null_mut(), + None, + ), + NemoRelayStatus::Ok + ); + assert_eq!( + nemo_relay_register_tool_execution_intercept( + legacy_name.as_ptr(), + 2, + tool_exec_intercept_cb, + ptr::null_mut(), + None, + ), + NemoRelayStatus::Ok + ); + + let tool_name = cstring("ffi_frame_tool"); + let args = cstring(r#"{"value":6}"#); + let mut out_json = ptr::null_mut(); + assert_eq!( + nemo_relay_tool_call_execute_frame( + tool_name.as_ptr(), + args.as_ptr(), + tool_exec_frame_cb, + ptr::null_mut(), + None, + ptr::null(), + 0, + ptr::null(), + ptr::null(), + &mut out_json, + ), + NemoRelayStatus::Ok + ); + let frame = returned_json(out_json); + assert_eq!(frame["result"]["value"], 6); + assert_eq!(frame["result"]["frame_seen"], true); + assert_eq!(frame["annotation"]["producer"], "ffi"); + assert_eq!(frame["annotation"]["observed_by"], "ffi_frame_outer"); + + assert_eq!( + nemo_relay_deregister_tool_execution_intercept(frame_name.as_ptr()), + NemoRelayStatus::Ok + ); + assert_eq!( + nemo_relay_deregister_tool_execution_intercept(legacy_name.as_ptr()), + NemoRelayStatus::Ok + ); + + // The generic deregistration removed the frame callback, so the same + // name can be registered again through the frame-aware surface. + assert_eq!( + nemo_relay_register_tool_execution_frame_intercept( + frame_name.as_ptr(), + 1, + tool_exec_frame_intercept_cb, + ptr::null_mut(), + None, + ), + NemoRelayStatus::Ok + ); + assert_eq!( + nemo_relay_deregister_tool_execution_intercept(frame_name.as_ptr()), + NemoRelayStatus::Ok + ); + nemo_relay_scope_stack_free(stack); + } +} + #[test] fn test_ffi_llm_execute_codec_parent_and_error_paths() { let _lock = TEST_MUTEX.lock().unwrap_or_else(|e| e.into_inner()); diff --git a/crates/ffi/tests/unit/api_tests.rs b/crates/ffi/tests/unit/api_tests.rs index 2a7f41246..59e607a51 100644 --- a/crates/ffi/tests/unit/api_tests.rs +++ b/crates/ffi/tests/unit/api_tests.rs @@ -17,7 +17,7 @@ use uuid::Uuid; use crate::callable::{ NemoRelayLlmExecNextFn, NemoRelayLlmSanitizeCodecKind, NemoRelayLlmSanitizeRequestContext, - NemoRelayLlmSanitizeResponseContext, NemoRelayToolExecNextFn, + NemoRelayLlmSanitizeResponseContext, NemoRelayToolExecFrameNextFn, NemoRelayToolExecNextFn, }; use crate::convert::nemo_relay_string_free; use crate::error::{NemoRelayStatus, nemo_relay_last_error}; diff --git a/crates/node/plugin.d.ts b/crates/node/plugin.d.ts index 21f8771dc..d94180d10 100644 --- a/crates/node/plugin.d.ts +++ b/crates/node/plugin.d.ts @@ -176,7 +176,7 @@ export interface LlmRequestInterceptOutcome { } /** - * Canonical result returned by a tool execution intercept. + * Relay-owned wrapper returned by a raw-result tool execution intercept. * * `result` is passed to the remaining middleware and application. `pendingMarks` * are Relay-owned lifecycle metadata emitted after the tool-end event and are @@ -187,6 +187,18 @@ export interface ToolExecutionInterceptOutcome { pendingMarks?: PendingMarkSpec[]; } +/** Raw tool result plus an optional annotation that Relay preserves opaquely. */ +export interface ToolExecutionFrame { + result: Json; + annotation?: Json; +} + +/** Outcome returned by annotation-aware tool execution middleware. */ +export interface ToolExecutionFrameOutcome { + frame: ToolExecutionFrame; + pendingMarks?: PendingMarkSpec[]; +} + /** Component-scoped registration context passed to plugin handlers. */ export interface PluginContext { /** Register an infallible event subscriber for this component. */ @@ -281,7 +293,7 @@ export interface PluginContext { callback: (name: string, args: Json) => Json | Promise, ): void; /** - * Register tool execution middleware that returns a canonical outcome. + * Register raw-result tool execution middleware that returns Relay's outcome wrapper. * The `next` callback resolves to the raw downstream result. */ registerToolExecutionIntercept( @@ -292,6 +304,15 @@ export interface PluginContext { next: (args: Json) => Json | Promise, ) => ToolExecutionInterceptOutcome | Promise, ): void; + /** Register annotation-aware middleware in the same tool execution chain. */ + registerToolExecutionFrameIntercept( + name: string, + priority: number, + callback: ( + args: Json, + next: (args: Json) => ToolExecutionFrame | Promise, + ) => ToolExecutionFrameOutcome | Promise, + ): void; } /** Plugin callback contract. */ diff --git a/crates/node/src/api/mod.rs b/crates/node/src/api/mod.rs index 12aff6e7f..d5887ee16 100644 --- a/crates/node/src/api/mod.rs +++ b/crates/node/src/api/mod.rs @@ -36,7 +36,7 @@ use nemo_relay::api::runtime::subscriber_dispatcher::{ }; use nemo_relay::api::runtime::{ EventSanitizeFn, LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, LlmStreamInner, - ScopeStackHandle as CoreScopeStackHandle, ToolExecutionNextFn, + ScopeStackHandle as CoreScopeStackHandle, ToolExecutionFrameNextFn, ToolExecutionNextFn, }; use nemo_relay::api::runtime::{ TASK_SCOPE_STACK, capture_propagation_context as capture_propagation_context_handle, @@ -86,7 +86,7 @@ use crate::convert::{ use crate::promise_call::PromiseAwareFn; use crate::promise_call::with_publication_callback_context; use crate::stream::LlmStream; -use crate::types::{LlmHandle, ScopeHandle, ScopeStack, ScopeType, ToolHandle}; +use crate::types::{LlmHandle, ScopeHandle, ScopeStack, ScopeType, ToolExecutionFrame, ToolHandle}; fn effective_scope_context( env: &Env, @@ -1186,7 +1186,7 @@ fn build_plugin_context( )?; let tool_regs = registrations.clone(); - let tool_exec_namespace = namespace_prefix; + let tool_exec_namespace = namespace_prefix.clone(); let register_tool_execution_intercept = env.create_function_from_closure( "__nemo_relay_adaptive_register_tool_execution_intercept", move |ctx| { @@ -1224,6 +1224,51 @@ fn build_plugin_context( register_tool_execution_intercept, )?; + let tool_frame_regs = registrations.clone(); + let register_tool_execution_frame_intercept = env.create_function_from_closure( + "__nemo_relay_adaptive_register_tool_execution_frame_intercept", + move |ctx| { + let name = format!("{}{}", namespace_prefix, ctx.get::(0)?); + let priority = ctx.get::(1)?; + let callback = ctx.get::(2)?; + let promise_fn = Arc::new(crate::promise_call::PromiseAwareFn::new( + ctx.env, &callback, + )?); + core_registry_api::register_tool_execution_frame_intercept( + &name, + priority, + callable::wrap_js_tool_exec_frame_intercept_fn(promise_fn.clone()), + ) + .map_err(to_napi_err)?; + + let name_clone = name.clone(); + tool_frame_regs + .lock() + .unwrap() + .push(PluginRegistration::new( + "plugin", + name_clone.clone(), + Box::new(move || { + let result = + core_registry_api::deregister_tool_execution_intercept(&name_clone) + .map(|_| ()) + .map_err(|e| { + PluginError::RegistrationFailed(format!( + "tool execution frame intercept deregistration failed: {e}" + )) + }); + promise_fn.close(); + result + }), + )); + ctx.env.get_undefined() + }, + )?; + context.set_named_property( + "registerToolExecutionFrameIntercept", + register_tool_execution_frame_intercept, + )?; + Ok(context) } @@ -2222,6 +2267,28 @@ pub fn tool_call_end( .map_err(to_napi_err) } +/// End a manual tool span with an optional opaque result annotation. +#[napi] +pub fn tool_call_end_frame( + handle: &ToolHandle, + frame: ToolExecutionFrame, + data: Option, + metadata: Option, + timestamp: Option, +) -> Result<()> { + let timestamp = parse_timestamp_micros(timestamp)?; + core_tool_api::tool_call_end_frame( + core_tool_api::ToolCallEndFrameParams::builder() + .handle(&handle.inner) + .frame(frame.into()) + .data_opt(opt_json(data)) + .metadata_opt(opt_json(metadata)) + .timestamp_opt(timestamp) + .build(), + ) + .map_err(to_napi_err) +} + /// Execute a tool call end-to-end with full lifecycle management. /// /// Runs conditional-execution guardrails (on raw args) → request intercepts → @@ -2283,6 +2350,63 @@ pub fn tool_call_execute( ) } +/// Execute a tool call with an optional opaque result annotation. +#[allow(clippy::too_many_arguments)] +#[napi(ts_return_type = "Promise")] +pub fn tool_call_execute_frame( + env: Env, + name: String, + args: Json, + #[napi(ts_arg_type = "(arg: Json) => ToolExecutionFrame")] func: JsFunction, + handle: Option<&ScopeHandle>, + attributes: Option, + data: Option, + metadata: Option, +) -> Result { + let attrs = ToolAttributes::from_bits_truncate(attributes.unwrap_or(0)); + let parent = handle + .map(|h| h.inner.clone()) + .unwrap_or_else(task_scope_top); + let callback = callable::safe_execution_callback(&env, &func)?; + let exec_fn = callable::wrap_js_tool_exec_fn(json_callback_tsfn(&env, &callback)?); + let default_fn: ToolExecutionFrameNextFn = Arc::new(move |args| { + let future = exec_fn(args); + Box::pin(async move { + let value = future.await?; + serde_json::from_value(value).map_err(|error| { + FlowError::Internal(format!( + "tool frame callback must return ToolExecutionFrame: {error}" + )) + }) + }) + }); + let scope_stack = current_scope_stack_handle(); + + env.execute_tokio_future( + async move { + TASK_SCOPE_STACK + .scope(scope_stack, async move { + core_tool_api::tool_call_execute_frame( + core_tool_api::ToolCallExecuteFrameParams::builder() + .name(name) + .args(args) + .func(default_fn) + .parent(parent) + .attributes(attrs) + .data_opt(opt_json(data)) + .metadata_opt(opt_json(metadata)) + .build(), + ) + .await + .map(ToolExecutionFrame::from) + .map_err(to_napi_err) + }) + .await + }, + |_env, frame| Ok(frame), + ) +} + /// Execute a tool call end-to-end, supporting both sync and async (Promise-returning) callbacks. /// /// Same lifecycle as `toolCallExecute` (guardrails → intercepts → func → response processing), @@ -2353,6 +2477,67 @@ pub fn tool_call_execute_async( ) } +/// Promise-aware form of `toolCallExecuteFrame`. +#[allow(clippy::too_many_arguments)] +#[napi(ts_return_type = "Promise")] +pub fn tool_call_execute_frame_async( + env: Env, + name: String, + args: Json, + #[napi(ts_arg_type = "(arg: Json) => ToolExecutionFrame | Promise")] + func: JsFunction, + handle: Option<&ScopeHandle>, + attributes: Option, + data: Option, + metadata: Option, +) -> Result { + let attrs = ToolAttributes::from_bits_truncate(attributes.unwrap_or(0)); + let parent = handle + .map(|h| h.inner.clone()) + .unwrap_or_else(task_scope_top); + let scope_stack = current_scope_stack_handle(); + let pa_fn = Arc::new( + crate::promise_call::PromiseAwareFn::new(&env, &func).map_err(|e| { + napi::Error::from_reason(format!("failed to create PromiseAwareFn: {e}")) + })?, + ); + let default_fn: ToolExecutionFrameNextFn = Arc::new(move |args| { + let pa_fn = pa_fn.clone(); + Box::pin(async move { + let value = pa_fn.call(args).await?; + serde_json::from_value(value).map_err(|error| { + FlowError::Internal(format!( + "tool frame callback must return ToolExecutionFrame: {error}" + )) + }) + }) + }); + + env.execute_tokio_future( + async move { + TASK_SCOPE_STACK + .scope(scope_stack, async move { + core_tool_api::tool_call_execute_frame( + core_tool_api::ToolCallExecuteFrameParams::builder() + .name(name) + .args(args) + .func(default_fn) + .parent(parent) + .attributes(attrs) + .data_opt(opt_json(data)) + .metadata_opt(opt_json(metadata)) + .build(), + ) + .await + .map(ToolExecutionFrame::from) + .map_err(to_napi_err) + }) + .await + }, + |_env, frame| Ok(frame), + ) +} + // --------------------------------------------------------------------------- // LLM lifecycle // --------------------------------------------------------------------------- @@ -3048,7 +3233,7 @@ pub fn register_tool_execution_intercept( Ok(()) } -/// Deregister a tool execution intercept by name. +/// Deregister a tool execution intercept of either callback form by name. /// /// Returns `true` if an intercept with that name was found and removed. #[napi] @@ -3056,6 +3241,31 @@ pub fn deregister_tool_execution_intercept(name: String) -> Result { core_registry_api::deregister_tool_execution_intercept(&name).map_err(to_napi_err) } +/// Register annotation-aware middleware in the existing tool execution chain. +#[napi] +pub fn register_tool_execution_frame_intercept( + env: Env, + name: String, + priority: i32, + #[napi( + ts_arg_type = "(args: Json, next: (args: Json) => ToolExecutionFrame | Promise) => { frame: ToolExecutionFrame; pendingMarks?: Array<{ name: string; category?: string | null; categoryProfile?: Json; data?: Json; metadata?: Json }> } | Promise<{ frame: ToolExecutionFrame; pendingMarks?: Array<{ name: string; category?: string | null; categoryProfile?: Json; data?: Json; metadata?: Json }> }>" + )] + callable: JsFunction, +) -> Result<()> { + let pa_fn = Arc::new( + crate::promise_call::PromiseAwareFn::new(&env, &callable).map_err(|e| { + napi::Error::from_reason(format!("failed to create PromiseAwareFn: {e}")) + })?, + ); + core_registry_api::register_tool_execution_frame_intercept( + &name, + priority, + callable::wrap_js_tool_exec_frame_intercept_fn(pa_fn.clone()), + ) + .map_err(to_napi_err)?; + Ok(()) +} + // --------------------------------------------------------------------------- // LLM guardrail registrations // --------------------------------------------------------------------------- @@ -3625,7 +3835,7 @@ pub fn scope_register_tool_execution_intercept( Ok(()) } -/// Deregister a scope-local tool execution intercept by name. +/// Deregister a scope-local tool execution intercept of either callback form by name. /// /// Returns `true` if an intercept with that name was found and removed from the specified scope. #[napi] @@ -3635,6 +3845,35 @@ pub fn scope_deregister_tool_execution_intercept(scope_uuid: String, name: Strin core_registry_api::scope_deregister_tool_execution_intercept(&uuid, &name).map_err(to_napi_err) } +/// Register annotation-aware scope-local middleware in the existing chain. +#[napi] +pub fn scope_register_tool_execution_frame_intercept( + env: Env, + scope_uuid: String, + name: String, + priority: i32, + #[napi( + ts_arg_type = "(args: Json, next: (args: Json) => ToolExecutionFrame | Promise) => { frame: ToolExecutionFrame; pendingMarks?: Array<{ name: string; category?: string | null; categoryProfile?: Json; data?: Json; metadata?: Json }> } | Promise<{ frame: ToolExecutionFrame; pendingMarks?: Array<{ name: string; category?: string | null; categoryProfile?: Json; data?: Json; metadata?: Json }> }>" + )] + callable: JsFunction, +) -> Result<()> { + let uuid = uuid::Uuid::parse_str(&scope_uuid) + .map_err(|e| napi::Error::from_reason(format!("invalid UUID: {e}")))?; + let pa_fn = Arc::new( + crate::promise_call::PromiseAwareFn::new(&env, &callable).map_err(|e| { + napi::Error::from_reason(format!("failed to create PromiseAwareFn: {e}")) + })?, + ); + core_registry_api::scope_register_tool_execution_frame_intercept( + &uuid, + &name, + priority, + callable::wrap_js_tool_exec_frame_intercept_fn(pa_fn.clone()), + ) + .map_err(to_napi_err)?; + Ok(()) +} + // --------------------------------------------------------------------------- // Scope-local guardrail registrations — LLM // --------------------------------------------------------------------------- diff --git a/crates/node/src/callable.rs b/crates/node/src/callable.rs index 731348696..c9b32b565 100644 --- a/crates/node/src/callable.rs +++ b/crates/node/src/callable.rs @@ -21,7 +21,8 @@ use nemo_relay::api::runtime::{ EventSanitizeFn, EventSubscriberFn, LlmCodecIdentity, LlmConditionalFn, LlmExecutionNextFn, LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn, LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, ToolConditionalFn, - ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, + ToolExecutionFrameFn, ToolExecutionFrameNextFn, ToolExecutionNextFn, ToolInterceptFn, + ToolSanitizeFn, }; use serde::{Deserialize, Serialize}; use serde_json::Value as Json; @@ -32,7 +33,9 @@ use nemo_relay::api::event::{ PendingMarkSpec, }; use nemo_relay::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; -use nemo_relay::api::tool::ToolExecutionInterceptOutcome; +use nemo_relay::api::tool::{ + ToolExecutionFrame, ToolExecutionFrameOutcome, ToolExecutionInterceptOutcome, +}; use nemo_relay::codec::optimization::LlmOptimizationContribution; use nemo_relay::codec::request::AnnotatedLlmRequest; use nemo_relay::codec::response::AnnotatedLlmResponse; @@ -1348,6 +1351,46 @@ pub fn wrap_js_tool_exec_intercept_fn( }) } +/// Wrap an annotation-aware JS tool execution intercept. +/// +/// The callback participates in the existing tool execution chain. Its +/// `next(args)` promise resolves to a serialized `ToolExecutionFrame`. +pub fn wrap_js_tool_exec_frame_intercept_fn(func: Arc) -> ToolExecutionFrameFn { + Arc::new( + move |_name: &str, args: Json, next: ToolExecutionFrameNextFn| { + let func = func.clone(); + let next_json: JsonNextFn = Arc::new(move |next_args| { + let next = next.clone(); + Box::pin(async move { + let frame = next(next_args).await?; + serde_json::to_value(frame).map_err(|error| { + FlowError::Internal(format!( + "failed to serialize downstream tool execution frame: {error}" + )) + }) + }) + }); + Box::pin(async move { + let result = func.call_with_json_next(args, next_json).await?; + #[derive(Deserialize)] + #[serde(rename_all = "camelCase")] + struct JsOutcome { + frame: ToolExecutionFrame, + #[serde(default)] + pending_marks: Vec, + } + let outcome: JsOutcome = serde_json::from_value(result).map_err(|error| { + FlowError::Internal(format!("invalid JS tool execution frame outcome: {error}")) + })?; + Ok(ToolExecutionFrameOutcome { + frame: outcome.frame, + pending_marks: outcome.pending_marks.into_iter().map(Into::into).collect(), + }) + }) + }, + ) +} + /// Wrap a JS function `(request, next) => result` for LLM execution intercept. /// /// The JS callback receives the `LlmRequest` serialized as a plain JSON object diff --git a/crates/node/src/types/mod.rs b/crates/node/src/types/mod.rs index 2a36c4dc9..350fcd2b2 100644 --- a/crates/node/src/types/mod.rs +++ b/crates/node/src/types/mod.rs @@ -315,6 +315,33 @@ pub struct EventSanitizeFields { pub metadata: Option, } +/// Raw tool result plus an optional opaque annotation. +#[napi(object)] +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ToolExecutionFrame { + pub result: Json, + pub annotation: Option, +} + +impl From for ToolExecutionFrame { + fn from(frame: nemo_relay::api::tool::ToolExecutionFrame) -> Self { + Self { + result: frame.result, + annotation: frame.annotation, + } + } +} + +impl From for nemo_relay::api::tool::ToolExecutionFrame { + fn from(frame: ToolExecutionFrame) -> Self { + Self { + result: frame.result, + annotation: frame.annotation, + } + } +} + pub(crate) fn event_sanitize_fields_from_json( value: Json, ) -> serde_json::Result { diff --git a/crates/node/tests/scope_local_tests.mjs b/crates/node/tests/scope_local_tests.mjs index 919bc2187..a9fd49561 100644 --- a/crates/node/tests/scope_local_tests.mjs +++ b/crates/node/tests/scope_local_tests.mjs @@ -32,6 +32,7 @@ const { popScope, event, toolCallExecute, + toolCallExecuteFrameAsync, llmCallExecute, llmStreamCallExecute, scopeRegisterToolSanitizeRequestGuardrail, @@ -43,6 +44,7 @@ const { scopeRegisterToolRequestIntercept, scopeDeregisterToolRequestIntercept, scopeRegisterToolExecutionIntercept, + scopeRegisterToolExecutionFrameIntercept, scopeDeregisterToolExecutionIntercept, scopeRegisterLlmSanitizeRequestGuardrail, scopeDeregisterLlmSanitizeRequestGuardrail, @@ -827,6 +829,68 @@ describe('Priority merge of global and scope-local middleware', () => { } }); + it('scope-local frame intercept uses the shared registry and deregistration', async () => { + const scope = pushScope('sl_frame_exec_scope', ScopeType.Agent, null, null); + scopeRegisterToolExecutionFrameIntercept(scope.uuid, 'sl_frame_exec', 10, async (args, next) => { + const frame = await next(args); + return { + frame: { + result: frame.result, + annotation: { + ...frame.annotation, + scope: 'node', + }, + }, + }; + }); + try { + const frame = await toolCallExecuteFrameAsync( + 'sl_frame_tool', + {}, + async () => ({ + result: { + ok: true, + }, + annotation: { + producer: 'node', + }, + }), + null, + null, + null, + null, + ); + assert.deepEqual(frame.annotation, { + producer: 'node', + scope: 'node', + }); + assert.equal(scopeDeregisterToolExecutionIntercept(scope.uuid, 'sl_frame_exec'), true); + + const unwrapped = await toolCallExecuteFrameAsync( + 'sl_frame_after_deregister', + {}, + async () => ({ + result: { + ok: true, + }, + annotation: { + producer: 'node', + }, + }), + null, + null, + null, + null, + ); + assert.deepEqual(unwrapped.annotation, { + producer: 'node', + }); + } finally { + scopeDeregisterToolExecutionIntercept(scope.uuid, 'sl_frame_exec'); + popScope(scope); + } + }); + it('global and scope-local llm request intercepts both run with priority ordering', async () => { const order = []; diff --git a/crates/node/tests/tools_tests.mjs b/crates/node/tests/tools_tests.mjs index 196dad76f..3ad2950ea 100644 --- a/crates/node/tests/tools_tests.mjs +++ b/crates/node/tests/tools_tests.mjs @@ -14,8 +14,11 @@ const { popScope, toolCall, toolCallEnd, + toolCallEndFrame, toolCallExecute, toolCallExecuteAsync, + toolCallExecuteFrame, + toolCallExecuteFrameAsync, toolRequestIntercepts, toolConditionalExecution, registerToolSanitizeRequestGuardrail, @@ -27,6 +30,7 @@ const { registerToolRequestIntercept, deregisterToolRequestIntercept, registerToolExecutionIntercept, + registerToolExecutionFrameIntercept, deregisterToolExecutionIntercept, clearLastCallbackError, getLastCallbackError, @@ -229,6 +233,137 @@ describe('Tool execute', () => { }); }); + it('frame execute preserves opaque annotation through the mixed chain', async () => { + const events = []; + registerSubscriber('node_frame_subscriber', (event) => events.push(event)); + registerToolExecutionFrameIntercept('node_frame_outer', 1, async (args, next) => { + const frame = await next(args); + assert.deepEqual(frame.annotation, { + producer: 'node', + status: 'failed', + }); + return { + frame: { + result: { + ...frame.result, + frameSeen: true, + }, + annotation: { + ...frame.annotation, + observedBy: 'node_frame_outer', + }, + }, + }; + }); + registerToolExecutionIntercept('node_frame_legacy', 2, async (args, next) => ({ + result: await next(args), + })); + try { + const frame = await toolCallExecuteFrameAsync( + 'node_frame_tool', + { + value: 3, + }, + async (args) => ({ + result: { + value: args.value * 2, + }, + annotation: { + producer: 'node', + status: 'failed', + }, + }), + null, + null, + null, + null, + ); + assert.deepEqual(frame, { + result: { + value: 6, + frameSeen: true, + }, + annotation: { + producer: 'node', + status: 'failed', + observedBy: 'node_frame_outer', + }, + }); + + await waitForSubscriberCallbacks(() => + events.some( + (event) => event.name === 'node_frame_tool' && event.kind === 'scope' && event.scope_category === 'end', + ), + ); + const end = events.find( + (event) => event.name === 'node_frame_tool' && event.kind === 'scope' && event.scope_category === 'end', + ); + assert.deepEqual(end.category_profile.tool_result_annotation, frame.annotation); + } finally { + assert.equal(deregisterToolExecutionIntercept('node_frame_outer'), true); + assert.equal(deregisterToolExecutionIntercept('node_frame_legacy'), true); + deregisterSubscriber('node_frame_subscriber'); + } + }); + + it('manual and synchronous frame APIs marshal the same frame shape', async () => { + const events = []; + registerSubscriber('node_manual_frame_subscriber', (event) => events.push(event)); + try { + const handle = toolCall('node_manual_frame_tool', {}, null, null, null, null); + toolCallEndFrame( + handle, + { + result: { + manual: true, + }, + annotation: { + producer: 'node-manual', + }, + }, + null, + null, + null, + ); + const frame = await toolCallExecuteFrame( + 'node_sync_frame_tool', + {}, + () => ({ + result: { + sync: true, + }, + annotation: { + producer: 'node-sync', + }, + }), + null, + null, + null, + null, + ); + assert.deepEqual(frame.annotation, { + producer: 'node-sync', + }); + await waitForSubscriberCallbacks( + () => + events.some( + (event) => + event.name === 'node_manual_frame_tool' && event.kind === 'scope' && event.scope_category === 'end', + ) && + events.some( + (event) => + event.name === 'node_sync_frame_tool' && event.kind === 'scope' && event.scope_category === 'end', + ), + ); + const manualEnd = events.find( + (event) => event.name === 'node_manual_frame_tool' && event.kind === 'scope' && event.scope_category === 'end', + ); + assert.equal(manualEnd.category_profile.tool_result_annotation.producer, 'node-manual'); + } finally { + deregisterSubscriber('node_manual_frame_subscriber'); + } + }); + it('treats implicit undefined tool results as null', async () => { const result = await toolCallExecute( 'exec_tool_undefined', diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index b6bf3da7d..c03f70160 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -22,7 +22,9 @@ pub use nemo_relay_types::api::event::{ }; pub use nemo_relay_types::api::llm::{LlmAttributes, LlmRequest, LlmRequestInterceptOutcome}; pub use nemo_relay_types::api::scope::{HandleAttributes, ScopeAttributes, ScopeType}; -pub use nemo_relay_types::api::tool::{ToolAttributes, ToolExecutionInterceptOutcome}; +pub use nemo_relay_types::api::tool::{ + ToolAttributes, ToolExecutionFrame, ToolExecutionFrameOutcome, ToolExecutionInterceptOutcome, +}; pub use nemo_relay_types::codec::optimization::{ LlmOptimizationContribution, LlmOptimizationEvidenceQuality, LlmOptimizationKind, LlmOptimizationModel, LlmOptimizationModelTransition, LlmOptimizationPayload, @@ -334,6 +336,17 @@ pub type NemoRelayNativeToolNextFn = unsafe extern "C" fn( out_json: *mut *mut NemoRelayNativeString, ) -> NemoRelayStatus; +/// Runtime-provided continuation for annotation-aware tool execution intercepts. +/// +/// The returned JSON serializes a [`ToolExecutionFrame`]. +/// The function and context are borrowed for one callback invocation and must +/// not be retained or used after that callback returns. +pub type NemoRelayNativeToolFrameNextFn = unsafe extern "C" fn( + args_json: *const NemoRelayNativeString, + next_ctx: *mut c_void, + out_frame_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus; + /// Runtime-provided continuation for LLM execution intercepts. pub type NemoRelayNativeLlmNextFn = unsafe extern "C" fn( request_json: *const NemoRelayNativeString, @@ -432,6 +445,19 @@ pub type NemoRelayNativeToolExecutionCb = unsafe extern "C" fn( out_outcome_json: *mut *mut NemoRelayNativeString, ) -> NemoRelayStatus; +/// Native annotation-aware tool execution intercept callback. +/// +/// `next_fn` returns serialized [`ToolExecutionFrame`] JSON and the callback +/// returns serialized [`ToolExecutionFrameOutcome`] JSON. +pub type NemoRelayNativeToolExecutionFrameCb = unsafe extern "C" fn( + user_data: *mut c_void, + name: *const NemoRelayNativeString, + args_json: *const NemoRelayNativeString, + next_fn: NemoRelayNativeToolFrameNextFn, + next_ctx: *mut c_void, + out_outcome_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus; + /// Native LLM request sanitizer callback. Return a successful null output to /// omit the observability payload and annotation. `request_json` is borrowed, /// but may be written directly to `out_request_json` as a pass-through; the @@ -1065,6 +1091,34 @@ pub struct NemoRelayNativeHostApiV3 { unsafe impl Send for NemoRelayNativeHostApiV3 {} unsafe impl Sync for NemoRelayNativeHostApiV3 {} +/// Optional tool-frame capabilities appended after the frozen ABI-v3 host +/// table. +/// +/// Hosts advertise this extension by setting +/// [`NemoRelayNativeHostApiV1::struct_size`] to at least the size of this +/// structure. Its prefix remains a complete [`NemoRelayNativeHostApiV3`], so +/// plugins compiled against the original ABI-v3 table remain compatible. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct NemoRelayNativeHostApiV3ToolFrames { + /// Compatibility prefix for ABI-v3 plugins. + pub v3: NemoRelayNativeHostApiV3, + /// Registers an annotation-aware tool execution intercept through the + /// plugin context. + pub plugin_context_register_tool_execution_frame_intercept: + unsafe extern "C" fn( + ctx: *mut NemoRelayNativePluginContext, + name: *const NemoRelayNativeString, + priority: i32, + cb: NemoRelayNativeToolExecutionFrameCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, + ) -> NemoRelayStatus, +} + +unsafe impl Send for NemoRelayNativeHostApiV3ToolFrames {} +unsafe impl Sync for NemoRelayNativeHostApiV3ToolFrames {} + // The host API table is immutable after construction. Function pointers and // the null-terminated version string pointer are safe to share across threads. unsafe impl Send for NemoRelayNativeHostApiV1 {} @@ -1321,6 +1375,35 @@ impl ToolNext<'_> { } } +/// Typed continuation passed to annotation-aware tool execution intercepts. +/// +/// Its lifetime is bounded by the intercept callback, preventing safe native +/// plugins from retaining the host continuation after the callback returns. +pub struct ToolFrameNext<'a> { + host: &'a NemoRelayNativeHostApiV1, + next_fn: NemoRelayNativeToolFrameNextFn, + next_ctx: *mut c_void, +} + +impl ToolFrameNext<'_> { + /// Continues the mixed tool execution chain with replacement arguments. + pub fn call(&self, args: Json) -> Result { + let args = HostString::from_json(self.host, &args) + .ok_or_else(|| "failed to allocate tool frame next args".to_string())?; + let mut out = ptr::null_mut(); + let status = unsafe { (self.next_fn)(args.as_ptr(), self.next_ctx, &mut out) }; + if status != NemoRelayStatus::Ok { + return Err(format!("tool frame next failed: {status:?}")); + } + if out.is_null() { + return Err("tool frame next returned null output".into()); + } + let result = read_json_value(self.host, out, "tool frame next result"); + unsafe { (self.host.string_free)(out) }; + result.map_err(|status| format!("tool frame next returned invalid JSON: {status:?}")) + } +} + /// Typed continuation passed to LLM execution intercepts. pub struct LlmNext<'a> { host: &'a NemoRelayNativeHostApiV1, @@ -2041,6 +2124,38 @@ impl<'a> PluginContext<'a> { finish_typed_registration::(self.host, status, user_data, "tool execution intercept") } + /// Registers an annotation-aware tool execution intercept in the same host + /// chain as legacy tool execution intercepts. + pub fn register_tool_execution_frame_intercept( + &mut self, + name: &str, + priority: i32, + callback: F, + ) -> Result<()> + where + F: for<'next> Fn(&str, Json, ToolFrameNext<'next>) -> Result + + Send + + Sync + + 'static, + { + let user_data = typed_callback_user_data(self.host, callback); + let status = unsafe { + self.register_tool_execution_frame_intercept_raw( + name, + priority, + typed_tool_execution_frame_trampoline::, + user_data, + Some(drop_typed_callback::), + ) + }; + finish_typed_registration::( + self.host, + status, + user_data, + "tool execution frame intercept", + ) + } + /// Registers a typed LLM sanitize-request guardrail. pub fn register_llm_sanitize_request_guardrail( &mut self, @@ -2409,6 +2524,34 @@ impl<'a> PluginContext<'a> { }) } + /// Registers a raw annotation-aware tool execution intercept callback. + /// + /// # Safety + /// `cb`, `user_data`, and `free_fn` must remain valid for every host + /// callback invocation until the host deregisters the callback or calls + /// `free_fn`. `free_fn` must match the allocation behind `user_data`. + pub unsafe fn register_tool_execution_frame_intercept_raw( + &mut self, + name: &str, + priority: i32, + cb: NemoRelayNativeToolExecutionFrameCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, + ) -> NemoRelayStatus { + if self.host.abi_version < NEMO_RELAY_NATIVE_ABI_VERSION_ASYNC_MIDDLEWARE + || self.host.struct_size < std::mem::size_of::() + { + return NemoRelayStatus::InvalidArg; + } + let host = + unsafe { &*(self.host as *const _ as *const NemoRelayNativeHostApiV3ToolFrames) }; + self.with_name(name, |_, name| unsafe { + (host.plugin_context_register_tool_execution_frame_intercept)( + self.raw, name, priority, cb, user_data, free_fn, + ) + }) + } + /// Registers a raw LLM sanitize-request guardrail callback. /// /// # Safety @@ -2875,6 +3018,56 @@ where } } +unsafe extern "C" fn typed_tool_execution_frame_trampoline( + user_data: *mut c_void, + name: *const NemoRelayNativeString, + args_json: *const NemoRelayNativeString, + next_fn: NemoRelayNativeToolFrameNextFn, + next_ctx: *mut c_void, + out_outcome_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus +where + F: for<'next> Fn(&str, Json, ToolFrameNext<'next>) -> Result + + Send + + Sync + + 'static, +{ + if user_data.is_null() || out_outcome_json.is_null() { + return NemoRelayStatus::NullPointer; + } + unsafe { *out_outcome_json = ptr::null_mut() }; + let state = unsafe { &*(user_data as *const TypedCallback) }; + let result = catch_unwind(AssertUnwindSafe(|| { + let name = read_required_host_string(&state.host, name, "tool name")?; + let args: Json = read_json_value(&state.host, args_json, "tool args")?; + let next = ToolFrameNext { + host: &state.host, + next_fn, + next_ctx, + }; + match (state.callback)(&name, args, next) { + Ok(outcome) => { + let Some(outcome) = HostString::from_json(&state.host, &outcome) else { + set_last_error( + &state.host, + "failed to allocate tool execution frame outcome", + ); + return Ok(NemoRelayStatus::Internal); + }; + unsafe { *out_outcome_json = outcome.ptr }; + std::mem::forget(outcome); + Ok(NemoRelayStatus::Ok) + } + Err(message) => Ok(callback_error(&state.host, message)), + } + })); + match result { + Ok(Ok(status)) => status, + Ok(Err(status)) => status, + Err(_) => callback_panic(&state.host, "tool execution frame callback"), + } +} + unsafe extern "C" fn typed_llm_sanitize_request_trampoline( user_data: *mut c_void, request_json: *const NemoRelayNativeString, @@ -3288,10 +3481,18 @@ impl<'a> OptionalHostJson<'a> { enum OwnedHostApi { V1(NemoRelayNativeHostApiV1), V3(NemoRelayNativeHostApiV3), + V3ToolFrames(NemoRelayNativeHostApiV3ToolFrames), } impl OwnedHostApi { unsafe fn copy_from(host: &NemoRelayNativeHostApiV1) -> Self { + if host.abi_version >= NEMO_RELAY_NATIVE_ABI_VERSION_ASYNC_MIDDLEWARE + && host.struct_size >= std::mem::size_of::() + { + return Self::V3ToolFrames(unsafe { + *(host as *const _ as *const NemoRelayNativeHostApiV3ToolFrames) + }); + } if host.abi_version >= NEMO_RELAY_NATIVE_ABI_VERSION_ASYNC_MIDDLEWARE && host.struct_size >= std::mem::size_of::() { @@ -3305,6 +3506,7 @@ impl OwnedHostApi { match self { Self::V1(host) => host, Self::V3(host) => &host.v1, + Self::V3ToolFrames(host) => &host.v3.v1, } } } diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index 39a6ece8b..c354db256 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -17,19 +17,24 @@ use nemo_relay_plugin::{ Event, EventCategory, EventSanitizeFields, Json, LlmCodecIdentity, LlmJsonStream, LlmNext, LlmRequest, LlmRequestInterceptOutcome, LlmStream, LlmStreamNext, NEMO_RELAY_NATIVE_ABI_VERSION, NativePlugin, NemoRelayNativeAsyncCallbackState, - NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeEventSanitizeCb, + NemoRelayNativeAsyncCompletion, NemoRelayNativeAsyncMiddlewareCb, + NemoRelayNativeAsyncMiddlewareKind, NemoRelayNativeAsyncNext, NemoRelayNativeAsyncNextResultCb, + NemoRelayNativeAsyncNextStreamCb, NemoRelayNativeAsyncStream, + NemoRelayNativeAsyncStreamMiddlewareCb, NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, - NemoRelayNativeHostApiV3, NemoRelayNativeLlmCodecKind, NemoRelayNativeLlmConditionalCb, - NemoRelayNativeLlmExecutionCb, NemoRelayNativeLlmRequestCodec, + NemoRelayNativeHostApiV3, NemoRelayNativeHostApiV3ToolFrames, NemoRelayNativeLlmCodecKind, + NemoRelayNativeLlmConditionalCb, NemoRelayNativeLlmExecutionCb, NemoRelayNativeLlmRequestCodec, NemoRelayNativeLlmRequestInterceptCb, NemoRelayNativeLlmResponseCodec, NemoRelayNativeLlmSanitizeRequestCb, NemoRelayNativeLlmSanitizeRequestContext, NemoRelayNativeLlmSanitizeResponseCb, NemoRelayNativeLlmSanitizeResponseContext, NemoRelayNativeLlmStreamExecutionCb, NemoRelayNativeLlmStreamV1, NemoRelayNativePluginContext, NemoRelayNativePluginV1, NemoRelayNativeScopeHandle, NemoRelayNativeScopeStack, NemoRelayNativeScopeStackBinding, NemoRelayNativeScopeType, NemoRelayNativeString, - NemoRelayNativeToolConditionalCb, NemoRelayNativeToolExecutionCb, NemoRelayNativeToolJsonCb, + NemoRelayNativeToolConditionalCb, NemoRelayNativeToolExecutionCb, + NemoRelayNativeToolExecutionFrameCb, NemoRelayNativeToolJsonCb, NemoRelayNativeWithScopeStackCb, NemoRelayStatus, PendingMarkSpec, PluginContext, - PluginRuntime, ScopeType, ToolExecutionInterceptOutcome, ToolNext, + PluginRuntime, ScopeType, ToolExecutionFrameOutcome, ToolExecutionInterceptOutcome, + ToolFrameNext, ToolNext, }; use serde_json::{Map, json}; @@ -139,6 +144,22 @@ struct RegisteredToolExecution { free_fn: NemoRelayNativeFreeFn, } +struct RegisteredToolExecutionFrame { + name: String, + priority: i32, + cb: NemoRelayNativeToolExecutionFrameCb, + user_data: usize, + free_fn: NemoRelayNativeFreeFn, +} + +impl RegisteredToolExecutionFrame { + unsafe fn free(self) { + if let Some(free_fn) = self.free_fn { + unsafe { free_fn(self.user_data as *mut c_void) }; + } + } +} + impl RegisteredToolExecution { unsafe fn free(self) { if let Some(free_fn) = self.free_fn { @@ -266,6 +287,7 @@ impl_captured_registration!( RegisteredToolJson, RegisteredToolConditional, RegisteredToolExecution, + RegisteredToolExecutionFrame, RegisteredLlmRequest, RegisteredLlmJson, RegisteredLlmConditional, @@ -323,6 +345,8 @@ static EVENT_SANITIZE_REGISTRATION: Mutex> = Mut static TOOL_JSON_REGISTRATION: Mutex> = Mutex::new(None); static TOOL_CONDITIONAL_REGISTRATION: Mutex> = Mutex::new(None); static TOOL_EXECUTION_REGISTRATION: Mutex> = Mutex::new(None); +static TOOL_EXECUTION_FRAME_REGISTRATION: Mutex> = + Mutex::new(None); static LLM_REQUEST_REGISTRATION: Mutex> = Mutex::new(None); static LLM_JSON_REGISTRATION: Mutex> = Mutex::new(None); static LLM_CONDITIONAL_REGISTRATION: Mutex> = Mutex::new(None); @@ -353,6 +377,15 @@ fn native_abi_v3_struct_sizes_are_self_describing() { { assert_eq!(align_of::(), 8); assert_eq!(size_of::(), 320); + assert_eq!( + offset_of!( + NemoRelayNativeHostApiV3ToolFrames, + plugin_context_register_tool_execution_frame_intercept + ), + 440, + "the frame registration hook must remain appended after the frozen ABI-v3 prefix" + ); + assert_eq!(size_of::(), 448); assert_eq!( host_api_offsets(), [ @@ -381,6 +414,15 @@ fn native_abi_v3_struct_sizes_are_self_describing() { { assert_eq!(align_of::(), 4); assert_eq!(size_of::(), 160); + assert_eq!( + offset_of!( + NemoRelayNativeHostApiV3ToolFrames, + plugin_context_register_tool_execution_frame_intercept + ), + 216, + "the frame registration hook must remain appended after the frozen ABI-v3 prefix" + ); + assert_eq!(size_of::(), 220); assert_eq!( host_api_offsets(), [ @@ -767,6 +809,35 @@ unsafe extern "C" fn capture_tool_execution( status } +unsafe extern "C" fn capture_tool_execution_frame( + _ctx: *mut NemoRelayNativePluginContext, + name: *const NemoRelayNativeString, + priority: i32, + cb: NemoRelayNativeToolExecutionFrameCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> NemoRelayStatus { + let status = *REGISTRATION_STATUS.lock().unwrap(); + if status == NemoRelayStatus::Ok { + let host = test_host(); + let name = match required_host_string(&host, name) { + Ok(name) => name, + Err(status) => return status, + }; + replace_registration( + &TOOL_EXECUTION_FRAME_REGISTRATION, + RegisteredToolExecutionFrame { + name, + priority, + cb, + user_data: user_data as usize, + free_fn, + }, + ); + } + status +} + unsafe extern "C" fn capture_llm_request( _ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, @@ -1305,6 +1376,141 @@ fn test_host() -> NemoRelayNativeHostApiV1 { } } +unsafe extern "C" fn unavailable_async_completion_resolve_json( + _completion: *const NemoRelayNativeAsyncCompletion, + _value_json: *const NemoRelayNativeString, +) -> NemoRelayStatus { + NemoRelayStatus::InvalidArg +} + +unsafe extern "C" fn unavailable_async_completion_reject( + _completion: *const NemoRelayNativeAsyncCompletion, + _message: *const NemoRelayNativeString, +) -> NemoRelayStatus { + NemoRelayStatus::InvalidArg +} + +unsafe extern "C" fn unavailable_async_completion_is_cancelled( + _completion: *const NemoRelayNativeAsyncCompletion, +) -> bool { + true +} + +unsafe extern "C" fn unavailable_async_completion_release( + _completion: *const NemoRelayNativeAsyncCompletion, +) { +} + +unsafe extern "C" fn unavailable_async_next_invoke( + _next: *const NemoRelayNativeAsyncNext, + _invocation_json: *const NemoRelayNativeString, + _completion: *const NemoRelayNativeAsyncCompletion, +) -> NemoRelayStatus { + NemoRelayStatus::InvalidArg +} + +unsafe extern "C" fn unavailable_async_next_release(_next: *const NemoRelayNativeAsyncNext) {} + +#[allow(clippy::too_many_arguments)] +unsafe extern "C" fn unavailable_async_middleware_registration( + _ctx: *mut NemoRelayNativePluginContext, + _kind: u32, + _name: *const NemoRelayNativeString, + _priority: i32, + _break_chain: bool, + _cb: NemoRelayNativeAsyncMiddlewareCb, + _user_data: *mut c_void, + _free_fn: NemoRelayNativeFreeFn, +) -> NemoRelayStatus { + NemoRelayStatus::InvalidArg +} + +unsafe extern "C" fn unavailable_async_stream_push_json( + _stream: *const NemoRelayNativeAsyncStream, + _chunk_json: *const NemoRelayNativeString, +) -> NemoRelayStatus { + NemoRelayStatus::InvalidArg +} + +unsafe extern "C" fn unavailable_async_stream_finish( + _stream: *const NemoRelayNativeAsyncStream, +) -> NemoRelayStatus { + NemoRelayStatus::InvalidArg +} + +unsafe extern "C" fn unavailable_async_stream_reject( + _stream: *const NemoRelayNativeAsyncStream, + _message: *const NemoRelayNativeString, +) -> NemoRelayStatus { + NemoRelayStatus::InvalidArg +} + +unsafe extern "C" fn unavailable_async_stream_is_cancelled( + _stream: *const NemoRelayNativeAsyncStream, +) -> bool { + true +} + +unsafe extern "C" fn unavailable_async_stream_release(_stream: *const NemoRelayNativeAsyncStream) {} + +unsafe extern "C" fn unavailable_async_next_invoke_stream( + _next: *const NemoRelayNativeAsyncNext, + _invocation_json: *const NemoRelayNativeString, + _stream: *const NemoRelayNativeAsyncStream, + _cb: NemoRelayNativeAsyncNextStreamCb, + _user_data: *mut c_void, +) -> NemoRelayStatus { + NemoRelayStatus::InvalidArg +} + +#[allow(clippy::too_many_arguments)] +unsafe extern "C" fn unavailable_async_stream_middleware_registration( + _ctx: *mut NemoRelayNativePluginContext, + _name: *const NemoRelayNativeString, + _priority: i32, + _cb: NemoRelayNativeAsyncStreamMiddlewareCb, + _user_data: *mut c_void, + _free_fn: NemoRelayNativeFreeFn, +) -> NemoRelayStatus { + NemoRelayStatus::InvalidArg +} + +unsafe extern "C" fn unavailable_async_next_invoke_result( + _next: *const NemoRelayNativeAsyncNext, + _invocation_json: *const NemoRelayNativeString, + _cb: NemoRelayNativeAsyncNextResultCb, + _user_data: *mut c_void, +) -> NemoRelayStatus { + NemoRelayStatus::InvalidArg +} + +fn test_host_v3() -> NemoRelayNativeHostApiV3ToolFrames { + let mut v1 = test_host(); + v1.struct_size = std::mem::size_of::(); + NemoRelayNativeHostApiV3ToolFrames { + v3: NemoRelayNativeHostApiV3 { + v1, + async_completion_resolve_json: unavailable_async_completion_resolve_json, + async_completion_reject: unavailable_async_completion_reject, + async_completion_is_cancelled: unavailable_async_completion_is_cancelled, + async_completion_release: unavailable_async_completion_release, + async_next_invoke: unavailable_async_next_invoke, + async_next_release: unavailable_async_next_release, + plugin_context_register_async_middleware: unavailable_async_middleware_registration, + async_stream_push_json: unavailable_async_stream_push_json, + async_stream_finish: unavailable_async_stream_finish, + async_stream_reject: unavailable_async_stream_reject, + async_stream_is_cancelled: unavailable_async_stream_is_cancelled, + async_stream_release: unavailable_async_stream_release, + async_next_invoke_stream: unavailable_async_next_invoke_stream, + plugin_context_register_async_stream_middleware: + unavailable_async_stream_middleware_registration, + async_next_invoke_result: unavailable_async_next_invoke_result, + }, + plugin_context_register_tool_execution_frame_intercept: capture_tool_execution_frame, + } +} + fn begin_test() -> MutexGuard<'static, ()> { let guard = TEST_LOCK .lock() @@ -1319,6 +1525,7 @@ fn reset_state() { clear_registration(&TOOL_JSON_REGISTRATION); clear_registration(&TOOL_CONDITIONAL_REGISTRATION); clear_registration(&TOOL_EXECUTION_REGISTRATION); + clear_registration(&TOOL_EXECUTION_FRAME_REGISTRATION); clear_registration(&LLM_REQUEST_REGISTRATION); clear_registration(&LLM_JSON_REGISTRATION); clear_registration(&LLM_CONDITIONAL_REGISTRATION); @@ -1547,6 +1754,14 @@ fn take_tool_execution_registration() -> RegisteredToolExecution { .expect("tool execution callback should be registered") } +fn take_tool_execution_frame_registration() -> RegisteredToolExecutionFrame { + TOOL_EXECUTION_FRAME_REGISTRATION + .lock() + .unwrap() + .take() + .expect("tool execution frame callback should be registered") +} + fn take_llm_request_registration() -> RegisteredLlmRequest { LLM_REQUEST_REGISTRATION .lock() @@ -3746,6 +3961,26 @@ unsafe extern "C" fn fake_tool_next( write_json(&state.host, &args, out_json) } +unsafe extern "C" fn fake_tool_frame_next( + args_json: *const NemoRelayNativeString, + next_ctx: *mut c_void, + out_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + let state = unsafe { &*(next_ctx as *const NextState) }; + state.called.fetch_add(1, Ordering::SeqCst); + let mut args: Json = + serde_json::from_str(&read_host_string(&state.host, args_json).unwrap()).unwrap(); + args["next_called"] = json!(true); + write_json( + &state.host, + &json!({ + "result": args, + "annotation": {"producer": "host"} + }), + out_json, + ) +} + unsafe extern "C" fn failing_tool_next( _args_json: *const NemoRelayNativeString, next_ctx: *mut c_void, @@ -4084,6 +4319,60 @@ fn typed_tool_execution_registration_calls_next() { } } +#[test] +fn typed_tool_execution_frame_registration_round_trips_opaque_annotation() { + let _guard = begin_test(); + let host_v3 = test_host_v3(); + let host = &host_v3.v3.v1; + let called = Arc::new(AtomicUsize::new(0)); + let mut ctx = test_context(host); + ctx.register_tool_execution_frame_intercept( + "tool-frame", + 17, + |_name, args, next: ToolFrameNext<'_>| { + let mut frame = next.call(args)?; + frame.annotation.as_mut().unwrap()["plugin"] = json!(true); + Ok(ToolExecutionFrameOutcome::new(frame) + .with_pending_mark(PendingMarkSpec::builder().name("plugin.tool.frame").build())) + }, + ) + .unwrap(); + + let registration = take_tool_execution_frame_registration(); + assert_eq!(registration.name, "tool-frame"); + assert_eq!(registration.priority, 17); + let next_state = Box::into_raw(Box::new(NextState { + host: *host, + called: called.clone(), + })); + let name = host_string(host, "tool-frame"); + let args = json_host_string(host, json!({ "input": true })); + let mut out = ptr::null_mut(); + let status = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + name, + args, + fake_tool_frame_next, + next_state.cast(), + &mut out, + ) + }; + assert_eq!(status, NemoRelayStatus::Ok); + assert_eq!(called.load(Ordering::SeqCst), 1); + let outcome = read_json_and_free(host, out); + assert_eq!(outcome["frame"]["result"]["next_called"], true); + assert_eq!(outcome["frame"]["annotation"]["producer"], "host"); + assert_eq!(outcome["frame"]["annotation"]["plugin"], true); + assert_eq!(outcome["pending_marks"][0]["name"], "plugin.tool.frame"); + unsafe { + (host.string_free)(name); + (host.string_free)(args); + drop(Box::from_raw(next_state)); + registration.free(); + } +} + #[test] fn typed_tool_execution_does_not_publish_partial_outcome() { let _guard = begin_test(); diff --git a/crates/python/src/py_api/mod.rs b/crates/python/src/py_api/mod.rs index 9d37efa45..e3d8386fc 100644 --- a/crates/python/src/py_api/mod.rs +++ b/crates/python/src/py_api/mod.rs @@ -20,7 +20,8 @@ use nemo_relay::api::runtime::subscriber_dispatcher::{ with_task_nested_publication_buffer, with_task_publication_context, }; use nemo_relay::api::runtime::{ - LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, ToolExecutionNextFn, + LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, ToolExecutionFrameNextFn, + ToolExecutionNextFn, }; use nemo_relay::api::runtime::{ TASK_SCOPE_STACK, capture_propagation_context as capture_propagation_context_handle, @@ -51,7 +52,7 @@ use crate::py_types::{ PyAnnotatedLLMResponse, PyAnthropicMessagesCodec, PyLLMAttributes, PyLLMHandle, PyLLMRequest, PyLlmStream, PyOpenAIChatCodec, PyOpenAIResponsesCodec, PyPropagationContext, PyScopeAttributes, PyScopeHandle, PyScopeStack, PyScopeType, PyThreadScopeStackBinding, - PyToolAttributes, PyToolHandle, + PyToolAttributes, PyToolExecutionFrame, PyToolHandle, }; pub(crate) type RustJsonStream = LlmJsonStream; @@ -658,6 +659,38 @@ fn tool_call_end( .map_err(to_py_err) } +/// End a manual tool call with an optional opaque result annotation. +#[pyfunction] +#[pyo3(signature = ( + handle: "ToolHandle", + frame: "ToolExecutionFrame", + *, + data: "object | None"=None, + metadata: "object | None"=None, + timestamp: "datetime.datetime | None"=None +) -> "None", text_signature = "(handle: ToolHandle, frame: ToolExecutionFrame, *, data: object | None = None, metadata: object | None = None, timestamp: datetime.datetime | None = None) -> None")] +fn tool_call_end_frame( + handle: &PyToolHandle, + frame: PyToolExecutionFrame, + data: Option<&Bound<'_, PyAny>>, + metadata: Option<&Bound<'_, PyAny>>, + timestamp: Option<&Bound<'_, PyAny>>, +) -> PyResult<()> { + let data = opt_py_to_json(data)?; + let metadata = opt_py_to_json(metadata)?; + let timestamp = opt_py_to_timestamp(timestamp)?; + core_tool_api::tool_call_end_frame( + core_tool_api::ToolCallEndFrameParams::builder() + .handle(&handle.inner) + .frame(frame.inner) + .data_opt(data) + .metadata_opt(metadata) + .timestamp_opt(timestamp) + .build(), + ) + .map_err(to_py_err) +} + /// Execute a tool call through the full middleware pipeline. /// /// Runs conditional-execution guardrails (on raw args) → request intercepts → @@ -743,6 +776,62 @@ fn tool_call_execute<'py>( }) } +/// Execute a tool call while carrying an optional opaque result annotation. +#[pyfunction] +#[pyo3(signature = ( + name: "str", + args: "object", + func: "object", + *, + handle: "ScopeHandle | None"=None, + attributes: "ToolAttributes | None"=None, + data: "object | None"=None, + metadata: "object | None"=None +) -> "object", text_signature = "(name: str, args: object, func: object, *, handle: ScopeHandle | None = None, attributes: ToolAttributes | None = None, data: object | None = None, metadata: object | None = None) -> object")] +#[allow(clippy::too_many_arguments)] +fn tool_call_execute_frame<'py>( + py: Python<'py>, + name: String, + args: &Bound<'py, PyAny>, + func: Py, + handle: Option, + attributes: Option, + data: Option<&Bound<'py, PyAny>>, + metadata: Option<&Bound<'py, PyAny>>, +) -> PyResult> { + let args_json = py_to_json(args)?; + let attrs = attributes + .map(|a| a.inner) + .unwrap_or(ToolAttributes::empty()); + let data_json = opt_py_to_json(data)?; + let metadata_json = opt_py_to_json(metadata)?; + let exec_fn = py_callable::wrap_py_tool_exec_frame_fn(func); + let default_fn: ToolExecutionFrameNextFn = Arc::new(move |args| exec_fn(args)); + let parent_handle = handle.map(|h| h.inner).unwrap_or_else(task_scope_top); + + let scope_stack = current_scope_stack_handle(); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + TASK_SCOPE_STACK + .scope(scope_stack, async move { + let frame = core_tool_api::tool_call_execute_frame( + core_tool_api::ToolCallExecuteFrameParams::builder() + .name(name) + .args(args_json) + .func(default_fn) + .parent(parent_handle) + .attributes(attrs) + .data_opt(data_json) + .metadata_opt(metadata_json) + .build(), + ) + .await + .map_err(to_py_err)?; + Ok(PyToolExecutionFrame { inner: frame }) + }) + .await + }) +} + // --------------------------------------------------------------------------- // LLM lifecycle // --------------------------------------------------------------------------- @@ -1274,7 +1363,22 @@ fn register_tool_execution_intercept( .map_err(to_py_err) } -/// Remove a previously registered tool execution intercept. +/// Register an annotation-aware callback in the existing tool execution chain. +#[pyfunction] +fn register_tool_execution_frame_intercept( + name: &str, + priority: i32, + callable: Py, +) -> PyResult<()> { + core_registry_api::register_tool_execution_frame_intercept( + name, + priority, + py_callable::wrap_py_tool_exec_frame_intercept_fn(callable), + ) + .map_err(to_py_err) +} + +/// Remove a previously registered tool execution intercept of either form. #[pyfunction] fn deregister_tool_execution_intercept(name: &str) -> PyResult { core_registry_api::deregister_tool_execution_intercept(name).map_err(to_py_err) @@ -1772,7 +1876,25 @@ fn scope_register_tool_execution_intercept( .map_err(to_py_err) } -/// Remove a previously registered scope-local tool execution intercept. +/// Register an annotation-aware scope-local callback in the existing chain. +#[pyfunction] +fn scope_register_tool_execution_frame_intercept( + scope_uuid: &str, + name: &str, + priority: i32, + callable: Py, +) -> PyResult<()> { + let uuid = parse_uuid(scope_uuid)?; + core_registry_api::scope_register_tool_execution_frame_intercept( + &uuid, + name, + priority, + py_callable::wrap_py_tool_exec_frame_intercept_fn(callable), + ) + .map_err(to_py_err) +} + +/// Remove a previously registered scope-local tool execution intercept of either form. #[pyfunction] fn scope_deregister_tool_execution_intercept(scope_uuid: &str, name: &str) -> PyResult { let uuid = parse_uuid(scope_uuid)?; @@ -1998,7 +2120,9 @@ pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { // Tool lifecycle m.add_function(wrap_pyfunction!(tool_call, m)?)?; m.add_function(wrap_pyfunction!(tool_call_end, m)?)?; + m.add_function(wrap_pyfunction!(tool_call_end_frame, m)?)?; m.add_function(wrap_pyfunction!(tool_call_execute, m)?)?; + m.add_function(wrap_pyfunction!(tool_call_execute_frame, m)?)?; // LLM lifecycle m.add_function(wrap_pyfunction!(llm_call, m)?)?; @@ -2054,6 +2178,10 @@ pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(deregister_tool_request_intercept, m)?)?; m.add_function(wrap_pyfunction!(register_tool_execution_intercept, m)?)?; m.add_function(wrap_pyfunction!(deregister_tool_execution_intercept, m)?)?; + m.add_function(wrap_pyfunction!( + register_tool_execution_frame_intercept, + m + )?)?; // LLM guardrails m.add_function(wrap_pyfunction!( @@ -2167,6 +2295,10 @@ pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { scope_deregister_tool_execution_intercept, m )?)?; + m.add_function(wrap_pyfunction!( + scope_register_tool_execution_frame_intercept, + m + )?)?; // Scope-local LLM guardrails m.add_function(wrap_pyfunction!( diff --git a/crates/python/src/py_callable.rs b/crates/python/src/py_callable.rs index 85d468d2b..9be24c3c1 100644 --- a/crates/python/src/py_callable.rs +++ b/crates/python/src/py_callable.rs @@ -32,8 +32,9 @@ use nemo_relay::api::runtime::{ EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn, LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, LlmStreamInner, - MiddlewareContinuationContext, ScopeStackHandle, ToolConditionalFn, ToolExecutionNextFn, - ToolInterceptFn, ToolSanitizeFn, capture_propagation_context, current_scope_stack, + MiddlewareContinuationContext, ScopeStackHandle, ToolConditionalFn, ToolExecutionFrameFn, + ToolExecutionFrameNextFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, + capture_propagation_context, current_scope_stack, }; use nemo_relay::error::{FlowError, Result as FlowResult}; use pyo3::exceptions::PyRuntimeError; @@ -46,6 +47,7 @@ use tokio_stream::wrappers::ReceiverStream; use nemo_relay::api::event::{Event, EventSanitizeFields}; use nemo_relay::api::llm::LlmRequest; +use nemo_relay::api::tool::ToolExecutionFrame; use nemo_relay::codec::request::AnnotatedLlmRequest as AnnotatedLLMRequest; use nemo_relay::codec::response::AnnotatedLlmResponse as AnnotatedLLMResponse; use nemo_relay::codec::traits::{LlmCodec, LlmResponseCodec}; @@ -53,8 +55,8 @@ use nemo_relay::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::convert::{json_to_py, py_to_json}; use crate::py_types::{ PyAnnotatedLLMRequest, PyAnnotatedLLMResponse, PyLLMRequest, PyLLMRequestInterceptOutcome, - PyLlmSanitizeRequestContext, PyLlmSanitizeResponseContext, PyScopeStack, - PyToolExecutionInterceptOutcome, + PyLlmSanitizeRequestContext, PyLlmSanitizeResponseContext, PyScopeStack, PyToolExecutionFrame, + PyToolExecutionFrameOutcome, PyToolExecutionInterceptOutcome, }; type PyValueFuture = Pin>> + Send>>; @@ -990,6 +992,57 @@ fn isolated_python_continuation_context( context.map_err(|error| PyRuntimeError::new_err(error.to_string())) } +/// Wrap a Python callable `(Json) -> ToolExecutionFrame`. +/// +/// Supports both synchronous and asynchronous Python callables. +pub fn wrap_py_tool_exec_frame_fn( + py_fn: Py, +) -> Box< + dyn Fn(Json) -> Pin> + Send>> + + Send + + Sync, +> { + let py_fn = Arc::new(py_fn); + let registered_task_locals = capture_python_task_locals(); + Box::new(move |args: Json| { + let py_fn = py_fn.clone(); + let task_locals = task_locals_with_running_loop(registered_task_locals.as_ref()); + Box::pin(async move { + let result = resolve_py_object_or_future(Python::attach(|py| { + let (invocation_context, task_locals) = copy_middleware_invocation(py, task_locals) + .map_err(|error| FlowError::Internal(error.to_string()))?; + let py_args = + json_to_py(py, &args).map_err(|e: PyErr| FlowError::Internal(e.to_string()))?; + let callback = + loop_affine_callback(py, py_fn.bind(py), task_locals.as_ref(), false) + .map_err(|error| FlowError::Internal(error.to_string()))?; + let result = match invocation_context.as_ref() { + Some(context) => context.call_method1("run", (callback.bind(py), py_args)), + None => callback.bind(py).call1((py_args,)), + } + .map_err(|error| FlowError::Internal(error.to_string()))?; + split_py_object_or_future_with_locals( + py, + result.unbind(), + task_locals.as_ref(), + invocation_context.as_ref(), + ) + })) + .await?; + Python::attach(|py| { + result + .extract::(py) + .map(|value| value.inner) + .map_err(|e| { + FlowError::Internal(format!( + "tool frame callback must return ToolExecutionFrame: {e}" + )) + }) + }) + }) + }) +} + /// Python-callable wrapper for the Rust `ToolExecutionNextFn`. /// /// The Python intercept calls `await next(args)` to invoke the next layer @@ -1021,6 +1074,33 @@ impl PyToolNextFn { } } +/// Python-callable wrapper for the annotation-aware tool continuation. +#[pyclass] +struct PyToolFrameNextFn { + inner: ToolExecutionFrameNextFn, + context: MiddlewareContinuationContext, +} + +#[pymethods] +impl PyToolFrameNextFn { + fn __call__<'py>( + &self, + py: Python<'py>, + args: &Bound<'py, PyAny>, + ) -> PyResult> { + let next = self.inner.clone(); + let context = isolated_python_continuation_context(py, &self.context)?; + let json_args = py_to_json(args)?; + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let frame = context + .invoke(move || next(json_args)) + .await + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; + Ok(PyToolExecutionFrame { inner: frame }) + }) + } +} + /// Python-callable wrapper for the Rust `LlmExecutionNextFn`. /// Reusable — calling `next` multiple times is supported (retry patterns). #[pyclass] @@ -1144,6 +1224,65 @@ pub fn wrap_py_tool_exec_intercept_fn( }) } +/// Wrap an annotation-aware Python tool execution intercept. +/// +/// The Python callback receives ``(name, args, next)`` and must return +/// ``ToolExecutionFrameOutcome`` synchronously or asynchronously. +pub fn wrap_py_tool_exec_frame_intercept_fn(py_fn: Py) -> ToolExecutionFrameFn { + let py_fn = Arc::new(py_fn); + let task_locals = capture_python_task_locals(); + Arc::new( + move |name: &str, args: Json, next: ToolExecutionFrameNextFn| { + let py_fn = py_fn.clone(); + let name = name.to_string(); + let task_locals = task_locals_with_running_loop(task_locals.as_ref()); + Box::pin(async move { + let result = resolve_py_object_or_future(Python::attach(|py| { + let (invocation_context, task_locals) = + copy_middleware_invocation(py, task_locals) + .map_err(|error| FlowError::Internal(error.to_string()))?; + let py_args = json_to_py(py, &args) + .map_err(|e: PyErr| FlowError::Internal(e.to_string()))?; + let py_next = PyToolFrameNextFn { + inner: next, + context: MiddlewareContinuationContext::capture(), + }; + let py_next = py_next + .into_pyobject(py) + .map_err(|e| FlowError::Internal(e.to_string()))? + .into_any(); + let callback = + loop_affine_callback(py, py_fn.bind(py), task_locals.as_ref(), false) + .map_err(|error| FlowError::Internal(error.to_string()))?; + let result = match invocation_context.as_ref() { + Some(context) => context + .call_method1("run", (callback.bind(py), &name, py_args, py_next)), + None => callback.bind(py).call1((&name, py_args, py_next)), + } + .map_err(|e: PyErr| FlowError::Internal(e.to_string()))?; + split_py_object_or_future_with_locals( + py, + result.unbind(), + task_locals.as_ref(), + invocation_context.as_ref(), + ) + })) + .await?; + Python::attach(|py| { + result + .extract::(py) + .map(|value| value.inner) + .map_err(|e| { + FlowError::Internal(format!( + "tool frame intercept must return ToolExecutionFrameOutcome: {e}" + )) + }) + }) + }) + }, + ) +} + /// Wrap a Python callable `(name, LlmRequest, next) -> dict` for LLM execution intercepts. pub fn wrap_py_llm_exec_intercept_fn( py_fn: Py, diff --git a/crates/python/src/py_plugin.rs b/crates/python/src/py_plugin.rs index 2cf76a4a8..268e8a092 100644 --- a/crates/python/src/py_plugin.rs +++ b/crates/python/src/py_plugin.rs @@ -25,8 +25,9 @@ use nemo_relay::api::registry::{ register_llm_sanitize_response_guardrail, register_llm_stream_execution_intercept, register_mark_sanitize_guardrail, register_scope_sanitize_end_guardrail, register_scope_sanitize_start_guardrail, register_tool_conditional_execution_guardrail, - register_tool_execution_intercept, register_tool_request_intercept, - register_tool_sanitize_request_guardrail, register_tool_sanitize_response_guardrail, + register_tool_execution_frame_intercept, register_tool_execution_intercept, + register_tool_request_intercept, register_tool_sanitize_request_guardrail, + register_tool_sanitize_response_guardrail, }; use nemo_relay::api::subscriber::{deregister_subscriber, register_subscriber}; use nemo_relay::error::Result as FlowResult; @@ -43,7 +44,8 @@ use crate::py_callable::{ wrap_py_llm_exec_intercept_fn, wrap_py_llm_request_intercept_fn, wrap_py_llm_sanitize_request_fn, wrap_py_llm_sanitize_response_fn, wrap_py_llm_stream_exec_intercept_fn, wrap_py_tool_conditional_fn, - wrap_py_tool_exec_intercept_fn, wrap_py_tool_fn, wrap_py_tool_request_intercept_fn, + wrap_py_tool_exec_frame_intercept_fn, wrap_py_tool_exec_intercept_fn, wrap_py_tool_fn, + wrap_py_tool_request_intercept_fn, }; #[cfg(test)] @@ -557,6 +559,27 @@ impl PyPluginContext { ) } + #[pyo3(signature = (name: "str", priority: "int", callback: "object") -> "None", text_signature = "(name: str, priority: int, callback: object) -> None")] + fn register_tool_execution_frame_intercept( + &self, + name: &str, + priority: i32, + callback: Py, + ) -> PyResult<()> { + self.register_callback( + name, + |qualified_name| { + register_tool_execution_frame_intercept( + qualified_name, + priority, + wrap_py_tool_exec_frame_intercept_fn(callback), + ) + }, + deregister_tool_execution_intercept, + "tool execution frame intercept", + ) + } + fn __repr__(&self) -> String { "".to_string() } diff --git a/crates/python/src/py_types/core.rs b/crates/python/src/py_types/core.rs index a798ad932..0b3f51535 100644 --- a/crates/python/src/py_types/core.rs +++ b/crates/python/src/py_types/core.rs @@ -12,6 +12,7 @@ use super::{ ScopeAttributes, ScopeHandle, ScopeStackHandle, ToolAttributes, ToolHandle, json_to_py, opt_json_to_py, py_to_json, }; +use crate::convert::opt_py_to_json; use nemo_relay::api::event::{CategoryProfile, EventCategory, PendingMarkSpec}; use nemo_relay::api::llm::LlmRequestInterceptOutcome; use nemo_relay::api::runtime::subscriber_dispatcher::PublicationBuffer; @@ -19,7 +20,9 @@ use nemo_relay::api::runtime::{ LlmSanitizeRequestContext, LlmSanitizeResponseContext, PropagationContext, ThreadScopeStackBinding, }; -use nemo_relay::api::tool::ToolExecutionInterceptOutcome; +use nemo_relay::api::tool::{ + ToolExecutionFrame, ToolExecutionFrameOutcome, ToolExecutionInterceptOutcome, +}; /// Structured identity of the codec active during LLM sanitization. #[pyclass(name = "LlmCodecIdentity", frozen)] @@ -942,7 +945,7 @@ impl PyLLMRequestInterceptOutcome { } } -/// Canonical result returned by Python tool execution intercepts. +/// Relay-owned wrapper returned by Python raw-result tool execution intercepts. #[pyclass(name = "ToolExecutionInterceptOutcome", from_py_object)] #[derive(Clone)] pub struct PyToolExecutionInterceptOutcome { @@ -977,3 +980,84 @@ impl PyToolExecutionInterceptOutcome { .collect() } } + +/// Raw tool result plus an optional opaque annotation. +#[pyclass(name = "ToolExecutionFrame", from_py_object)] +#[derive(Clone)] +pub struct PyToolExecutionFrame { + pub inner: ToolExecutionFrame, +} + +#[pymethods] +impl PyToolExecutionFrame { + #[new] + #[pyo3(signature = (result, annotation=None))] + fn new(result: &Bound<'_, PyAny>, annotation: Option<&Bound<'_, PyAny>>) -> PyResult { + Ok(Self { + inner: ToolExecutionFrame { + result: py_to_json(result)?, + annotation: opt_py_to_json(annotation)?, + }, + }) + } + + #[getter] + fn result(&self, py: Python<'_>) -> PyResult> { + json_to_py(py, &self.inner.result) + } + + #[setter] + fn set_result(&mut self, result: &Bound<'_, PyAny>) -> PyResult<()> { + self.inner.result = py_to_json(result)?; + Ok(()) + } + + #[getter] + fn annotation(&self, py: Python<'_>) -> PyResult> { + opt_json_to_py(py, &self.inner.annotation) + } + + #[setter] + fn set_annotation(&mut self, annotation: Option<&Bound<'_, PyAny>>) -> PyResult<()> { + self.inner.annotation = opt_py_to_json(annotation)?; + Ok(()) + } +} + +/// Result returned by annotation-aware Python tool execution intercepts. +#[pyclass(name = "ToolExecutionFrameOutcome", from_py_object)] +#[derive(Clone)] +pub struct PyToolExecutionFrameOutcome { + pub inner: ToolExecutionFrameOutcome, +} + +#[pymethods] +impl PyToolExecutionFrameOutcome { + #[new] + #[pyo3(signature = (frame, pending_marks=Vec::new()))] + fn new(frame: PyToolExecutionFrame, pending_marks: Vec) -> Self { + Self { + inner: ToolExecutionFrameOutcome { + frame: frame.inner, + pending_marks: pending_marks.into_iter().map(|value| value.inner).collect(), + }, + } + } + + #[getter] + fn frame(&self) -> PyToolExecutionFrame { + PyToolExecutionFrame { + inner: self.inner.frame.clone(), + } + } + + #[getter] + fn pending_marks(&self) -> Vec { + self.inner + .pending_marks + .iter() + .cloned() + .map(|inner| PyPendingMarkSpec { inner }) + .collect() + } +} diff --git a/crates/python/src/py_types/mod.rs b/crates/python/src/py_types/mod.rs index c1679f4ef..3509dbee0 100644 --- a/crates/python/src/py_types/mod.rs +++ b/crates/python/src/py_types/mod.rs @@ -160,6 +160,8 @@ fn register_llm_types(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; Ok(()) diff --git a/crates/types/src/api/tool.rs b/crates/types/src/api/tool.rs index 372425072..5ddb303b4 100644 --- a/crates/types/src/api/tool.rs +++ b/crates/types/src/api/tool.rs @@ -9,6 +9,13 @@ use serde::{Deserialize, Serialize}; use crate::Json; use crate::api::event::PendingMarkSpec; +/// Category-profile key used to expose a tool-result annotation on lifecycle events. +pub const TOOL_RESULT_ANNOTATION_PROFILE_KEY: &str = "tool_result_annotation"; +/// JSON-envelope schema for an opaque tool execution frame. +pub const TOOL_EXECUTION_FRAME_SCHEMA: &str = "nemo.relay.ToolExecutionFrame@1"; +/// JSON-envelope schema for a frame-aware tool execution intercept outcome. +pub const TOOL_EXECUTION_FRAME_OUTCOME_SCHEMA: &str = "nemo.relay.ToolExecutionFrameOutcome@1"; + bitflags! { /// Bitflags that modify tool-call behavior and observability. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] @@ -18,7 +25,7 @@ bitflags! { } } -/// Canonical result returned by a tool execution intercept. +/// Relay-owned wrapper returned by a raw-result tool execution intercept. /// /// `result` is passed to the remaining middleware and application. `pending_marks` /// are Relay-owned lifecycle metadata retained separately and emitted after the @@ -54,3 +61,126 @@ impl From for ToolExecutionInterceptOutcome { Self::new(result) } } + +/// Tool result plus an optional opaque annotation for Relay interception. +/// +/// The raw [`Self::result`] remains the application-visible value. The +/// annotation is carried only as adjacent middleware and lifecycle context; +/// Relay does not define or interpret the schema of either value. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ToolExecutionFrame { + /// Raw application-owned tool result. + pub result: Json, + /// Optional application-supplied result annotation, opaque to Relay. + #[serde( + default, + deserialize_with = "deserialize_annotation", + skip_serializing_if = "annotation_is_absent" + )] + pub annotation: Option, +} + +fn deserialize_annotation<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + Option::::deserialize(deserializer).map(normalize_annotation) +} + +fn annotation_is_absent(annotation: &Option) -> bool { + annotation.as_ref().is_none_or(Json::is_null) +} + +fn normalize_annotation(annotation: Option) -> Option { + annotation.filter(|value| !value.is_null()) +} + +impl ToolExecutionFrame { + /// Create a frame without an annotation. + #[must_use] + pub fn new(result: Json) -> Self { + Self { + result, + annotation: None, + } + } + + /// Create a frame carrying an opaque annotation. + #[must_use] + pub fn annotated(result: Json, annotation: Json) -> Self { + Self { + result, + annotation: normalize_annotation(Some(annotation)), + } + } + + /// Attach or replace the opaque annotation. + #[must_use] + pub fn with_annotation(mut self, annotation: Json) -> Self { + self.annotation = normalize_annotation(Some(annotation)); + self + } + + /// Remove any annotation. + #[must_use] + pub fn without_annotation(mut self) -> Self { + self.annotation = None; + self + } + + /// Normalize JSON `null` to the frame's absent-annotation representation. + /// + /// Direct struct construction can produce `Some(Json::Null)` even though + /// JSON deserialization maps both a missing field and `null` to `None`. + /// Relay calls this at middleware boundaries to keep in-memory and wire + /// semantics stable. + #[must_use] + pub fn normalized(mut self) -> Self { + self.annotation = normalize_annotation(self.annotation); + self + } +} + +impl From for ToolExecutionFrame { + fn from(result: Json) -> Self { + Self::new(result) + } +} + +/// Result returned by an annotation-aware tool execution intercept. +/// +/// Pending marks remain Relay-owned lifecycle metadata and are not exposed +/// through the annotation-aware continuation, matching the existing v1 +/// execution-intercept behavior. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ToolExecutionFrameOutcome { + /// Raw result and optional opaque annotation returned to upstream middleware. + pub frame: ToolExecutionFrame, + /// Ordered marks for the managed tool lifecycle owner to emit. + #[serde(default)] + pub pending_marks: Vec, +} + +impl ToolExecutionFrameOutcome { + /// Create an outcome without pending marks. + #[must_use] + pub fn new(frame: ToolExecutionFrame) -> Self { + Self { + frame: frame.normalized(), + pending_marks: Vec::new(), + } + } + + /// Append one pending mark while preserving callback order. + #[must_use] + pub fn with_pending_mark(mut self, mark: PendingMarkSpec) -> Self { + self.pending_marks.push(mark); + self + } +} + +impl From for ToolExecutionFrameOutcome { + fn from(frame: ToolExecutionFrame) -> Self { + Self::new(frame) + } +} diff --git a/crates/types/tests/serialization_tests.rs b/crates/types/tests/serialization_tests.rs index ee0819947..53b74e68d 100644 --- a/crates/types/tests/serialization_tests.rs +++ b/crates/types/tests/serialization_tests.rs @@ -10,7 +10,9 @@ use nemo_relay_types::api::event::{ llm_attributes_to_strings, }; use nemo_relay_types::api::llm::{LlmAttributes, LlmRequest, LlmRequestInterceptOutcome}; -use nemo_relay_types::api::tool::ToolExecutionInterceptOutcome; +use nemo_relay_types::api::tool::{ + ToolExecutionFrame, ToolExecutionFrameOutcome, ToolExecutionInterceptOutcome, +}; use nemo_relay_types::codec::request::{AnnotatedLlmRequest, ContentPart, Message, MessageContent}; use nemo_relay_types::codec::response::AnnotatedLlmResponse; use serde_json::{Map, json}; @@ -226,6 +228,90 @@ fn tool_execution_intercept_outcome_converts_from_json() { assert_eq!(outcome, ToolExecutionInterceptOutcome::new(result)); } +#[test] +fn tool_execution_frame_round_trips_opaque_annotation() { + let annotation = json!({ + "producer_status": "failed", + "representation": { + "media_type": "application/json", + "data_schema": "example.tool.failure@7" + }, + "tool_call_id": "call-42", + "producer_detail": { + "retryable": true + } + }); + let outcome = ToolExecutionFrameOutcome::new(ToolExecutionFrame::annotated( + json!({"error": "timeout"}), + annotation.clone(), + )) + .with_pending_mark(PendingMarkSpec::builder().name("tool.failed").build()); + + let encoded = serde_json::to_value(&outcome).expect("frame outcome should serialize"); + assert_eq!(encoded["frame"]["result"]["error"], "timeout"); + assert_eq!(encoded["frame"]["annotation"]["producer_status"], "failed"); + assert_eq!( + encoded["frame"]["annotation"]["representation"]["data_schema"], + "example.tool.failure@7" + ); + assert_eq!( + encoded["frame"]["annotation"]["producer_detail"]["retryable"], + true + ); + + let decoded: ToolExecutionFrameOutcome = + serde_json::from_value(encoded).expect("frame outcome should deserialize"); + assert_eq!(decoded, outcome); +} + +#[test] +fn tool_execution_frame_requires_raw_result_and_defaults_to_no_annotation() { + let frame: ToolExecutionFrame = serde_json::from_value(json!({ + "result": ["plain", "json"] + })) + .expect("annotation should be optional"); + assert_eq!(frame.result, json!(["plain", "json"])); + assert!(frame.annotation.is_none()); + + assert!( + serde_json::from_value::(json!({ + "annotation": { + "version": 1, + "status": "succeeded" + } + })) + .is_err(), + "raw result remains required" + ); +} + +#[test] +fn tool_execution_frame_normalizes_null_annotation_to_absence() { + let decoded: ToolExecutionFrame = serde_json::from_value(json!({ + "result": {"ok": true}, + "annotation": null + })) + .expect("null annotation should decode as absence"); + assert!(decoded.annotation.is_none()); + + let constructed = ToolExecutionFrame::annotated(json!({"ok": true}), json!(null)); + assert!(constructed.annotation.is_none()); + assert_eq!( + serde_json::to_value(&constructed).unwrap(), + json!({"result": {"ok": true}}) + ); + + let direct = ToolExecutionFrame { + result: json!({"ok": true}), + annotation: Some(json!(null)), + }; + assert_eq!( + serde_json::to_value(&direct).unwrap(), + json!({"result": {"ok": true}}) + ); + assert!(direct.normalized().annotation.is_none()); +} + #[test] fn annotated_request_helpers_cover_portable_and_native_components() { let mut request = AnnotatedLlmRequest { diff --git a/crates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.proto b/crates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.proto index 5c441912a..c50854570 100644 --- a/crates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.proto +++ b/crates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.proto @@ -23,6 +23,7 @@ service RelayHostRuntime { rpc CreateScopeStack(CreateScopeStackRequest) returns (CreateScopeStackResponse); rpc DropScopeStack(DropScopeStackRequest) returns (HostAck); rpc ToolNext(ToolNextRequest) returns (JsonResult); + rpc ToolFrameNext(ToolFrameNextRequest) returns (JsonResult); rpc LlmNext(LlmNextRequest) returns (JsonResult); rpc LlmStreamNext(LlmStreamNextRequest) returns (stream StreamChunk); rpc DecodeLlmCodecRequest(LlmCodecDecodeRequest) returns (JsonResult); @@ -43,6 +44,7 @@ enum RegistrationSurface { TOOL_CONDITIONAL_EXECUTION_GUARDRAIL = 12; TOOL_REQUEST_INTERCEPT = 13; TOOL_EXECUTION_INTERCEPT = 14; + TOOL_EXECUTION_FRAME_INTERCEPT = 15; LLM_SANITIZE_REQUEST_GUARDRAIL = 20; LLM_SANITIZE_RESPONSE_GUARDRAIL = 21; LLM_CONDITIONAL_EXECUTION_GUARDRAIL = 22; @@ -203,6 +205,7 @@ message InvokeResponse { LlmRequestInterceptResult llm_request = 4; WorkerError error = 5; ToolExecutionInterceptResult tool_execution = 6; + ToolExecutionFrameInterceptResult tool_execution_frame = 7; } } @@ -225,6 +228,10 @@ message ToolExecutionInterceptResult { JsonEnvelope outcome = 1; } +message ToolExecutionFrameInterceptResult { + JsonEnvelope outcome = 1; +} + message StreamChunk { oneof item { JsonEnvelope value = 1; @@ -323,6 +330,14 @@ message ToolNextRequest { ScopeContext scope = 5; } +message ToolFrameNextRequest { + string activation_id = 1; + string auth_token = 2; + string continuation_id = 3; + JsonEnvelope value = 4; + ScopeContext scope = 5; +} + message LlmNextRequest { string activation_id = 1; string auth_token = 2; diff --git a/crates/worker-proto/tests/proto_tests.rs b/crates/worker-proto/tests/proto_tests.rs index 3b34ad9d1..8e08f8167 100644 --- a/crates/worker-proto/tests/proto_tests.rs +++ b/crates/worker-proto/tests/proto_tests.rs @@ -29,6 +29,7 @@ fn registration_surface_values_are_stable() { ); assert_eq!(RegistrationSurface::ToolRequestIntercept as i32, 13); assert_eq!(RegistrationSurface::ToolExecutionIntercept as i32, 14); + assert_eq!(RegistrationSurface::ToolExecutionFrameIntercept as i32, 15); assert_eq!(RegistrationSurface::LlmSanitizeRequestGuardrail as i32, 20); assert_eq!(RegistrationSurface::LlmSanitizeResponseGuardrail as i32, 21); assert_eq!( diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index fb15f886f..506b760da 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -37,7 +37,10 @@ pub use nemo_relay_types::Json; pub use nemo_relay_types::api::event::{DataSchema, Event, EventSanitizeFields, PendingMarkSpec}; pub use nemo_relay_types::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; pub use nemo_relay_types::api::scope::ScopeType; -pub use nemo_relay_types::api::tool::ToolExecutionInterceptOutcome; +pub use nemo_relay_types::api::tool::{ + TOOL_EXECUTION_FRAME_OUTCOME_SCHEMA, TOOL_EXECUTION_FRAME_SCHEMA, ToolExecutionFrame, + ToolExecutionFrameOutcome, ToolExecutionInterceptOutcome, +}; pub use nemo_relay_types::codec::optimization::{ LlmOptimizationContribution, LlmOptimizationEvidenceQuality, LlmOptimizationKind, LlmOptimizationModel, LlmOptimizationModelTransition, LlmOptimizationPayload, @@ -56,8 +59,9 @@ use nemo_relay_worker_proto::v1::{ LlmCodecDecodeResponse, LlmCodecEncodeRequest, LlmCodecKind, LlmNextRequest, LlmRequestInterceptResult, LlmStreamNextRequest, PopScopeRequest, PushScopeRequest, RegisterRequest, RegisterResponse, Registration, RegistrationSurface, ScopeContext, - ShutdownRequest, StreamChunk, ToolExecutionInterceptResult, ToolNextRequest, ValidateRequest, - ValidateResponse, WorkerAck, WorkerError, + ShutdownRequest, StreamChunk, ToolExecutionFrameInterceptResult, ToolExecutionInterceptResult, + ToolFrameNextRequest, ToolNextRequest, ValidateRequest, ValidateResponse, WorkerAck, + WorkerError, }; use nemo_relay_worker_proto::{WORKER_PROTOCOL_GRPC_V1, decode_json_envelope, json_envelope}; use tokio::net::TcpListener; @@ -137,6 +141,9 @@ type ToolRequestFn = Arc BoxFutureResult + Send + type ToolExecutionFn = Arc< dyn Fn(&str, Json, ToolNext) -> BoxFutureResult + Send + Sync, >; +type ToolExecutionFrameFn = Arc< + dyn Fn(&str, Json, ToolFrameNext) -> BoxFutureResult + Send + Sync, +>; type LlmSanitizeRequestFn = Arc< dyn Fn(LlmRequest, LlmSanitizeRequestContext) -> BoxFutureResult> + Send @@ -287,6 +294,7 @@ struct WorkerHandlers { tool_conditionals: HashMap, tool_requests: HashMap, tool_executions: HashMap, + tool_execution_frames: HashMap, llm_sanitize_requests: HashMap, llm_sanitize_responses: HashMap, llm_conditionals: HashMap, @@ -533,6 +541,32 @@ impl PluginContext { ); } + /// Registers an annotation-aware tool execution intercept. + /// + /// Calling [`ToolFrameNext::call`] continues the same execution-intercept + /// chain and returns the raw harness-owned result together with its optional + /// opaque annotation. Relay does not interpret the annotation schema. + pub fn register_tool_execution_frame_intercept( + &mut self, + name: &str, + priority: i32, + callback: F, + ) where + F: Fn(&str, Json, ToolFrameNext) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + self.push_registration( + name, + RegistrationSurface::ToolExecutionFrameIntercept, + priority, + false, + ); + self.handlers.tool_execution_frames.insert( + name.into(), + Arc::new(move |tool, value, next| Box::pin(callback(tool, value, next))), + ); + } + /// Registers an LLM sanitize-request guardrail. pub fn register_llm_sanitize_request_guardrail( &mut self, @@ -959,6 +993,35 @@ impl ToolNext { } } +/// Continuation handle for annotation-aware tool execution intercepts. +/// +/// The host invalidates this capability when the worker callback invocation +/// finishes; retained handles fail instead of invoking the chain later. +#[derive(Clone)] +pub struct ToolFrameNext { + runtime: PluginRuntime, + continuation_id: String, +} + +impl ToolFrameNext { + /// Calls the remaining mixed tool execution chain. + pub async fn call(&self, value: Json) -> Result { + let mut client = self.runtime.host_client().await?; + let response = client + .tool_frame_next(Request::new(ToolFrameNextRequest { + activation_id: self.runtime.activation_id.clone(), + auth_token: self.runtime.auth_token.clone(), + continuation_id: self.continuation_id.clone(), + value: Some(json_envelope(JSON_SCHEMA, &value)?), + scope: self.runtime.current_scope_context(), + })) + .await + .map_err(|err| WorkerSdkError::Transport(err.to_string()))? + .into_inner(); + decode_typed_json_result(response, TOOL_EXECUTION_FRAME_SCHEMA) + } +} + /// Continuation handle for LLM execution intercepts. #[derive(Clone)] pub struct LlmNext { @@ -1690,7 +1753,8 @@ impl WorkerService { | RegistrationSurface::ToolSanitizeResponseGuardrail | RegistrationSurface::ToolConditionalExecutionGuardrail | RegistrationSurface::ToolRequestIntercept - | RegistrationSurface::ToolExecutionIntercept => { + | RegistrationSurface::ToolExecutionIntercept + | RegistrationSurface::ToolExecutionFrameIntercept => { self.invoke_tool_response(request, &scope, surface).await } RegistrationSurface::LlmSanitizeRequestGuardrail @@ -1758,6 +1822,10 @@ impl WorkerService { RegistrationSurface::ToolExecutionIntercept => { self.invoke_tool_execution_response(request, scope).await } + RegistrationSurface::ToolExecutionFrameIntercept => { + self.invoke_tool_execution_frame_response(request, scope) + .await + } _ => unreachable!("tool surface was pre-filtered"), } } @@ -1823,6 +1891,21 @@ impl WorkerService { tool_execution_response(future.await?) } + async fn invoke_tool_execution_frame_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = tool_payload(request.payload)?; + let handler = self.tool_execution_frame(&request.registration_name)?; + let next = ToolFrameNext { + runtime: self.runtime.clone(), + continuation_id: request.continuation_id, + }; + let future = with_thread_scope(scope, || handler(&payload.tool_name, payload.value, next)); + tool_execution_frame_response(future.await?) + } + async fn invoke_llm_response( &self, request: InvokeRequest, @@ -2030,6 +2113,20 @@ impl WorkerService { }) } + fn tool_execution_frame(&self, name: &str) -> Result { + self.handlers + .lock() + .map_err(|err| WorkerSdkError::Callback(format!("handler lock poisoned: {err}")))? + .tool_execution_frames + .get(name) + .cloned() + .ok_or_else(|| { + WorkerSdkError::InvalidInput(format!( + "tool execution frame '{name}' not registered" + )) + }) + } + fn llm_sanitize_request(&self, name: &str) -> Result { self.handlers .lock() @@ -2290,6 +2387,21 @@ fn tool_execution_response(outcome: ToolExecutionInterceptOutcome) -> Result Result { + Ok(InvokeResponse { + result: Some( + nemo_relay_worker_proto::v1::invoke_response::Result::ToolExecutionFrame( + ToolExecutionFrameInterceptResult { + outcome: Some(json_envelope( + TOOL_EXECUTION_FRAME_OUTCOME_SCHEMA, + &outcome, + )?), + }, + ), + ), + }) +} + fn stream_chunk_to_json(chunk: StreamChunk) -> Result { match chunk.item { Some(nemo_relay_worker_proto::v1::stream_chunk::Item::Value(value)) => { @@ -2523,6 +2635,7 @@ fn all_surfaces() -> Vec { RegistrationSurface::ToolConditionalExecutionGuardrail, RegistrationSurface::ToolRequestIntercept, RegistrationSurface::ToolExecutionIntercept, + RegistrationSurface::ToolExecutionFrameIntercept, RegistrationSurface::LlmSanitizeRequestGuardrail, RegistrationSurface::LlmSanitizeResponseGuardrail, RegistrationSurface::LlmConditionalExecutionGuardrail, diff --git a/crates/worker/tests/worker_sdk_tests.rs b/crates/worker/tests/worker_sdk_tests.rs index d120425b6..52efb2666 100644 --- a/crates/worker/tests/worker_sdk_tests.rs +++ b/crates/worker/tests/worker_sdk_tests.rs @@ -20,9 +20,10 @@ use hyper_util::rt::TokioIo; use nemo_relay_types::api::event::{BaseEvent, Event, MarkEvent, PendingMarkSpec}; use nemo_relay_worker::{ ANNOTATED_LLM_REQUEST_SCHEMA, Json, JsonStream, LlmNext, LlmRequest, LlmStreamNext, - PluginContext, PluginRuntime, Result, ScopeType, ToolExecutionInterceptOutcome, ToolNext, - WorkerPlugin, WorkerSdkError, WorkerServerConfig, serve_plugin, serve_plugin_arc, - serve_plugin_arc_with_config, + PluginContext, PluginRuntime, Result, ScopeType, TOOL_EXECUTION_FRAME_OUTCOME_SCHEMA, + TOOL_EXECUTION_FRAME_SCHEMA, ToolExecutionFrame, ToolExecutionFrameOutcome, + ToolExecutionInterceptOutcome, ToolFrameNext, ToolNext, WorkerPlugin, WorkerSdkError, + WorkerServerConfig, serve_plugin, serve_plugin_arc, serve_plugin_arc_with_config, }; use nemo_relay_worker_proto::v1::plugin_worker_client::PluginWorkerClient; use nemo_relay_worker_proto::v1::relay_host_runtime_server::{ @@ -33,8 +34,8 @@ use nemo_relay_worker_proto::v1::{ DropScopeStackRequest, EmitMarkRequest, HandshakeRequest, HealthRequest, HostAck, InvokeRequest, InvokeResponse, JsonEnvelope, JsonResult, LlmInvocation, LlmNextRequest, LlmStreamNextRequest, PopScopeRequest, PushScopeRequest, PushScopeResponse, RegisterRequest, - RegistrationSurface, ScopeContext, ShutdownRequest, StreamChunk, ToolInvocation, - ToolNextRequest, ValidateRequest, WorkerError, + RegistrationSurface, ScopeContext, ShutdownRequest, StreamChunk, ToolFrameNextRequest, + ToolInvocation, ToolNextRequest, ValidateRequest, WorkerError, }; use nemo_relay_worker_proto::{WORKER_PROTOCOL_GRPC_V1, decode_json_envelope, json_envelope}; use serde_json::json; @@ -200,12 +201,13 @@ async fn worker_service_enforces_auth_and_reports_registrations() { assert_eq!(invalid_register_config.code(), tonic::Code::InvalidArgument); let registrations = register_plugin(&mut client).await; - assert_eq!(registrations.len(), 21); + assert_eq!(registrations.len(), 22); for local_name in [ "llm-sanitize-request", "llm-sanitize-response", "llm-sanitize-omit-request", "llm-sanitize-omit-response", + "tool-exec-frame", ] { assert_eq!( registrations @@ -615,6 +617,28 @@ async fn worker_service_invokes_every_registration_surface() { let tool_exec = tool_outcome.result; assert_json_field(tool_exec.clone(), "next", "tool"); assert_json_field(tool_exec, "phase", "tool_exec"); + let frame_outcome = invoke_tool_execution_frame( + &mut client, + tool_invoke( + "tool-exec-frame", + RegistrationSurface::ToolExecutionFrameIntercept, + json!({}), + ), + ) + .await; + assert_eq!(frame_outcome.pending_marks.len(), 1); + assert_eq!( + frame_outcome.pending_marks[0].name, + "worker.tool.execution.frame" + ); + assert_json_field(frame_outcome.frame.result, "phase", "tool_exec_frame"); + assert_eq!( + frame_outcome.frame.annotation, + Some(json!({ + "producer": "host", + "observed_by": "worker", + })) + ); assert!( invoke_json( &mut client, @@ -785,6 +809,7 @@ async fn worker_service_invokes_every_registration_surface() { assert!(calls.contains(&"pop:scope-handle-1".into())); assert!(calls.contains(&"drop:isolated-stack".into())); assert!(calls.contains(&"tool_next:next-1".into())); + assert!(calls.contains(&"tool_frame_next:next-1".into())); assert!(calls.contains(&"llm_next:next-1".into())); assert!(calls.contains(&"llm_stream_next:next-1".into())); assert!(calls.contains(&"codec_request_decode:request-capability".into())); @@ -1492,6 +1517,23 @@ async fn worker_service_propagates_host_runtime_errors() { ); } + host.set_failures(MockHostFailures { + tool_frame_next: true, + ..Default::default() + }); + assert_worker_error( + client + .invoke(Request::new(tool_invoke( + "tool-exec-frame", + RegistrationSurface::ToolExecutionFrameIntercept, + json!({}), + ))) + .await + .expect("tool frame next failure returns structured error") + .into_inner(), + "tool frame next failed", + ); + host.set_failures(MockHostFailures { llm_next: true, ..Default::default() @@ -1823,6 +1865,22 @@ impl WorkerPlugin for SurfacePlugin { )) } }); + ctx.register_tool_execution_frame_intercept( + "tool-exec-frame", + 1, + |_, value, next: ToolFrameNext| async move { + let mut frame = next.call(value).await?; + frame.result = set_json_field(frame.result, "phase", "tool_exec_frame"); + let mut annotation = frame.annotation.unwrap_or_else(|| json!({})); + annotation["observed_by"] = json!("worker"); + frame.annotation = Some(annotation); + Ok(ToolExecutionFrameOutcome::new(frame).with_pending_mark( + PendingMarkSpec::builder() + .name("worker.tool.execution.frame") + .build(), + )) + }, + ); let scope_runtime = runtime.clone(); ctx.register_tool_execution_intercept("tool-scope-types", 1, move |_, _, _| { let runtime = scope_runtime.clone(); @@ -2018,6 +2076,7 @@ struct MockHostFailures { pop_scope: HostFailure, drop_scope_stack: HostFailure, tool_next: bool, + tool_frame_next: bool, llm_next: bool, llm_stream_mode: MockStreamMode, codec_request_decode: bool, @@ -2182,6 +2241,34 @@ impl RelayHostRuntime for MockHost { })) } + async fn tool_frame_next( + &self, + request: Request, + ) -> std::result::Result, Status> { + let request = request.into_inner(); + authorize_host(&request.activation_id, &request.auth_token)?; + self.record(format!("tool_frame_next:{}", request.continuation_id)); + if self.failures().tool_frame_next { + return Ok(Response::new(JsonResult { + value: None, + error: Some(worker_error("tool frame next failed")), + })); + } + Ok(Response::new(JsonResult { + value: Some( + json_envelope( + TOOL_EXECUTION_FRAME_SCHEMA, + &ToolExecutionFrame::annotated( + json!({"next": "tool_frame"}), + json!({"producer": "host"}), + ), + ) + .expect("encode tool execution frame"), + ), + error: None, + })) + } + async fn llm_next( &self, request: Request, @@ -2759,6 +2846,25 @@ async fn invoke_tool_execution( } } +async fn invoke_tool_execution_frame( + client: &mut PluginWorkerClient, + request: InvokeRequest, +) -> ToolExecutionFrameOutcome { + let response = client + .invoke(Request::new(request)) + .await + .expect("invoke succeeds") + .into_inner(); + match response.result.expect("invoke result") { + nemo_relay_worker_proto::v1::invoke_response::Result::ToolExecutionFrame(result) => { + let outcome = result.outcome.expect("tool execution frame outcome"); + assert_eq!(outcome.schema, TOOL_EXECUTION_FRAME_OUTCOME_SCHEMA); + decode_json_envelope(&outcome).expect("decode tool execution frame outcome") + } + other => panic!("unexpected invoke result: {other:?}"), + } +} + async fn invoke_guardrail( client: &mut PluginWorkerClient, request: InvokeRequest, diff --git a/docs/reference/tool-execution-intercept-outcomes.mdx b/docs/reference/tool-execution-intercept-outcomes.mdx index 7567f6a23..86d8d8c97 100644 --- a/docs/reference/tool-execution-intercept-outcomes.mdx +++ b/docs/reference/tool-execution-intercept-outcomes.mdx @@ -1,13 +1,14 @@ --- title: "Tool Execution Intercept Outcomes" -description: "Canonical result returned by tool execution intercepts and its managed lifecycle behavior." +description: "Relay middleware wrappers for tool results, opaque annotations, and managed lifecycle behavior." --- {/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 */} A tool execution intercept wraps or short-circuits a managed tool callback. -Every intercept returns one canonical outcome so Relay can keep lifecycle -control data separate from the application-visible tool result. +The harness still owns tool dispatch, receipt, and the schema of the tool +result. Relay wraps that result only to keep Relay-owned lifecycle control data +separate from the application-visible value. ```json { @@ -30,12 +31,61 @@ them in effective middleware order, and materializes them only after the final outcome succeeds. Use the existing global, scope-local, or plugin-context tool execution -registration APIs to produce pending marks, and return the canonical outcome -from every registered callback. Use a +registration APIs to produce pending marks, and return the Relay outcome +wrapper from every registered callback. Use a [mark event sanitizer](/reference/event-sanitizers) when the emitted pending mark's `data`, `category_profile`, or `metadata` must be sanitized. Legacy raw intercept returns are rejected at public and dynamic-plugin boundaries. +## Opaque Result Annotations + +An integration that needs to carry metadata adjacent to a tool result can use +the frame-aware execution API: + +```json +{ + "result": {}, + "annotation": {} +} +``` + +`result` is the harness-owned value. `annotation` is optional and opaque to +Relay; a missing value or JSON `null` means no annotation. Relay does not +define fields inside either value, convert one tool-result schema into +another, or participate in tool dispatch. + +For lifecycle emission, Relay places the annotation in the tool event's +category profile before applying the existing event-sanitizer chain. Tool +response sanitizers continue to receive only the harness-owned raw result; +Relay does not infer how a result rewrite should transform opaque annotation +fields. + +Frame-aware callbacks use `ToolExecutionFrame` for `next(args)` and place the +returned frame inside `ToolExecutionFrameOutcome` so Relay can continue to +carry pending marks separately. They share the same registry, name namespace, +priority ordering, and chain as raw-result callbacks. The existing generic +tool-execution deregistration operation removes either callback form. + +Compatibility with raw-result middleware is conservative: + +- A raw callback that invokes `next(args)` exactly once and returns the + downstream result unchanged preserves the downstream annotation. At this + JSON boundary, "unchanged" means structural JSON equality; Relay cannot + distinguish a reconstructed equal value from the original allocation. +- A raw callback that changes or replaces the result drops the annotation so + stale metadata cannot become attached to a different result. +- A raw callback that calls `next` more than once or short-circuits execution + also returns without an implicit annotation. +- A frame-aware callback may explicitly preserve, replace, or remove the + annotation. + +Frame-aware continuations are invocation-scoped. A frame-aware callback must +await every `next` call before it returns; Relay rejects frame continuation +calls that begin after the callback finishes, but it cannot cancel arbitrary +downstream work that a callback already detached after the continuation +started. The existing raw callback and continuation API retains its prior +behavior. + ## Managed Lifecycle On successful managed execution, Relay emits the tool end event before any @@ -59,20 +109,24 @@ or object shape: - Go callbacks return `ToolExecutionInterceptOutcome`. - Node.js callbacks return `{ result, pendingMarks? }`, where JavaScript pending-mark DTOs use `categoryProfile`. -- Public C callbacks return canonical JSON with `result` and optional +- Public C callbacks return Relay outcome JSON with `result` and optional `pending_marks`. - `grpc-v1` worker SDKs return a `ToolExecutionInterceptOutcome` in a `JsonEnvelope` with schema `nemo.relay.ToolExecutionInterceptOutcome@1`. -Canonical JSON uses `pending_marks` and `category_profile` across bindings. +Relay's outcome JSON uses `pending_marks` and `category_profile` across +bindings. Frame-aware worker and C envelopes additionally carry the opaque +`frame.result` and optional `frame.annotation` values without interpreting +them. ## Migration This finalizes the unpublished tool execution intercept contract. Update every -registered tool execution intercept to return the canonical outcome, while -leaving the default tool callback and `next(args)` continuation as raw JSON. -Rebuild development native plugins and workers against the same NeMo Relay -release that hosts them. +registered raw-result tool execution intercept to return Relay's outcome +wrapper, while leaving the default raw tool callback and `next(args)` +continuation unchanged. Use the frame-aware registration only when an +integration must carry an adjacent opaque annotation. Rebuild development +native plugins and workers against the same NeMo Relay release that hosts them. ## Related Topics diff --git a/go/nemo_relay/callbacks.go b/go/nemo_relay/callbacks.go index 8096c7447..4b3987e34 100644 --- a/go/nemo_relay/callbacks.go +++ b/go/nemo_relay/callbacks.go @@ -54,6 +54,8 @@ typedef struct FfiPluginContext FfiPluginContext; // Middleware chain next function types typedef char* (*NemoRelayToolExecNextFn)(const char* args_json, void* next_ctx); typedef char* (*NemoRelayToolExecInterceptCb)(void* user_data, const char* args_json, NemoRelayToolExecNextFn next_fn, void* next_ctx); +typedef char* (*NemoRelayToolExecFrameNextFn)(const char* args_json, void* next_ctx); +typedef char* (*NemoRelayToolExecFrameInterceptCb)(void* user_data, const char* args_json, NemoRelayToolExecFrameNextFn next_fn, void* next_ctx); typedef char* (*NemoRelayLlmExecNextFn)(const char* native_json, void* next_ctx); typedef char* (*NemoRelayLlmExecInterceptCb)(void* user_data, const char* native_json, NemoRelayLlmExecNextFn next_fn, void* next_ctx); @@ -62,6 +64,11 @@ static inline char* callToolExecNext(NemoRelayToolExecNextFn next_fn, const char return next_fn(args_json, next_ctx); } +// Helper to call the annotation-aware tool exec next function pointer from Go +static inline char* callToolExecFrameNext(NemoRelayToolExecFrameNextFn next_fn, const char* args_json, void* next_ctx) { + return next_fn(args_json, next_ctx); +} + // Helper to call the LLM exec next function pointer from Go static inline char* callLlmExecNext(NemoRelayLlmExecNextFn next_fn, const char* native_json, void* next_ctx) { return next_fn(native_json, next_ctx); @@ -171,14 +178,26 @@ type ToolConditionalFunc func(name string, args json.RawMessage) *string // arguments as JSON and returning the result JSON or an error. type ToolExecutionFunc func(args json.RawMessage) (json.RawMessage, error) +// ToolExecutionFrameFunc executes a tool and returns its raw result together +// with an optional opaque annotation. +type ToolExecutionFrameFunc func(args json.RawMessage) (ToolExecutionFrame, error) + // ToolExecutionInterceptFunc is a callback for tool execution intercepts // following the middleware chain pattern. It receives the tool arguments and // a `next` function. Call `next` to invoke the next intercept in the chain // (or the original tool implementation if this is the innermost intercept). // Skip calling `next` to short-circuit the chain entirely. The callback returns -// the canonical outcome containing the tool result and any pending marks. +// Relay's outcome wrapper containing the tool result and any pending marks. +// The `next` function is invocation-scoped and must not be retained after the +// callback returns. type ToolExecutionInterceptFunc func(args json.RawMessage, next func(json.RawMessage) (json.RawMessage, error)) (ToolExecutionInterceptOutcome, error) +// ToolExecutionFrameInterceptFunc is the annotation-aware form of a tool +// execution intercept. It participates in the same priority-ordered chain as +// ToolExecutionInterceptFunc. The `next` function is invocation-scoped and +// must not be retained after the callback returns. +type ToolExecutionFrameInterceptFunc func(args json.RawMessage, next func(json.RawMessage) (ToolExecutionFrame, error)) (ToolExecutionFrameOutcome, error) + // LLMCodecKind identifies the active codec state supplied to a sanitizer. type LLMCodecKind string @@ -226,28 +245,48 @@ func (context LLMSanitizeResponseContext) ResolveCodec() *LLMResponseSanitizeCod // capability is used after its sanitizer callback has returned. var ErrLLMSanitizeCodecExpired = errors.New("LLM sanitizer codec capability is no longer active") -type llmSanitizeCodecInvocation struct { +// ErrToolExecutionContinuationExpired is returned when an execution-intercept +// continuation is retained and called after its callback has returned. +var ErrToolExecutionContinuationExpired = errors.New("tool execution continuation is no longer active") + +type callbackLifetime struct { mu sync.RWMutex active bool } +func newCallbackLifetime() *callbackLifetime { + return &callbackLifetime{active: true} +} + +func (lifetime *callbackLifetime) acquire(expired error) (func(), error) { + lifetime.mu.RLock() + if !lifetime.active { + lifetime.mu.RUnlock() + return nil, expired + } + return lifetime.mu.RUnlock, nil +} + +func (lifetime *callbackLifetime) invalidate() { + lifetime.mu.Lock() + lifetime.active = false + lifetime.mu.Unlock() +} + +type llmSanitizeCodecInvocation struct { + lifetime *callbackLifetime +} + func newLLMSanitizeCodecInvocation() *llmSanitizeCodecInvocation { - return &llmSanitizeCodecInvocation{active: true} + return &llmSanitizeCodecInvocation{lifetime: newCallbackLifetime()} } func (invocation *llmSanitizeCodecInvocation) acquire() (func(), error) { - invocation.mu.RLock() - if !invocation.active { - invocation.mu.RUnlock() - return nil, ErrLLMSanitizeCodecExpired - } - return invocation.mu.RUnlock, nil + return invocation.lifetime.acquire(ErrLLMSanitizeCodecExpired) } func (invocation *llmSanitizeCodecInvocation) invalidate() { - invocation.mu.Lock() - invocation.active = false - invocation.mu.Unlock() + invocation.lifetime.invalidate() } // LLMRequestSanitizeCodec is a callback-scoped request codec capability. @@ -444,8 +483,8 @@ type LLMRequestInterceptOutcome struct { OptimizationContributions []LLMOptimizationContribution `json:"optimization_contributions"` } -// ToolExecutionInterceptOutcome is the canonical result of a tool execution -// intercept. Result is passed to the remaining middleware and application; +// ToolExecutionInterceptOutcome is Relay's wrapper for a tool execution +// intercept result. Result is passed to the remaining middleware and application; // PendingMarks are Relay-owned lifecycle metadata emitted after the tool-end // event and are not included in the application-visible result. type ToolExecutionInterceptOutcome struct { @@ -453,6 +492,20 @@ type ToolExecutionInterceptOutcome struct { PendingMarks []PendingMarkSpec `json:"pending_marks"` } +// ToolExecutionFrame carries a raw application result and an optional opaque +// annotation. Relay does not interpret Annotation. +type ToolExecutionFrame struct { + Result json.RawMessage `json:"result"` + Annotation json.RawMessage `json:"annotation,omitempty"` +} + +// ToolExecutionFrameOutcome is returned by annotation-aware execution +// intercepts. Pending marks retain their existing Relay-owned semantics. +type ToolExecutionFrameOutcome struct { + Frame ToolExecutionFrame `json:"frame"` + PendingMarks []PendingMarkSpec `json:"pending_marks"` +} + // LLMRequestInterceptFunc is a callback for LLM request intercepts. When // annotatedJSON is non-nil, request.Content is read-only, request.Headers may // be changed, and the returned annotation is authoritative for provider body @@ -614,6 +667,22 @@ func goToolExecTrampoline(userData unsafe.Pointer, argsJSON *C.char) *C.char { return C.CString(string(result)) } +//export goToolExecFrameTrampoline +func goToolExecFrameTrampoline(userData unsafe.Pointer, argsJSON *C.char) *C.char { + fn := lookupClosure(userData).(ToolExecutionFrameFunc) + frame, err := fn(json.RawMessage(C.GoString(argsJSON))) + if err != nil { + setLastErrorMessage(err.Error()) + return nil + } + value, err := json.Marshal(frame) + if err != nil { + setLastErrorMessage(err.Error()) + return nil + } + return C.CString(string(value)) +} + //export goEventSubscriberTrampoline func goEventSubscriberTrampoline(userData unsafe.Pointer, event *C.FfiEvent) { fn := lookupClosure(userData).(EventSubscriberFunc) @@ -772,7 +841,14 @@ func goLlmExecTrampoline(userData unsafe.Pointer, nativeJSON *C.char) *C.char { func goToolExecInterceptTrampoline(userData unsafe.Pointer, argsJSON *C.char, nextFn C.NemoRelayToolExecNextFn, nextCtx unsafe.Pointer) *C.char { fn := lookupClosure(userData).(ToolExecutionInterceptFunc) goArgs := json.RawMessage(C.GoString(argsJSON)) + invocation := newCallbackLifetime() + defer invocation.invalidate() goNext := func(args json.RawMessage) (json.RawMessage, error) { + release, err := invocation.acquire(ErrToolExecutionContinuationExpired) + if err != nil { + return nil, err + } + defer release() cArgs := C.CString(string(args)) defer C.free(unsafe.Pointer(cArgs)) result := C.callToolExecNext(nextFn, cArgs, nextCtx) @@ -798,6 +874,46 @@ func goToolExecInterceptTrampoline(userData unsafe.Pointer, argsJSON *C.char, ne return C.CString(string(outcomeJSON)) } +//export goToolExecFrameInterceptTrampoline +func goToolExecFrameInterceptTrampoline(userData unsafe.Pointer, argsJSON *C.char, nextFn C.NemoRelayToolExecFrameNextFn, nextCtx unsafe.Pointer) *C.char { + fn := lookupClosure(userData).(ToolExecutionFrameInterceptFunc) + invocation := newCallbackLifetime() + defer invocation.invalidate() + goNext := func(args json.RawMessage) (ToolExecutionFrame, error) { + release, err := invocation.acquire(ErrToolExecutionContinuationExpired) + if err != nil { + return ToolExecutionFrame{}, err + } + defer release() + cArgs := C.CString(string(args)) + defer C.free(unsafe.Pointer(cArgs)) + result := C.callToolExecFrameNext(nextFn, cArgs, nextCtx) + if result == nil { + return ToolExecutionFrame{}, lastError() + } + defer C.nemo_relay_string_free(result) + var frame ToolExecutionFrame + if err := json.Unmarshal([]byte(C.GoString(result)), &frame); err != nil { + return ToolExecutionFrame{}, err + } + return frame, nil + } + outcome, err := fn(json.RawMessage(C.GoString(argsJSON)), goNext) + if err != nil { + setLastErrorMessage(err.Error()) + return nil + } + if outcome.PendingMarks == nil { + outcome.PendingMarks = []PendingMarkSpec{} + } + outcomeJSON, err := json.Marshal(outcome) + if err != nil { + setLastErrorMessage(err.Error()) + return nil + } + return C.CString(string(outcomeJSON)) +} + //export goLlmExecInterceptTrampoline func goLlmExecInterceptTrampoline(userData unsafe.Pointer, nativeJSON *C.char, nextFn C.NemoRelayLlmExecNextFn, nextCtx unsafe.Pointer) *C.char { fn := lookupClosure(userData).(LLMExecutionInterceptFunc) diff --git a/go/nemo_relay/deregister_test.go b/go/nemo_relay/deregister_test.go index e243bed6b..8e451adff 100644 --- a/go/nemo_relay/deregister_test.go +++ b/go/nemo_relay/deregister_test.go @@ -157,6 +157,25 @@ func TestRegisterDeregisterReregisterToolExecutionIntercept(t *testing.T) { DeregisterToolExecutionIntercept(name) } +func TestRegisterDeregisterReregisterToolExecutionFrameIntercept(t *testing.T) { + name := "go_reregister_frame_exec_int" + fn := func(args json.RawMessage, next func(json.RawMessage) (ToolExecutionFrame, error)) (ToolExecutionFrameOutcome, error) { + frame, err := next(args) + return ToolExecutionFrameOutcome{Frame: frame}, err + } + + if err := RegisterToolExecutionFrameIntercept(name, 1, fn); err != nil { + t.Fatalf(firstRegisterFailed, err) + } + if err := DeregisterToolExecutionIntercept(name); err != nil { + t.Fatalf(deregisterFailed, err) + } + if err := RegisterToolExecutionFrameIntercept(name, 1, fn); err != nil { + t.Fatalf(reregisterFailed, err) + } + DeregisterToolExecutionIntercept(name) +} + func TestRegisterDeregisterReregisterLlmSanitizeRequestGuardrail(t *testing.T) { name := "go_reregister_llm_san_req" fn := func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { diff --git a/go/nemo_relay/nemo_relay.go b/go/nemo_relay/nemo_relay.go index a3e80fd18..2c813a75d 100644 --- a/go/nemo_relay/nemo_relay.go +++ b/go/nemo_relay/nemo_relay.go @@ -54,15 +54,23 @@ extern int32_t nemo_relay_event(const char* name, const FfiScopeHandle* parent, // Tool lifecycle extern int32_t nemo_relay_tool_call(const char* name, const char* args_json, const FfiScopeHandle* parent, uint32_t attributes, const char* data_json, const char* metadata_json, const char* tool_call_id, const int64_t* timestamp_unix_micros, FfiToolHandle** out); extern int32_t nemo_relay_tool_call_end(const FfiToolHandle* handle, const char* result_json, const char* data_json, const char* metadata_json, const int64_t* timestamp_unix_micros); +extern int32_t nemo_relay_tool_call_end_frame(const FfiToolHandle* handle, const char* frame_json, const char* data_json, const char* metadata_json, const int64_t* timestamp_unix_micros); // Tool call execute (with C function pointer callbacks) typedef char* (*NemoRelayToolExecFn)(void* user_data, const char* args_json); +typedef char* (*NemoRelayToolExecFrameFn)(void* user_data, const char* args_json); extern int32_t nemo_relay_tool_call_execute( const char* name, const char* args_json, NemoRelayToolExecFn func_cb, void* func_user_data, NemoRelayFreeFn func_free, const FfiScopeHandle* parent, uint32_t attributes, const char* data_json, const char* metadata_json, char** out); +extern int32_t nemo_relay_tool_call_execute_frame( + const char* name, const char* args_json, + NemoRelayToolExecFrameFn func_cb, void* func_user_data, NemoRelayFreeFn func_free, + const FfiScopeHandle* parent, uint32_t attributes, + const char* data_json, const char* metadata_json, + char** out); // LLM lifecycle typedef void (*NemoRelayCollectorCb)(const char* chunk_json); @@ -135,7 +143,10 @@ extern int32_t nemo_relay_deregister_tool_request_intercept(const char* name); // Middleware chain intercept callback types (must be declared before use in externs) typedef char* (*NemoRelayToolExecNextFn)(const char* args_json, void* next_ctx); typedef char* (*NemoRelayToolExecInterceptCb)(void* user_data, const char* args_json, NemoRelayToolExecNextFn next_fn, void* next_ctx); +typedef char* (*NemoRelayToolExecFrameNextFn)(const char* args_json, void* next_ctx); +typedef char* (*NemoRelayToolExecFrameInterceptCb)(void* user_data, const char* args_json, NemoRelayToolExecFrameNextFn next_fn, void* next_ctx); extern int32_t nemo_relay_register_tool_execution_intercept(const char* name, int32_t priority, NemoRelayToolExecInterceptCb exec_cb, void* exec_user_data, NemoRelayFreeFn exec_free); +extern int32_t nemo_relay_register_tool_execution_frame_intercept(const char* name, int32_t priority, NemoRelayToolExecFrameInterceptCb exec_cb, void* exec_user_data, NemoRelayFreeFn exec_free); extern int32_t nemo_relay_deregister_tool_execution_intercept(const char* name); // LLM guardrails @@ -194,6 +205,7 @@ extern int32_t nemo_relay_scope_deregister_tool_conditional_execution_guardrail( extern int32_t nemo_relay_scope_register_tool_request_intercept(const char* scope_uuid, const char* name, int32_t priority, _Bool break_chain, NemoRelayToolSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_scope_deregister_tool_request_intercept(const char* scope_uuid, const char* name); extern int32_t nemo_relay_scope_register_tool_execution_intercept(const char* scope_uuid, const char* name, int32_t priority, NemoRelayToolExecInterceptCb exec_cb, void* exec_user_data, NemoRelayFreeFn exec_free); +extern int32_t nemo_relay_scope_register_tool_execution_frame_intercept(const char* scope_uuid, const char* name, int32_t priority, NemoRelayToolExecFrameInterceptCb exec_cb, void* exec_user_data, NemoRelayFreeFn exec_free); extern int32_t nemo_relay_scope_deregister_tool_execution_intercept(const char* scope_uuid, const char* name); // Scope-local LLM guardrails @@ -270,6 +282,7 @@ extern char* goToolSanitizeTrampoline(void*, const char*, const char*); extern char* goEventSanitizeTrampoline(void*, const FfiEvent*, const char*); extern char* goToolConditionalTrampoline(void*, const char*, const char*); extern char* goToolExecTrampoline(void*, const char*); +extern char* goToolExecFrameTrampoline(void*, const char*); extern void goEventSubscriberTrampoline(void*, const FfiEvent*); extern void goFreeTrampoline(void*); extern FfiLLMRequest* goLlmRequestTrampoline(void*, const FfiLLMRequest*, NemoRelayLlmSanitizeRequestContext); @@ -277,6 +290,7 @@ extern char* goLlmResponseTrampoline(void*, const char*, NemoRelayLlmSanitizeRes extern char* goLlmConditionalTrampoline(void*, const FfiLLMRequest*); extern char* goLlmExecTrampoline(void*, const char*); extern char* goToolExecInterceptTrampoline(void*, const char*, NemoRelayToolExecNextFn, void*); +extern char* goToolExecFrameInterceptTrampoline(void*, const char*, NemoRelayToolExecFrameNextFn, void*); extern char* goLlmExecInterceptTrampoline(void*, const char*, NemoRelayLlmExecNextFn, void*); // Codec trampolines (used at execute time, not registration) @@ -779,6 +793,26 @@ func ToolCallEnd(handle *ToolHandle, result json.RawMessage, opts ...ToolCallOpt return checkStatus(C.nemo_relay_tool_call_end(handle.ptr, cResult, o.data, o.metadata, o.timestamp)) } +// ToolCallEndFrame completes a manually started tool call with a raw result +// and an optional opaque annotation. The annotation is carried on Relay's +// lifecycle event; Relay does not interpret it. +func ToolCallEndFrame(handle *ToolHandle, frame ToolExecutionFrame, opts ...ToolCallOption) error { + o := &toolCallOptions{} + for _, opt := range opts { + opt(o) + } + defer freeToolOpts(o) + + frameJSON, err := json.Marshal(frame) + if err != nil { + return err + } + cFrame := C.CString(string(frameJSON)) + defer C.free(unsafe.Pointer(cFrame)) + + return checkStatus(C.nemo_relay_tool_call_end_frame(handle.ptr, cFrame, o.data, o.metadata, o.timestamp)) +} + // ToolCallExecute runs a complete tool call lifecycle through the full // middleware pipeline: conditional-execution guardrails (on raw args), // request intercepts, sanitize-request guardrails for the emitted Start event @@ -820,6 +854,45 @@ func ToolCallExecute(name string, args json.RawMessage, fn ToolExecutionFunc, op return result, nil } +// ToolCallExecuteFrame runs the existing tool middleware and lifecycle +// pipeline while allowing the producer and frame-aware intercepts to carry an +// optional opaque annotation beside the raw result. +func ToolCallExecuteFrame(name string, args json.RawMessage, fn ToolExecutionFrameFunc, opts ...ToolCallOption) (ToolExecutionFrame, error) { + o := &toolCallOptions{} + for _, opt := range opts { + opt(o) + } + defer freeToolOpts(o) + + id := registerClosure(fn) + + cName := C.CString(name) + cArgs := C.CString(string(args)) + defer C.free(unsafe.Pointer(cName)) + defer C.free(unsafe.Pointer(cArgs)) + + var out *C.char + status := C.nemo_relay_tool_call_execute_frame( + cName, cArgs, + C.NemoRelayToolExecFrameFn(C.goToolExecFrameTrampoline), + id, + C.NemoRelayFreeFn(C.goFreeTrampoline), + o.parent, C.uint32_t(o.attributes), + o.data, o.metadata, + &out, + ) + if err := checkStatus(status); err != nil { + return ToolExecutionFrame{}, err + } + defer C.nemo_relay_string_free(out) + + var frame ToolExecutionFrame + if err := json.Unmarshal([]byte(C.GoString(out)), &frame); err != nil { + return ToolExecutionFrame{}, err + } + return frame, nil +} + // --------------------------------------------------------------------------- // LLM lifecycle // --------------------------------------------------------------------------- @@ -1354,8 +1427,22 @@ func RegisterToolExecutionIntercept(name string, priority int32, execFn ToolExec )) } -// DeregisterToolExecutionIntercept removes a previously registered tool -// execution intercept by name. +// RegisterToolExecutionFrameIntercept registers an annotation-aware callback +// in the existing tool execution registry and priority order. +func RegisterToolExecutionFrameIntercept(name string, priority int32, execFn ToolExecutionFrameInterceptFunc) error { + execID := registerClosure(execFn) + cName := C.CString(name) + defer C.free(unsafe.Pointer(cName)) + return checkStatus(C.nemo_relay_register_tool_execution_frame_intercept( + cName, C.int32_t(priority), + C.NemoRelayToolExecFrameInterceptCb(C.goToolExecFrameInterceptTrampoline), + execID, + C.NemoRelayFreeFn(C.goFreeTrampoline), + )) +} + +// DeregisterToolExecutionIntercept removes a previously registered raw or +// frame-aware tool execution intercept by name. func DeregisterToolExecutionIntercept(name string) error { cName := C.CString(name) defer C.free(unsafe.Pointer(cName)) @@ -2350,6 +2437,23 @@ func ScopeRegisterToolExecutionIntercept(scopeUUID, name string, priority int32, )) } +// ScopeRegisterToolExecutionFrameIntercept registers an annotation-aware +// scope-local callback in the existing tool execution registry and priority +// order. +func ScopeRegisterToolExecutionFrameIntercept(scopeUUID, name string, priority int32, execFn ToolExecutionFrameInterceptFunc) error { + execID := registerClosure(execFn) + cScopeUUID := C.CString(scopeUUID) + defer C.free(unsafe.Pointer(cScopeUUID)) + cName := C.CString(name) + defer C.free(unsafe.Pointer(cName)) + return checkStatus(C.nemo_relay_scope_register_tool_execution_frame_intercept( + cScopeUUID, cName, C.int32_t(priority), + C.NemoRelayToolExecFrameInterceptCb(C.goToolExecFrameInterceptTrampoline), + execID, + C.NemoRelayFreeFn(C.goFreeTrampoline), + )) +} + // ScopeDeregisterToolExecutionIntercept removes a scope-local tool execution // intercept by name. func ScopeDeregisterToolExecutionIntercept(scopeUUID, name string) error { diff --git a/go/nemo_relay/plugin.go b/go/nemo_relay/plugin.go index 8a7f1942e..03be8d37b 100644 --- a/go/nemo_relay/plugin.go +++ b/go/nemo_relay/plugin.go @@ -30,6 +30,8 @@ typedef char* (*NemoRelayLlmExecNextFn)(const char* native_json, void* next_ctx) typedef char* (*NemoRelayLlmExecInterceptCb)(void* user_data, const char* native_json, NemoRelayLlmExecNextFn next_fn, void* next_ctx); typedef char* (*NemoRelayToolExecNextFn)(const char* args_json, void* next_ctx); typedef char* (*NemoRelayToolExecInterceptCb)(void* user_data, const char* args_json, NemoRelayToolExecNextFn next_fn, void* next_ctx); +typedef char* (*NemoRelayToolExecFrameNextFn)(const char* args_json, void* next_ctx); +typedef char* (*NemoRelayToolExecFrameInterceptCb)(void* user_data, const char* args_json, NemoRelayToolExecFrameNextFn next_fn, void* next_ctx); extern int32_t nemo_relay_validate_plugin_config(const char* config_json, char** out_json); extern int32_t nemo_relay_initialize_plugins(const char* config_json, char** out_json); @@ -58,6 +60,7 @@ extern int32_t nemo_relay_plugin_context_register_tool_request_intercept(FfiPlug extern int32_t nemo_relay_plugin_context_register_llm_execution_intercept(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayLlmExecInterceptCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_plugin_context_register_llm_stream_execution_intercept(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayLlmExecInterceptCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_plugin_context_register_tool_execution_intercept(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayToolExecInterceptCb cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_plugin_context_register_tool_execution_frame_intercept(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayToolExecFrameInterceptCb cb, void* user_data, NemoRelayFreeFn free_fn); extern char* goPluginValidateTrampoline(void*, const char*); extern int32_t goPluginRegisterTrampoline(void*, const char*, FfiPluginContext*); @@ -72,6 +75,7 @@ extern char* goLlmConditionalTrampoline(void*, const void*); extern char* goLlmExecInterceptTrampoline(void*, const char*, NemoRelayLlmExecNextFn, void*); extern int32_t goLlmRequestInterceptTrampoline(void*, const char*, const void*, const char*, char**); extern char* goToolExecInterceptTrampoline(void*, const char*, NemoRelayToolExecNextFn, void*); +extern char* goToolExecFrameInterceptTrampoline(void*, const char*, NemoRelayToolExecFrameNextFn, void*); */ import "C" @@ -813,6 +817,25 @@ func (ctx *PluginContext) RegisterToolExecutionIntercept(name string, priority i )) } +// RegisterToolExecutionFrameIntercept registers an annotation-aware tool +// execution intercept for this component in the existing execution chain. +func (ctx *PluginContext) RegisterToolExecutionFrameIntercept(name string, priority int32, fn ToolExecutionFrameInterceptFunc) error { + if ctx == nil || ctx.ptr == nil { + return errors.New(errPluginContextClosed) + } + cName := C.CString(name) + defer C.free(unsafe.Pointer(cName)) + userData := registerClosure(fn) + return checkStatus(C.nemo_relay_plugin_context_register_tool_execution_frame_intercept( + ctx.ptr, + cName, + C.int32_t(priority), + (C.NemoRelayToolExecFrameInterceptCb)(C.goToolExecFrameInterceptTrampoline), + userData, + (C.NemoRelayFreeFn)(C.goFreeTrampoline), + )) +} + func pluginConfigCString(config PluginConfig) (*C.char, error) { payload, err := jsonMarshal(config) if err != nil { diff --git a/go/nemo_relay/tools_test.go b/go/nemo_relay/tools_test.go index e8c6b784b..048aa3137 100644 --- a/go/nemo_relay/tools_test.go +++ b/go/nemo_relay/tools_test.go @@ -144,6 +144,158 @@ func TestToolCallExecuteBasic(t *testing.T) { } } +func TestToolCallExecuteFramePreservesOpaqueAnnotationThroughMixedChain(t *testing.T) { + const ( + frameName = "go_frame_outer" + legacyName = "go_frame_legacy" + ) + _ = DeregisterToolExecutionIntercept(frameName) + _ = DeregisterToolExecutionIntercept(legacyName) + + err := RegisterToolExecutionFrameIntercept( + frameName, + 1, + func(args json.RawMessage, next func(json.RawMessage) (ToolExecutionFrame, error)) (ToolExecutionFrameOutcome, error) { + frame, err := next(args) + if err != nil { + return ToolExecutionFrameOutcome{}, err + } + var result map[string]interface{} + if err := json.Unmarshal(frame.Result, &result); err != nil { + return ToolExecutionFrameOutcome{}, err + } + var annotation map[string]interface{} + if err := json.Unmarshal(frame.Annotation, &annotation); err != nil { + return ToolExecutionFrameOutcome{}, err + } + result["frame_seen"] = true + annotation["observed_by"] = frameName + frame.Result, _ = json.Marshal(result) + frame.Annotation, _ = json.Marshal(annotation) + return ToolExecutionFrameOutcome{Frame: frame}, nil + }, + ) + if err != nil { + t.Fatalf("RegisterToolExecutionFrameIntercept failed: %v", err) + } + defer DeregisterToolExecutionIntercept(frameName) + + err = RegisterToolExecutionIntercept( + legacyName, + 2, + func(args json.RawMessage, next func(json.RawMessage) (json.RawMessage, error)) (ToolExecutionInterceptOutcome, error) { + result, err := next(args) + return ToolExecutionInterceptOutcome{Result: result}, err + }, + ) + if err != nil { + t.Fatalf(registerFailed, err) + } + defer DeregisterToolExecutionIntercept(legacyName) + + events := []Event{} + var mu sync.Mutex + const subscriberName = "go_frame_subscriber" + _ = DeregisterSubscriber(subscriberName) + if err := RegisterSubscriber(subscriberName, func(event Event) { + mu.Lock() + events = append(events, event) + mu.Unlock() + }); err != nil { + t.Fatalf(registerFailed, err) + } + defer DeregisterSubscriber(subscriberName) + + frame, err := ToolCallExecuteFrame( + "go_frame_tool", + json.RawMessage(`{"value":3}`), + func(args json.RawMessage) (ToolExecutionFrame, error) { + return ToolExecutionFrame{ + Result: json.RawMessage(`{"value":6}`), + Annotation: json.RawMessage(`{"producer":"go","status":"failed"}`), + }, nil + }, + ) + if err != nil { + t.Fatalf("ToolCallExecuteFrame failed: %v", err) + } + + var result map[string]interface{} + if err := json.Unmarshal(frame.Result, &result); err != nil { + t.Fatalf("unmarshal frame result: %v", err) + } + if result["frame_seen"] != true { + t.Fatalf("expected frame middleware mutation, got %v", result) + } + var annotation map[string]interface{} + if err := json.Unmarshal(frame.Annotation, &annotation); err != nil { + t.Fatalf("unmarshal frame annotation: %v", err) + } + if annotation["producer"] != "go" || annotation["observed_by"] != frameName { + t.Fatalf("unexpected annotation: %v", annotation) + } + + if err := FlushSubscribers(); err != nil { + t.Fatalf(toolFlushSubscribersFailed, err) + } + mu.Lock() + defer mu.Unlock() + for _, event := range events { + if event.Name() != "go_frame_tool" || event.Kind() != "scope" || event.ScopeCategory() != "end" { + continue + } + var profile map[string]interface{} + if err := json.Unmarshal(event.CategoryProfile(), &profile); err != nil { + t.Fatalf("unmarshal category profile: %v", err) + } + if profileAnnotation, ok := profile["tool_result_annotation"].(map[string]interface{}); ok && + profileAnnotation["producer"] == "go" { + return + } + } + t.Fatal("tool end event did not carry the opaque result annotation") +} + +func TestToolExecutionFrameContinuationExpiresWithCallback(t *testing.T) { + const interceptName = "go_frame_retained_next" + _ = DeregisterToolExecutionIntercept(interceptName) + + var retainedNext func(json.RawMessage) (ToolExecutionFrame, error) + err := RegisterToolExecutionFrameIntercept( + interceptName, + 1, + func(_ json.RawMessage, next func(json.RawMessage) (ToolExecutionFrame, error)) (ToolExecutionFrameOutcome, error) { + retainedNext = next + return ToolExecutionFrameOutcome{ + Frame: ToolExecutionFrame{Result: json.RawMessage(`{"short_circuit":true}`)}, + }, nil + }, + ) + if err != nil { + t.Fatalf("RegisterToolExecutionFrameIntercept failed: %v", err) + } + defer DeregisterToolExecutionIntercept(interceptName) + + _, err = ToolCallExecuteFrame( + "go_frame_retained_next_tool", + json.RawMessage(`{}`), + func(args json.RawMessage) (ToolExecutionFrame, error) { + return ToolExecutionFrame{Result: args}, nil + }, + ) + if err != nil { + t.Fatalf("ToolCallExecuteFrame failed: %v", err) + } + if retainedNext == nil { + t.Fatal("frame intercept did not retain the continuation") + } + + _, err = retainedNext(json.RawMessage(`{}`)) + if !errors.Is(err, ErrToolExecutionContinuationExpired) { + t.Fatalf("retained frame continuation returned %v", err) + } +} + func TestToolCallExecuteWithAttributes(t *testing.T) { fn := func(args json.RawMessage) (json.RawMessage, error) { return args, nil diff --git a/python/nemo_relay/__init__.py b/python/nemo_relay/__init__.py index fe3b9f43f..e13cd7f59 100644 --- a/python/nemo_relay/__init__.py +++ b/python/nemo_relay/__init__.py @@ -117,6 +117,8 @@ async def main(): ScopeStack, ScopeType, ToolAttributes, + ToolExecutionFrame, + ToolExecutionFrameOutcome, ToolExecutionInterceptOutcome, ToolHandle, ) @@ -202,6 +204,11 @@ class EventSanitizeFields(TypedDict): [str, Json, Callable[[Json], Awaitable[Json]]], ToolExecutionInterceptOutcome | Awaitable[ToolExecutionInterceptOutcome], ] +#: Annotation-aware execution intercept in the same tool middleware chain. +ToolExecutionFrameIntercept: TypeAlias = Callable[ + [str, Json, Callable[[Json], Awaitable[ToolExecutionFrame]]], + ToolExecutionFrameOutcome | Awaitable[ToolExecutionFrameOutcome], +] #: Request intercept callback that returns the canonical request, annotation, #: and pending-mark outcome passed to later intercepts and managed execution. LlmRequestIntercept: TypeAlias = Callable[ @@ -586,6 +593,8 @@ def worker() -> None: "ScopeHandle", "ToolHandle", "ToolExecutionInterceptOutcome", + "ToolExecutionFrame", + "ToolExecutionFrameOutcome", "LLMHandle", "LLMRequest", "LLMRequestInterceptOutcome", @@ -619,6 +628,7 @@ def worker() -> None: "LlmConditionalExecutionGuardrail", "ToolRequestIntercept", "ToolExecutionIntercept", + "ToolExecutionFrameIntercept", "LlmRequestIntercept", "LlmExecutionIntercept", "LlmStreamExecutionIntercept", diff --git a/python/nemo_relay/__init__.pyi b/python/nemo_relay/__init__.pyi index 31dee7ea2..5b9d2c721 100644 --- a/python/nemo_relay/__init__.pyi +++ b/python/nemo_relay/__init__.pyi @@ -120,6 +120,12 @@ from nemo_relay._native import ( from nemo_relay._native import ( ToolAttributes as ToolAttributes, ) +from nemo_relay._native import ( + ToolExecutionFrame as ToolExecutionFrame, +) +from nemo_relay._native import ( + ToolExecutionFrameOutcome as ToolExecutionFrameOutcome, +) from nemo_relay._native import ( ToolExecutionInterceptOutcome as ToolExecutionInterceptOutcome, ) @@ -257,12 +263,17 @@ Arguments: The tool name, current JSON arguments, and next callable. Return: - A canonical tool execution outcome, either directly or as an awaitable. + Relay's raw-result execution wrapper, either directly or as an awaitable. Exceptional flow: The callback may short-circuit by not invoking ``next``. Exceptions propagate through the managed tool call. """ +ToolExecutionFrameIntercept: TypeAlias = Callable[ + [str, Json, Callable[[Json], Awaitable[ToolExecutionFrame]]], + ToolExecutionFrameOutcome | Awaitable[ToolExecutionFrameOutcome], +] +"""Annotation-aware execution intercept in the existing tool chain.""" LlmRequestIntercept: TypeAlias = Callable[ [str, LLMRequest, AnnotatedLLMRequest | None], LLMRequestInterceptOutcome | Awaitable[LLMRequestInterceptOutcome], diff --git a/python/nemo_relay/_native.pyi b/python/nemo_relay/_native.pyi index 512d311c8..4232e458b 100644 --- a/python/nemo_relay/_native.pyi +++ b/python/nemo_relay/_native.pyi @@ -57,6 +57,10 @@ _ToolExecutionIntercept: TypeAlias = Callable[ [str, _Json, Callable[[_Json], Awaitable[_Json]]], "ToolExecutionInterceptOutcome | Awaitable[ToolExecutionInterceptOutcome]", ] +_ToolExecutionFrameIntercept: TypeAlias = Callable[ + [str, _Json, Callable[[_Json], Awaitable["ToolExecutionFrame"]]], + "ToolExecutionFrameOutcome | Awaitable[ToolExecutionFrameOutcome]", +] _LlmRequestIntercept: TypeAlias = Callable[ [str, "LLMRequest", "AnnotatedLLMRequest | None"], "LLMRequestInterceptOutcome | Awaitable[LLMRequestInterceptOutcome]", @@ -470,7 +474,7 @@ class LLMRequestInterceptOutcome: ... class ToolExecutionInterceptOutcome: - """Canonical result returned by a tool execution intercept. + """Relay-owned wrapper returned by a raw-result tool execution intercept. ``result`` is passed to the remaining middleware and application. ``pending_marks`` are Relay-owned lifecycle metadata emitted after the @@ -486,6 +490,30 @@ class ToolExecutionInterceptOutcome: @property def pending_marks(self) -> list[PendingMarkSpec]: ... +class ToolExecutionFrame: + """Raw tool result plus an optional opaque annotation.""" + def __init__(self, result: _Json, annotation: _Json | None = None) -> None: ... + @property + def result(self) -> _Json: ... + @result.setter + def result(self, value: _Json) -> None: ... + @property + def annotation(self) -> _Json | None: ... + @annotation.setter + def annotation(self, value: _Json | None) -> None: ... + +class ToolExecutionFrameOutcome: + """Result returned by an annotation-aware tool execution intercept.""" + def __init__( + self, + frame: ToolExecutionFrame, + pending_marks: list[PendingMarkSpec] = ..., + ) -> None: ... + @property + def frame(self) -> ToolExecutionFrame: ... + @property + def pending_marks(self) -> list[PendingMarkSpec]: ... + class AnnotatedLLMRequest: """Structured view of an LLM request produced by a codec. @@ -1474,6 +1502,17 @@ def tool_call_end( """ ... +def tool_call_end_frame( + handle: ToolHandle, + frame: ToolExecutionFrame, + *, + data: _Json | None = None, + metadata: _Json | None = None, + timestamp: datetime | None = None, +) -> None: + """End a manual tool span with an optional opaque result annotation.""" + ... + def tool_call_execute( name: str, args: _Json, @@ -1497,6 +1536,15 @@ def tool_call_execute( """ ... +def tool_call_execute_frame( + name: str, + args: _Json, + func: Callable[[_Json], ToolExecutionFrame | Awaitable[ToolExecutionFrame]], + **kwargs: object, +) -> Awaitable[ToolExecutionFrame]: + """Execute a tool through the existing chain with an opaque annotation.""" + ... + def llm_call( name: str, request: LLMRequest, @@ -1877,6 +1925,14 @@ def deregister_tool_execution_intercept(name: str) -> bool: """ ... +def register_tool_execution_frame_intercept( + name: str, + priority: int, + callable: _ToolExecutionFrameIntercept, +) -> None: + """Register annotation-aware middleware in the existing tool chain.""" + ... + def register_llm_request_intercept( name: str, priority: int, @@ -2127,6 +2183,15 @@ def scope_deregister_tool_execution_intercept(scope_uuid: str, name: str) -> boo """ ... +def scope_register_tool_execution_frame_intercept( + scope_uuid: str, + name: str, + priority: int, + callable: _ToolExecutionFrameIntercept, +) -> None: + """Register annotation-aware scope-local middleware in the existing chain.""" + ... + def scope_register_llm_sanitize_request_guardrail( scope_uuid: str, name: str, priority: int, guardrail: _LlmSanitizeRequestGuardrail ) -> None: diff --git a/python/nemo_relay/intercepts.py b/python/nemo_relay/intercepts.py index b0181c84a..e61e046e7 100644 --- a/python/nemo_relay/intercepts.py +++ b/python/nemo_relay/intercepts.py @@ -29,6 +29,7 @@ def add_header( LlmExecutionIntercept, LlmRequestIntercept, LlmStreamExecutionIntercept, + ToolExecutionFrameIntercept, ToolExecutionIntercept, ToolRequestIntercept, ) @@ -56,6 +57,9 @@ def add_header( from nemo_relay._native import ( register_llm_stream_execution_intercept as _native_register_llm_stream_execution, ) +from nemo_relay._native import ( + register_tool_execution_frame_intercept as _native_register_tool_execution_frame, +) from nemo_relay._native import ( register_tool_execution_intercept as _native_register_tool_execution, ) @@ -144,7 +148,7 @@ def register_tool_execution(name: str, priority: int, fn: ToolExecutionIntercept def deregister_tool_execution(name: str) -> bool: - """Remove a previously registered tool execution intercept. + """Remove a previously registered tool execution intercept of either form. Args: name: Intercept name previously passed to @@ -160,6 +164,11 @@ def deregister_tool_execution(name: str) -> bool: return _native_deregister_tool_execution(name) +def register_tool_execution_frame(name: str, priority: int, fn: ToolExecutionFrameIntercept) -> None: + """Register annotation-aware middleware in the existing execution chain.""" + return _native_register_tool_execution_frame(name, priority, fn) + + # --------------------------------------------------------------------------- # LLM intercepts # --------------------------------------------------------------------------- @@ -315,6 +324,7 @@ def deregister_llm_stream_execution(name: str) -> bool: __all__ = [ "ToolRequestIntercept", "ToolExecutionIntercept", + "ToolExecutionFrameIntercept", "LlmRequestIntercept", "LlmExecutionIntercept", "LlmStreamExecutionIntercept", @@ -322,6 +332,7 @@ def deregister_llm_stream_execution(name: str) -> bool: "deregister_tool_request", "register_tool_execution", "deregister_tool_execution", + "register_tool_execution_frame", "register_llm_request", "deregister_llm_request", "register_llm_execution", diff --git a/python/nemo_relay/plugin.py b/python/nemo_relay/plugin.py index 8cc927f36..75b675c44 100644 --- a/python/nemo_relay/plugin.py +++ b/python/nemo_relay/plugin.py @@ -25,6 +25,7 @@ LlmSanitizeResponseGuardrail, LlmStreamExecutionIntercept, ToolConditionalExecutionGuardrail, + ToolExecutionFrameIntercept, ToolExecutionIntercept, ToolRequestIntercept, ToolSanitizeGuardrail, @@ -169,6 +170,12 @@ def register_tool_execution_intercept(self, name: str, priority: int, callback: """Register a tool execution intercept for this component.""" ... + def register_tool_execution_frame_intercept( + self, name: str, priority: int, callback: ToolExecutionFrameIntercept + ) -> None: + """Register an annotation-aware tool execution intercept for this component.""" + ... + class Plugin(Protocol): """Custom plugin callback contract.""" diff --git a/python/nemo_relay/plugin.pyi b/python/nemo_relay/plugin.pyi index d5d387ff5..769f43161 100644 --- a/python/nemo_relay/plugin.pyi +++ b/python/nemo_relay/plugin.pyi @@ -16,6 +16,7 @@ from nemo_relay import ( LlmSanitizeResponseGuardrail, LlmStreamExecutionIntercept, ToolConditionalExecutionGuardrail, + ToolExecutionFrameIntercept, ToolExecutionIntercept, ToolRequestIntercept, ToolSanitizeGuardrail, @@ -74,6 +75,9 @@ class PluginContext(Protocol): self, name: str, priority: int, break_chain: bool, callback: ToolRequestIntercept ) -> None: ... def register_tool_execution_intercept(self, name: str, priority: int, callback: ToolExecutionIntercept) -> None: ... + def register_tool_execution_frame_intercept( + self, name: str, priority: int, callback: ToolExecutionFrameIntercept + ) -> None: ... class Plugin(Protocol): def validate(self, plugin_config: JsonObject) -> list[ConfigDiagnostic] | None: ... diff --git a/python/nemo_relay/scope_local.py b/python/nemo_relay/scope_local.py index 7b4da72a7..c92d338a5 100644 --- a/python/nemo_relay/scope_local.py +++ b/python/nemo_relay/scope_local.py @@ -97,6 +97,9 @@ def redact(tool_name, args): from nemo_relay._native import ( scope_register_tool_conditional_execution_guardrail as _register_tool_conditional_execution, ) +from nemo_relay._native import ( + scope_register_tool_execution_frame_intercept as _register_tool_execution_frame, +) from nemo_relay._native import ( scope_register_tool_execution_intercept as _register_tool_execution, ) @@ -346,8 +349,8 @@ def deregister_tool_execution(scope_handle, name): Args: scope_handle: Scope handle that owns the registration. - name: Intercept name previously passed to - ``register_tool_execution()``. + name: Intercept name previously passed to ``register_tool_execution()`` + or ``register_tool_execution_frame()``. Returns: bool: ``True`` if an intercept was removed, otherwise ``False``. @@ -359,6 +362,11 @@ def deregister_tool_execution(scope_handle, name): return _deregister_tool_execution(scope_handle.uuid, name) +def register_tool_execution_frame(scope_handle, name, priority, fn): + """Register annotation-aware scope-local middleware in the existing chain.""" + return _register_tool_execution_frame(scope_handle.uuid, name, priority, fn) + + # --------------------------------------------------------------------------- # LLM guardrails (scope-local) # --------------------------------------------------------------------------- @@ -690,6 +698,7 @@ def deregister_subscriber(scope_handle, name): "deregister_tool_request", "register_tool_execution", "deregister_tool_execution", + "register_tool_execution_frame", # LLM guardrails "register_llm_sanitize_request", "deregister_llm_sanitize_request", diff --git a/python/nemo_relay/tools.py b/python/nemo_relay/tools.py index 1f53127c2..690721cb3 100644 --- a/python/nemo_relay/tools.py +++ b/python/nemo_relay/tools.py @@ -29,9 +29,15 @@ async def search(args): from nemo_relay._native import ( tool_call_end as _native_tool_call_end, ) +from nemo_relay._native import ( + tool_call_end_frame as _native_tool_call_end_frame, +) from nemo_relay._native import ( tool_call_execute as _native_tool_call_execute, ) +from nemo_relay._native import ( + tool_call_execute_frame as _native_tool_call_execute_frame, +) from nemo_relay._native import ( tool_conditional_execution as _native_tool_conditional_execution, ) @@ -137,6 +143,17 @@ def call_end(handle, result, *, data=None, metadata=None, timestamp: datetime | return _native_tool_call_end(handle, result, data=data, metadata=metadata, timestamp=timestamp) +def call_end_frame(handle, frame, *, data=None, metadata=None, timestamp: datetime | None = None): + """Finish a manual tool span with an optional opaque result annotation. + + ``frame.result`` follows the existing response-sanitization path. Relay + carries ``frame.annotation`` on the lifecycle event without interpreting + its schema. + """ + ensure_scope_stack() + return _native_tool_call_end_frame(handle, frame, data=data, metadata=metadata, timestamp=timestamp) + + def execute(name, args, func, *, handle=None, attributes=None, data=None, metadata=None): """Run a tool through the managed middleware pipeline. @@ -198,6 +215,18 @@ async def local_tool(args): ) +def execute_frame(name, args, func, *, handle=None, attributes=None, data=None, metadata=None): + """Run the managed tool pipeline with an optional opaque annotation. + + ``func`` must return ``ToolExecutionFrame``. Frame-aware and legacy + execution intercepts share the same priority-ordered chain. + """ + ensure_scope_stack() + return _native_tool_call_execute_frame( + name, args, func, handle=handle, attributes=attributes, data=data, metadata=metadata + ) + + def request_intercepts(name, args): """Apply global tool request intercepts to ``args``. @@ -241,4 +270,12 @@ def conditional_execution(name, args): return _native_tool_conditional_execution(name, args) -__all__ = ["call", "call_end", "execute", "request_intercepts", "conditional_execution"] +__all__ = [ + "call", + "call_end", + "call_end_frame", + "execute", + "execute_frame", + "request_intercepts", + "conditional_execution", +] diff --git a/python/plugin/src/nemo_relay_plugin/__init__.py b/python/plugin/src/nemo_relay_plugin/__init__.py index 54466a5b5..f9763dbf8 100644 --- a/python/plugin/src/nemo_relay_plugin/__init__.py +++ b/python/plugin/src/nemo_relay_plugin/__init__.py @@ -36,7 +36,9 @@ LlmOptimizationTokens: Explicit token evidence by category. LlmOptimizationTokenImpact: Baseline, effective, and saved token evidence. LlmRequestInterceptOutcome: Canonical LLM request-intercept result. - ToolExecutionInterceptOutcome: Canonical tool execution-intercept result. + ToolExecutionInterceptOutcome: Relay wrapper for a raw-result tool execution intercept. + ToolExecutionFrame: Raw tool result with optional opaque annotation. + ToolExecutionFrameOutcome: Relay wrapper for a frame-aware tool execution intercept. DiagnosticLevel: Severity of a configuration diagnostic. ConfigDiagnostic: Structured configuration warning or error. ScopeType: Semantic category for a Relay execution scope. @@ -49,6 +51,7 @@ ToolConditionalCallback: Tool execution guardrail callback. ToolRequestCallback: Tool request intercept callback. ToolExecutionCallback: Tool execution intercept callback. + ToolExecutionFrameCallback: Frame-aware tool execution intercept callback. LlmSanitizeRequestCallback: LLM request sanitizer callback. LlmSanitizeResponseCallback: LLM response sanitizer callback. LlmConditionalCallback: LLM execution guardrail callback. @@ -61,6 +64,7 @@ PluginContext: Component-scoped callback registration context. PluginRuntime: Host runtime handle for event and scope operations. ToolNext: Continuation for a tool execution intercept. + ToolFrameNext: Continuation for a frame-aware tool execution intercept. LlmNext: Continuation for a unary LLM execution intercept. LlmStreamNext: Continuation for a streaming LLM execution intercept. @@ -103,7 +107,11 @@ SubscriberCallback, ToolConditionalCallback, ToolExecutionCallback, + ToolExecutionFrame, + ToolExecutionFrameCallback, + ToolExecutionFrameOutcome, ToolExecutionInterceptOutcome, + ToolFrameNext, ToolNext, ToolRequestCallback, ToolSanitizeCallback, @@ -149,7 +157,11 @@ "SubscriberCallback", "ToolConditionalCallback", "ToolExecutionCallback", + "ToolExecutionFrame", + "ToolExecutionFrameCallback", + "ToolExecutionFrameOutcome", "ToolExecutionInterceptOutcome", + "ToolFrameNext", "ToolNext", "ToolRequestCallback", "ToolSanitizeCallback", diff --git a/python/plugin/src/nemo_relay_plugin/_api.py b/python/plugin/src/nemo_relay_plugin/_api.py index bea81181a..ed6b2335b 100644 --- a/python/plugin/src/nemo_relay_plugin/_api.py +++ b/python/plugin/src/nemo_relay_plugin/_api.py @@ -208,6 +208,8 @@ def _llm_codec_capability(invocation: pb.LlmInvocation) -> str | None: ANNOTATED_LLM_REQUEST_SCHEMA = "nemo.relay.AnnotatedLlmRequest@2" LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA = "nemo.relay.LlmRequestInterceptOutcome@2" TOOL_EXECUTION_INTERCEPT_OUTCOME_SCHEMA = "nemo.relay.ToolExecutionInterceptOutcome@1" +TOOL_EXECUTION_FRAME_SCHEMA = "nemo.relay.ToolExecutionFrame@1" +TOOL_EXECUTION_FRAME_OUTCOME_SCHEMA = "nemo.relay.ToolExecutionFrameOutcome@1" PLUGIN_DIAGNOSTICS_SCHEMA = "nemo.relay.PluginDiagnostics@1" _OBJECT_SCHEMAS = frozenset( { @@ -216,6 +218,8 @@ def _llm_codec_capability(invocation: pb.LlmInvocation) -> str | None: ANNOTATED_LLM_REQUEST_SCHEMA, LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA, TOOL_EXECUTION_INTERCEPT_OUTCOME_SCHEMA, + TOOL_EXECUTION_FRAME_SCHEMA, + TOOL_EXECUTION_FRAME_OUTCOME_SCHEMA, } ) _UNREGISTERED = object() @@ -653,13 +657,13 @@ def to_json(self) -> dict[str, Json]: @dataclass(slots=True) class ToolExecutionInterceptOutcome: - """Canonical result returned by a Python worker tool execution intercept.""" + """Relay wrapper returned by a Python worker tool execution intercept.""" result: Json pending_marks: list[PendingMarkSpec] = field(default_factory=list) def to_json(self) -> dict[str, Json]: - """Convert this outcome to the canonical worker-envelope payload.""" + """Convert this outcome to the worker-envelope payload.""" marks = [] for mark in self.pending_marks: if not isinstance(mark, PendingMarkSpec): @@ -673,6 +677,53 @@ def to_json(self) -> dict[str, Json]: } +@dataclass(slots=True) +class ToolExecutionFrame: + """Raw tool result plus optional metadata that remains opaque to Relay.""" + + result: Json + annotation: Json | None = None + + @classmethod + def from_json(cls, value: Json) -> "ToolExecutionFrame": + """Decode the Relay frame envelope without interpreting its annotation.""" + if not isinstance(value, Mapping) or "result" not in value: + raise WorkerSdkError("tool execution frame must be an object containing result") + return cls( + result=value["result"], + annotation=value.get("annotation"), + ) + + def to_json(self) -> dict[str, Json]: + """Convert this frame to its Relay JSON-envelope payload.""" + value: dict[str, Json] = {"result": self.result} + if self.annotation is not None: + value["annotation"] = self.annotation + return value + + +@dataclass(slots=True) +class ToolExecutionFrameOutcome: + """Result returned by a Python worker frame-aware tool intercept.""" + + frame: ToolExecutionFrame + pending_marks: list[PendingMarkSpec] = field(default_factory=list) + + def to_json(self) -> dict[str, Json]: + """Convert this outcome to the canonical worker-envelope payload.""" + if not isinstance(self.frame, ToolExecutionFrame): + raise WorkerSdkError("tool execution frame outcome frame must be ToolExecutionFrame") + marks = [] + for mark in self.pending_marks: + if not isinstance(mark, PendingMarkSpec): + raise WorkerSdkError("tool execution frame outcome pending_marks must contain PendingMarkSpec values") + marks.append(mark.to_json()) + return { + "frame": self.frame.to_json(), + "pending_marks": marks, + } + + def _normalize_diagnostic(value: Mapping[str, Any]) -> dict[str, Any]: try: level = DiagnosticLevel(value.get("level")).value @@ -829,6 +880,10 @@ def register(self, ctx: PluginContext, config: Json) -> None | Awaitable[None]: [str, Json, "ToolNext"], ToolExecutionInterceptOutcome | Awaitable[ToolExecutionInterceptOutcome], ] +ToolExecutionFrameCallback: TypeAlias = Callable[ + [str, Json, "ToolFrameNext"], + ToolExecutionFrameOutcome | Awaitable[ToolExecutionFrameOutcome], +] LlmSanitizeRequestCallback: TypeAlias = Callable[ [LlmRequest, LlmSanitizeRequestContext], LlmRequest | None | Awaitable[LlmRequest | None] ] @@ -859,6 +914,7 @@ class _Handlers: tool_conditionals: dict[str, ToolConditionalCallback] tool_requests: dict[str, ToolRequestCallback] tool_executions: dict[str, ToolExecutionCallback] + tool_execution_frames: dict[str, ToolExecutionFrameCallback] llm_sanitize_requests: dict[str, LlmSanitizeRequestCallback] llm_sanitize_responses: dict[str, LlmSanitizeResponseCallback] llm_conditionals: dict[str, LlmConditionalCallback] @@ -879,6 +935,7 @@ def empty(cls) -> _Handlers: tool_conditionals={}, tool_requests={}, tool_executions={}, + tool_execution_frames={}, llm_sanitize_requests={}, llm_sanitize_responses={}, llm_conditionals={}, @@ -1131,6 +1188,21 @@ def register_tool_execution_intercept( self._push_registration(name, pb.TOOL_EXECUTION_INTERCEPT, priority, False) self._handlers.tool_executions[name] = callback + def register_tool_execution_frame_intercept( + self, + name: str, + callback: ToolExecutionFrameCallback, + *, + priority: int = 0, + ) -> None: + """Register middleware that carries an opaque annotation with the tool result. + + The callback participates in the same host registry and priority order + as legacy tool execution intercepts. + """ + self._push_registration(name, pb.TOOL_EXECUTION_FRAME_INTERCEPT, priority, False) + self._handlers.tool_execution_frames[name] = callback + def register_llm_sanitize_request_guardrail( self, name: str, @@ -1633,6 +1705,46 @@ async def call(self, value: Json) -> Json: return _json_result_to_value(response) +class ToolFrameNext: + """Continue the remaining tool chain and receive its result frame. + + The SDK creates this handle for one frame-aware execution-intercept + invocation. The Relay host invalidates it when that invocation finishes. + + Args: + runtime: Runtime used to call the Relay host. + continuation_id: Opaque host-issued continuation identifier. + """ + + def __init__(self, runtime: PluginRuntime, continuation_id: str) -> None: + self._runtime = runtime + self._continuation_id = continuation_id + + async def call(self, value: Json) -> ToolExecutionFrame: + """Call the remaining chain with replacement arguments. + + Args: + value: JSON arguments passed to the next intercept or real tool. + + Returns: + Raw downstream result plus its optional opaque annotation. + + Raises: + WorkerSdkError: The continuation is invalid, complete, cancelled, + or fails in the host. + TypeError: ``value`` is not JSON-serializable. + """ + response = await self._runtime._host_stub.ToolFrameNext( + pb.ToolFrameNextRequest( + activation_id=self._runtime._activation_id, + auth_token=self._runtime._auth_token, + continuation_id=self._continuation_id, + value=_json_envelope(JSON_SCHEMA, value), + ) + ) + return ToolExecutionFrame.from_json(_json_result_to_value(response, TOOL_EXECUTION_FRAME_SCHEMA)) + + class LlmNext: """Continue the remaining unary LLM execution chain. @@ -2071,6 +2183,27 @@ async def _invoke_result(self, request: Any) -> Any: ), ) ) + if request.surface == pb.TOOL_EXECUTION_FRAME_INTERCEPT: + result = await _maybe_await( + self._handler( + self._handlers.tool_execution_frames, + request.registration_name, + )( + request.tool.tool_name, + _decode_required_envelope(request.tool.value, "tool value"), + ToolFrameNext(self._runtime, request.continuation_id), + ) + ) + if not isinstance(result, ToolExecutionFrameOutcome): + raise WorkerSdkError("tool execution frame intercept must return ToolExecutionFrameOutcome") + return pb.InvokeResponse( + tool_execution_frame=pb.ToolExecutionFrameInterceptResult( + outcome=_json_envelope( + TOOL_EXECUTION_FRAME_OUTCOME_SCHEMA, + result.to_json(), + ), + ) + ) raise WorkerSdkError(f"unsupported registration surface {request.surface}") async def _invoke_llm_result(self, request: Any) -> Any: @@ -2200,6 +2333,7 @@ def _all_surfaces() -> list[int]: pb.TOOL_CONDITIONAL_EXECUTION_GUARDRAIL, pb.TOOL_REQUEST_INTERCEPT, pb.TOOL_EXECUTION_INTERCEPT, + pb.TOOL_EXECUTION_FRAME_INTERCEPT, pb.LLM_SANITIZE_REQUEST_GUARDRAIL, pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL, diff --git a/python/tests/plugin/test_public_api_docstrings.py b/python/tests/plugin/test_public_api_docstrings.py index 31abf7427..b7ceaf291 100644 --- a/python/tests/plugin/test_public_api_docstrings.py +++ b/python/tests/plugin/test_public_api_docstrings.py @@ -29,6 +29,7 @@ "ToolConditionalCallback", "ToolRequestCallback", "ToolExecutionCallback", + "ToolExecutionFrameCallback", "LlmSanitizeRequestCallback", "LlmSanitizeResponseCallback", "LlmConditionalCallback", diff --git a/python/tests/plugin/test_worker_sdk.py b/python/tests/plugin/test_worker_sdk.py index 7ab918d20..0c1439970 100644 --- a/python/tests/plugin/test_worker_sdk.py +++ b/python/tests/plugin/test_worker_sdk.py @@ -35,7 +35,10 @@ PluginContext, PluginRuntime, ScopeType, + ToolExecutionFrame, + ToolExecutionFrameOutcome, ToolExecutionInterceptOutcome, + ToolFrameNext, ToolNext, WorkerPlugin, WorkerSdkError, @@ -48,6 +51,8 @@ JSON_SCHEMA, LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA, LLM_REQUEST_SCHEMA, + TOOL_EXECUTION_FRAME_OUTCOME_SCHEMA, + TOOL_EXECUTION_FRAME_SCHEMA, TOOL_EXECUTION_INTERCEPT_OUTCOME_SCHEMA, WORKER_PROTOCOL, _announced_worker_endpoint, @@ -157,6 +162,34 @@ def test_optimization_contribution_drops_known_fields_from_extra(): } +def test_tool_execution_frame_dtos_preserve_opaque_annotation(): + annotation = { + "producer_status": "failed", + "representation": {"schema": "example.failure@1"}, + } + frame = ToolExecutionFrame.from_json({"result": {"raw": True}, "annotation": annotation}) + outcome = ToolExecutionFrameOutcome( + frame=frame, + pending_marks=[PendingMarkSpec("worker.frame")], + ) + + assert outcome.to_json() == { + "frame": { + "result": {"raw": True}, + "annotation": annotation, + }, + "pending_marks": [ + { + "name": "worker.frame", + "category": None, + "category_profile": None, + "data": None, + "metadata": None, + } + ], + } + + class GrpcAbort(Exception): def __init__(self, code: object, details: str) -> None: super().__init__(f"{code}: {details}") @@ -225,6 +258,21 @@ async def ToolNext(self, request: Any) -> Any: value = json.loads(request.value.json.decode("utf-8")) return pb.JsonResult(value=_json_envelope(JSON_SCHEMA, {"next_tool": value})) + async def ToolFrameNext(self, request: Any) -> Any: + self.requests.append(request) + if self.failures.get("ToolFrameNext") == "error": + return pb.JsonResult(error=_worker_error("ToolFrameNext failed")) + value = json.loads(request.value.json.decode("utf-8")) + return pb.JsonResult( + value=_json_envelope( + TOOL_EXECUTION_FRAME_SCHEMA, + { + "result": {"next_frame": value}, + "annotation": {"host_annotation": True}, + }, + ) + ) + async def LlmNext(self, request: Any) -> Any: self.requests.append(request) if self.failures.get("LlmNext") == "error": @@ -334,6 +382,21 @@ async def tool_execution(name: str, value: Json, next_call: ToolNext) -> ToolExe pending_marks=[PendingMarkSpec("worker.tool.execution")], ) + async def tool_execution_frame( + name: str, + value: Json, + next_call: ToolFrameNext, + ) -> ToolExecutionFrameOutcome: + frame = await next_call.call(_tag(value, f"execute_frame_{name}")) + frame.result = _tag(frame.result, "tool_execution_frame") + annotation = dict(frame.annotation or {}) + annotation["worker_annotation"] = True + frame.annotation = annotation + return ToolExecutionFrameOutcome( + frame=frame, + pending_marks=[PendingMarkSpec("worker.tool.execution.frame")], + ) + def llm_sanitize_request(request: Json, context: LlmSanitizeRequestContext) -> Json: del context return _tag_llm_request(request, "llm_sanitize_request") @@ -372,6 +435,11 @@ async def llm_stream_execution(name: str, request: Json, next_call: Any) -> Asyn ctx.register_tool_conditional_execution_guardrail("tool_conditional", tool_block, priority=3) ctx.register_tool_request_intercept("tool_request", tool_request, priority=4, break_chain=True) ctx.register_tool_execution_intercept("tool_execution", tool_execution, priority=5) + ctx.register_tool_execution_frame_intercept( + "tool_execution_frame", + tool_execution_frame, + priority=6, + ) ctx.register_llm_sanitize_request_guardrail("llm_sanitize_request", llm_sanitize_request, priority=6) ctx.register_llm_sanitize_response_guardrail("llm_sanitize_response", llm_sanitize_response, priority=7) ctx.register_llm_conditional_execution_guardrail("llm_conditional", llm_block, priority=8) @@ -457,6 +525,7 @@ async def test_health_handshake_validate_register_and_all_surfaces(service: _Wor ("tool_conditional", pb.TOOL_CONDITIONAL_EXECUTION_GUARDRAIL, 3, False), ("tool_request", pb.TOOL_REQUEST_INTERCEPT, 4, True), ("tool_execution", pb.TOOL_EXECUTION_INTERCEPT, 5, False), + ("tool_execution_frame", pb.TOOL_EXECUTION_FRAME_INTERCEPT, 6, False), ("llm_sanitize_request", pb.LLM_SANITIZE_REQUEST_GUARDRAIL, 6, False), ("llm_sanitize_response", pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, 7, False), ("llm_conditional", pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL, 8, False), @@ -1364,6 +1433,17 @@ async def test_unary_invoke_success_paths(service: _WorkerService, host_stub: Re assert tool_execution["result"]["tag"] == "tool_execution" assert tool_execution["result"]["next_tool"]["tag"] == "execute_lookup" assert tool_execution["pending_marks"][0]["name"] == "worker.tool.execution" + tool_execution_frame = await _invoke_tool_execution_frame_async( + service, + "tool_execution_frame", + ) + assert tool_execution_frame["frame"]["result"]["tag"] == "tool_execution_frame" + assert tool_execution_frame["frame"]["result"]["next_frame"]["tag"] == "execute_frame_lookup" + assert tool_execution_frame["frame"]["annotation"] == { + "host_annotation": True, + "worker_annotation": True, + } + assert tool_execution_frame["pending_marks"][0]["name"] == "worker.tool.execution.frame" llm_sanitize_request = await _invoke_json_async( service, @@ -1882,6 +1962,7 @@ async def test_runtime_host_calls_and_scope_context(host_stub: RecordingHostStub scope_id = await runtime.push_scope("scope", scope_type=ScopeType.TOOL, input={"in": True}) await runtime.pop_scope(scope_id, output={"out": True}) tool_next = await ToolNext(runtime, "tool-next").call({"value": 1}) + tool_frame_next = await ToolFrameNext(runtime, "tool-frame-next").call({"value": 2}) llm_next = await _llm_next(runtime, {"content": {"prompt": "hello"}}) stream_next = [chunk async for chunk in _llm_stream_next(runtime, {"content": {"prompt": "hello"}})] with runtime.clear_scope_stack(): @@ -1890,6 +1971,8 @@ async def test_runtime_host_calls_and_scope_context(host_stub: RecordingHostStub assert runtime.current_scope_stack_id() == stack_id assert runtime.current_scope_stack_id() is None assert tool_next["next_tool"]["value"] == 1 + assert tool_frame_next.result["next_frame"]["value"] == 2 + assert tool_frame_next.annotation == {"host_annotation": True} assert llm_next["next_llm"]["content"]["prompt"] == "hello" assert stream_next[0]["next_stream"]["content"]["prompt"] == "hello" @@ -2023,6 +2106,10 @@ async def test_runtime_host_call_error_paths(host_stub: RecordingHostStub): with pytest.raises(WorkerSdkError, match="ToolNext failed"): await ToolNext(runtime, "tool-next").call({"value": 1}) + host_stub.failures["ToolFrameNext"] = "error" + with pytest.raises(WorkerSdkError, match="ToolFrameNext failed"): + await ToolFrameNext(runtime, "tool-frame-next").call({"value": 1}) + host_stub.failures["LlmNext"] = "error" with pytest.raises(WorkerSdkError, match="LlmNext failed"): await _llm_next(runtime, {"content": {}}) @@ -2784,6 +2871,23 @@ async def _invoke_tool_execution_async( return _envelope_value(response.tool_execution.outcome) +async def _invoke_tool_execution_frame_async( + service: _WorkerService, + registration_name: str, +) -> Json: + response = await service.Invoke( + _tool_request( + registration_name, + pb.TOOL_EXECUTION_FRAME_INTERCEPT, + {"query": "relay"}, + ), + AbortContext(), + ) + assert response.WhichOneof("result") == "tool_execution_frame", response + assert response.tool_execution_frame.outcome.schema == TOOL_EXECUTION_FRAME_OUTCOME_SCHEMA + return _envelope_value(response.tool_execution_frame.outcome) + + def _envelope_value(envelope: Any) -> Json: return json.loads(envelope.json.decode("utf-8")) @@ -2820,6 +2924,7 @@ def _all_expected_surfaces() -> list[int]: pb.TOOL_CONDITIONAL_EXECUTION_GUARDRAIL, pb.TOOL_REQUEST_INTERCEPT, pb.TOOL_EXECUTION_INTERCEPT, + pb.TOOL_EXECUTION_FRAME_INTERCEPT, pb.LLM_SANITIZE_REQUEST_GUARDRAIL, pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL, diff --git a/python/tests/test_scope_local.py b/python/tests/test_scope_local.py index e0601dd35..389f13838 100644 --- a/python/tests/test_scope_local.py +++ b/python/tests/test_scope_local.py @@ -20,6 +20,8 @@ MarkEvent, ScopeEvent, ScopeType, + ToolExecutionFrame, + ToolExecutionFrameOutcome, ToolExecutionInterceptOutcome, guardrails, llm, @@ -528,6 +530,46 @@ def intercept_fn(name, args, next_fn): assert result["value"] == 12 assert result["intercepted"] is True + async def test_frame_intercept_uses_shared_scope_registry_and_deregistration(self): + async def frame_intercept(name, args, next_fn): + frame = await next_fn(args) + annotation = frame.annotation + assert isinstance(annotation, dict) + annotation["scope"] = name + frame.annotation = annotation + return ToolExecutionFrameOutcome(frame) + + with scope.scope("frame_exec_scope", ScopeType.Agent) as handle: + scope_local.register_tool_execution_frame( + handle, + "sl_frame_exec", + 1, + frame_intercept, + ) + frame = await tools.execute_frame( + "scope_frame_tool", + {}, + lambda _args: ToolExecutionFrame( + {"scope": True}, + {"producer": "python"}, + ), + ) + assert frame.annotation == { + "producer": "python", + "scope": "scope_frame_tool", + } + assert scope_local.deregister_tool_execution(handle, "sl_frame_exec") + + unwrapped = await tools.execute_frame( + "scope_frame_after_deregister", + {}, + lambda _args: ToolExecutionFrame( + {"scope": True}, + {"producer": "python"}, + ), + ) + assert unwrapped.annotation == {"producer": "python"} + # --------------------------------------------------------------------------- # Deregistration within scope diff --git a/python/tests/test_tools.py b/python/tests/test_tools.py index a65ea43aa..aca33f071 100644 --- a/python/tests/test_tools.py +++ b/python/tests/test_tools.py @@ -20,6 +20,8 @@ ScopeEvent, ScopeType, ToolAttributes, + ToolExecutionFrame, + ToolExecutionFrameOutcome, ToolExecutionInterceptOutcome, ToolHandle, create_scope_stack, @@ -99,6 +101,75 @@ def my_func(args): result = await tools.execute("double", {"x": 5}, my_func) assert result == {"result": 10} + async def test_execute_frame_round_trips_annotation_through_mixed_chain(self): + events = [] + + async def frame_middleware(name, args, next_fn): + assert name == "python_frame_tool" + frame = await next_fn(args) + assert frame.annotation == {"producer": "python", "status": "failed"} + result = frame.result + result["frame_seen"] = True + frame.result = result + annotation = frame.annotation + assert isinstance(annotation, dict) + annotation["observed_by"] = "frame_middleware" + frame.annotation = annotation + return ToolExecutionFrameOutcome(frame) + + async def legacy_passthrough(name, args, next_fn): + assert name == "python_frame_tool" + return ToolExecutionInterceptOutcome(await next_fn(args)) + + intercepts.register_tool_execution_frame("python_frame", 1, frame_middleware) + intercepts.register_tool_execution("python_legacy_passthrough", 2, legacy_passthrough) + subscribers.register("python_frame_subscriber", lambda event: events.append(event)) + try: + frame = await tools.execute_frame( + "python_frame_tool", + {"value": 3}, + lambda args: ToolExecutionFrame( + {"value": args["value"] * 2}, + {"producer": "python", "status": "failed"}, + ), + ) + assert frame.result == {"value": 6, "frame_seen": True} + assert frame.annotation == { + "producer": "python", + "status": "failed", + "observed_by": "frame_middleware", + } + + await subscribers.flush_async() + end = _tool_event(events, "python_frame_tool", "end") + assert end.category_profile is not None + assert end.category_profile["tool_result_annotation"] == frame.annotation + finally: + assert intercepts.deregister_tool_execution("python_frame") + assert intercepts.deregister_tool_execution("python_legacy_passthrough") + subscribers.deregister("python_frame_subscriber") + + async def test_legacy_mutation_drops_frame_annotation(self): + async def legacy_mutator(name, args, next_fn): + result = await next_fn(args) + result["mutated"] = True + return ToolExecutionInterceptOutcome(result) + + intercepts.register_tool_execution("python_frame_legacy_mutator", 1, legacy_mutator) + try: + frame = await tools.execute_frame( + "python_frame_mutated", + {}, + lambda _args: ToolExecutionFrame( + {"original": True}, + {"producer": "python"}, + ), + ) + assert frame.result == {"original": True, "mutated": True} + assert frame.annotation is None + finally: + intercepts.deregister_tool_execution("python_frame_legacy_mutator") + async def test_execute_rejects_cyclic_results_and_remains_usable(self): def cyclic_result(_args): result = {}