diff --git a/crates/cli/src/gateway/mod.rs b/crates/cli/src/gateway/mod.rs index a42444553..b4ebdbe9e 100644 --- a/crates/cli/src/gateway/mod.rs +++ b/crates/cli/src/gateway/mod.rs @@ -26,7 +26,7 @@ use nemo_relay::api::llm::{ llm_stream_call_execute, }; use nemo_relay::api::runtime::{ - LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, TASK_SCOPE_STACK, + LlmCodecIdentity, LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, TASK_SCOPE_STACK, }; use nemo_relay::codec::request::AnnotatedLlmRequest; use nemo_relay::codec::resolve::{ @@ -160,6 +160,10 @@ struct GatewayRequestCodec { } impl LlmCodec for GatewayRequestCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + self.inner.codec_identity() + } + fn decode(&self, request: &LlmRequest) -> nemo_relay::error::Result { self.inner.decode(request) } diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index 8484d18e3..29e4655a9 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -15,11 +15,13 @@ use crate::api::event::{ use crate::api::optimization::{ LlmOptimizationRecorder, finalize_optimization_summary, scope_llm_optimization_recorder, }; +#[cfg(test)] +use crate::api::runtime::LlmCodecIdentity; use crate::api::runtime::NemoRelayContextState; use crate::api::runtime::global_context; use crate::api::runtime::{ EventSubscriberFn, LlmCollectorFn, LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream, - LlmStreamExecutionNextFn, + LlmSanitizeRequestContext, LlmSanitizeResponseContext, LlmStreamExecutionNextFn, }; use crate::api::runtime::{ScopeStackHandle, current_scope_stack}; use crate::api::scope::event; @@ -391,7 +393,7 @@ fn emit_llm_start( handle: &LlmHandle, request: &LlmRequest, annotated_request: Option>, - request_codec: Option<&dyn LlmCodec>, + request_codec: Option>, ) -> Result<()> { ensure_runtime_owner()?; let subscribers = { @@ -412,7 +414,7 @@ fn emit_llm_start_with_subscribers( handle: &LlmHandle, request: &LlmRequest, annotated_request: Option>, - request_codec: Option<&dyn LlmCodec>, + request_codec: Option>, subscribers: &[EventSubscriberFn], ) -> Result<()> { ensure_runtime_owner()?; @@ -428,36 +430,43 @@ fn emit_llm_start_with_subscribers( .map_err(|error| FlowError::Internal(error.to_string()))?; state.llm_sanitize_request_entries(&scope_locals) }; - let mut sanitized_request = - NemoRelayContextState::llm_sanitize_request_snapshot_chain(request.clone(), &entries); - let mut annotated_request = match request_codec { - Some(codec) - if sanitized_request.headers != request.headers - || sanitized_request.content != request.content => - { - codec.decode(&sanitized_request).ok().map(Arc::new) + let mut sanitized_request = NemoRelayContextState::llm_sanitize_request_snapshot_chain( + request.clone(), + LlmSanitizeRequestContext::for_request_codec(request_codec.clone()), + &entries, + ); + let request_changed = sanitized_request + .as_ref() + .is_some_and(|sanitized_request| sanitized_request != request); + let mut annotated_request = match (sanitized_request.as_ref(), request_codec.as_deref()) { + (Some(sanitized_request), Some(codec)) if request_changed => { + codec.decode(sanitized_request).ok().map(Arc::new) } - _ => annotated_request, + (Some(_), _) if !request_changed => annotated_request, + (None, _) => None, + (Some(_), _) => None, }; let scope_stack = handle.captured_scope_stack(); let agent_is_fresh = { let mut scope_guard = scope_stack.write().expect("scope stack lock poisoned"); scope_guard.take_agent_freshness(handle.parent_uuid) }; - if !agent_is_fresh { + if !agent_is_fresh && let Some(sanitized_request) = sanitized_request.as_mut() { project_llm_request_to_current_user_turn( - &mut sanitized_request, + sanitized_request, &mut annotated_request, - request_codec, + request_codec.as_deref(), ); } - let input = serde_json::to_value(&sanitized_request).unwrap_or(Json::Null); + let input = sanitized_request + .as_ref() + .and_then(|sanitized_request| serde_json::to_value(sanitized_request).ok()); let event = { let context = global_context(); let state = context .read() .map_err(|error| FlowError::Internal(error.to_string()))?; - state.build_llm_start_event(handle, Some(input), annotated_request) + state.build_llm_start_event(handle, input, annotated_request) }; if let Some(event) = sanitize_event_with_scope_stack(event, scope_stack) { NemoRelayContextState::emit_event(&event, subscribers); @@ -626,9 +635,8 @@ struct LlmCallEndBehavior { /// # Parameters /// - `handle`: LLM handle to close. /// - `response`: Raw provider response associated with the end event. -/// - `data`: Optional application payload retained for compatibility. The -/// emitted end event data is the sanitized `response` unless it sanitizes to -/// JSON null, in which case this payload is used. +/// - `data`: Optional application payload retained for compatibility. When the +/// raw `response` is JSON null, this payload is sanitized in its place. /// - `metadata`: Optional JSON metadata recorded on the end event. /// - `annotated_response`: Optional normalized response annotation produced by /// a response codec. When omitted and `response_codec` is supplied, the @@ -693,20 +701,36 @@ fn llm_call_end_with_behavior( let entries = state.llm_sanitize_response_entries(&scope_locals); (entries, subscribers) }; - let sanitized_response = - NemoRelayContextState::llm_sanitize_response_snapshot_chain(response, &entries); - let data = if sanitized_response.is_null() { - data + let response_was_null_without_fallback = response.is_null() && data.is_none(); + let response = if response.is_null() { + data.unwrap_or(response) } else { - Some(sanitized_response) + response }; - let (mut annotated_response, decode_error) = resolve_llm_end_annotation( - annotated_response, - response_codec, - data.as_ref(), - &behavior, - &handle.name, + let sanitized_response = NemoRelayContextState::llm_sanitize_response_snapshot_chain( + response.clone(), + LlmSanitizeResponseContext::for_response_codec(response_codec.clone()), + &entries, ); + let response_changed = sanitized_response + .as_ref() + .is_some_and(|sanitized_response| sanitized_response != &response); + let data = match sanitized_response { + Some(response) if response_was_null_without_fallback && response.is_null() => None, + response => response, + }; + let annotation_omitted = data.as_ref().is_none_or(Json::is_null); + let (mut annotated_response, decode_error) = if annotation_omitted { + (None, None) + } else { + resolve_llm_end_annotation( + (!response_changed).then_some(annotated_response).flatten(), + response_codec, + data.as_ref(), + &behavior, + &handle.name, + ) + }; handle.optimization_recorder.close_for_finalization(None); emit_optimization_marks(handle, &subscribers); let pricing = crate::codec::response::active_pricing_resolver(); @@ -716,7 +740,8 @@ fn llm_call_end_with_behavior( handle.model_name.as_deref(), &pricing, ); - if annotated_response.is_none() + if !annotation_omitted + && annotated_response.is_none() && let Some(summary) = summary { annotated_response = Some(AnnotatedLlmResponse { @@ -753,6 +778,22 @@ fn llm_call_end_with_behavior( } } +#[cfg(test)] +fn sanitize_context_for_request_codec(codec: Option<&dyn LlmCodec>) -> LlmSanitizeRequestContext { + LlmSanitizeRequestContext::with_identity( + codec.map_or(LlmCodecIdentity::None, LlmCodec::codec_identity), + ) +} + +#[cfg(test)] +pub(crate) fn sanitize_context_for_response_codec( + codec: Option<&dyn LlmResponseCodec>, +) -> LlmSanitizeResponseContext { + LlmSanitizeResponseContext::with_identity( + codec.map_or(LlmCodecIdentity::None, LlmResponseCodec::codec_identity), + ) +} + fn resolve_llm_end_annotation( annotated_response: Option>, response_codec: Option>, @@ -780,39 +821,63 @@ fn resolve_llm_end_annotation( fn emit_llm_end_without_output( handle: &LlmHandle, metadata: Option, + response_codec: Option>, lifecycle_subscribers: Option<&[EventSubscriberFn]>, ) -> Result<()> { ensure_runtime_owner()?; - let subscribers = { + let (entries, subscribers) = { let scope_stack = handle.captured_scope_stack(); let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); + let scope_locals = scope_guard.collect_scope_local_registries(|registries| { + ®istries.llm_sanitize_response_guardrails + }); let scope_subscribers = scope_guard.collect_scope_local_subscribers(); - match lifecycle_subscribers { + let subscribers = match lifecycle_subscribers { Some(subscribers) => subscribers.to_vec(), None => snapshot_event_subscribers(scope_subscribers)?, - } + }; + let context = global_context(); + let state = context + .read() + .map_err(|error| FlowError::Internal(error.to_string()))?; + let entries = state.llm_sanitize_response_entries(&scope_locals); + (entries, subscribers) }; + let had_fallback_data = handle.data.is_some(); + let data = handle.data.clone().and_then(|data| { + NemoRelayContextState::llm_sanitize_response_snapshot_chain( + data, + LlmSanitizeResponseContext::for_response_codec(response_codec), + &entries, + ) + }); + let annotation_omitted = + (had_fallback_data && data.is_none()) || data.as_ref().is_some_and(Json::is_null); handle.optimization_recorder.close_for_finalization(None); emit_optimization_marks(handle, &subscribers); let pricing = crate::codec::response::active_pricing_resolver(); - let annotated_response = finalize_optimization_summary( - &handle.optimization_recorder, - None, - handle.model_name.as_deref(), - &pricing, - ) - .map(|summary| { - Arc::new(AnnotatedLlmResponse { - optimization_summary: Some(summary), - ..AnnotatedLlmResponse::default() + let annotated_response = (!annotation_omitted) + .then(|| { + finalize_optimization_summary( + &handle.optimization_recorder, + None, + handle.model_name.as_deref(), + &pricing, + ) }) - }); + .flatten() + .map(|summary| { + Arc::new(AnnotatedLlmResponse { + optimization_summary: Some(summary), + ..AnnotatedLlmResponse::default() + }) + }); let event = { let context = global_context(); let state = context .read() .map_err(|error| FlowError::Internal(error.to_string()))?; - state.end_llm_handle(handle, handle.data.clone(), metadata, annotated_response) + state.end_llm_handle(handle, data, metadata, annotated_response) }; if let Some(event) = sanitize_event_with_scope_stack(event, handle.captured_scope_stack()) { NemoRelayContextState::emit_event(&event, &subscribers); @@ -950,7 +1015,7 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { &handle, &intercepted_request, annotated_request.clone(), - request_codec.as_deref(), + request_codec.clone(), &lifecycle_subscribers, )?; emit_pending_request_marks(&handle, pending_marks, &lifecycle_subscribers)?; @@ -999,8 +1064,12 @@ pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result { Err(error) => { let end_metadata = metadata_with_otel_status(metadata, "ERROR", Some(error.to_string())); - let _ = - emit_llm_end_without_output(&handle, end_metadata, Some(&lifecycle_subscribers)); + let _ = emit_llm_end_without_output( + &handle, + end_metadata, + response_codec, + Some(&lifecycle_subscribers), + ); Err(error) } } @@ -1136,7 +1205,7 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu &handle, &intercepted_request, annotated_request, - request_codec.as_deref(), + request_codec.clone(), &lifecycle_subscribers, )?; emit_pending_request_marks(&handle, pending_marks, &lifecycle_subscribers)?; @@ -1180,8 +1249,12 @@ pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Resu Err(error) => { let end_metadata = metadata_with_otel_status(metadata, "ERROR", Some(error.to_string())); - let _ = - emit_llm_end_without_output(&handle, end_metadata, Some(&lifecycle_subscribers)); + let _ = emit_llm_end_without_output( + &handle, + end_metadata, + response_codec, + Some(&lifecycle_subscribers), + ); Err(error) } } diff --git a/crates/core/src/api/registry.rs b/crates/core/src/api/registry.rs index 72199095c..984ac43d8 100644 --- a/crates/core/src/api/registry.rs +++ b/crates/core/src/api/registry.rs @@ -111,7 +111,7 @@ macro_rules! global_guardrail_registry_api { state .$field .register( - Guardrail::new(name, priority, guardrail), + Guardrail::new(name, priority, guardrail.into()), ) .map_err(FlowError::AlreadyExists) } @@ -298,7 +298,7 @@ macro_rules! scope_guardrail_registry_api { registries .$field .register( - Guardrail::new(name, priority, guardrail), + Guardrail::new(name, priority, guardrail.into()), ) .map_err(FlowError::AlreadyExists) } @@ -475,6 +475,7 @@ global_guardrail_registry_api!( mark_sanitize_guardrails, EventSanitizeFn ); + global_guardrail_registry_api!( /// Register a global scope-start event sanitizer. register_scope_sanitize_start_guardrail, diff --git a/crates/core/src/api/runtime.rs b/crates/core/src/api/runtime.rs index 8b02d998b..0261657ba 100644 --- a/crates/core/src/api/runtime.rs +++ b/crates/core/src/api/runtime.rs @@ -10,10 +10,12 @@ pub mod state; pub mod subscriber_dispatcher; pub use callbacks::{ - EventSanitizeFn, EventSubscriberFn, LlmCollectorFn, LlmConditionalFn, LlmExecutionFn, - LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestFn, - LlmSanitizeResponseFn, LlmStreamExecutionFn, LlmStreamExecutionNextFn, LlmStreamInner, - ToolConditionalFn, ToolExecutionFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, + BuiltinLlmCodec, EventSanitizeFn, EventSubscriberFn, LlmCodecIdentity, LlmCollectorFn, + LlmConditionalFn, LlmExecutionFn, LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream, + LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn, + LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionFn, + LlmStreamExecutionNextFn, LlmStreamInner, ToolConditionalFn, ToolExecutionFn, + ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; pub use global::global_context; pub use scope_stack::{ diff --git a/crates/core/src/api/runtime/callbacks.rs b/crates/core/src/api/runtime/callbacks.rs index 21fe853f1..e40558da1 100644 --- a/crates/core/src/api/runtime/callbacks.rs +++ b/crates/core/src/api/runtime/callbacks.rs @@ -19,6 +19,7 @@ use crate::api::event::{Event, EventSanitizeFields}; use crate::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; use crate::api::tool::ToolExecutionInterceptOutcome; use crate::codec::request::AnnotatedLlmRequest; +use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::error::Result; use crate::json::Json; @@ -132,6 +133,165 @@ pub(crate) type ToolExecutionOutcomeNextFn = Arc< + Sync, >; +/// Relay's built-in LLM codec identities. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BuiltinLlmCodec { + /// OpenAI Chat Completions request and response payloads. + OpenAiChat, + /// OpenAI Responses request and response payloads. + OpenAiResponses, + /// Anthropic Messages request and response payloads. + AnthropicMessages, +} + +impl BuiltinLlmCodec { + /// Stable identifier used in configuration and language bindings. + #[must_use] + pub const fn id(self) -> &'static str { + match self { + Self::OpenAiChat => "openai_chat", + Self::OpenAiResponses => "openai_responses", + Self::AnthropicMessages => "anthropic_messages", + } + } +} + +/// Per-call LLM codec identity supplied to sanitize guardrails. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub enum LlmCodecIdentity { + /// No codec was active for this payload direction. + #[default] + None, + /// A Relay built-in codec was active. + BuiltIn(BuiltinLlmCodec), + /// A runtime-registered codec was active, identified by its stable ID. + Runtime(String), + /// A codec was active but does not expose a registered identity. + Opaque, +} + +/// Per-call codec context for LLM request sanitize guardrails. +/// +/// The context distinguishes no codec, Relay built-ins, runtime-registered +/// codecs, and active codecs with no stable identity. +#[derive(Clone, Default)] +pub struct LlmSanitizeRequestContext { + /// Identity of the codec active for this payload direction. + codec: LlmCodecIdentity, + request_codec: Option>, +} + +impl std::fmt::Debug for LlmSanitizeRequestContext { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("LlmSanitizeRequestContext") + .field("codec", &self.codec) + .finish_non_exhaustive() + } +} + +impl LlmSanitizeRequestContext { + /// Construct a context that carries only a codec identity. + /// + /// Identity-only contexts do not carry a codec handle, so + /// [`Self::resolve_codec`] returns `None` even when the identity describes + /// an active codec. + #[must_use] + pub fn with_identity(codec: LlmCodecIdentity) -> Self { + Self { + codec, + ..Self::default() + } + } + + /// Construct request-sanitizer context from the active request codec. + #[must_use] + pub fn for_request_codec(codec: Option>) -> Self { + let identity = codec + .as_deref() + .map_or(LlmCodecIdentity::None, LlmCodec::codec_identity); + Self { + codec: identity, + request_codec: codec, + } + } + + /// Return the identity of the codec active for this payload direction. + #[must_use] + pub fn codec(&self) -> &LlmCodecIdentity { + &self.codec + } + + /// Resolve the active request codec. + /// + /// Returns `None` for contexts constructed with [`Self::with_identity`]. + #[must_use] + pub fn resolve_codec(&self) -> Option> { + self.request_codec.clone() + } +} + +/// Per-call codec context for LLM response sanitize guardrails. +/// +/// The context distinguishes no codec, Relay built-ins, runtime-registered +/// codecs, and active codecs with no stable identity. +#[derive(Clone, Default)] +pub struct LlmSanitizeResponseContext { + /// Identity of the codec active for this payload direction. + codec: LlmCodecIdentity, + response_codec: Option>, +} + +impl std::fmt::Debug for LlmSanitizeResponseContext { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("LlmSanitizeResponseContext") + .field("codec", &self.codec) + .finish_non_exhaustive() + } +} + +impl LlmSanitizeResponseContext { + /// Construct a context that carries only a codec identity. + /// + /// Identity-only contexts do not carry a codec handle, so + /// [`Self::resolve_codec`] returns `None` even when the identity describes + /// an active codec. + #[must_use] + pub fn with_identity(codec: LlmCodecIdentity) -> Self { + Self { + codec, + ..Self::default() + } + } + + /// Construct response-sanitizer context from the active response codec. + #[must_use] + pub fn for_response_codec(codec: Option>) -> Self { + let identity = codec + .as_deref() + .map_or(LlmCodecIdentity::None, LlmResponseCodec::codec_identity); + Self { + codec: identity, + response_codec: codec, + } + } + + /// Return the identity of the codec active for this payload direction. + #[must_use] + pub fn codec(&self) -> &LlmCodecIdentity { + &self.codec + } + + /// Resolve the active response codec. + /// + /// Returns `None` for contexts constructed with [`Self::with_identity`]. + #[must_use] + pub fn resolve_codec(&self) -> Option> { + self.response_codec.clone() + } +} + /// Sanitize an LLM request before the runtime records it. /// /// LLM request sanitizers affect the serialized request payload emitted on @@ -140,10 +300,16 @@ pub(crate) type ToolExecutionOutcomeNextFn = Arc< /// /// # Parameters /// - First argument: LLM request payload to sanitize for observability. +/// - Second argument: Per-call request codec identity and capability. /// /// # Returns -/// Sanitized [`LlmRequest`] for the emitted event. -pub type LlmSanitizeRequestFn = Arc LlmRequest + Send + Sync>; +/// `Some` contains the sanitized request for the emitted event. `None` omits +/// both the raw request payload and its annotation from that event. +/// +/// The context is always supplied and distinguishes no codec, built-in codecs, +/// runtime-registered codecs, and opaque active codecs. +pub type LlmSanitizeRequestFn = + Arc Option + Send + Sync>; /// Sanitize an LLM response before the runtime records it. /// /// These callbacks rewrite the JSON response payload captured on LLM-end @@ -151,10 +317,16 @@ pub type LlmSanitizeRequestFn = Arc LlmRequest + Send + Sy /// /// # Parameters /// - First argument: JSON response payload to sanitize for observability. +/// - Second argument: Per-call response codec identity and capability. /// /// # Returns -/// Sanitized JSON response payload for the emitted event. -pub type LlmSanitizeResponseFn = Arc Json + Send + Sync>; +/// `Some` contains the sanitized response for the emitted event. `None` omits +/// both the raw response payload and its annotation from that event. +/// +/// The context is always supplied and distinguishes no codec, built-in codecs, +/// runtime-registered codecs, and opaque active codecs. +pub type LlmSanitizeResponseFn = + Arc Option + Send + Sync>; /// Decide whether an LLM call is allowed to continue. /// /// The callback receives the current [`LlmRequest`] and can allow execution, diff --git a/crates/core/src/api/runtime/state.rs b/crates/core/src/api/runtime/state.rs index 60a01aa8e..eb3dbb950 100644 --- a/crates/core/src/api/runtime/state.rs +++ b/crates/core/src/api/runtime/state.rs @@ -22,7 +22,8 @@ use crate::api::llm::{LlmHandle, LlmRequest}; use crate::api::registry::{ExecutionIntercept, Guardrail, Intercept}; use crate::api::runtime::callbacks::{ EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionFn, LlmExecutionNextFn, - LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionFn, + LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn, + LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionFn, LlmStreamExecutionNextFn, LlmStreamExecutionRegistryRefs, ToolConditionalFn, ToolExecutionFn, ToolExecutionNextFn, ToolExecutionOutcomeNextFn, ToolInterceptFn, ToolSanitizeFn, }; @@ -965,11 +966,12 @@ impl NemoRelayContextState { /// The sanitized [`LlmRequest`] after every provided guardrail has run. pub(crate) fn llm_sanitize_request_snapshot_chain( request: LlmRequest, + context: LlmSanitizeRequestContext, entries: &[Guardrail], - ) -> LlmRequest { - let mut value = request; + ) -> Option { + let mut value = Some(request); for entry in entries { - value = (entry.payload)(value); + value = value.and_then(|value| (entry.payload)(value, context.clone())); } value } @@ -1003,11 +1005,12 @@ impl NemoRelayContextState { /// The sanitized response payload after every provided guardrail has run. pub(crate) fn llm_sanitize_response_snapshot_chain( response: Json, + context: LlmSanitizeResponseContext, entries: &[Guardrail], - ) -> Json { - let mut value = response; + ) -> Option { + let mut value = Some(response); for entry in entries { - value = (entry.payload)(value); + value = value.and_then(|value| (entry.payload)(value, context.clone())); } value } diff --git a/crates/core/src/codec/anthropic.rs b/crates/core/src/codec/anthropic.rs index 9272aa062..88272d748 100644 --- a/crates/core/src/codec/anthropic.rs +++ b/crates/core/src/codec/anthropic.rs @@ -19,6 +19,7 @@ use serde::Deserialize; use crate::api::llm::LlmRequest; +use crate::api::runtime::{BuiltinLlmCodec, LlmCodecIdentity}; use crate::error::{FlowError, Result}; use crate::json::Json; @@ -922,6 +923,10 @@ fn anthropic_usage( // --------------------------------------------------------------------------- impl LlmResponseCodec for AnthropicMessagesCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) + } + fn decode_response(&self, response: &Json) -> Result { let raw: RawAnthropicResponse = serde_json::from_value(response.clone()) .map_err(|e| FlowError::Internal(format!("Anthropic Messages response decode: {e}")))?; @@ -968,6 +973,10 @@ impl LlmResponseCodec for AnthropicMessagesCodec { // --------------------------------------------------------------------------- impl LlmCodec for AnthropicMessagesCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) + } + fn decode(&self, request: &LlmRequest) -> Result { let obj = request .content diff --git a/crates/core/src/codec/openai_chat.rs b/crates/core/src/codec/openai_chat.rs index 4717426ac..dc3c46dca 100644 --- a/crates/core/src/codec/openai_chat.rs +++ b/crates/core/src/codec/openai_chat.rs @@ -9,6 +9,7 @@ use serde::Deserialize; use crate::api::llm::LlmRequest; +use crate::api::runtime::{BuiltinLlmCodec, LlmCodecIdentity}; use crate::error::{FlowError, Result}; use crate::json::Json; @@ -1116,6 +1117,10 @@ fn patch_chat_api_fields( // --------------------------------------------------------------------------- impl LlmResponseCodec for OpenAIChatCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) + } + fn decode_response(&self, response: &Json) -> Result { let raw: RawChatCompletion = serde_json::from_value(response.clone()) .map_err(|e| FlowError::Internal(format!("OpenAI Chat response decode: {e}")))?; @@ -1205,6 +1210,10 @@ impl LlmResponseCodec for OpenAIChatCodec { // --------------------------------------------------------------------------- impl LlmCodec for OpenAIChatCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) + } + fn decode(&self, request: &LlmRequest) -> Result { let obj = request .content diff --git a/crates/core/src/codec/openai_responses.rs b/crates/core/src/codec/openai_responses.rs index a36e5e1cd..279cf3d0f 100644 --- a/crates/core/src/codec/openai_responses.rs +++ b/crates/core/src/codec/openai_responses.rs @@ -18,6 +18,7 @@ use serde::Deserialize; use crate::api::llm::LlmRequest; +use crate::api::runtime::{BuiltinLlmCodec, LlmCodecIdentity}; use crate::error::{FlowError, Result}; use crate::json::Json; @@ -1194,6 +1195,10 @@ fn patch_responses_common_fields( // --------------------------------------------------------------------------- impl LlmResponseCodec for OpenAIResponsesCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiResponses) + } + fn decode_response(&self, response: &Json) -> Result { let raw: RawResponsesResponse = serde_json::from_value(response.clone()) .map_err(|e| FlowError::Internal(format!("OpenAI Responses response decode: {e}")))?; @@ -1275,6 +1280,10 @@ impl LlmResponseCodec for OpenAIResponsesCodec { // --------------------------------------------------------------------------- impl LlmCodec for OpenAIResponsesCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiResponses) + } + fn decode(&self, request: &LlmRequest) -> Result { let obj = request .content diff --git a/crates/core/src/codec/traits.rs b/crates/core/src/codec/traits.rs index 1c62647c7..379f3c263 100644 --- a/crates/core/src/codec/traits.rs +++ b/crates/core/src/codec/traits.rs @@ -4,6 +4,7 @@ //! LLM codec traits for bidirectional request translation. use crate::api::llm::LlmRequest; +use crate::api::runtime::LlmCodecIdentity; use crate::error::Result; use crate::json::Json; @@ -33,6 +34,15 @@ use super::response::AnnotatedLlmResponse; /// so the Rust core cannot know concrete types at compile time. /// Store as `Arc`. pub trait LlmCodec: Send + Sync { + /// Return this codec's identity for LLM sanitizer context. + /// + /// Custom codecs should keep the default [`LlmCodecIdentity::Opaque`] unless + /// they have a stable runtime registration ID. Callers must not infer a + /// provider surface from an opaque implementation or request shape. + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::Opaque + } + /// Parse opaque request content into structured form. fn decode(&self, request: &LlmRequest) -> Result; @@ -72,6 +82,14 @@ pub trait LlmCodec: Send + Sync { /// 1. Deserialize raw JSON into API-specific intermediate structs /// 2. Map intermediate structs into the normalized `AnnotatedLlmResponse` pub trait LlmResponseCodec: Send + Sync { + /// Return this codec's identity for LLM sanitizer context. + /// + /// Custom codecs should keep the default [`LlmCodecIdentity::Opaque`] unless + /// they have a stable runtime registration ID. + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::Opaque + } + /// Parse a raw JSON response into normalized structured form. /// /// Implementations should return `Err` only for genuinely unparseable input. diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index be1b658a9..f1217c746 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -3,9 +3,14 @@ //! Native dynamic plugin loader and host-side ABI adapter. +#[cfg(test)] +use std::cell::Cell; use std::cell::RefCell; +#[cfg(test)] +use std::collections::HashSet; use std::ffi::c_void; use std::future::Future; +use std::panic::{AssertUnwindSafe, catch_unwind}; use std::path::{Path, PathBuf}; use std::pin::Pin; use std::ptr; @@ -17,13 +22,16 @@ use libloading::{Library, Symbol}; use nemo_relay_plugin::{ NEMO_RELAY_NATIVE_ABI_VERSION, NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, - NemoRelayNativeJsonCb, NemoRelayNativeLlmConditionalCb, NemoRelayNativeLlmExecutionCb, - NemoRelayNativeLlmRequestCb, NemoRelayNativeLlmRequestInterceptCb, - NemoRelayNativeLlmStreamExecutionCb, NemoRelayNativeLlmStreamV1, NemoRelayNativePluginContext, - NemoRelayNativePluginEntry, NemoRelayNativePluginV1, NemoRelayNativeScopeHandle, - NemoRelayNativeScopeStack, NemoRelayNativeScopeStackBinding, NemoRelayNativeScopeType, - NemoRelayNativeString, NemoRelayNativeToolConditionalCb, NemoRelayNativeToolExecutionCb, - NemoRelayNativeToolJsonCb, NemoRelayNativeWithScopeStackCb, NemoRelayStatus, + NemoRelayNativeLlmCodecKind, NemoRelayNativeLlmConditionalCb, NemoRelayNativeLlmExecutionCb, + NemoRelayNativeLlmRequestCodec, NemoRelayNativeLlmRequestInterceptCb, + NemoRelayNativeLlmResponseCodec, NemoRelayNativeLlmSanitizeRequestCb, + NemoRelayNativeLlmSanitizeRequestContext, NemoRelayNativeLlmSanitizeResponseCb, + NemoRelayNativeLlmSanitizeResponseContext, NemoRelayNativeLlmStreamExecutionCb, + NemoRelayNativeLlmStreamV1, NemoRelayNativePluginContext, NemoRelayNativePluginEntry, + NemoRelayNativePluginV1, NemoRelayNativeScopeHandle, NemoRelayNativeScopeStack, + NemoRelayNativeScopeStackBinding, NemoRelayNativeScopeType, NemoRelayNativeString, + NemoRelayNativeToolConditionalCb, NemoRelayNativeToolExecutionCb, NemoRelayNativeToolJsonCb, + NemoRelayNativeWithScopeStackCb, NemoRelayStatus, }; use semver::{Version, VersionReq}; use serde_json::{Map, Value as Json}; @@ -34,10 +42,11 @@ use tokio_stream::{Stream, StreamExt}; use crate::api::event::{Event, EventSanitizeFields}; use crate::api::llm::{LlmRequest, LlmRequestInterceptOutcome}; use crate::api::runtime::{ - EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionFn, LlmExecutionNextFn, - LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, - LlmStreamExecutionFn, LlmStreamExecutionNextFn, ToolConditionalFn, ToolExecutionFn, - ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, + EventSanitizeFn, EventSubscriberFn, LlmCodecIdentity, LlmConditionalFn, LlmExecutionFn, + LlmExecutionNextFn, LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestContext, + LlmSanitizeRequestFn, LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionFn, + LlmStreamExecutionNextFn, ToolConditionalFn, ToolExecutionFn, ToolExecutionNextFn, + ToolInterceptFn, ToolSanitizeFn, }; use crate::api::runtime::{ ScopeStackHandle, ThreadScopeStackBinding, capture_thread_scope_stack, create_scope_stack, @@ -49,6 +58,8 @@ use crate::api::scope::{ }; use crate::api::scope::{event as emit_scope_mark, get_handle, pop_scope, push_scope}; use crate::api::tool::ToolExecutionInterceptOutcome; +use crate::codec::request::AnnotatedLlmRequest; +use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::error::{FlowError, Result as FlowResult}; use crate::plugin::{ ConfigDiagnostic, DiagnosticLevel, Plugin, PluginError, PluginRegistrationContext, @@ -491,6 +502,9 @@ struct NativeHostPluginContext { struct NativeHostString(Vec); +struct NativeHostLlmRequestCodec(Arc); +struct NativeHostLlmResponseCodec(Arc); + struct NativeHostScopeHandle(ScopeHandle); struct NativeHostScopeStack(ScopeStackHandle); @@ -499,6 +513,10 @@ struct NativeHostScopeStackBinding(ThreadScopeStackBinding); thread_local! { static NATIVE_LAST_ERROR: RefCell> = const { RefCell::new(None) }; + #[cfg(test)] + static NATIVE_STRING_LIVE_ALLOCATIONS: RefCell> = RefCell::new(HashSet::new()); + #[cfg(test)] + static NATIVE_STRING_FAIL_AFTER: Cell> = const { Cell::new(None) }; } fn set_native_last_error(message: impl Into) { @@ -513,6 +531,16 @@ fn native_last_error_message() -> Option { NATIVE_LAST_ERROR.with(|cell| cell.borrow().clone()) } +#[cfg(test)] +fn fail_native_string_allocation_after(successful_allocations: usize) { + NATIVE_STRING_FAIL_AFTER.with(|cell| cell.set(Some(successful_allocations))); +} + +#[cfg(test)] +fn native_string_live_allocations() -> usize { + NATIVE_STRING_LIVE_ALLOCATIONS.with(|allocations| allocations.borrow().len()) +} + unsafe extern "C" fn native_string_new( data: *const u8, len: usize, @@ -536,8 +564,29 @@ unsafe extern "C" fn native_string_new( set_native_last_error(format!("string data is not valid UTF-8: {err}")); return NemoRelayStatus::InvalidUtf8; } + #[cfg(test)] + let should_fail = NATIVE_STRING_FAIL_AFTER.with(|cell| match cell.get() { + Some(0) => { + cell.set(None); + true + } + Some(remaining) => { + cell.set(Some(remaining - 1)); + false + } + None => false, + }); + #[cfg(test)] + if should_fail { + set_native_last_error("injected native string allocation failure"); + return NemoRelayStatus::Internal; + } let handle = Box::new(NativeHostString(bytes.to_vec())); unsafe { *out = Box::into_raw(handle) as *mut NemoRelayNativeString }; + #[cfg(test)] + NATIVE_STRING_LIVE_ALLOCATIONS.with(|allocations| { + allocations.borrow_mut().insert(unsafe { *out } as usize); + }); NemoRelayStatus::Ok } @@ -560,6 +609,10 @@ unsafe extern "C" fn native_string_len(value: *const NemoRelayNativeString) -> u unsafe extern "C" fn native_string_free(value: *mut NemoRelayNativeString) { if !value.is_null() { drop(unsafe { Box::from_raw(value as *mut NativeHostString) }); + #[cfg(test)] + NATIVE_STRING_LIVE_ALLOCATIONS.with(|allocations| { + allocations.borrow_mut().remove(&(value as usize)); + }); } } @@ -574,6 +627,162 @@ unsafe extern "C" fn native_last_error_set(message: *const NemoRelayNativeString } } +unsafe extern "C" fn native_llm_request_codec_decode( + codec: *const NemoRelayNativeLlmRequestCodec, + request_json: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + clear_native_last_error(); + if out.is_null() { + set_native_last_error("request codec decode output pointer is null"); + return NemoRelayStatus::NullPointer; + } + unsafe { *out = ptr::null_mut() }; + if codec.is_null() { + set_native_last_error("request codec decode capability is null"); + return NemoRelayStatus::NullPointer; + } + if request_json.is_null() { + set_native_last_error("request codec decode request is null"); + return NemoRelayStatus::NullPointer; + } + let result = catch_unwind(AssertUnwindSafe(|| -> std::result::Result<_, String> { + let request: LlmRequest = serde_json::from_str( + &read_native_string(request_json).map_err(|error| error.to_string())?, + ) + .map_err(|error| format!("invalid request JSON: {error}"))?; + let codec = unsafe { &*(codec as *const NativeHostLlmRequestCodec) }; + let annotated = codec + .0 + .decode(&request) + .map_err(|error| error.to_string())?; + let annotated = serde_json::to_value(annotated).map_err(|error| error.to_string())?; + native_string_from_json(&annotated) + .ok_or_else(|| "failed to allocate decoded request".to_string()) + })); + match result { + Ok(Ok(value)) => { + unsafe { *out = value }; + NemoRelayStatus::Ok + } + Ok(Err(error)) => { + set_native_last_error(format!("request codec decode failed: {error}")); + NemoRelayStatus::Internal + } + Err(_) => { + set_native_last_error("request codec decode panicked"); + NemoRelayStatus::Internal + } + } +} + +unsafe extern "C" fn native_llm_request_codec_encode( + codec: *const NemoRelayNativeLlmRequestCodec, + annotated_json: *const NemoRelayNativeString, + original_json: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + clear_native_last_error(); + if out.is_null() { + set_native_last_error("request codec encode output pointer is null"); + return NemoRelayStatus::NullPointer; + } + unsafe { *out = ptr::null_mut() }; + if codec.is_null() { + set_native_last_error("request codec encode capability is null"); + return NemoRelayStatus::NullPointer; + } + if annotated_json.is_null() { + set_native_last_error("request codec encode annotated request is null"); + return NemoRelayStatus::NullPointer; + } + if original_json.is_null() { + set_native_last_error("request codec encode original request is null"); + return NemoRelayStatus::NullPointer; + } + let result = catch_unwind(AssertUnwindSafe(|| -> std::result::Result<_, String> { + let annotated: AnnotatedLlmRequest = serde_json::from_str( + &read_native_string(annotated_json).map_err(|error| error.to_string())?, + ) + .map_err(|error| format!("invalid annotated request JSON: {error}"))?; + let original: LlmRequest = serde_json::from_str( + &read_native_string(original_json).map_err(|error| error.to_string())?, + ) + .map_err(|error| format!("invalid original request JSON: {error}"))?; + let codec = unsafe { &*(codec as *const NativeHostLlmRequestCodec) }; + let request = codec + .0 + .encode(&annotated, &original) + .map_err(|error| error.to_string())?; + let request = serde_json::to_value(request).map_err(|error| error.to_string())?; + native_string_from_json(&request) + .ok_or_else(|| "failed to allocate encoded request".to_string()) + })); + match result { + Ok(Ok(value)) => { + unsafe { *out = value }; + NemoRelayStatus::Ok + } + Ok(Err(error)) => { + set_native_last_error(format!("request codec encode failed: {error}")); + NemoRelayStatus::Internal + } + Err(_) => { + set_native_last_error("request codec encode panicked"); + NemoRelayStatus::Internal + } + } +} + +unsafe extern "C" fn native_llm_response_codec_decode( + codec: *const NemoRelayNativeLlmResponseCodec, + response_json: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + clear_native_last_error(); + if out.is_null() { + set_native_last_error("response codec decode output pointer is null"); + return NemoRelayStatus::NullPointer; + } + unsafe { *out = ptr::null_mut() }; + if codec.is_null() { + set_native_last_error("response codec decode capability is null"); + return NemoRelayStatus::NullPointer; + } + if response_json.is_null() { + set_native_last_error("response codec decode response is null"); + return NemoRelayStatus::NullPointer; + } + let result = catch_unwind(AssertUnwindSafe(|| -> std::result::Result<_, String> { + let response: Json = serde_json::from_str( + &read_native_string(response_json).map_err(|error| error.to_string())?, + ) + .map_err(|error| format!("invalid response JSON: {error}"))?; + let codec = unsafe { &*(codec as *const NativeHostLlmResponseCodec) }; + let annotated = codec + .0 + .decode_response(&response) + .map_err(|error| error.to_string())?; + let annotated = serde_json::to_value(annotated).map_err(|error| error.to_string())?; + native_string_from_json(&annotated) + .ok_or_else(|| "failed to allocate decoded response".to_string()) + })); + match result { + Ok(Ok(value)) => { + unsafe { *out = value }; + NemoRelayStatus::Ok + } + Ok(Err(error)) => { + set_native_last_error(format!("response codec decode failed: {error}")); + NemoRelayStatus::Internal + } + Err(_) => { + set_native_last_error("response codec decode panicked"); + NemoRelayStatus::Internal + } + } +} + fn native_host_api() -> *const NemoRelayNativeHostApiV1 { static HOST_API: OnceLock = OnceLock::new(); static RELAY_VERSION: &[u8] = concat!(env!("CARGO_PKG_VERSION"), "\0").as_bytes(); @@ -587,6 +796,9 @@ fn native_host_api() -> *const NemoRelayNativeHostApiV1 { string_free: native_string_free, last_error_clear: native_last_error_clear, last_error_set: native_last_error_set, + llm_request_codec_decode: native_llm_request_codec_decode, + llm_request_codec_encode: native_llm_request_codec_encode, + llm_response_codec_decode: native_llm_response_codec_decode, plugin_context_register_subscriber: native_plugin_context_register_subscriber, plugin_context_register_tool_sanitize_request_guardrail: native_plugin_context_register_tool_sanitize_request_guardrail, @@ -681,6 +893,22 @@ fn take_json_from_native_string( result } +unsafe fn free_native_sanitizer_strings( + input: *mut NemoRelayNativeString, + codec_id: Option<*mut NemoRelayNativeString>, + output: *mut NemoRelayNativeString, +) { + unsafe { native_string_free(input) }; + if let Some(codec_id) = codec_id + && codec_id != input + { + unsafe { native_string_free(codec_id) }; + } + if !output.is_null() && output != input && Some(output) != codec_id { + unsafe { native_string_free(output) }; + } +} + fn optional_json_from_native_string( value: *const NemoRelayNativeString, field: &str, @@ -1261,7 +1489,7 @@ unsafe extern "C" fn native_plugin_context_register_llm_sanitize_request_guardra ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, priority: i32, - cb: NemoRelayNativeLlmRequestCb, + cb: NemoRelayNativeLlmSanitizeRequestCb, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus { @@ -1279,7 +1507,7 @@ unsafe extern "C" fn native_plugin_context_register_llm_sanitize_request_guardra match ctx.register_llm_sanitize_request_guardrail( &name, priority, - wrap_llm_request_fn(instance, cb, user_data, free_fn), + wrap_llm_sanitize_request_fn(instance, cb, user_data, free_fn), ) { Ok(()) => NemoRelayStatus::Ok, Err(err) => status_from_plugin_error(err), @@ -1290,7 +1518,7 @@ unsafe extern "C" fn native_plugin_context_register_llm_sanitize_response_guardr ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, priority: i32, - cb: NemoRelayNativeJsonCb, + cb: NemoRelayNativeLlmSanitizeResponseCb, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus { @@ -1308,7 +1536,7 @@ unsafe extern "C" fn native_plugin_context_register_llm_sanitize_response_guardr match ctx.register_llm_sanitize_response_guardrail( &name, priority, - wrap_json_fn(instance, cb, user_data, free_fn), + wrap_llm_sanitize_response_fn(instance, cb, user_data, free_fn), ) { Ok(()) => NemoRelayStatus::Ok, Err(err) => status_from_plugin_error(err), @@ -1725,74 +1953,161 @@ unsafe extern "C" fn native_tool_next( } } -fn wrap_llm_request_fn( +fn wrap_llm_sanitize_request_fn( instance: Arc, - cb: NemoRelayNativeLlmRequestCb, + cb: NemoRelayNativeLlmSanitizeRequestCb, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> LlmSanitizeRequestFn { let user_data = make_user_data(instance, user_data, free_fn); - Arc::new(move |request| { - call_llm_request_callback(cb, user_data.ptr, &request).unwrap_or_else(|_| LlmRequest { - headers: Map::new(), - content: Json::Null, - }) + Arc::new(move |request, context| { + call_llm_sanitize_request_callback(cb, user_data.ptr, &request, context) + .ok() + .flatten() }) } -fn call_llm_request_callback( - cb: NemoRelayNativeLlmRequestCb, +fn wrap_llm_sanitize_response_fn( + instance: Arc, + cb: NemoRelayNativeLlmSanitizeResponseCb, + user_data: *mut c_void, + free_fn: NemoRelayNativeFreeFn, +) -> LlmSanitizeResponseFn { + let user_data = make_user_data(instance, user_data, free_fn); + Arc::new(move |payload, context| { + call_llm_sanitize_response_callback(cb, user_data.ptr, &payload, context) + .ok() + .flatten() + }) +} + +fn call_llm_sanitize_request_callback( + cb: NemoRelayNativeLlmSanitizeRequestCb, user_data: *mut c_void, request: &LlmRequest, -) -> FlowResult { + context: LlmSanitizeRequestContext, +) -> FlowResult> { clear_native_last_error(); - let request_json = serde_json::to_value(request) - .map_err(|err| FlowError::Internal(format!("failed to serialize LLM request: {err}")))?; - let request_string = native_string_from_json(&request_json) - .ok_or_else(|| FlowError::Internal("failed to allocate native LLM request".into()))?; + let codec = context.resolve_codec().map(NativeHostLlmRequestCodec); + let (codec_kind, context_id) = native_llm_codec_identity(context.codec())?; + let request_json = match serde_json::to_value(request) { + Ok(request_json) => request_json, + Err(err) => { + if let Some(context_id) = context_id { + unsafe { native_string_free(context_id) }; + } + return Err(FlowError::Internal(format!( + "failed to serialize LLM request: {err}" + ))); + } + }; + let request_string = match native_string_from_json(&request_json) { + Some(request_string) => request_string, + None => { + if let Some(context_id) = context_id { + unsafe { native_string_free(context_id) }; + } + return Err(FlowError::Internal( + "failed to allocate native LLM request".into(), + )); + } + }; + let context = NemoRelayNativeLlmSanitizeRequestContext { + codec_kind, + codec_id: context_id.map_or(ptr::null(), |value| value.cast_const()), + codec: codec + .as_ref() + .map_or(ptr::null(), |value| std::ptr::from_ref(value).cast()), + }; let mut out = ptr::null_mut(); - let status = unsafe { cb(user_data, request_string, &mut out) }; - unsafe { native_string_free(request_string) }; + let status = unsafe { cb(user_data, request_string, context, &mut out) }; if status != NemoRelayStatus::Ok { - if !out.is_null() { - unsafe { native_string_free(out) }; - } + unsafe { free_native_sanitizer_strings(request_string, context_id, out) }; return Err(flow_error_from_status( status, - "native LLM request callback failed", + "native LLM sanitize-request callback failed", )); } + if out.is_null() { + unsafe { free_native_sanitizer_strings(request_string, context_id, out) }; + return Ok(None); + } let result_json = - take_json_from_native_string(out, "native LLM request callback returned null")?; + json_from_native_string(out, "native LLM sanitize-request returned invalid JSON"); + unsafe { free_native_sanitizer_strings(request_string, context_id, out) }; + let result_json = result_json?; serde_json::from_value(result_json) + .map(Some) .map_err(|err| FlowError::Internal(format!("invalid LLM request JSON: {err}"))) } -fn wrap_json_fn( - instance: Arc, - cb: NemoRelayNativeJsonCb, +fn call_llm_sanitize_response_callback( + cb: NemoRelayNativeLlmSanitizeResponseCb, user_data: *mut c_void, - free_fn: NemoRelayNativeFreeFn, -) -> LlmSanitizeResponseFn { - let user_data = make_user_data(instance, user_data, free_fn); - Arc::new(move |payload| { - clear_native_last_error(); - let payload_string = native_string_from_json(&payload); - let Some(payload_string) = payload_string else { - return Json::Null; - }; - let mut out = ptr::null_mut(); - let status = unsafe { cb(user_data.ptr, payload_string, &mut out) }; - unsafe { native_string_free(payload_string) }; - if status != NemoRelayStatus::Ok { - if !out.is_null() { - unsafe { native_string_free(out) }; + payload: &Json, + context: LlmSanitizeResponseContext, +) -> FlowResult> { + clear_native_last_error(); + let codec = context.resolve_codec().map(NativeHostLlmResponseCodec); + let (codec_kind, context_id) = native_llm_codec_identity(context.codec())?; + let payload_string = match native_string_from_json(payload) { + Some(payload_string) => payload_string, + None => { + if let Some(context_id) = context_id { + unsafe { native_string_free(context_id) }; } - return Json::Null; + return Err(FlowError::Internal( + "failed to allocate native LLM response".into(), + )); } - take_json_from_native_string(out, "native JSON callback returned null") - .unwrap_or(Json::Null) - }) + }; + let context = NemoRelayNativeLlmSanitizeResponseContext { + codec_kind, + codec_id: context_id.map_or(ptr::null(), |value| value.cast_const()), + codec: codec + .as_ref() + .map_or(ptr::null(), |value| std::ptr::from_ref(value).cast()), + }; + let mut out = ptr::null_mut(); + let status = unsafe { cb(user_data, payload_string, context, &mut out) }; + if status != NemoRelayStatus::Ok { + unsafe { free_native_sanitizer_strings(payload_string, context_id, out) }; + return Err(flow_error_from_status( + status, + "native LLM sanitize-response callback failed", + )); + } + if out.is_null() { + unsafe { free_native_sanitizer_strings(payload_string, context_id, out) }; + return Ok(None); + } + let result = json_from_native_string(out, "native LLM sanitize-response returned invalid JSON"); + unsafe { free_native_sanitizer_strings(payload_string, context_id, out) }; + result.map(Some) +} + +fn native_llm_codec_identity( + context: &LlmCodecIdentity, +) -> FlowResult<( + NemoRelayNativeLlmCodecKind, + Option<*mut NemoRelayNativeString>, +)> { + let (codec_kind, codec_id) = match context { + LlmCodecIdentity::None => (NemoRelayNativeLlmCodecKind::None, None), + LlmCodecIdentity::BuiltIn(codec) => { + (NemoRelayNativeLlmCodecKind::BuiltIn, Some(codec.id())) + } + LlmCodecIdentity::Runtime(id) => (NemoRelayNativeLlmCodecKind::Runtime, Some(id.as_str())), + LlmCodecIdentity::Opaque => (NemoRelayNativeLlmCodecKind::Opaque, None), + }; + let codec_id = + match codec_id { + Some(codec_id) => Some(native_string_from_str(codec_id).ok_or_else(|| { + FlowError::Internal("failed to allocate native LLM codec ID".into()) + })?), + None => None, + }; + Ok((codec_kind, codec_id)) } fn wrap_llm_conditional_fn( diff --git a/crates/core/src/plugin/dynamic/worker.rs b/crates/core/src/plugin/dynamic/worker.rs index ce25e4935..79e2e70ce 100644 --- a/crates/core/src/plugin/dynamic/worker.rs +++ b/crates/core/src/plugin/dynamic/worker.rs @@ -19,10 +19,14 @@ use nemo_relay_worker_proto::v1::relay_host_runtime_server::{ use nemo_relay_worker_proto::v1::{ CancelInvocationRequest, CreateScopeStackRequest, CreateScopeStackResponse, DropScopeStackRequest, EmitMarkRequest, GuardrailResult, HandshakeRequest, HealthRequest, - HostAck, InvokeRequest, InvokeResponse, JsonEnvelope, JsonResult, LlmInvocation, - LlmNextRequest, LlmStreamNextRequest, PopScopeRequest, PushScopeRequest, PushScopeResponse, - RegisterRequest, RegisterResponse, Registration, RegistrationSurface, ScopeContext, - ShutdownRequest, StreamChunk, ToolInvocation, ToolNextRequest, ValidateRequest, WorkerError, + HostAck, InvokeRequest, InvokeResponse, JsonEnvelope, JsonResult, LlmCodecDecodeRequest, + LlmCodecDecodeResponse, LlmCodecEncodeRequest, LlmCodecIdentity as ProtoLlmCodecIdentity, + LlmCodecKind, LlmInvocation, LlmNextRequest, + LlmSanitizeRequestContext as ProtoLlmSanitizeRequestContext, + LlmSanitizeResponseContext as ProtoLlmSanitizeResponseContext, LlmStreamNextRequest, + PopScopeRequest, PushScopeRequest, PushScopeResponse, RegisterRequest, RegisterResponse, + Registration, RegistrationSurface, ScopeContext, ShutdownRequest, StreamChunk, ToolInvocation, + ToolNextRequest, ValidateRequest, WorkerError, }; use nemo_relay_worker_proto::{WORKER_PROTOCOL_GRPC_V1, decode_json_envelope, json_envelope}; use semver::{Version, VersionReq}; @@ -55,8 +59,9 @@ use tower::service_fn; use crate::api::event::{Event, EventSanitizeFields}; use crate::api::llm::{LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA, LlmRequest}; use crate::api::runtime::{ - LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, ToolExecutionNextFn, - current_scope_stack, with_scope_stack, + LlmCodecIdentity, LlmExecutionNextFn, LlmJsonStream, LlmSanitizeRequestContext, + LlmSanitizeResponseContext, LlmStreamExecutionNextFn, ToolExecutionNextFn, current_scope_stack, + with_scope_stack, }; use crate::api::scope::{ EmitMarkEventParams, PopScopeParams, PushScopeParams, ScopeAttributes, ScopeHandle, ScopeType, @@ -64,6 +69,7 @@ use crate::api::scope::{ }; use crate::api::tool::ToolExecutionInterceptOutcome; use crate::codec::request::{ANNOTATED_LLM_REQUEST_SCHEMA, AnnotatedLlmRequest}; +use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::error::{FlowError, Result as FlowResult}; use crate::plugin::{ ConfigDiagnostic, DiagnosticLevel, Plugin, PluginError, PluginRegistrationContext, @@ -1048,271 +1054,278 @@ impl WorkerPluginInstance { self.plugin_kind, registration.surface )) })?; - let name = registration.local_name.clone(); - let priority = registration.priority; - let break_chain = registration.break_chain; match surface { RegistrationSurface::Subscriber => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_subscriber( - &name, - Arc::new(move |event| { - if instance.invoke_subscriber(&callback_name, event).is_err() { - instance.log_callback_fallback( - &callback_name, - RegistrationSurface::Subscriber, - ); - } - }), - )?; + self.install_subscriber_registration(ctx, ®istration.local_name)? } RegistrationSurface::MarkSanitizeGuardrail | RegistrationSurface::ScopeSanitizeStartGuardrail - | RegistrationSurface::ScopeSanitizeEndGuardrail => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - let callback = Arc::new(move |event: &Event, _fields: EventSanitizeFields| { - instance - .invoke_event_sanitize(&callback_name, surface, event) - .unwrap_or_else(|_| { - instance.log_callback_fallback(&callback_name, surface); - EventSanitizeFields::default() - }) - }); - match surface { - RegistrationSurface::MarkSanitizeGuardrail => { - ctx.register_mark_sanitize_guardrail(&name, priority, callback)? - } - RegistrationSurface::ScopeSanitizeStartGuardrail => { - ctx.register_scope_sanitize_start_guardrail(&name, priority, callback)? - } - RegistrationSurface::ScopeSanitizeEndGuardrail => { - ctx.register_scope_sanitize_end_guardrail(&name, priority, callback)? - } - _ => unreachable!(), - } + | RegistrationSurface::ScopeSanitizeEndGuardrail => self + .install_event_sanitize_registration( + ctx, + ®istration.local_name, + registration.priority, + surface, + )?, + RegistrationSurface::ToolSanitizeRequestGuardrail + | RegistrationSurface::ToolSanitizeResponseGuardrail + | RegistrationSurface::ToolConditionalExecutionGuardrail + | RegistrationSurface::ToolRequestIntercept + | RegistrationSurface::ToolExecutionIntercept => { + self.install_tool_registration(ctx, registration, surface)? } - RegistrationSurface::ToolSanitizeRequestGuardrail => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_tool_sanitize_request_guardrail( - &name, - priority, - Arc::new(move |tool_name, value| { - instance - .invoke_tool_json( - &callback_name, - RegistrationSurface::ToolSanitizeRequestGuardrail, - tool_name, - value.clone(), - None, - ) - .unwrap_or_else(|_| { - instance.log_callback_fallback( - &callback_name, - RegistrationSurface::ToolSanitizeRequestGuardrail, - ); - value - }) - }), - )?; + RegistrationSurface::LlmSanitizeRequestGuardrail + | RegistrationSurface::LlmSanitizeResponseGuardrail + | RegistrationSurface::LlmConditionalExecutionGuardrail + | RegistrationSurface::LlmRequestIntercept + | RegistrationSurface::LlmExecutionIntercept + | RegistrationSurface::LlmStreamExecutionIntercept => { + self.install_llm_registration(ctx, registration, surface)? } - RegistrationSurface::ToolSanitizeResponseGuardrail => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_tool_sanitize_response_guardrail( - &name, - priority, - Arc::new(move |tool_name, value| { - instance - .invoke_tool_json( - &callback_name, - RegistrationSurface::ToolSanitizeResponseGuardrail, - tool_name, - value.clone(), - None, - ) - .unwrap_or_else(|_| { - instance.log_callback_fallback( - &callback_name, - RegistrationSurface::ToolSanitizeResponseGuardrail, - ); - value - }) - }), - )?; + RegistrationSurface::Unspecified => { + return Err(PluginError::RegistrationFailed(format!( + "worker plugin '{}' returned unspecified registration surface", + self.plugin_kind + ))); } - RegistrationSurface::ToolConditionalExecutionGuardrail => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_tool_conditional_execution_guardrail( - &name, - priority, - Arc::new(move |tool_name, value| { - instance.invoke_tool_guardrail(&callback_name, tool_name, value.clone()) - }), - )?; + } + } + Ok(()) + } + + fn install_subscriber_registration( + &self, + ctx: &mut PluginRegistrationContext, + name: &str, + ) -> crate::plugin::Result<()> { + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.to_owned(); + ctx.register_subscriber( + name, + Arc::new(move |event| { + if instance.invoke_subscriber(&callback_name, event).is_err() { + instance.log_callback_fallback(&callback_name, RegistrationSurface::Subscriber); } - RegistrationSurface::ToolRequestIntercept => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_tool_request_intercept( - &name, - priority, - break_chain, - Arc::new(move |tool_name, value| { - instance.invoke_tool_json( + }), + ) + } + + fn install_event_sanitize_registration( + &self, + ctx: &mut PluginRegistrationContext, + name: &str, + priority: i32, + surface: RegistrationSurface, + ) -> crate::plugin::Result<()> { + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.to_owned(); + let callback = Arc::new(move |event: &Event, _fields: EventSanitizeFields| { + instance + .invoke_event_sanitize(&callback_name, surface, event) + .unwrap_or_else(|_| { + instance.log_callback_fallback(&callback_name, surface); + EventSanitizeFields::default() + }) + }); + match surface { + RegistrationSurface::MarkSanitizeGuardrail => { + ctx.register_mark_sanitize_guardrail(name, priority, callback) + } + RegistrationSurface::ScopeSanitizeStartGuardrail => { + ctx.register_scope_sanitize_start_guardrail(name, priority, callback) + } + RegistrationSurface::ScopeSanitizeEndGuardrail => { + ctx.register_scope_sanitize_end_guardrail(name, priority, callback) + } + _ => unreachable!("event sanitizer surface was pre-filtered"), + } + } + + fn install_tool_registration( + &self, + ctx: &mut PluginRegistrationContext, + registration: &Registration, + surface: RegistrationSurface, + ) -> crate::plugin::Result<()> { + let name = registration.local_name.as_str(); + let priority = registration.priority; + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.to_owned(); + match surface { + RegistrationSurface::ToolSanitizeRequestGuardrail => ctx + .register_tool_sanitize_request_guardrail( + name, + priority, + Arc::new(move |tool_name, value| { + instance + .invoke_tool_json( &callback_name, - RegistrationSurface::ToolRequestIntercept, + surface, tool_name, - value, + value.clone(), None, ) - }), - )?; - } - RegistrationSurface::ToolExecutionIntercept => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_tool_execution_intercept( - &name, - priority, - Arc::new(move |tool_name, value, next| { - let instance = instance.clone(); - let name = callback_name.clone(); - let tool_name = tool_name.to_string(); - Box::pin(async move { - instance - .invoke_tool_execution(&name, &tool_name, value, next) - .await + .unwrap_or_else(|_| { + instance.log_callback_fallback(&callback_name, surface); + value }) - }), - )?; - } - RegistrationSurface::LlmSanitizeRequestGuardrail => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_llm_sanitize_request_guardrail( - &name, - priority, - Arc::new(move |request| { - instance - .invoke_llm_request_json( - &callback_name, - RegistrationSurface::LlmSanitizeRequestGuardrail, - "", - request.clone(), - None, - None, - ) - .unwrap_or_else(|_| { - instance.log_callback_fallback( - &callback_name, - RegistrationSurface::LlmSanitizeRequestGuardrail, - ); - request - }) - }), - )?; - } - RegistrationSurface::LlmSanitizeResponseGuardrail => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_llm_sanitize_response_guardrail( - &name, - priority, - Arc::new(move |value| { - instance - .invoke_llm_response_json( - &callback_name, - RegistrationSurface::LlmSanitizeResponseGuardrail, - "", - value.clone(), - ) - .unwrap_or_else(|_| { - instance.log_callback_fallback( - &callback_name, - RegistrationSurface::LlmSanitizeResponseGuardrail, - ); - value - }) - }), - )?; - } - RegistrationSurface::LlmConditionalExecutionGuardrail => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_llm_conditional_execution_guardrail( - &name, - priority, - Arc::new(move |request| { - instance.invoke_llm_guardrail(&callback_name, request.clone()) - }), - )?; - } - RegistrationSurface::LlmRequestIntercept => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_llm_request_intercept( - &name, - priority, - break_chain, - Arc::new(move |model_name, request, annotated| { - instance.invoke_llm_request_intercept( + }), + ), + RegistrationSurface::ToolSanitizeResponseGuardrail => ctx + .register_tool_sanitize_response_guardrail( + name, + priority, + Arc::new(move |tool_name, value| { + instance + .invoke_tool_json( &callback_name, - model_name, - request, - annotated, + surface, + tool_name, + value.clone(), + None, ) - }), - )?; - } - RegistrationSurface::LlmExecutionIntercept => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_llm_execution_intercept( - &name, - priority, - Arc::new(move |model_name, request, next| { - let instance = instance.clone(); - let name = callback_name.clone(); - let model_name = model_name.to_string(); - Box::pin(async move { - instance - .invoke_llm_execution(&name, &model_name, request, next) - .await + .unwrap_or_else(|_| { + instance.log_callback_fallback(&callback_name, surface); + value }) - }), - )?; - } - RegistrationSurface::LlmStreamExecutionIntercept => { - let instance = Arc::new(self.clone_for_callback()); - let callback_name = name.clone(); - ctx.register_llm_stream_execution_intercept( - &name, - priority, - Arc::new(move |model_name, request, next| { - let instance = instance.clone(); - let name = callback_name.clone(); - let model_name = model_name.to_string(); - Box::pin(async move { - instance - .invoke_llm_stream_execution(&name, &model_name, request, next) - .await + }), + ), + RegistrationSurface::ToolConditionalExecutionGuardrail => ctx + .register_tool_conditional_execution_guardrail( + name, + priority, + Arc::new(move |tool_name, value| { + instance.invoke_tool_guardrail(&callback_name, tool_name, value.clone()) + }), + ), + RegistrationSurface::ToolRequestIntercept => ctx.register_tool_request_intercept( + name, + priority, + registration.break_chain, + Arc::new(move |tool_name, value| { + instance.invoke_tool_json(&callback_name, surface, tool_name, value, None) + }), + ), + RegistrationSurface::ToolExecutionIntercept => ctx.register_tool_execution_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(&callback_name, &tool_name, value, next) + .await + }) + }), + ), + _ => Err(PluginError::RegistrationFailed(format!( + "worker plugin '{}' cannot install registration surface {} as a tool callback", + self.plugin_kind, + surface.as_str_name() + ))), + } + } + + fn install_llm_registration( + &self, + ctx: &mut PluginRegistrationContext, + registration: &Registration, + surface: RegistrationSurface, + ) -> crate::plugin::Result<()> { + let name = registration.local_name.as_str(); + let priority = registration.priority; + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.to_owned(); + match surface { + RegistrationSurface::LlmSanitizeRequestGuardrail => ctx + .register_llm_sanitize_request_guardrail( + name, + priority, + Arc::new(move |request, context| { + instance + .invoke_llm_sanitize_request(&callback_name, request.clone(), context) + .unwrap_or_else(|_| { + instance.log_callback_fallback(&callback_name, surface); + None }) - }), - )?; - } - RegistrationSurface::Unspecified => { - return Err(PluginError::RegistrationFailed(format!( - "worker plugin '{}' returned unspecified registration surface", - self.plugin_kind - ))); - } - } + }), + ), + RegistrationSurface::LlmSanitizeResponseGuardrail => ctx + .register_llm_sanitize_response_guardrail( + name, + priority, + Arc::new(move |value, context| { + instance + .invoke_llm_sanitize_response(&callback_name, value.clone(), context) + .unwrap_or_else(|_| { + instance.log_callback_fallback(&callback_name, surface); + None + }) + }), + ), + RegistrationSurface::LlmConditionalExecutionGuardrail => ctx + .register_llm_conditional_execution_guardrail( + name, + priority, + Arc::new(move |request| { + instance.invoke_llm_guardrail(&callback_name, request.clone()) + }), + ), + RegistrationSurface::LlmRequestIntercept => ctx.register_llm_request_intercept( + name, + priority, + registration.break_chain, + Arc::new(move |model_name, request, annotated| { + instance.invoke_llm_request_intercept( + &callback_name, + model_name, + request, + annotated, + ) + }), + ), + RegistrationSurface::LlmExecutionIntercept => ctx.register_llm_execution_intercept( + name, + priority, + Arc::new(move |model_name, request, next| { + let instance = instance.clone(); + let callback_name = callback_name.clone(); + let model_name = model_name.to_owned(); + Box::pin(async move { + instance + .invoke_llm_execution(&callback_name, &model_name, request, next) + .await + }) + }), + ), + RegistrationSurface::LlmStreamExecutionIntercept => ctx + .register_llm_stream_execution_intercept( + name, + priority, + Arc::new(move |model_name, request, next| { + let instance = instance.clone(); + let callback_name = callback_name.clone(); + let model_name = model_name.to_owned(); + Box::pin(async move { + instance + .invoke_llm_stream_execution( + &callback_name, + &model_name, + request, + next, + ) + .await + }) + }), + ), + _ => Err(PluginError::RegistrationFailed(format!( + "worker plugin '{}' cannot install registration surface {} as an LLM callback", + self.plugin_kind, + surface.as_str_name() + ))), } - Ok(()) } fn clone_for_callback(&self) -> WorkerPluginCallback { @@ -1555,51 +1568,99 @@ impl WorkerPluginCallback { } } - fn invoke_llm_request_json( + fn invoke_llm_sanitize_request( &self, registration_name: &str, - surface: RegistrationSurface, - model_name: &str, request: LlmRequest, - annotated: Option, - continuation_id: Option, - ) -> FlowResult { - let invoke = self.base_request( + context: LlmSanitizeRequestContext, + ) -> FlowResult> { + let mut invoke = self.base_request( registration_name, - surface, - continuation_id, - Some(invoke_request_payload_llm( - model_name, + RegistrationSurface::LlmSanitizeRequestGuardrail, + None, + Some(invoke_request_payload_llm_context( + "", Some(request), - annotated, None, + None, + llm_invocation::SanitizeContext::RequestSanitizeContext( + ProtoLlmSanitizeRequestContext { + codec: Some(codec_identity_to_proto(context.codec())), + codec_capability_id: None, + }, + ), )), ); - let value = json_from_invoke_response(self.invoke_blocking(invoke)?)?; - serde_json::from_value(value).map_err(|err| { - FlowError::Internal(format!("worker returned invalid LLM request: {err}")) - }) + let capability_id = context.resolve_codec().map(|codec| { + let capability_id = self + .host_state + .insert_request_codec(&invoke.invocation_id, codec); + let Some(invoke_request_payload::Payload::Llm(llm)) = invoke.payload.as_mut() else { + unreachable!("LLM sanitizer invocation must have an LLM payload"); + }; + let Some(llm_invocation::SanitizeContext::RequestSanitizeContext(context)) = + llm.sanitize_context.as_mut() + else { + unreachable!("request sanitizer invocation must have a request context"); + }; + context.codec_capability_id = Some(capability_id.clone()); + capability_id + }); + let response = self.invoke_blocking(invoke); + if let Some(capability_id) = capability_id { + self.host_state.remove_codec(&capability_id); + } + optional_json_from_invoke_response(response?)? + .map(serde_json::from_value) + .transpose() + .map_err(|err| { + FlowError::Internal(format!("worker returned invalid LLM request: {err}")) + }) } - fn invoke_llm_response_json( + fn invoke_llm_sanitize_response( &self, registration_name: &str, - surface: RegistrationSurface, - model_name: &str, response: Json, - ) -> FlowResult { - let invoke = self.base_request( + context: LlmSanitizeResponseContext, + ) -> FlowResult> { + let mut invoke = self.base_request( registration_name, - surface, + RegistrationSurface::LlmSanitizeResponseGuardrail, None, - Some(invoke_request_payload_llm( - model_name, + Some(invoke_request_payload_llm_context( + "", None, None, Some(response), + llm_invocation::SanitizeContext::ResponseSanitizeContext( + ProtoLlmSanitizeResponseContext { + codec: Some(codec_identity_to_proto(context.codec())), + codec_capability_id: None, + }, + ), )), ); - json_from_invoke_response(self.invoke_blocking(invoke)?) + let capability_id = context.resolve_codec().map(|codec| { + let capability_id = self + .host_state + .insert_response_codec(&invoke.invocation_id, codec); + let Some(invoke_request_payload::Payload::Llm(llm)) = invoke.payload.as_mut() else { + unreachable!("LLM sanitizer invocation must have an LLM payload"); + }; + let Some(llm_invocation::SanitizeContext::ResponseSanitizeContext(context)) = + llm.sanitize_context.as_mut() + else { + unreachable!("response sanitizer invocation must have a response context"); + }; + context.codec_capability_id = Some(capability_id.clone()); + capability_id + }); + let response = self.invoke_blocking(invoke); + if let Some(capability_id) = capability_id { + self.host_state.remove_codec(&capability_id); + } + optional_json_from_invoke_response(response?) } fn invoke_llm_guardrail( @@ -1973,6 +2034,17 @@ struct WorkerHostRuntimeState { scope_stack_cleanup_complete: Condvar, scope_handles: Mutex>, continuations: Mutex>, + codecs: Mutex>, +} + +struct WorkerCodecCapability { + invocation_id: String, + direction: WorkerCodecDirection, +} + +enum WorkerCodecDirection { + Request(Arc), + Response(Arc), } struct StoredScopeStack { @@ -2024,6 +2096,85 @@ impl WorkerHostRuntimeState { scope_stack_cleanup_complete: Condvar::new(), scope_handles: Mutex::new(HashMap::new()), continuations: Mutex::new(HashMap::new()), + codecs: Mutex::new(HashMap::new()), + } + } + + fn insert_request_codec(&self, invocation_id: &str, codec: Arc) -> String { + self.insert_codec(invocation_id, WorkerCodecDirection::Request(codec)) + } + + fn insert_response_codec( + &self, + invocation_id: &str, + codec: Arc, + ) -> String { + self.insert_codec(invocation_id, WorkerCodecDirection::Response(codec)) + } + + fn insert_codec(&self, invocation_id: &str, direction: WorkerCodecDirection) -> String { + let id = format!("codec-{}", Uuid::now_v7()); + if let Ok(mut codecs) = self.codecs.lock() { + codecs.insert( + id.clone(), + WorkerCodecCapability { + invocation_id: invocation_id.to_owned(), + direction, + }, + ); + } + id + } + + fn remove_codec(&self, id: &str) { + if let Ok(mut codecs) = self.codecs.lock() { + codecs.remove(id); + } + } + + fn request_codec(&self, id: &str, invocation_id: &str) -> Result, Status> { + let codecs = self + .codecs + .lock() + .map_err(|err| Status::internal(format!("codec lock poisoned: {err}")))?; + let capability = codecs + .get(id) + .ok_or_else(|| Status::not_found("codec capability is unavailable"))?; + if capability.invocation_id != invocation_id { + return Err(Status::permission_denied( + "codec capability does not belong to this invocation", + )); + } + match &capability.direction { + WorkerCodecDirection::Request(codec) => Ok(codec.clone()), + WorkerCodecDirection::Response(_) => Err(Status::invalid_argument( + "codec capability is response-only", + )), + } + } + + fn response_codec( + &self, + id: &str, + invocation_id: &str, + ) -> Result, Status> { + let codecs = self + .codecs + .lock() + .map_err(|err| Status::internal(format!("codec lock poisoned: {err}")))?; + let capability = codecs + .get(id) + .ok_or_else(|| Status::not_found("codec capability is unavailable"))?; + if capability.invocation_id != invocation_id { + return Err(Status::permission_denied( + "codec capability does not belong to this invocation", + )); + } + match &capability.direction { + WorkerCodecDirection::Response(codec) => Ok(codec.clone()), + WorkerCodecDirection::Request(_) => { + Err(Status::invalid_argument("codec capability is request-only")) + } } } @@ -2425,6 +2576,79 @@ impl RelayHostRuntime for WorkerHostRuntimeService { }); Ok(Response::new(Box::pin(mapped))) } + + async fn decode_llm_codec_request( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + self.state + .authorize(&request.activation_id, &request.auth_token)?; + let codec = self + .state + .request_codec(&request.codec_capability_id, &request.invocation_id)?; + let request = required_envelope(request.request, "codec request") + .and_then(|value| { + decode_json_envelope::(&value) + .map_err(|err| FlowError::Internal(err.to_string())) + }) + .map_err(status_from_flow)?; + Ok(Response::new(typed_json_result( + ANNOTATED_LLM_REQUEST_SCHEMA, + codec.decode(&request), + ))) + } + + async fn encode_llm_codec_request( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + self.state + .authorize(&request.activation_id, &request.auth_token)?; + let codec = self + .state + .request_codec(&request.codec_capability_id, &request.invocation_id)?; + let annotated = required_envelope(request.annotated_request, "annotated codec request") + .and_then(|value| { + decode_json_envelope::(&value) + .map_err(|err| FlowError::Internal(err.to_string())) + }) + .map_err(status_from_flow)?; + let original = required_envelope(request.original_request, "original codec request") + .and_then(|value| { + decode_json_envelope::(&value) + .map_err(|err| FlowError::Internal(err.to_string())) + }) + .map_err(status_from_flow)?; + Ok(Response::new(typed_json_result( + LLM_REQUEST_SCHEMA, + codec.encode(&annotated, &original), + ))) + } + + async fn decode_llm_codec_response( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + self.state + .authorize(&request.activation_id, &request.auth_token)?; + let codec = self + .state + .response_codec(&request.codec_capability_id, &request.invocation_id)?; + let response = required_envelope(request.response, "codec response") + .and_then(|value| { + decode_json_envelope::(&value) + .map_err(|err| FlowError::Internal(err.to_string())) + }) + .map_err(status_from_flow)?; + Ok(Response::new(json_result( + codec.decode_response(&response).and_then(|value| { + serde_json::to_value(value).map_err(|err| FlowError::Internal(err.to_string())) + }), + ))) + } } impl WorkerHostRuntimeService { @@ -2451,6 +2675,10 @@ mod invoke_request_payload { pub(crate) use nemo_relay_worker_proto::v1::invoke_request::Payload; } +mod llm_invocation { + pub(crate) use nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext; +} + mod invoke_response_result { pub(crate) use nemo_relay_worker_proto::v1::invoke_response::Result; } @@ -2487,9 +2715,44 @@ fn invoke_request_payload_llm( response: response .as_ref() .map(|response| json_envelope_infallible(JSON_SCHEMA, response)), + sanitize_context: None, + }) +} + +fn invoke_request_payload_llm_context( + model_name: &str, + request: Option, + annotated_request: Option, + response: Option, + context: impl Into, +) -> invoke_request_payload::Payload { + invoke_request_payload::Payload::Llm(LlmInvocation { + model_name: model_name.into(), + request: request + .as_ref() + .map(|request| json_envelope_infallible(LLM_REQUEST_SCHEMA, request)), + annotated_request: annotated_request + .as_ref() + .map(|request| json_envelope_infallible(ANNOTATED_LLM_REQUEST_SCHEMA, request)), + response: response + .as_ref() + .map(|response| json_envelope_infallible(JSON_SCHEMA, response)), + sanitize_context: Some(context.into()), }) } +fn codec_identity_to_proto(identity: &LlmCodecIdentity) -> ProtoLlmCodecIdentity { + let (kind, id) = match identity { + LlmCodecIdentity::None => (LlmCodecKind::Unspecified as i32, None), + LlmCodecIdentity::BuiltIn(codec) => { + (LlmCodecKind::Builtin as i32, Some(codec.id().to_owned())) + } + LlmCodecIdentity::Runtime(id) => (LlmCodecKind::Runtime as i32, Some(id.clone())), + LlmCodecIdentity::Opaque => (LlmCodecKind::Opaque as i32, None), + }; + ProtoLlmCodecIdentity { kind, id } +} + fn json_envelope_infallible(schema: &str, value: &T) -> JsonEnvelope { json_envelope(schema, value).expect("Relay DTO JSON serialization should be infallible") } @@ -2512,6 +2775,27 @@ fn json_from_invoke_response(response: InvokeResponse) -> FlowResult { } } +fn optional_json_from_invoke_response(response: InvokeResponse) -> FlowResult> { + match response.result { + Some(invoke_response_result::Result::Empty(_)) => Ok(None), + Some(invoke_response_result::Result::Json(result)) => { + if let Some(error) = result.error { + return Err(worker_error_to_flow(error)); + } + let envelope = required_envelope(result.value, "worker JSON result")?; + decode_json_envelope::(&envelope) + .map(Some) + .map_err(|err| { + FlowError::Internal(format!("worker returned invalid JSON result: {err}")) + }) + } + Some(invoke_response_result::Result::Error(error)) => Err(worker_error_to_flow(error)), + _ => Err(FlowError::Internal( + "worker returned unexpected LLM sanitizer result".into(), + )), + } +} + fn guardrail_from_invoke_response(response: InvokeResponse) -> FlowResult> { match response.result { Some(invoke_response_result::Result::Guardrail(GuardrailResult { block_reason })) => { @@ -2572,6 +2856,19 @@ fn json_result(result: FlowResult) -> JsonResult { } } +fn typed_json_result(schema: &str, result: FlowResult) -> JsonResult { + match result { + Ok(value) => JsonResult { + value: Some(json_envelope_infallible(schema, &value)), + error: None, + }, + Err(err) => JsonResult { + value: None, + error: Some(flow_error_to_worker(err)), + }, + } +} + fn flow_error_to_worker(err: FlowError) -> WorkerError { WorkerError { code: "host.runtime_error".into(), diff --git a/crates/core/src/stream.rs b/crates/core/src/stream.rs index 1124b3d0d..3fee8dd3e 100644 --- a/crates/core/src/stream.rs +++ b/crates/core/src/stream.rs @@ -36,6 +36,7 @@ use crate::api::event::{BaseEvent, MarkEvent}; use crate::api::llm::LlmHandle; use crate::api::llm::emit_optimization_marks; use crate::api::optimization::finalize_optimization_summary; +use crate::api::runtime::LlmSanitizeResponseContext; use crate::api::runtime::NemoRelayContextState; use crate::api::runtime::global_context; use crate::api::runtime::{ @@ -71,6 +72,7 @@ pub struct LlmStreamWrapper { collector: Box Result<()> + Send>, finalizer: Option Json + Send>>, response_codec: Option>, + sanitize_context: LlmSanitizeResponseContext, metadata: Option, subscribers: Vec, chunk_index: u64, @@ -140,6 +142,8 @@ impl LlmStreamWrapper { subscribers: Vec, ) -> Self { let scope_stack = handle.captured_scope_stack().clone(); + let sanitize_context = + LlmSanitizeResponseContext::for_response_codec(response_codec.clone()); Self { inner, handle, @@ -147,6 +151,7 @@ impl LlmStreamWrapper { collector, finalizer: Some(finalizer), response_codec, + sanitize_context, metadata, subscribers, chunk_index: 0, @@ -204,6 +209,12 @@ impl LlmStreamWrapper { Some(finalizer) => finalizer(), None => Json::Null, }; + let response_was_null_without_fallback = aggregated.is_null() && self.handle.data.is_none(); + let response = if aggregated.is_null() { + self.handle.data.clone().unwrap_or(aggregated) + } else { + aggregated + }; let snapshot = { let ss_guard = self.scope_stack.read().expect("scope stack lock poisoned"); @@ -222,19 +233,27 @@ impl LlmStreamWrapper { let Some(entries) = snapshot else { return; }; - let sanitized = - NemoRelayContextState::llm_sanitize_response_snapshot_chain(aggregated, &entries); - let data = if sanitized.is_null() { - self.handle.data.clone() - } else { - Some(sanitized) + let sanitized = NemoRelayContextState::llm_sanitize_response_snapshot_chain( + response, + self.sanitize_context.clone(), + &entries, + ); + let data = match sanitized { + Some(response) if response_was_null_without_fallback && response.is_null() => None, + response => response, }; - let mut annotated_response: Option = - self.response_codec.as_ref().and_then(|codec| { - let mut decoded = codec.decode_response(data.as_ref()?).ok()?; - attach_estimated_cost_for_provider(&mut decoded, Some(&self.handle.name)); - Some(decoded) - }); + let annotation_omitted = data.as_ref().is_none_or(Json::is_null); + let mut annotated_response: Option = (!annotation_omitted) + .then(|| { + data.as_ref().and_then(|response| { + self.response_codec.as_ref().and_then(|codec| { + let mut decoded = codec.decode_response(response).ok()?; + attach_estimated_cost_for_provider(&mut decoded, Some(&self.handle.name)); + Some(decoded) + }) + }) + }) + .flatten(); let interruption = (interrupted && !has_authoritative_final_usage(annotated_response.as_ref())) .then_some("stream_interrupted"); @@ -249,7 +268,8 @@ impl LlmStreamWrapper { self.handle.model_name.as_deref(), &pricing, ); - if annotated_response.is_none() + if !annotation_omitted + && annotated_response.is_none() && let Some(summary) = summary { annotated_response = Some(AnnotatedLlmResponse { diff --git a/crates/core/tests/fixtures/native_plugin/src/lib.rs b/crates/core/tests/fixtures/native_plugin/src/lib.rs index 08b3191d8..350cf1414 100644 --- a/crates/core/tests/fixtures/native_plugin/src/lib.rs +++ b/crates/core/tests/fixtures/native_plugin/src/lib.rs @@ -130,12 +130,12 @@ impl NativePlugin for FixtureNativePlugin { ctx.register_llm_sanitize_request_guardrail( "fixture_llm_sanitize_request", 0, - |request| mark_llm_request(request, "native_plugin_llm_sanitize_request"), + |request, _context| Some(mark_llm_request(request, "native_plugin_llm_sanitize_request")), )?; ctx.register_llm_sanitize_response_guardrail( "fixture_llm_sanitize_response", 0, - |response| mark_json(response, "native_plugin_llm_sanitize_response"), + |response, _context| Some(mark_json(response, "native_plugin_llm_sanitize_response")), )?; ctx.register_llm_conditional_execution_guardrail( "fixture_llm_conditional", diff --git a/crates/core/tests/fixtures/worker_plugin/src/main.rs b/crates/core/tests/fixtures/worker_plugin/src/main.rs index c0320f2d5..a9dd41a67 100644 --- a/crates/core/tests/fixtures/worker_plugin/src/main.rs +++ b/crates/core/tests/fixtures/worker_plugin/src/main.rs @@ -2,11 +2,11 @@ // SPDX-License-Identifier: Apache-2.0 use nemo_relay_worker::{ - JsonStream, LlmNext, LlmStreamNext, PluginContext, ScopeType, ToolNext, WorkerPlugin, - ToolExecutionInterceptOutcome, WorkerSdkError, serve_plugin, + ConfigDiagnostic, DiagnosticLevel, EventSanitizeFields, Json, LlmRequest, PendingMarkSpec, }; use nemo_relay_worker::{ - ConfigDiagnostic, DiagnosticLevel, EventSanitizeFields, Json, LlmRequest, PendingMarkSpec, + JsonStream, LlmNext, LlmStreamNext, PluginContext, ScopeType, ToolExecutionInterceptOutcome, + ToolNext, WorkerPlugin, WorkerSdkError, serve_plugin, }; use serde_json::json; @@ -61,22 +61,26 @@ impl WorkerPlugin for FixtureWorkerPlugin { let runtime = ctx .runtime() .ok_or_else(|| WorkerSdkError::Callback("runtime handle missing".into()))?; - ctx.register_mark_sanitize_guardrail("fixture_mark_sanitize", 0, |_, fields| { - mark_event_fields(fields, "worker_plugin_mark") - }); - ctx.register_mark_sanitize_guardrail("fixture_mark_sanitize_data", 1, |_, mut fields| { - fields.data = Some(json!({"worker_plugin_mark_data": true})); - fields + ctx.register_mark_sanitize_guardrail("fixture_mark_sanitize", 0, |_, fields| async move { + Ok(mark_event_fields(fields, "worker_plugin_mark")) }); + ctx.register_mark_sanitize_guardrail( + "fixture_mark_sanitize_data", + 1, + |_, mut fields| async move { + fields.data = Some(json!({"worker_plugin_mark_data": true})); + Ok(fields) + }, + ); ctx.register_scope_sanitize_start_guardrail( "fixture_scope_start_sanitize", 0, - |_, fields| mark_event_fields(fields, "worker_plugin_scope_start"), + |_, fields| async move { Ok(mark_event_fields(fields, "worker_plugin_scope_start")) }, ); ctx.register_scope_sanitize_end_guardrail( "fixture_scope_end_sanitize", 0, - |_, fields| mark_event_fields(fields, "worker_plugin_scope_end"), + |_, fields| async move { Ok(mark_event_fields(fields, "worker_plugin_scope_end")) }, ); register_fixture_subscriber(ctx, runtime.clone()); register_fixture_tool_hooks( @@ -135,12 +139,14 @@ fn register_fixture_tool_hooks( ctx.register_tool_sanitize_request_guardrail( "fixture_tool_sanitize_request", 0, - |_name, args| mark_json(args, "worker_plugin_tool_sanitize_request"), + |_name, args| async move { Ok(mark_json(args, "worker_plugin_tool_sanitize_request")) }, ); ctx.register_tool_sanitize_response_guardrail( "fixture_tool_sanitize_response", 0, - |_name, result| mark_json(result, "worker_plugin_tool_sanitize_response"), + |_name, result| async move { + Ok(mark_json(result, "worker_plugin_tool_sanitize_response")) + }, ); ctx.register_tool_conditional_execution_guardrail( "fixture_tool_conditional", @@ -175,17 +181,15 @@ fn register_fixture_tool_hooks( let result = next .call(mark_json(args, "worker_plugin_tool_execution_request")) .await?; - Ok( - ToolExecutionInterceptOutcome::new(mark_json( - result, - "worker_plugin_tool_execution", - )) - .with_pending_mark( - PendingMarkSpec::builder() - .name("fixture.worker.tool_execution.mark") - .build(), - ), - ) + Ok(ToolExecutionInterceptOutcome::new(mark_json( + result, + "worker_plugin_tool_execution", + )) + .with_pending_mark( + PendingMarkSpec::builder() + .name("fixture.worker.tool_execution.mark") + .build(), + )) }, ); } @@ -198,18 +202,26 @@ fn register_fixture_llm_hooks( ctx.register_llm_sanitize_request_guardrail( "fixture_llm_sanitize_request", 0, - |request| mark_llm_request(request, "worker_plugin_llm_sanitize_request"), + |request, _context| async move { + Ok(Some(mark_llm_request( + request, + "worker_plugin_llm_sanitize_request", + ))) + }, ); ctx.register_llm_sanitize_response_guardrail( "fixture_llm_sanitize_response", 0, - |response| mark_json(response, "worker_plugin_llm_sanitize_response"), - ); - ctx.register_llm_conditional_execution_guardrail( - "fixture_llm_conditional", - 0, - |_request| Ok(None), + |response, _context| async move { + Ok(Some(mark_json( + response, + "worker_plugin_llm_sanitize_response", + ))) + }, ); + ctx.register_llm_conditional_execution_guardrail("fixture_llm_conditional", 0, |_request| { + Ok(None) + }); ctx.register_llm_request_intercept( "fixture_llm_request_intercept", 0, @@ -232,14 +244,16 @@ fn register_fixture_llm_hooks( None, ), }; - Ok(nemo_relay_worker::LlmRequestInterceptOutcome::new(request, annotated) - .with_pending_mark( - PendingMarkSpec::builder() - .name("fixture.worker.llm_request.mark") - .data(json!({ "source": "worker_request_intercept" })) - .metadata(json!({ "fixture": true })) - .build(), - )) + Ok( + nemo_relay_worker::LlmRequestInterceptOutcome::new(request, annotated) + .with_pending_mark( + PendingMarkSpec::builder() + .name("fixture.worker.llm_request.mark") + .data(json!({ "source": "worker_request_intercept" })) + .metadata(json!({ "fixture": true })) + .build(), + ), + ) }, ); ctx.register_llm_execution_intercept( @@ -278,7 +292,9 @@ fn register_fixture_llm_hooks( ); } -async fn emit_runtime_events(runtime: nemo_relay_worker::PluginRuntime) -> nemo_relay_worker::Result<()> { +async fn emit_runtime_events( + runtime: nemo_relay_worker::PluginRuntime, +) -> nemo_relay_worker::Result<()> { runtime .emit_mark("fixture.worker.mark", Some(json!("current")), None) .await?; @@ -311,7 +327,11 @@ async fn emit_runtime_events(runtime: nemo_relay_worker::PluginRuntime) -> nemo_ runtime .with_scope_stack(&isolated, || async move { isolated_runtime - .emit_mark("fixture.worker.isolated.mark", Some(json!("isolated")), None) + .emit_mark( + "fixture.worker.isolated.mark", + Some(json!("isolated")), + None, + ) .await }) .await?; diff --git a/crates/core/tests/integration/api_surface_tests.rs b/crates/core/tests/integration/api_surface_tests.rs index 6b3d83bd9..7a6045ac0 100644 --- a/crates/core/tests/integration/api_surface_tests.rs +++ b/crates/core/tests/integration/api_surface_tests.rs @@ -903,14 +903,18 @@ fn test_global_registry_and_subscriber_wrappers_cover_success_and_duplicates() { .unwrap(); assert!(deregister_tool_execution_intercept("tool-execution").unwrap()); - register_llm_sanitize_request_guardrail("llm-sanitize-request", 1, Arc::new(|request| request)) - .unwrap(); + register_llm_sanitize_request_guardrail( + "llm-sanitize-request", + 1, + Arc::new(|request, _context| Some(request)), + ) + .unwrap(); assert!(deregister_llm_sanitize_request_guardrail("llm-sanitize-request").unwrap()); register_llm_sanitize_response_guardrail( "llm-sanitize-response", 1, - Arc::new(|response| response), + Arc::new(|response, _context| Some(response)), ) .unwrap(); assert!(deregister_llm_sanitize_response_guardrail("llm-sanitize-response").unwrap()); @@ -1134,7 +1138,7 @@ fn test_scope_registry_and_subscriber_wrappers_cover_success_duplicates_and_miss &scope.uuid, "llm-sanitize-request", 1, - Arc::new(|request| request), + Arc::new(|request, _context| Some(request)), ) .unwrap(); assert!( @@ -1146,7 +1150,7 @@ fn test_scope_registry_and_subscriber_wrappers_cover_success_duplicates_and_miss &scope.uuid, "llm-sanitize-response", 1, - Arc::new(|response| response), + Arc::new(|response, _context| Some(response)), ) .unwrap(); assert!( @@ -1486,21 +1490,21 @@ async fn test_llm_api_emits_sanitized_events_and_covers_error_paths() { register_llm_sanitize_request_guardrail( "llm-sanitize-request", 1, - Arc::new(|mut request| { + Arc::new(|mut request, _context| { request.headers.insert("x-sanitized".into(), json!(true)); - request + Some(request) }), ) .unwrap(); register_llm_sanitize_response_guardrail( "llm-sanitize-response", 1, - Arc::new(|mut response| { + Arc::new(|mut response, _context| { response .as_object_mut() .unwrap() .insert("sanitized_response".into(), json!(true)); - response + Some(response) }), ) .unwrap(); diff --git a/crates/core/tests/integration/middleware_tests.rs b/crates/core/tests/integration/middleware_tests.rs index 5acc08b63..41093cdcb 100644 --- a/crates/core/tests/integration/middleware_tests.rs +++ b/crates/core/tests/integration/middleware_tests.rs @@ -2638,10 +2638,10 @@ async fn test_llm_middleware_callbacks_run_without_registry_or_scope_locks() { register_llm_sanitize_request_guardrail( "lock_global_llm_sanitize_request", 1, - Arc::new(move |request| { + Arc::new(move |request, _context| { record_middleware_callback(&tracked, "llm_sanitize_request_global"); assert_middleware_callback_locks_are_free(); - request + Some(request) }), ) .unwrap(); @@ -2650,10 +2650,10 @@ async fn test_llm_middleware_callbacks_run_without_registry_or_scope_locks() { &scope.uuid, "lock_scope_llm_sanitize_request", 2, - Arc::new(move |request| { + Arc::new(move |request, _context| { record_middleware_callback(&tracked, "llm_sanitize_request_scope"); assert_middleware_callback_locks_are_free(); - request + Some(request) }), ) .unwrap(); @@ -2707,10 +2707,10 @@ async fn test_llm_middleware_callbacks_run_without_registry_or_scope_locks() { register_llm_sanitize_response_guardrail( "lock_global_llm_sanitize_response", 1, - Arc::new(move |response| { + Arc::new(move |response, _context| { record_middleware_callback(&tracked, "llm_sanitize_response_global"); assert_middleware_callback_locks_are_free(); - response + Some(response) }), ) .unwrap(); @@ -2719,10 +2719,10 @@ async fn test_llm_middleware_callbacks_run_without_registry_or_scope_locks() { &scope.uuid, "lock_scope_llm_sanitize_response", 2, - Arc::new(move |response| { + Arc::new(move |response, _context| { record_middleware_callback(&tracked, "llm_sanitize_response_scope"); assert_middleware_callback_locks_are_free(); - response + Some(response) }), ) .unwrap(); diff --git a/crates/core/tests/integration/native_plugin_tests.rs b/crates/core/tests/integration/native_plugin_tests.rs index ef2604f58..ee8a860be 100644 --- a/crates/core/tests/integration/native_plugin_tests.rs +++ b/crates/core/tests/integration/native_plugin_tests.rs @@ -640,12 +640,17 @@ async fn sdk_cdylib_registers_tool_request_intercept() { flush_subscribers().expect("llm response annotation event should flush"); let llm_events = events.lock().unwrap().clone(); let llm_end = find_event(&llm_events, "native-fixture-llm", Some(ScopeCategory::End)); - let annotated = llm_end - .annotated_response() - .expect("native plugin should preserve response annotation"); - assert_eq!(annotated.id.as_deref(), Some("annotation-before-plugin")); - assert_eq!(annotated.extra["preexisting_annotation"], json!("kept")); - assert!(annotated.extra.get("native_plugin_annotation").is_none()); + assert_eq!( + llm_end.output().unwrap()["native_plugin_llm_sanitize_response"], + true + ); + assert!( + llm_end.annotated_response().is_none(), + "a changed response without an active codec must discard the stale caller annotation" + ); + let serialized = serde_json::to_string(llm_end).expect("LLM end event should serialize"); + assert!(!serialized.contains("annotation-before-plugin")); + assert!(!serialized.contains("preexisting_annotation")); drop(cleanup); activation.clear(); diff --git a/crates/core/tests/integration/pipeline_tests.rs b/crates/core/tests/integration/pipeline_tests.rs index 918407663..1bc192b4b 100644 --- a/crates/core/tests/integration/pipeline_tests.rs +++ b/crates/core/tests/integration/pipeline_tests.rs @@ -1448,7 +1448,7 @@ async fn test_response_codec_annotation_uses_sanitized_managed_response() { register_llm_sanitize_response_guardrail( "sanitize_resp_codec_annotation", 1, - Arc::new(|_response| make_openai_chat_response("Sanitized")), + Arc::new(|_response, _context| Some(make_openai_chat_response("Sanitized"))), ) .unwrap(); @@ -1648,9 +1648,11 @@ async fn test_request_codec_annotation_uses_sanitized_start_payload() { register_llm_sanitize_request_guardrail( "sanitize_req_codec_annotation", 1, - Arc::new(|request| LlmRequest { - headers: request.headers, - content: make_openai_chat_request("Sanitized").content, + Arc::new(|request, _context| { + Some(LlmRequest { + headers: request.headers, + content: make_openai_chat_request("Sanitized").content, + }) }), ) .unwrap(); @@ -1914,7 +1916,7 @@ async fn test_stream_response_codec_annotation_uses_sanitized_aggregated_respons register_llm_sanitize_response_guardrail( "stream_sanitize_resp_codec_annotation", 1, - Arc::new(|_response| make_openai_chat_response("Sanitized")), + Arc::new(|_response, _context| Some(make_openai_chat_response("Sanitized"))), ) .unwrap(); diff --git a/crates/core/tests/unit/context_tests.rs b/crates/core/tests/unit/context_tests.rs index 07e4ea6b6..49e57b2cb 100644 --- a/crates/core/tests/unit/context_tests.rs +++ b/crates/core/tests/unit/context_tests.rs @@ -320,8 +320,12 @@ fn context_state_supports_extensions_events_and_builders() { content: json!({"messages": []}), }; let entries = state.llm_sanitize_request_entries(&[]); - let sanitized = - NemoRelayContextState::llm_sanitize_request_snapshot_chain(request.clone(), &entries); + let sanitized = NemoRelayContextState::llm_sanitize_request_snapshot_chain( + request.clone(), + crate::api::runtime::LlmSanitizeRequestContext::default(), + &entries, + ) + .expect("an empty sanitizer chain must retain the request"); assert!(sanitized.headers.is_empty()); let events = Arc::new(Mutex::new(Vec::::new())); diff --git a/crates/core/tests/unit/dynamic_worker_tests.rs b/crates/core/tests/unit/dynamic_worker_tests.rs index 1e28358ad..7754238dc 100644 --- a/crates/core/tests/unit/dynamic_worker_tests.rs +++ b/crates/core/tests/unit/dynamic_worker_tests.rs @@ -4,7 +4,12 @@ use std::sync::{Arc, Mutex}; use crate::api::event::{BaseEvent, MarkEvent}; -use crate::api::runtime::NemoRelayContextState; +use crate::api::runtime::{ + BuiltinLlmCodec, LlmCodecIdentity, LlmSanitizeRequestContext, LlmSanitizeResponseContext, + NemoRelayContextState, +}; +use crate::codec::openai_chat::OpenAIChatCodec; +use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use nemo_relay_worker_proto::json_envelope; use nemo_relay_worker_proto::v1::invoke_response::Result as InvokeResult; use nemo_relay_worker_proto::v1::plugin_worker_server::{PluginWorker, PluginWorkerServer}; @@ -12,10 +17,10 @@ use nemo_relay_worker_proto::v1::stream_chunk::Item as StreamItem; use nemo_relay_worker_proto::v1::{ CancelInvocationRequest, CreateScopeStackRequest, DropScopeStackRequest, EmitMarkRequest, EmptyResult, GuardrailResult, HandshakeRequest, HandshakeResponse, HealthRequest, - HealthResponse, JsonEnvelope, JsonResult, LlmNextRequest, LlmRequestInterceptResult, - LlmStreamNextRequest, PopScopeRequest, PushScopeRequest, Registration, ScopeContext, - ScopeType as ProtoScopeType, ShutdownRequest, StreamChunk, ToolNextRequest, ValidateRequest, - ValidateResponse, WorkerAck, + HealthResponse, JsonEnvelope, JsonResult, LlmCodecDecodeRequest, LlmCodecDecodeResponse, + LlmCodecEncodeRequest, LlmNextRequest, LlmRequestInterceptResult, LlmStreamNextRequest, + PopScopeRequest, PushScopeRequest, Registration, ScopeContext, ScopeType as ProtoScopeType, + ShutdownRequest, StreamChunk, ToolNextRequest, ValidateRequest, ValidateResponse, WorkerAck, }; use serde_json::json; use tokio_stream::StreamExt; @@ -464,13 +469,10 @@ async fn callback_helpers_cover_worker_response_edges() { ); let error = callback - .invoke_llm_request_json( + .invoke_llm_sanitize_request( "llm_json_invalid", - RegistrationSurface::LlmSanitizeRequestGuardrail, - "model", valid_llm_request(), - None, - None, + LlmSanitizeRequestContext::default(), ) .expect_err("invalid LLM JSON result should fail"); assert!(error.to_string().contains("invalid type")); @@ -537,6 +539,287 @@ async fn callback_helpers_cover_worker_response_edges() { ); } +#[tokio::test(flavor = "multi_thread")] +async fn llm_worker_sanitizers_forward_codec_context_and_omission() { + enable_operational_logs(); + let seen = Arc::new(Mutex::new(Vec::new())); + let (callback, _shutdown) = fake_callback_service({ + let seen = seen.clone(); + move |request| { + let Some(nemo_relay_worker_proto::v1::invoke_request::Payload::Llm(invocation)) = + request.payload + else { + panic!("LLM sanitizer must receive an LLM invocation"); + }; + let codec = match invocation.sanitize_context.as_ref() { + Some(nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::RequestSanitizeContext(context)) => context.codec.as_ref(), + Some(nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::ResponseSanitizeContext(context)) => context.codec.as_ref(), + None => None, + }; + seen.lock().unwrap().push(( + request.registration_name, + codec + .map(|codec| codec.kind) + .unwrap_or(LlmCodecKind::Unspecified as i32), + codec.and_then(|codec| codec.id.clone()), + invocation.request.is_some(), + invocation.response.is_some(), + )); + InvokeResponse { + result: Some(InvokeResult::Empty(EmptyResult {})), + } + } + }) + .await; + + let identities = [ + LlmCodecIdentity::None, + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiResponses), + LlmCodecIdentity::Runtime("com.example.chat.v1".into()), + LlmCodecIdentity::Opaque, + ]; + for identity in identities { + assert!( + callback + .invoke_llm_sanitize_request( + "request", + valid_llm_request(), + LlmSanitizeRequestContext::with_identity(identity.clone()), + ) + .expect("empty worker result must represent request omission") + .is_none() + ); + assert!( + callback + .invoke_llm_sanitize_response( + "response", + json!({"secret": "value"}), + LlmSanitizeResponseContext::with_identity(identity), + ) + .expect("empty worker result must represent response omission") + .is_none() + ); + } + + let seen = seen.lock().unwrap().clone(); + assert_eq!( + seen, + [ + ( + "request".into(), + LlmCodecKind::Unspecified as i32, + None, + true, + false, + ), + ( + "response".into(), + LlmCodecKind::Unspecified as i32, + None, + false, + true, + ), + ( + "request".into(), + LlmCodecKind::Builtin as i32, + Some("openai_responses".into()), + true, + false, + ), + ( + "response".into(), + LlmCodecKind::Builtin as i32, + Some("openai_responses".into()), + false, + true, + ), + ( + "request".into(), + LlmCodecKind::Runtime as i32, + Some("com.example.chat.v1".into()), + true, + false, + ), + ( + "response".into(), + LlmCodecKind::Runtime as i32, + Some("com.example.chat.v1".into()), + false, + true, + ), + ( + "request".into(), + LlmCodecKind::Opaque as i32, + None, + true, + false, + ), + ( + "response".into(), + LlmCodecKind::Opaque as i32, + None, + false, + true, + ), + ] + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn llm_worker_codec_capabilities_are_active_only_during_sanitizer_invocation() { + enable_operational_logs(); + type SeenCapability = (String, String, String); + + let host_state = Arc::new(Mutex::new(None::>)); + let seen = Arc::new(Mutex::new(Vec::::new())); + let (callback, _shutdown) = fake_callback_service({ + let host_state = host_state.clone(); + let seen = seen.clone(); + move |request| { + let invocation_id = request.invocation_id.clone(); + let registration_name = request.registration_name.clone(); + let Some(nemo_relay_worker_proto::v1::invoke_request::Payload::Llm(invocation)) = + request.payload + else { + panic!("LLM sanitizer must receive an LLM invocation"); + }; + let state = host_state + .lock() + .unwrap() + .clone() + .expect("host state must be installed before invoking the worker"); + let capability_id = match invocation + .sanitize_context + .as_ref() + .expect("codec context must be forwarded") + { + nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::RequestSanitizeContext( + context, + ) => { + let id = context + .codec_capability_id + .as_deref() + .expect("active request codec must receive a capability"); + let codec = state + .request_codec(id, &invocation_id) + .expect("request capability must resolve during invocation"); + let request: LlmRequest = decode_json_envelope( + invocation.request.as_ref().expect("request payload"), + ) + .expect("request payload must decode"); + let annotated = codec + .decode(&request) + .expect("resolved request codec must be usable"); + assert_eq!(annotated.model.as_deref(), Some("gpt-test")); + id.to_owned() + } + nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::ResponseSanitizeContext( + context, + ) => { + let id = context + .codec_capability_id + .as_deref() + .expect("active response codec must receive a capability"); + let codec = state + .response_codec(id, &invocation_id) + .expect("response capability must resolve during invocation"); + let response: Json = decode_json_envelope( + invocation.response.as_ref().expect("response payload"), + ) + .expect("response payload must decode"); + let annotated = codec + .decode_response(&response) + .expect("resolved response codec must be usable"); + assert_eq!(annotated.model.as_deref(), Some("gpt-test")); + id.to_owned() + } + }; + seen.lock().unwrap().push(( + registration_name.clone(), + capability_id, + invocation_id, + )); + if registration_name == "response-error" { + InvokeResponse { + result: Some(InvokeResult::Error( + nemo_relay_worker_proto::v1::WorkerError { + code: "worker.failed".into(), + message: "boom".into(), + retryable: false, + }, + )), + } + } else { + InvokeResponse { + result: Some(InvokeResult::Empty(EmptyResult {})), + } + } + } + }) + .await; + *host_state.lock().unwrap() = Some(callback.host_state.clone()); + + let codec = Arc::new(OpenAIChatCodec); + let request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-test", + "messages": [{"role": "user", "content": "secret"}] + }), + }; + assert!( + callback + .invoke_llm_sanitize_request( + "request-success", + request, + LlmSanitizeRequestContext::for_request_codec(Some(codec.clone())), + ) + .expect("request sanitizer must succeed") + .is_none() + ); + + let response = json!({ + "id": "chatcmpl-test", + "model": "gpt-test", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "secret"}, + "finish_reason": "stop" + }] + }); + let error = callback + .invoke_llm_sanitize_response( + "response-error", + response, + LlmSanitizeResponseContext::for_response_codec(Some(codec)), + ) + .expect_err("worker sanitizer error must surface"); + assert!(error.to_string().contains("worker.failed: boom")); + + let seen = seen.lock().unwrap().clone(); + assert_eq!(seen.len(), 2); + let (_, request_capability, request_invocation) = &seen[0]; + assert_eq!( + callback + .host_state + .request_codec(request_capability, request_invocation) + .err() + .expect("successful invocation must expire its request capability") + .code(), + tonic::Code::NotFound + ); + let (_, response_capability, response_invocation) = &seen[1]; + assert_eq!( + callback + .host_state + .response_codec(response_capability, response_invocation) + .err() + .expect("failed invocation must expire its response capability") + .code(), + tonic::Code::NotFound + ); +} + #[tokio::test(flavor = "multi_thread")] async fn callback_stream_transport_error_surfaces_to_host_stream() { enable_operational_logs(); @@ -1164,20 +1447,24 @@ async fn installed_callbacks_apply_surface_specific_fallbacks() { tool_response ); let entries = state.llm_sanitize_request_entries(&[]); - assert_eq!( + assert!( NemoRelayContextState::llm_sanitize_request_snapshot_chain( llm_request.clone(), + crate::api::runtime::LlmSanitizeRequestContext::default(), &entries, - ), - llm_request + ) + .is_none(), + "a worker request sanitizer failure must omit the observability payload" ); let entries = state.llm_sanitize_response_entries(&[]); - assert_eq!( + assert!( NemoRelayContextState::llm_sanitize_response_snapshot_chain( llm_response.clone(), + crate::api::runtime::LlmSanitizeResponseContext::default(), &entries, - ), - llm_response + ) + .is_none(), + "a worker response sanitizer failure must omit the observability payload" ); } crate::api::subscriber::flush_subscribers().expect("subscriber callback should flush"); @@ -1346,6 +1633,152 @@ async fn host_runtime_service_covers_auth_scope_and_ack_errors() { ); } +#[tokio::test] +async fn host_runtime_codec_capabilities_are_directional_authorized_and_ephemeral() { + let state = Arc::new(WorkerHostRuntimeState::new( + ACTIVATION_ID.into(), + AUTH_TOKEN.into(), + )); + let service = WorkerHostRuntimeService { + state: state.clone(), + }; + let codec = Arc::new(OpenAIChatCodec); + let request_codec: Arc = codec.clone(); + let response_codec: Arc = codec.clone(); + let invocation_id = "sanitize-invocation"; + let request_capability = state.insert_request_codec(invocation_id, request_codec); + let response_capability = state.insert_response_codec(invocation_id, response_codec); + let request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-test", + "messages": [{"role": "user", "content": "secret"}], + "preserve": true + }), + }; + + let unauthorized = service + .decode_llm_codec_request(Request::new(LlmCodecDecodeRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: "wrong".into(), + codec_capability_id: request_capability.clone(), + invocation_id: invocation_id.into(), + request: Some(json_envelope("nemo.relay.LlmRequest@1", &request).unwrap()), + })) + .await + .expect_err("wrong activation credentials must be rejected"); + assert_eq!(unauthorized.code(), tonic::Code::PermissionDenied); + + let forged = service + .decode_llm_codec_request(Request::new(LlmCodecDecodeRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + codec_capability_id: "codec-forged".into(), + invocation_id: invocation_id.into(), + request: Some(json_envelope("nemo.relay.LlmRequest@1", &request).unwrap()), + })) + .await + .expect_err("forged capability must be rejected"); + assert_eq!(forged.code(), tonic::Code::NotFound); + + let wrong_direction = service + .decode_llm_codec_request(Request::new(LlmCodecDecodeRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + codec_capability_id: response_capability.clone(), + invocation_id: invocation_id.into(), + request: Some(json_envelope("nemo.relay.LlmRequest@1", &request).unwrap()), + })) + .await + .expect_err("response capability cannot decode requests"); + assert_eq!(wrong_direction.code(), tonic::Code::InvalidArgument); + + let wrong_invocation = service + .decode_llm_codec_request(Request::new(LlmCodecDecodeRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + codec_capability_id: request_capability.clone(), + invocation_id: "another-invocation".into(), + request: Some(json_envelope("nemo.relay.LlmRequest@1", &request).unwrap()), + })) + .await + .expect_err("a capability cannot be reused by another invocation"); + assert_eq!(wrong_invocation.code(), tonic::Code::PermissionDenied); + + let decoded = service + .decode_llm_codec_request(Request::new(LlmCodecDecodeRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + codec_capability_id: request_capability.clone(), + invocation_id: invocation_id.into(), + request: Some(json_envelope("nemo.relay.LlmRequest@1", &request).unwrap()), + })) + .await + .expect("active request capability decodes") + .into_inner(); + assert!(decoded.error.is_none()); + let decoded = decoded.value.expect("decoded request value"); + assert_eq!(decoded.schema, "nemo.relay.AnnotatedLlmRequest@2"); + let annotated: AnnotatedLlmRequest = decode_json_envelope(&decoded).unwrap(); + assert_eq!(annotated.model.as_deref(), Some("gpt-test")); + + let encoded = service + .encode_llm_codec_request(Request::new(LlmCodecEncodeRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + codec_capability_id: request_capability.clone(), + invocation_id: invocation_id.into(), + annotated_request: Some( + json_envelope("nemo.relay.AnnotatedLlmRequest@2", &annotated).unwrap(), + ), + original_request: Some(json_envelope("nemo.relay.LlmRequest@1", &request).unwrap()), + })) + .await + .expect("active request capability encodes") + .into_inner(); + assert!(encoded.error.is_none()); + let encoded = encoded.value.expect("encoded request value"); + assert_eq!(encoded.schema, "nemo.relay.LlmRequest@1"); + let encoded: LlmRequest = decode_json_envelope(&encoded).unwrap(); + assert_eq!(encoded.content["preserve"], json!(true)); + + let response = json!({ + "id": "chatcmpl-test", + "model": "gpt-test", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "secret"}, + "finish_reason": "stop" + }] + }); + let decoded = service + .decode_llm_codec_response(Request::new(LlmCodecDecodeResponse { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + codec_capability_id: response_capability.clone(), + invocation_id: invocation_id.into(), + response: Some(json_envelope(JSON_SCHEMA, &response).unwrap()), + })) + .await + .expect("active response capability decodes") + .into_inner(); + assert!(decoded.error.is_none()); + + state.remove_codec(&request_capability); + state.remove_codec(&response_capability); + let expired = service + .decode_llm_codec_request(Request::new(LlmCodecDecodeRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + codec_capability_id: request_capability, + invocation_id: invocation_id.into(), + request: Some(json_envelope("nemo.relay.LlmRequest@1", &request).unwrap()), + })) + .await + .expect_err("removed capability must expire"); + assert_eq!(expired.code(), tonic::Code::NotFound); +} + #[tokio::test] async fn host_runtime_service_reports_poisoned_internal_locks() { enable_operational_logs(); diff --git a/crates/core/tests/unit/llm_api_tests.rs b/crates/core/tests/unit/llm_api_tests.rs index 5dd2e3914..733f5b2b7 100644 --- a/crates/core/tests/unit/llm_api_tests.rs +++ b/crates/core/tests/unit/llm_api_tests.rs @@ -12,25 +12,36 @@ use serde_json::json; use tokio_stream::StreamExt; use super::{ - LlmCallExecuteParams, LlmCallParams, LlmHandle, LlmRequest, LlmStreamCallExecuteParams, - emit_optimization_marks_with, llm_call, llm_call_execute, llm_stream_call_execute, - project_llm_request_to_current_user_turn, + CreateLlmHandleParams, LlmCallEndParams, LlmCallExecuteParams, LlmCallParams, LlmHandle, + LlmRequest, LlmStreamCallExecuteParams, create_llm_handle, emit_llm_start, + emit_optimization_marks_with, llm_call, llm_call_end, llm_call_execute, + llm_stream_call_execute, project_llm_request_to_current_user_turn, + sanitize_context_for_request_codec, sanitize_context_for_response_codec, }; use crate::api::event::{Event, ScopeCategory}; use crate::api::optimization::finalize_optimization_summary; -use crate::api::runtime::LlmJsonStream; +use crate::api::registry::{ + deregister_llm_sanitize_request_guardrail, deregister_llm_sanitize_response_guardrail, + register_llm_sanitize_request_guardrail, register_llm_sanitize_response_guardrail, +}; +use crate::api::runtime::{BuiltinLlmCodec, LlmCodecIdentity, LlmJsonStream}; use crate::api::runtime::{ NemoRelayContextState, create_scope_stack, global_context, set_thread_scope_stack, }; use crate::api::scope::{COMPACTION_EVENT_NAME, EmitMarkEventParams, event}; use crate::api::scope::{PopScopeParams, PushScopeParams, ScopeType, pop_scope, push_scope}; use crate::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; +use crate::codec::anthropic::AnthropicMessagesCodec; use crate::codec::openai_chat::OpenAIChatCodec; -use crate::codec::request::{AnnotatedLlmRequest, Message}; -use crate::codec::traits::LlmCodec; +use crate::codec::openai_responses::OpenAIResponsesCodec; +use crate::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; +use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::error::FlowError; use crate::json::Json; -use crate::{codec::optimization::LlmOptimizationContribution, codec::response::PricingResolver}; +use crate::{ + codec::optimization::LlmOptimizationContribution, + codec::response::{AnnotatedLlmResponse, PricingResolver}, +}; fn reset_global() { let _ = spdlog::init_log_crate_proxy(); @@ -76,6 +87,115 @@ struct ProjectionFailingCodec { projection_attempts: Arc, } +struct RuntimeIdentityCodec; + +impl LlmCodec for RuntimeIdentityCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::Runtime("com.example.chat.v1".into()) + } + + fn decode(&self, request: &LlmRequest) -> crate::error::Result { + OpenAIChatCodec.decode(request) + } + + fn encode( + &self, + annotated: &AnnotatedLlmRequest, + original: &LlmRequest, + ) -> crate::error::Result { + OpenAIChatCodec.encode(annotated, original) + } +} + +impl LlmResponseCodec for RuntimeIdentityCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + LlmCodecIdentity::Runtime("com.example.chat.v1".into()) + } + + fn decode_response(&self, response: &Json) -> crate::error::Result { + OpenAIChatCodec.decode_response(response) + } +} + +#[test] +fn sanitizer_context_preserves_all_codec_identity_states() { + let identity_only_request = crate::api::runtime::LlmSanitizeRequestContext::with_identity( + LlmCodecIdentity::Runtime("identity-only.request.v1".into()), + ); + assert_eq!( + identity_only_request.codec(), + &LlmCodecIdentity::Runtime("identity-only.request.v1".into()) + ); + assert!(identity_only_request.resolve_codec().is_none()); + + let identity_only_response = crate::api::runtime::LlmSanitizeResponseContext::with_identity( + LlmCodecIdentity::Runtime("identity-only.response.v1".into()), + ); + assert_eq!( + identity_only_response.codec(), + &LlmCodecIdentity::Runtime("identity-only.response.v1".into()) + ); + assert!(identity_only_response.resolve_codec().is_none()); + + assert_eq!( + sanitize_context_for_request_codec(None).codec(), + &LlmCodecIdentity::None + ); + assert_eq!( + sanitize_context_for_request_codec(Some(&OpenAIChatCodec)).codec(), + &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) + ); + assert_eq!( + sanitize_context_for_request_codec(Some(&OpenAIResponsesCodec)).codec(), + &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiResponses) + ); + assert_eq!( + sanitize_context_for_request_codec(Some(&AnthropicMessagesCodec)).codec(), + &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) + ); + assert_eq!( + sanitize_context_for_request_codec(Some(&RuntimeIdentityCodec)).codec(), + &LlmCodecIdentity::Runtime("com.example.chat.v1".into()) + ); + assert_eq!( + sanitize_context_for_request_codec(Some(&ProjectionFailingCodec { + projection_attempts: Arc::new(AtomicUsize::new(0)), + })) + .codec(), + &LlmCodecIdentity::Opaque + ); + assert_eq!( + sanitize_context_for_response_codec(Some(&OpenAIChatCodec as &dyn LlmResponseCodec)) + .codec(), + &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) + ); + assert_eq!( + sanitize_context_for_response_codec(Some(&OpenAIResponsesCodec as &dyn LlmResponseCodec)) + .codec(), + &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiResponses) + ); + assert_eq!( + sanitize_context_for_response_codec(Some(&AnthropicMessagesCodec as &dyn LlmResponseCodec)) + .codec(), + &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) + ); + assert_eq!( + sanitize_context_for_response_codec(Some(&RuntimeIdentityCodec)).codec(), + &LlmCodecIdentity::Runtime("com.example.chat.v1".into()) + ); + assert_eq!( + sanitize_context_for_response_codec(None).codec(), + &LlmCodecIdentity::None + ); + assert_eq!( + sanitize_context_for_response_codec(Some(&ProjectionFailingCodec { + projection_attempts: Arc::new(AtomicUsize::new(0)), + })) + .codec(), + &LlmCodecIdentity::Opaque + ); +} + impl LlmCodec for ProjectionFailingCodec { fn decode(&self, request: &LlmRequest) -> crate::error::Result { OpenAIChatCodec.decode(request) @@ -95,6 +215,15 @@ impl LlmCodec for ProjectionFailingCodec { } } +impl LlmResponseCodec for ProjectionFailingCodec { + fn decode_response( + &self, + response: &Json, + ) -> crate::error::Result { + OpenAIChatCodec.decode_response(response) + } +} + fn emit_compaction() { event( EmitMarkEventParams::builder() @@ -104,6 +233,486 @@ fn emit_compaction() { .unwrap(); } +fn secret_request() -> LlmRequest { + LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "SECRET"}] + }), + } +} + +fn redacted_request() -> LlmRequest { + LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "[REDACTED]"}] + }), + } +} + +fn secret_response() -> Json { + json!({ + "id": "chatcmpl-test", + "model": "gpt-4o-mini", + "choices": [{"message": {"role": "assistant", "content": "SECRET"}}] + }) +} + +fn redacted_response() -> Json { + json!({ + "id": "chatcmpl-test", + "model": "gpt-4o-mini", + "choices": [{"message": {"role": "assistant", "content": "[REDACTED]"}}] + }) +} + +#[test] +fn sanitization_invalidates_manual_annotations_without_a_codec() { + let _guard = lock_global_runtime(); + reset_global(); + set_thread_scope_stack(create_scope_stack()); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "manual-annotation-invalidation", + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + register_llm_sanitize_request_guardrail( + "manual-annotation-invalidation-request", + 1, + Arc::new(|_request, _context| Some(redacted_request())), + ) + .unwrap(); + register_llm_sanitize_response_guardrail( + "manual-annotation-invalidation-response", + 1, + Arc::new(|_response, _context| Some(redacted_response())), + ) + .unwrap(); + + let request = secret_request(); + let handle = llm_call( + LlmCallParams::builder() + .name("manual") + .request(&request) + .annotated_request(Arc::new(OpenAIChatCodec.decode(&request).unwrap())) + .build(), + ) + .unwrap(); + let response = secret_response(); + llm_call_end( + LlmCallEndParams::builder() + .handle(&handle) + .response(response.clone()) + .annotated_response(Arc::new( + OpenAIChatCodec.decode_response(&response).unwrap(), + )) + .build(), + ) + .unwrap(); + + flush_subscribers().unwrap(); + let captured = events.lock().unwrap(); + assert_eq!(captured.len(), 2); + assert!( + captured + .iter() + .all(|event| !serde_json::to_string(event).unwrap().contains("SECRET")) + ); + assert!(captured[0].annotated_request().is_none()); + assert!(captured[1].annotated_response().is_none()); + + assert!( + deregister_llm_sanitize_request_guardrail("manual-annotation-invalidation-request") + .unwrap() + ); + assert!( + deregister_llm_sanitize_response_guardrail("manual-annotation-invalidation-response") + .unwrap() + ); + assert!(deregister_subscriber("manual-annotation-invalidation").unwrap()); +} + +#[test] +fn no_op_sanitizers_keep_manual_annotations() { + let _guard = lock_global_runtime(); + reset_global(); + set_thread_scope_stack(create_scope_stack()); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "manual-annotation-noop", + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + register_llm_sanitize_request_guardrail( + "manual-annotation-noop-request", + 1, + Arc::new(|request, _context| Some(request)), + ) + .unwrap(); + register_llm_sanitize_response_guardrail( + "manual-annotation-noop-response", + 1, + Arc::new(|response, _context| Some(response)), + ) + .unwrap(); + + let request = secret_request(); + let handle = llm_call( + LlmCallParams::builder() + .name("manual") + .request(&request) + .annotated_request(Arc::new(OpenAIChatCodec.decode(&request).unwrap())) + .build(), + ) + .unwrap(); + let response = secret_response(); + llm_call_end( + LlmCallEndParams::builder() + .handle(&handle) + .response(response.clone()) + .annotated_response(Arc::new( + OpenAIChatCodec.decode_response(&response).unwrap(), + )) + .build(), + ) + .unwrap(); + + flush_subscribers().unwrap(); + let captured = events.lock().unwrap(); + assert_eq!(captured.len(), 2); + assert!(captured[0].annotated_request().is_some()); + assert!(captured[1].annotated_response().is_some()); + + assert!(deregister_llm_sanitize_request_guardrail("manual-annotation-noop-request").unwrap()); + assert!(deregister_llm_sanitize_response_guardrail("manual-annotation-noop-response").unwrap()); + assert!(deregister_subscriber("manual-annotation-noop").unwrap()); +} + +#[test] +fn sanitization_regenerates_annotations_with_active_codecs() { + let _guard = lock_global_runtime(); + reset_global(); + set_thread_scope_stack(create_scope_stack()); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "active-codec-annotation-regeneration", + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + register_llm_sanitize_request_guardrail( + "active-codec-annotation-regeneration-request", + 1, + Arc::new(|_request, _context| Some(redacted_request())), + ) + .unwrap(); + register_llm_sanitize_response_guardrail( + "active-codec-annotation-regeneration-response", + 1, + Arc::new(|_response, _context| Some(redacted_response())), + ) + .unwrap(); + + let request = secret_request(); + let handle = create_llm_handle( + CreateLlmHandleParams::builder() + .name("manual-with-codec") + .build(), + ) + .unwrap(); + emit_llm_start( + &handle, + &request, + Some(Arc::new(OpenAIChatCodec.decode(&request).unwrap())), + Some(Arc::new(OpenAIChatCodec)), + ) + .unwrap(); + let response = secret_response(); + llm_call_end( + LlmCallEndParams::builder() + .handle(&handle) + .response(response.clone()) + .annotated_response(Arc::new( + OpenAIChatCodec.decode_response(&response).unwrap(), + )) + .response_codec(Arc::new(OpenAIChatCodec)) + .build(), + ) + .unwrap(); + + flush_subscribers().unwrap(); + let captured = events.lock().unwrap(); + assert_eq!(captured.len(), 2); + assert!( + captured + .iter() + .all(|event| !serde_json::to_string(event).unwrap().contains("SECRET")) + ); + assert_eq!( + captured[0].annotated_request().unwrap().messages, + vec![Message::User { + content: MessageContent::Text("[REDACTED]".to_string()), + name: None, + }] + ); + assert_eq!( + captured[1].annotated_response().unwrap().message, + Some(MessageContent::Text("[REDACTED]".to_string())) + ); + + assert!( + deregister_llm_sanitize_request_guardrail("active-codec-annotation-regeneration-request") + .unwrap() + ); + assert!( + deregister_llm_sanitize_response_guardrail("active-codec-annotation-regeneration-response") + .unwrap() + ); + assert!(deregister_subscriber("active-codec-annotation-regeneration").unwrap()); +} + +#[test] +fn buffered_null_fallback_is_sanitized_before_emission() { + let _guard = lock_global_runtime(); + reset_global(); + set_thread_scope_stack(create_scope_stack()); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = Arc::clone(&events); + register_subscriber( + "buffered-null-fallback", + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + + let seen = Arc::new(Mutex::new(Vec::::new())); + let sanitizer_inputs = Arc::clone(&seen); + register_llm_sanitize_response_guardrail( + "buffered-null-fallback-null", + 1, + Arc::new(move |response, _context| { + sanitizer_inputs.lock().unwrap().push(response); + Some(Json::Null) + }), + ) + .unwrap(); + + let handle = create_llm_handle( + CreateLlmHandleParams::builder() + .name("buffered-null-fallback-null") + .build(), + ) + .unwrap(); + let fallback = secret_response(); + llm_call_end( + LlmCallEndParams::builder() + .handle(&handle) + .response(Json::Null) + .data(fallback.clone()) + .annotated_response(Arc::new( + OpenAIChatCodec.decode_response(&fallback).unwrap(), + )) + .response_codec(Arc::new(OpenAIChatCodec)) + .build(), + ) + .unwrap(); + assert!(deregister_llm_sanitize_response_guardrail("buffered-null-fallback-null").unwrap()); + + register_llm_sanitize_response_guardrail( + "buffered-null-fallback-redacted", + 1, + Arc::new(|_response, _context| Some(redacted_response())), + ) + .unwrap(); + let handle = create_llm_handle( + CreateLlmHandleParams::builder() + .name("buffered-null-fallback-redacted") + .build(), + ) + .unwrap(); + llm_call_end( + LlmCallEndParams::builder() + .handle(&handle) + .response(Json::Null) + .data(fallback.clone()) + .annotated_response(Arc::new( + OpenAIChatCodec.decode_response(&fallback).unwrap(), + )) + .response_codec(Arc::new(OpenAIChatCodec)) + .build(), + ) + .unwrap(); + assert!(deregister_llm_sanitize_response_guardrail("buffered-null-fallback-redacted").unwrap()); + + let handle = create_llm_handle( + CreateLlmHandleParams::builder() + .name("buffered-null-without-fallback") + .build(), + ) + .unwrap(); + llm_call_end( + LlmCallEndParams::builder() + .handle(&handle) + .response(Json::Null) + .build(), + ) + .unwrap(); + + flush_subscribers().unwrap(); + let captured = events.lock().unwrap(); + assert_eq!(*seen.lock().unwrap(), vec![fallback]); + assert_eq!(captured.len(), 3); + assert_eq!(captured[0].output(), Some(&Json::Null)); + assert!(captured[0].annotated_response().is_none()); + assert_eq!(captured[1].output(), Some(&redacted_response())); + assert_eq!( + captured[1].annotated_response().unwrap().message, + Some(MessageContent::Text("[REDACTED]".to_string())) + ); + assert!(captured[2].output().is_none()); + assert!(captured[2].annotated_response().is_none()); + assert!( + captured + .iter() + .all(|event| !serde_json::to_string(event).unwrap().contains("SECRET")) + ); + + assert!(deregister_subscriber("buffered-null-fallback").unwrap()); +} + +#[test] +fn streaming_null_fallback_is_sanitized_before_emission() { + let _guard = lock_global_runtime(); + reset_global(); + set_thread_scope_stack(create_scope_stack()); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = Arc::clone(&events); + register_subscriber( + "streaming-null-fallback", + Arc::new(move |event| { + if event.scope_category() == Some(ScopeCategory::End) { + captured.lock().unwrap().push(event.clone()); + } + }), + ) + .unwrap(); + + let seen = Arc::new(Mutex::new(Vec::::new())); + let sanitizer_inputs = Arc::clone(&seen); + register_llm_sanitize_response_guardrail( + "streaming-null-fallback-null", + 1, + Arc::new(move |response, _context| { + sanitizer_inputs.lock().unwrap().push(response); + Some(Json::Null) + }), + ) + .unwrap(); + + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("streaming-null-fallback-null") + .request(request()) + .func(Arc::new(|_request| { + Box::pin(async { Ok(LlmJsonStream::new(tokio_stream::empty())) }) + })) + .collector(Box::new(|_chunk| Ok(()))) + .finalizer(Box::new(|| Json::Null)) + .data(secret_response()) + .response_codec(Arc::new(OpenAIChatCodec)) + .build(), + ) + .await + .unwrap(); + while let Some(chunk) = stream.next().await { + chunk.unwrap(); + } + }); + assert!(deregister_llm_sanitize_response_guardrail("streaming-null-fallback-null").unwrap()); + + register_llm_sanitize_response_guardrail( + "streaming-null-fallback-redacted", + 1, + Arc::new(|_response, _context| Some(redacted_response())), + ) + .unwrap(); + runtime.block_on(async { + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("streaming-null-fallback-redacted") + .request(request()) + .func(Arc::new(|_request| { + Box::pin(async { Ok(LlmJsonStream::new(tokio_stream::empty())) }) + })) + .collector(Box::new(|_chunk| Ok(()))) + .finalizer(Box::new(|| Json::Null)) + .data(secret_response()) + .response_codec(Arc::new(OpenAIChatCodec)) + .build(), + ) + .await + .unwrap(); + while let Some(chunk) = stream.next().await { + chunk.unwrap(); + } + }); + assert!( + deregister_llm_sanitize_response_guardrail("streaming-null-fallback-redacted").unwrap() + ); + + runtime.block_on(async { + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("streaming-null-without-fallback") + .request(request()) + .func(Arc::new(|_request| { + Box::pin(async { Ok(LlmJsonStream::new(tokio_stream::empty())) }) + })) + .collector(Box::new(|_chunk| Ok(()))) + .finalizer(Box::new(|| Json::Null)) + .build(), + ) + .await + .unwrap(); + while let Some(chunk) = stream.next().await { + chunk.unwrap(); + } + }); + + flush_subscribers().unwrap(); + let captured = events.lock().unwrap(); + assert_eq!(*seen.lock().unwrap(), vec![secret_response()]); + assert_eq!(captured.len(), 3); + assert_eq!(captured[0].output(), Some(&Json::Null)); + assert!(captured[0].annotated_response().is_none()); + assert_eq!(captured[1].output(), Some(&redacted_response())); + assert_eq!( + captured[1].annotated_response().unwrap().message, + Some(MessageContent::Text("[REDACTED]".to_string())) + ); + assert!(captured[2].output().is_none()); + assert!(captured[2].annotated_response().is_none()); + assert!( + captured + .iter() + .all(|event| !serde_json::to_string(event).unwrap().contains("SECRET")) + ); + + assert!(deregister_subscriber("streaming-null-fallback").unwrap()); +} + #[test] fn freshness_culls_annotations_and_repeated_compactions_are_idempotent() { let _guard = lock_global_runtime(); @@ -746,6 +1355,99 @@ fn close_boundary_freezes_identical_mark_and_summary_contributions() { ); } +#[test] +fn failed_managed_calls_sanitize_fallback_end_data() { + let _guard = lock_global_runtime(); + reset_global(); + set_thread_scope_stack(create_scope_stack()); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = Arc::clone(&events); + register_subscriber( + "failed-managed-call-sanitization", + Arc::new(move |event| { + if event.scope_category() == Some(ScopeCategory::End) { + captured.lock().unwrap().push(event.clone()); + } + }), + ) + .unwrap(); + + let seen = Arc::new(Mutex::new(Vec::new())); + let sanitizer_inputs = Arc::clone(&seen); + register_llm_sanitize_response_guardrail( + "failed-managed-call-sanitization", + 1, + Arc::new(move |response, context| { + sanitizer_inputs + .lock() + .unwrap() + .push((response, context.codec().clone())); + Some(redacted_response()) + }), + ) + .unwrap(); + + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let buffered_error = llm_call_execute( + LlmCallExecuteParams::builder() + .name("failed-buffered-call") + .request(request()) + .func(Arc::new(|_request| { + Box::pin(async { Err(FlowError::Internal("buffered boom".to_string())) }) + })) + .data(secret_response()) + .response_codec(Arc::new(OpenAIChatCodec)) + .build(), + ) + .await + .unwrap_err(); + assert!(buffered_error.to_string().contains("buffered boom")); + + let stream_error = match llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("failed-stream-call") + .request(request()) + .func(Arc::new(|_request| { + Box::pin(async { Err(FlowError::Internal("stream setup boom".to_string())) }) + })) + .collector(Box::new(|_chunk| Ok(()))) + .finalizer(Box::new(|| Json::Null)) + .data(secret_response()) + .response_codec(Arc::new(OpenAIChatCodec)) + .build(), + ) + .await + { + Ok(_) => panic!("stream setup should fail"), + Err(error) => error, + }; + assert!(stream_error.to_string().contains("stream setup boom")); + }); + + flush_subscribers().unwrap(); + assert!( + deregister_llm_sanitize_response_guardrail("failed-managed-call-sanitization").unwrap() + ); + assert!(deregister_subscriber("failed-managed-call-sanitization").unwrap()); + + let seen = seen.lock().unwrap(); + assert_eq!(seen.len(), 2); + assert!(seen.iter().all(|(response, codec)| { + response == &secret_response() + && codec == &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) + })); + + let events = events.lock().unwrap(); + assert_eq!(events.len(), 2); + assert!(events.iter().all(|event| { + event.output() == Some(&redacted_response()) + && event.annotated_response().is_none() + && !serde_json::to_string(event).unwrap().contains("SECRET") + })); +} + #[test] fn llm_call_execute_adds_otel_status_metadata_to_end_events() { let _guard = lock_global_runtime(); diff --git a/crates/core/tests/unit/native_plugin_tests.rs b/crates/core/tests/unit/native_plugin_tests.rs index 3e3243b16..beec626af 100644 --- a/crates/core/tests/unit/native_plugin_tests.rs +++ b/crates/core/tests/unit/native_plugin_tests.rs @@ -10,11 +10,18 @@ use std::panic::{AssertUnwindSafe, catch_unwind}; use std::sync::atomic::{AtomicUsize, Ordering}; use nemo_relay_plugin::{ - NemoRelayNativeLlmNextFn, NemoRelayNativeLlmStreamNextFn, NemoRelayNativeToolNextFn, + NemoRelayNativeLlmNextFn, NemoRelayNativeLlmSanitizeRequestContext, + NemoRelayNativeLlmSanitizeResponseContext, NemoRelayNativeLlmStreamNextFn, + NemoRelayNativeToolNextFn, }; use serde_json::json; -use crate::api::runtime::{NemoRelayContextState, global_context}; +use crate::api::runtime::{ + BuiltinLlmCodec, LlmSanitizeRequestContext, LlmSanitizeResponseContext, NemoRelayContextState, + global_context, +}; +use crate::codec::openai_chat::OpenAIChatCodec; +use crate::codec::response::AnnotatedLlmResponse; struct ThreadScopeStackRestore(Option); @@ -65,6 +72,50 @@ fn assert_last_error_contains(expected: &str) { ); } +struct FailingNativeCodec; + +impl LlmCodec for FailingNativeCodec { + fn decode(&self, _request: &LlmRequest) -> FlowResult { + Err(FlowError::Internal("request decode rejected".into())) + } + + fn encode( + &self, + _annotated: &AnnotatedLlmRequest, + _original: &LlmRequest, + ) -> FlowResult { + Err(FlowError::Internal("request encode rejected".into())) + } +} + +impl LlmResponseCodec for FailingNativeCodec { + fn decode_response(&self, _response: &Json) -> FlowResult { + Err(FlowError::Internal("response decode rejected".into())) + } +} + +struct PanickingNativeCodec; + +impl LlmCodec for PanickingNativeCodec { + fn decode(&self, _request: &LlmRequest) -> FlowResult { + panic!("request decode panic") + } + + fn encode( + &self, + _annotated: &AnnotatedLlmRequest, + _original: &LlmRequest, + ) -> FlowResult { + panic!("request encode panic") + } +} + +impl LlmResponseCodec for PanickingNativeCodec { + fn decode_response(&self, _response: &Json) -> FlowResult { + panic!("response decode panic") + } +} + #[test] fn native_string_and_json_helpers_cover_abi_boundaries() { clear_native_last_error(); @@ -512,6 +563,7 @@ unsafe extern "C" fn noop_tool_execution( unsafe extern "C" fn noop_llm_request( _user_data: *mut c_void, _request_json: *const NemoRelayNativeString, + _context: NemoRelayNativeLlmSanitizeRequestContext, _out_request_json: *mut *mut NemoRelayNativeString, ) -> NemoRelayStatus { NemoRelayStatus::Ok @@ -520,6 +572,7 @@ unsafe extern "C" fn noop_llm_request( unsafe extern "C" fn noop_json( _user_data: *mut c_void, _payload_json: *const NemoRelayNativeString, + _context: NemoRelayNativeLlmSanitizeResponseContext, _out_json: *mut *mut NemoRelayNativeString, ) -> NemoRelayStatus { NemoRelayStatus::Ok @@ -705,6 +758,239 @@ fn native_registration_entrypoints_reject_null_contexts() { assert_last_error_contains("plugin context is null"); } +#[test] +fn native_codec_operations_report_json_and_codec_failures() { + let openai_request_codec = + NativeHostLlmRequestCodec(Arc::new(OpenAIChatCodec) as Arc); + let openai_response_codec = + NativeHostLlmResponseCodec(Arc::new(OpenAIChatCodec) as Arc); + let failing_request_codec = + NativeHostLlmRequestCodec(Arc::new(FailingNativeCodec) as Arc); + let failing_response_codec = + NativeHostLlmResponseCodec(Arc::new(FailingNativeCodec) as Arc); + let invalid_json = native_string("not-json"); + let mut output = ptr::null_mut(); + + assert_eq!( + unsafe { + native_llm_request_codec_decode( + ptr::from_ref(&openai_request_codec).cast(), + invalid_json, + &mut output, + ) + }, + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("invalid request JSON"); + + assert_eq!( + unsafe { + native_llm_response_codec_decode( + ptr::from_ref(&openai_response_codec).cast(), + invalid_json, + &mut output, + ) + }, + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("invalid response JSON"); + unsafe { native_string_free(invalid_json) }; + + let request = LlmRequest { + headers: Map::new(), + content: json!({ + "model": "gpt-test", + "messages": [{"role": "user", "content": "secret"}] + }), + }; + let request_json = native_string(&serde_json::to_string(&request).unwrap()); + assert_eq!( + unsafe { + native_llm_request_codec_decode( + ptr::from_ref(&failing_request_codec).cast(), + request_json, + &mut output, + ) + }, + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("request decode rejected"); + + let annotated = OpenAIChatCodec.decode(&request).unwrap(); + let annotated_json = native_string(&serde_json::to_string(&annotated).unwrap()); + assert_eq!( + unsafe { + native_llm_request_codec_encode( + ptr::from_ref(&failing_request_codec).cast(), + annotated_json, + request_json, + &mut output, + ) + }, + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("request encode rejected"); + + let response_json = native_string( + r#"{"id":"chatcmpl-test","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"secret"},"finish_reason":"stop"}]}"#, + ); + assert_eq!( + unsafe { + native_llm_response_codec_decode( + ptr::from_ref(&failing_response_codec).cast(), + response_json, + &mut output, + ) + }, + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("response decode rejected"); + + unsafe { + native_string_free(request_json); + native_string_free(annotated_json); + native_string_free(response_json); + } +} + +#[test] +fn native_codec_operations_contain_codec_panics() { + let request_codec = + NativeHostLlmRequestCodec(Arc::new(PanickingNativeCodec) as Arc); + let response_codec = + NativeHostLlmResponseCodec(Arc::new(PanickingNativeCodec) as Arc); + let request = LlmRequest { + headers: Map::new(), + content: json!({ + "model": "gpt-test", + "messages": [{"role": "user", "content": "secret"}] + }), + }; + let request_json = native_string(&serde_json::to_string(&request).unwrap()); + let annotated = OpenAIChatCodec.decode(&request).unwrap(); + let annotated_json = native_string(&serde_json::to_string(&annotated).unwrap()); + let response_json = native_string( + r#"{"id":"chatcmpl-test","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"secret"},"finish_reason":"stop"}]}"#, + ); + let mut output = ptr::null_mut(); + + assert_eq!( + unsafe { + native_llm_request_codec_decode( + ptr::from_ref(&request_codec).cast(), + request_json, + &mut output, + ) + }, + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("request codec decode panicked"); + + assert_eq!( + unsafe { + native_llm_request_codec_encode( + ptr::from_ref(&request_codec).cast(), + annotated_json, + request_json, + &mut output, + ) + }, + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("request codec encode panicked"); + + assert_eq!( + unsafe { + native_llm_response_codec_decode( + ptr::from_ref(&response_codec).cast(), + response_json, + &mut output, + ) + }, + NemoRelayStatus::Internal + ); + assert!(output.is_null()); + assert_last_error_contains("response codec decode panicked"); + + unsafe { + native_string_free(request_json); + native_string_free(annotated_json); + native_string_free(response_json); + } +} + +#[test] +fn native_codec_operations_clear_output_slots_on_null_arguments() { + let request_codec = NativeHostLlmRequestCodec(Arc::new(OpenAIChatCodec) as Arc); + let response_codec = + NativeHostLlmResponseCodec(Arc::new(OpenAIChatCodec) as Arc); + let request_json = native_string( + r#"{"headers":{},"content":{"model":"gpt-test","messages":[{"role":"user","content":"hello"}]}}"#, + ); + + let request_decode_sentinel = native_string("request-decode-sentinel"); + let mut output = request_decode_sentinel; + set_native_last_error("stale request decode error"); + assert_eq!( + unsafe { + native_llm_request_codec_decode( + ptr::from_ref(&request_codec).cast(), + ptr::null(), + &mut output, + ) + }, + NemoRelayStatus::NullPointer + ); + assert!(output.is_null()); + assert_last_error_contains("request codec decode request is null"); + unsafe { native_string_free(request_decode_sentinel) }; + + let request_encode_sentinel = native_string("request-encode-sentinel"); + output = request_encode_sentinel; + set_native_last_error("stale request encode error"); + assert_eq!( + unsafe { + native_llm_request_codec_encode( + ptr::from_ref(&request_codec).cast(), + ptr::null(), + request_json, + &mut output, + ) + }, + NemoRelayStatus::NullPointer + ); + assert!(output.is_null()); + assert_last_error_contains("request codec encode annotated request is null"); + unsafe { native_string_free(request_encode_sentinel) }; + + let response_decode_sentinel = native_string("response-decode-sentinel"); + output = response_decode_sentinel; + set_native_last_error("stale response decode error"); + assert_eq!( + unsafe { + native_llm_response_codec_decode( + ptr::from_ref(&response_codec).cast(), + ptr::null(), + &mut output, + ) + }, + NemoRelayStatus::NullPointer + ); + assert!(output.is_null()); + assert_last_error_contains("response codec decode response is null"); + unsafe { + native_string_free(response_decode_sentinel); + native_string_free(request_json); + } +} + unsafe extern "C" fn tool_json_echo( _user_data: *mut c_void, _name: *const NemoRelayNativeString, @@ -730,6 +1016,7 @@ unsafe extern "C" fn tool_json_error( unsafe extern "C" fn llm_request_echo( _user_data: *mut c_void, request_json: *const NemoRelayNativeString, + _context: NemoRelayNativeLlmSanitizeRequestContext, out_request_json: *mut *mut NemoRelayNativeString, ) -> NemoRelayStatus { let request = read_native_string(request_json).unwrap(); @@ -737,6 +1024,65 @@ unsafe extern "C" fn llm_request_echo( NemoRelayStatus::Ok } +unsafe extern "C" fn llm_request_alias( + _user_data: *mut c_void, + request_json: *const NemoRelayNativeString, + _context: NemoRelayNativeLlmSanitizeRequestContext, + out_request_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { *out_request_json = request_json.cast_mut() }; + NemoRelayStatus::Ok +} + +unsafe extern "C" fn llm_request_codec_round_trip( + _user_data: *mut c_void, + request_json: *const NemoRelayNativeString, + context: NemoRelayNativeLlmSanitizeRequestContext, + out_request_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + assert!(!context.codec.is_null()); + let mut annotated = ptr::null_mut(); + let status = + unsafe { native_llm_request_codec_decode(context.codec, request_json, &mut annotated) }; + if status != NemoRelayStatus::Ok { + return status; + } + let status = unsafe { + native_llm_request_codec_encode(context.codec, annotated, request_json, out_request_json) + }; + unsafe { native_string_free(annotated) }; + status +} + +unsafe extern "C" fn llm_response_codec_decode_and_echo( + _user_data: *mut c_void, + response_json: *const NemoRelayNativeString, + context: NemoRelayNativeLlmSanitizeResponseContext, + out_response_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + assert!(!context.codec.is_null()); + let mut annotated = ptr::null_mut(); + let status = + unsafe { native_llm_response_codec_decode(context.codec, response_json, &mut annotated) }; + if status != NemoRelayStatus::Ok { + return status; + } + unsafe { native_string_free(annotated) }; + let response = read_native_string(response_json).unwrap(); + unsafe { *out_response_json = native_string(&response) }; + NemoRelayStatus::Ok +} + +unsafe extern "C" fn llm_response_alias( + _user_data: *mut c_void, + response_json: *const NemoRelayNativeString, + _context: NemoRelayNativeLlmSanitizeResponseContext, + out_response_json: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { *out_response_json = response_json.cast_mut() }; + NemoRelayStatus::Ok +} + #[test] fn native_callback_helpers_cover_success_error_and_invalid_output() { assert_eq!( @@ -755,9 +1101,172 @@ fn native_callback_helpers_cover_success_error_and_invalid_output() { content: json!({"model": "test"}), }; assert_eq!( - call_llm_request_callback(llm_request_echo, ptr::null_mut(), &request).unwrap(), - request + call_llm_sanitize_request_callback( + llm_request_echo, + ptr::null_mut(), + &request, + LlmSanitizeRequestContext::default(), + ) + .unwrap(), + Some(request) + ); + + let request = LlmRequest { + headers: Map::new(), + content: json!({"model": "alias"}), + }; + assert_eq!( + call_llm_sanitize_request_callback( + llm_request_alias, + ptr::null_mut(), + &request, + LlmSanitizeRequestContext::default(), + ) + .unwrap(), + Some(request) + ); + + let response = json!({"message": "alias"}); + assert_eq!( + call_llm_sanitize_response_callback( + llm_response_alias, + ptr::null_mut(), + &response, + LlmSanitizeResponseContext::default(), + ) + .unwrap(), + Some(response) + ); +} + +#[test] +fn native_sanitizer_context_resolves_directional_codecs() { + let codec = Arc::new(OpenAIChatCodec); + let request = LlmRequest { + headers: Map::new(), + content: json!({ + "model": "gpt-test", + "messages": [{"role": "user", "content": "secret"}], + "preserve": true + }), + }; + let sanitized = call_llm_sanitize_request_callback( + llm_request_codec_round_trip, + ptr::null_mut(), + &request, + LlmSanitizeRequestContext::for_request_codec(Some(codec.clone())), + ) + .unwrap() + .expect("native request sanitizer returns a request"); + assert_eq!(sanitized, request); + + let response = json!({ + "id": "chatcmpl-test", + "model": "gpt-test", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "secret"}, + "finish_reason": "stop" + }] + }); + let sanitized = call_llm_sanitize_response_callback( + llm_response_codec_decode_and_echo, + ptr::null_mut(), + &response, + LlmSanitizeResponseContext::for_response_codec(Some(codec)), + ) + .unwrap() + .expect("native response sanitizer returns a response"); + assert_eq!(sanitized, response); +} + +#[test] +fn native_llm_sanitize_context_preserves_all_codec_identity_states() { + let cases = [ + ( + LlmCodecIdentity::None, + NemoRelayNativeLlmCodecKind::None, + None, + ), + ( + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat), + NemoRelayNativeLlmCodecKind::BuiltIn, + Some("openai_chat"), + ), + ( + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiResponses), + NemoRelayNativeLlmCodecKind::BuiltIn, + Some("openai_responses"), + ), + ( + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages), + NemoRelayNativeLlmCodecKind::BuiltIn, + Some("anthropic_messages"), + ), + ( + LlmCodecIdentity::Runtime("com.example.chat.v1".into()), + NemoRelayNativeLlmCodecKind::Runtime, + Some("com.example.chat.v1"), + ), + ( + LlmCodecIdentity::Opaque, + NemoRelayNativeLlmCodecKind::Opaque, + None, + ), + ]; + + for (codec, expected_kind, expected_id) in cases { + let (codec_kind, context_id) = + native_llm_codec_identity(&codec).expect("native context conversion should succeed"); + assert_eq!(codec_kind, expected_kind); + assert_eq!( + context_id.map(|id| read_native_string(id).unwrap()), + expected_id.map(str::to_owned), + ); + if let Some(context_id) = context_id { + unsafe { native_string_free(context_id) }; + } + } +} + +#[test] +fn native_llm_sanitizer_input_allocation_failures_release_codec_ids() { + let request = LlmRequest { + headers: Map::new(), + content: json!({"model": "test"}), + }; + let identity = LlmCodecIdentity::Runtime("com.example.chat.v1".into()); + let live_before = native_string_live_allocations(); + + fail_native_string_allocation_after(1); + let request_error = call_llm_sanitize_request_callback( + llm_request_alias, + ptr::null_mut(), + &request, + LlmSanitizeRequestContext::with_identity(identity.clone()), + ) + .unwrap_err(); + assert!( + request_error + .to_string() + .contains("failed to allocate native LLM request") + ); + assert_eq!(native_string_live_allocations(), live_before); + + fail_native_string_allocation_after(1); + let response_error = call_llm_sanitize_response_callback( + llm_response_alias, + ptr::null_mut(), + &json!({"model": "test"}), + LlmSanitizeResponseContext::with_identity(identity), + ) + .unwrap_err(); + assert!( + response_error + .to_string() + .contains("failed to allocate native LLM response") ); + assert_eq!(native_string_live_allocations(), live_before); } fn tool_next(output: FlowResult) -> ToolExecutionNextFn { diff --git a/crates/core/tests/unit/plugin_tests.rs b/crates/core/tests/unit/plugin_tests.rs index f3d1824a9..b39abb9d2 100644 --- a/crates/core/tests/unit/plugin_tests.rs +++ b/crates/core/tests/unit/plugin_tests.rs @@ -1511,13 +1511,13 @@ fn test_plugin_registration_context_supports_guardrail_helpers() { ctx.register_llm_sanitize_request_guardrail( "llm_sanitize_request", 1, - Arc::new(|request| request), + Arc::new(|request, _context| Some(request)), ) .unwrap(); ctx.register_llm_sanitize_response_guardrail( "llm_sanitize_response", 1, - Arc::new(|response| response), + Arc::new(|response, _context| Some(response)), ) .unwrap(); ctx.register_llm_conditional_execution_guardrail( @@ -1651,14 +1651,14 @@ fn test_plugin_registration_context_maps_duplicate_registration_errors() { ctx.register_llm_sanitize_request_guardrail( "llm-sanitize-request", 1, - Arc::new(|request| request), + Arc::new(|request, _context| Some(request)), ) .unwrap(); expect_registration_failed( ctx.register_llm_sanitize_request_guardrail( "llm-sanitize-request", 1, - Arc::new(|request| request), + Arc::new(|request, _context| Some(request)), ), "llm sanitize request guardrail:", ); @@ -1666,14 +1666,14 @@ fn test_plugin_registration_context_maps_duplicate_registration_errors() { ctx.register_llm_sanitize_response_guardrail( "llm-sanitize-response", 1, - Arc::new(|response| response), + Arc::new(|response, _context| Some(response)), ) .unwrap(); expect_registration_failed( ctx.register_llm_sanitize_response_guardrail( "llm-sanitize-response", 1, - Arc::new(|response| response), + Arc::new(|response, _context| Some(response)), ), "llm sanitize response guardrail:", ); @@ -1815,13 +1815,13 @@ fn test_plugin_registration_context_maps_deregistration_errors() { ctx.register_llm_sanitize_request_guardrail( "llm-sanitize-request", 1, - Arc::new(|request| request), + Arc::new(|request, _context| Some(request)), ) .unwrap(); ctx.register_llm_sanitize_response_guardrail( "llm-sanitize-response", 1, - Arc::new(|response| response), + Arc::new(|response, _context| Some(response)), ) .unwrap(); ctx.register_llm_conditional_execution_guardrail("llm-conditional", 1, Arc::new(|_| Ok(None))) diff --git a/crates/ffi/nemo_relay.h b/crates/ffi/nemo_relay.h index f5acfa076..3c381015a 100644 --- a/crates/ffi/nemo_relay.h +++ b/crates/ffi/nemo_relay.h @@ -58,6 +58,29 @@ enum NemoRelayStatus { }; typedef int32_t NemoRelayStatus; +/** + * Codec identity kind supplied to an LLM sanitizer. + */ +enum NemoRelayLlmSanitizeCodecKind { + /** + * No codec was active. + */ + NEMO_RELAY_LLM_SANITIZE_CODEC_KIND_NONE = 0, + /** + * A Relay built-in codec was active. + */ + NEMO_RELAY_LLM_SANITIZE_CODEC_KIND_BUILT_IN = 1, + /** + * A runtime-registered codec was active. + */ + NEMO_RELAY_LLM_SANITIZE_CODEC_KIND_RUNTIME = 2, + /** + * A codec was active but has no registered identity. + */ + NEMO_RELAY_LLM_SANITIZE_CODEC_KIND_OPAQUE = 3, +}; +typedef uint32_t NemoRelayLlmSanitizeCodecKind; + /** * The type of scope in the agent execution hierarchy. */ @@ -149,6 +172,16 @@ typedef struct FfiLLMHandle FfiLLMHandle; */ typedef struct FfiLLMRequest FfiLLMRequest; +/** + * Borrowed, callback-scoped request codec capability supplied to an LLM sanitizer. + */ +typedef struct FfiLlmSanitizeRequestCodec FfiLlmSanitizeRequestCodec; + +/** + * Borrowed, callback-scoped response codec capability supplied to an LLM sanitizer. + */ +typedef struct FfiLlmSanitizeResponseCodec FfiLlmSanitizeResponseCodec; + /** * Opaque OpenInference subscriber handle. */ @@ -248,17 +281,61 @@ typedef char *(*NemoRelayCodecEncodeFn)(void *user_data, const struct FfiLLMRequest *original_request); /** - * Callback for LLM request sanitization. Receives an `FfiLLMRequest` and returns - * a new (possibly modified) `FfiLLMRequest`. Return null to use defaults. + * Codec identity supplied to an LLM sanitizer. `codec_id` is null for + * `None` and `Opaque`, and is valid only for the duration of the callback. */ -typedef struct FfiLLMRequest *(*NemoRelayLlmRequestCb)(void *user_data, - const struct FfiLLMRequest *request); +typedef struct NemoRelayLlmSanitizeRequestContext { + /** + * Kind of active codec identity. + */ + NemoRelayLlmSanitizeCodecKind codec_kind; + /** + * Built-in or runtime codec ID, when applicable. + */ + const char *codec_id; + /** + * Borrowed request codec capability, or null when no codec is active. + */ + const struct FfiLlmSanitizeRequestCodec *codec; +} NemoRelayLlmSanitizeRequestContext; /** - * Generic JSON-to-JSON callback, used for LLM response sanitization and intercepts. - * The returned string must be allocated with `malloc` or equivalent. + * LLM request sanitizer. It receives the request first and its codec context + * second. Return null to omit the observability payload. The request is + * borrowed, but returning that same pointer is supported as a pass-through. + * Any other non-null result transfers ownership to Relay. */ -typedef char *(*NemoRelayJsonCb)(void *user_data, const char *json); +typedef struct FfiLLMRequest *(*NemoRelayLlmSanitizeRequestCb)(void *user_data, + const struct FfiLLMRequest *request, + struct NemoRelayLlmSanitizeRequestContext context); + +/** + * Directional codec context supplied to an LLM response sanitizer. + */ +typedef struct NemoRelayLlmSanitizeResponseContext { + /** + * Kind of active codec identity. + */ + NemoRelayLlmSanitizeCodecKind codec_kind; + /** + * Built-in or runtime codec ID, when applicable. + */ + const char *codec_id; + /** + * Borrowed response codec capability, or null when no codec is active. + */ + const struct FfiLlmSanitizeResponseCodec *codec; +} NemoRelayLlmSanitizeResponseContext; + +/** + * LLM response sanitizer. It receives response JSON first and its codec + * context second. Return null to omit the observability payload. The response + * is borrowed, but returning that same pointer is supported as a pass-through. + * Any other non-null result transfers ownership to Relay. + */ +typedef char *(*NemoRelayLlmSanitizeResponseCb)(void *user_data, + const char *response_json, + struct NemoRelayLlmSanitizeResponseContext context); /** * Callback for LLM conditional execution guardrails. @@ -703,6 +780,41 @@ NemoRelayStatus nemo_relay_scope_register_scope_sanitize_end_guardrail(const cha NemoRelayStatus nemo_relay_scope_deregister_scope_sanitize_end_guardrail(const char *scope_uuid, const char *name); +/** + * Decode a request through a callback-scoped sanitizer codec capability. + * + * The returned JSON string must be freed with `nemo_relay_string_free`. + * + * # Safety + * Both pointers must be non-null and valid only during the sanitizer callback. + */ +char *nemo_relay_llm_sanitize_request_codec_decode(const struct FfiLlmSanitizeRequestCodec *codec, + const struct FfiLLMRequest *request); + +/** + * Encode normalized request changes through a callback-scoped codec capability. + * + * The returned request is owned by the caller and must be freed with + * `nemo_relay_llm_request_free`. Returns null on failure. + * + * # Safety + * All pointers must be non-null and valid only during the sanitizer callback. + */ +struct FfiLLMRequest *nemo_relay_llm_sanitize_request_codec_encode(const struct FfiLlmSanitizeRequestCodec *codec, + const char *annotated_json, + const struct FfiLLMRequest *original); + +/** + * Decode a response through a callback-scoped sanitizer codec capability. + * + * The returned JSON string must be freed with `nemo_relay_string_free`. + * + * # Safety + * All pointers must be non-null and valid only during the sanitizer callback. + */ +char *nemo_relay_llm_sanitize_response_codec_decode(const struct FfiLlmSanitizeResponseCodec *codec, + const char *response_json); + /** * Begin a manual LLM call lifecycle span. * @@ -942,8 +1054,8 @@ int32_t nemo_relay_stream_next(struct FfiStream *stream, char **out_chunk); void nemo_relay_stream_free(struct FfiStream *stream); /** - * Register an LLM request sanitization guardrail. The callback can modify or - * replace the LLM request before it is sent. + * Register an LLM request sanitizer. The callback receives the emitted + * request first and per-call codec context second; null omits observability. * * # Parameters * - `name`: Unique guardrail name. @@ -957,7 +1069,7 @@ void nemo_relay_stream_free(struct FfiStream *stream); */ NemoRelayStatus nemo_relay_register_llm_sanitize_request_guardrail(const char *name, int32_t priority, - NemoRelayLlmRequestCb cb, + NemoRelayLlmSanitizeRequestCb cb, void *user_data, NemoRelayFreeFn free_fn); @@ -985,7 +1097,7 @@ NemoRelayStatus nemo_relay_deregister_llm_sanitize_request_guardrail(const char */ NemoRelayStatus nemo_relay_register_llm_sanitize_response_guardrail(const char *name, int32_t priority, - NemoRelayJsonCb cb, + NemoRelayLlmSanitizeResponseCb cb, void *user_data, NemoRelayFreeFn free_fn); @@ -1704,7 +1816,7 @@ NemoRelayStatus nemo_relay_plugin_context_register_tool_conditional_execution_gu NemoRelayStatus nemo_relay_plugin_context_register_llm_sanitize_request_guardrail(struct FfiPluginContext *ctx, const char *name, int32_t priority, - NemoRelayLlmRequestCb cb, + NemoRelayLlmSanitizeRequestCb cb, void *user_data, NemoRelayFreeFn free_fn); @@ -1718,7 +1830,7 @@ NemoRelayStatus nemo_relay_plugin_context_register_llm_sanitize_request_guardrai NemoRelayStatus nemo_relay_plugin_context_register_llm_sanitize_response_guardrail(struct FfiPluginContext *ctx, const char *name, int32_t priority, - NemoRelayJsonCb cb, + NemoRelayLlmSanitizeResponseCb cb, void *user_data, NemoRelayFreeFn free_fn); @@ -2007,7 +2119,7 @@ NemoRelayStatus nemo_relay_scope_deregister_tool_execution_intercept(const char NemoRelayStatus nemo_relay_scope_register_llm_sanitize_request_guardrail(const char *scope_uuid, const char *name, int32_t priority, - NemoRelayLlmRequestCb cb, + NemoRelayLlmSanitizeRequestCb cb, void *user_data, NemoRelayFreeFn free_fn); @@ -2037,7 +2149,7 @@ NemoRelayStatus nemo_relay_scope_deregister_llm_sanitize_request_guardrail(const NemoRelayStatus nemo_relay_scope_register_llm_sanitize_response_guardrail(const char *scope_uuid, const char *name, int32_t priority, - NemoRelayJsonCb cb, + NemoRelayLlmSanitizeResponseCb cb, void *user_data, NemoRelayFreeFn free_fn); diff --git a/crates/ffi/src/api/llm.rs b/crates/ffi/src/api/llm.rs index 2cacc1c7b..62060d846 100644 --- a/crates/ffi/src/api/llm.rs +++ b/crates/ffi/src/api/llm.rs @@ -2,17 +2,141 @@ // SPDX-License-Identifier: Apache-2.0 use super::{ - Arc, FfiCodecHandle, FfiLLMHandle, FfiScopeHandle, FlowResult, LlmAttributes, - LlmExecutionNextFn, LlmRequest, LlmStreamExecutionNextFn, NemoRelayCodecDecodeFn, - NemoRelayCodecEncodeFn, NemoRelayCollectorCb, NemoRelayFinalizerCb, NemoRelayFreeFn, - NemoRelayLlmExecCb, NemoRelayStatus, TASK_SCOPE_STACK, c_char, c_str_to_json, - c_str_to_opt_json, c_str_to_string, clear_last_error, core_llm_api, current_scope_stack, - json_to_c_string, set_last_error, status_from_error, tokio_runtime, - unix_micros_to_opt_timestamp, wrap_codec_fn, wrap_collector_fn, wrap_finalizer_fn, - wrap_llm_exec_fn, wrap_llm_stream_exec_fn, + Arc, FfiCodecHandle, FfiLLMHandle, FfiLLMRequest, FfiLlmSanitizeRequestCodec, + FfiLlmSanitizeResponseCodec, FfiScopeHandle, FlowResult, LlmAttributes, LlmExecutionNextFn, + LlmRequest, LlmStreamExecutionNextFn, NemoRelayCodecDecodeFn, NemoRelayCodecEncodeFn, + NemoRelayCollectorCb, NemoRelayFinalizerCb, NemoRelayFreeFn, NemoRelayLlmExecCb, + NemoRelayStatus, TASK_SCOPE_STACK, c_char, c_str_to_json, c_str_to_opt_json, c_str_to_string, + clear_last_error, core_llm_api, current_scope_stack, json_to_c_string, set_last_error, + status_from_error, tokio_runtime, unix_micros_to_opt_timestamp, wrap_codec_fn, + wrap_collector_fn, wrap_finalizer_fn, wrap_llm_exec_fn, wrap_llm_stream_exec_fn, }; +use std::panic::{AssertUnwindSafe, catch_unwind}; use tokio_stream::StreamExt; +/// Decode a request through a callback-scoped sanitizer codec capability. +/// +/// The returned JSON string must be freed with `nemo_relay_string_free`. +/// +/// # Safety +/// Both pointers must be non-null and valid only during the sanitizer callback. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_llm_sanitize_request_codec_decode( + codec: *const FfiLlmSanitizeRequestCodec, + request: *const FfiLLMRequest, +) -> *mut c_char { + clear_last_error(); + if codec.is_null() || request.is_null() { + set_last_error("null sanitizer request codec argument"); + return std::ptr::null_mut(); + } + let result = catch_unwind(AssertUnwindSafe( + || -> std::result::Result<*mut c_char, String> { + let annotated = unsafe { &*codec } + .0 + .decode(&unsafe { &*request }.0) + .map_err(|error| error.to_string())?; + let value = serde_json::to_value(annotated).map_err(|error| error.to_string())?; + Ok(json_to_c_string(&value)) + }, + )); + match result { + Ok(Ok(value)) => value, + Ok(Err(error)) => { + set_last_error(&error); + std::ptr::null_mut() + } + Err(_) => { + set_last_error("sanitizer request codec decode panicked"); + std::ptr::null_mut() + } + } +} + +/// Encode normalized request changes through a callback-scoped codec capability. +/// +/// The returned request is owned by the caller and must be freed with +/// `nemo_relay_llm_request_free`. Returns null on failure. +/// +/// # Safety +/// All pointers must be non-null and valid only during the sanitizer callback. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_llm_sanitize_request_codec_encode( + codec: *const FfiLlmSanitizeRequestCodec, + annotated_json: *const c_char, + original: *const FfiLLMRequest, +) -> *mut FfiLLMRequest { + clear_last_error(); + if codec.is_null() || annotated_json.is_null() || original.is_null() { + set_last_error("null sanitizer request codec argument"); + return std::ptr::null_mut(); + } + let result = catch_unwind(AssertUnwindSafe( + || -> std::result::Result<*mut FfiLLMRequest, String> { + let annotated = c_str_to_json(annotated_json) + .and_then(|value| serde_json::from_value(value).ok()) + .ok_or_else(|| "invalid annotated request JSON".to_string())?; + let request = unsafe { &*codec } + .0 + .encode(&annotated, &unsafe { &*original }.0) + .map_err(|error| error.to_string())?; + Ok(Box::into_raw(Box::new(FfiLLMRequest(request)))) + }, + )); + match result { + Ok(Ok(value)) => value, + Ok(Err(error)) => { + set_last_error(&error); + std::ptr::null_mut() + } + Err(_) => { + set_last_error("sanitizer request codec encode panicked"); + std::ptr::null_mut() + } + } +} + +/// Decode a response through a callback-scoped sanitizer codec capability. +/// +/// The returned JSON string must be freed with `nemo_relay_string_free`. +/// +/// # Safety +/// All pointers must be non-null and valid only during the sanitizer callback. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nemo_relay_llm_sanitize_response_codec_decode( + codec: *const FfiLlmSanitizeResponseCodec, + response_json: *const c_char, +) -> *mut c_char { + clear_last_error(); + if codec.is_null() || response_json.is_null() { + set_last_error("null sanitizer response codec argument"); + return std::ptr::null_mut(); + } + let result = catch_unwind(AssertUnwindSafe( + || -> std::result::Result<*mut c_char, String> { + let response = + c_str_to_json(response_json).ok_or_else(|| "invalid response JSON".to_string())?; + let annotated = unsafe { &*codec } + .0 + .decode_response(&response) + .map_err(|error| error.to_string())?; + let value = serde_json::to_value(annotated).map_err(|error| error.to_string())?; + Ok(json_to_c_string(&value)) + }, + )); + match result { + Ok(Ok(value)) => value, + Ok(Err(error)) => { + set_last_error(&error); + std::ptr::null_mut() + } + Err(_) => { + set_last_error("sanitizer response codec decode panicked"); + std::ptr::null_mut() + } + } +} + // --------------------------------------------------------------------------- // LLM lifecycle // --------------------------------------------------------------------------- diff --git a/crates/ffi/src/api/llm_registry.rs b/crates/ffi/src/api/llm_registry.rs index 1d53e4ce4..38883ecee 100644 --- a/crates/ffi/src/api/llm_registry.rs +++ b/crates/ffi/src/api/llm_registry.rs @@ -2,20 +2,20 @@ // SPDX-License-Identifier: Apache-2.0 use super::{ - NemoRelayEventSubscriberCb, NemoRelayFreeFn, NemoRelayJsonCb, NemoRelayLlmConditionalCb, - NemoRelayLlmExecInterceptCb, NemoRelayLlmRequestCb, NemoRelayLlmRequestInterceptCb, - NemoRelayStatus, c_char, c_str_to_string, clear_last_error, core_registry_api, - core_subscriber_api, status_from_error, wrap_event_subscriber, wrap_llm_conditional_fn, - wrap_llm_exec_intercept_fn, wrap_llm_request_intercept_fn, wrap_llm_response_fn, - wrap_llm_sanitize_request_fn, wrap_llm_stream_exec_intercept_fn, + NemoRelayEventSubscriberCb, NemoRelayFreeFn, NemoRelayLlmConditionalCb, + NemoRelayLlmExecInterceptCb, NemoRelayLlmRequestInterceptCb, NemoRelayLlmSanitizeRequestCb, + NemoRelayLlmSanitizeResponseCb, NemoRelayStatus, c_char, c_str_to_string, clear_last_error, + core_registry_api, core_subscriber_api, 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, }; // --------------------------------------------------------------------------- // LLM guardrail registrations // --------------------------------------------------------------------------- -/// Register an LLM request sanitization guardrail. The callback can modify or -/// replace the LLM request before it is sent. +/// Register an LLM request sanitizer. The callback receives the emitted +/// request first and per-call codec context second; null omits observability. /// /// # Parameters /// - `name`: Unique guardrail name. @@ -30,7 +30,7 @@ use super::{ pub unsafe extern "C" fn nemo_relay_register_llm_sanitize_request_guardrail( name: *const c_char, priority: i32, - cb: NemoRelayLlmRequestCb, + cb: NemoRelayLlmSanitizeRequestCb, user_data: *mut libc::c_void, free_fn: NemoRelayFreeFn, ) -> NemoRelayStatus { @@ -81,7 +81,7 @@ pub unsafe extern "C" fn nemo_relay_deregister_llm_sanitize_request_guardrail( pub unsafe extern "C" fn nemo_relay_register_llm_sanitize_response_guardrail( name: *const c_char, priority: i32, - cb: NemoRelayJsonCb, + cb: NemoRelayLlmSanitizeResponseCb, user_data: *mut libc::c_void, free_fn: NemoRelayFreeFn, ) -> NemoRelayStatus { @@ -90,7 +90,7 @@ pub unsafe extern "C" fn nemo_relay_register_llm_sanitize_response_guardrail( Ok(s) => s, Err(status) => return status, }; - let wrapped = wrap_llm_response_fn(cb, user_data, free_fn); + let wrapped = wrap_llm_sanitize_response_fn(cb, user_data, free_fn); match core_registry_api::register_llm_sanitize_response_guardrail(&name, priority, wrapped) { Ok(()) => NemoRelayStatus::Ok, Err(e) => status_from_error(&e), diff --git a/crates/ffi/src/api/mod.rs b/crates/ffi/src/api/mod.rs index 7759b6e3b..42f40bdca 100644 --- a/crates/ffi/src/api/mod.rs +++ b/crates/ffi/src/api/mod.rs @@ -15,14 +15,14 @@ use std::time::Duration; use crate::callable::{ NemoRelayCodecDecodeFn, NemoRelayCodecEncodeFn, NemoRelayCollectorCb, NemoRelayEventSanitizeCb, - NemoRelayEventSubscriberCb, NemoRelayFinalizerCb, NemoRelayFreeFn, NemoRelayJsonCb, - NemoRelayLlmConditionalCb, NemoRelayLlmExecCb, NemoRelayLlmExecInterceptCb, - NemoRelayLlmRequestCb, NemoRelayLlmRequestInterceptCb, NemoRelayPluginRegisterCb, + NemoRelayEventSubscriberCb, NemoRelayFinalizerCb, NemoRelayFreeFn, NemoRelayLlmConditionalCb, + 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_response_fn, wrap_llm_sanitize_request_fn, wrap_llm_stream_exec_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, }; @@ -36,8 +36,9 @@ use crate::error::{ }; use crate::types::{ FfiAtifExporter, FfiAtofExporter, FfiCodecHandle, FfiLLMHandle, FfiLLMRequest, - FfiOpenInferenceSubscriber, FfiOpenTelemetrySubscriber, FfiPluginActivation, FfiPluginContext, - FfiScopeHandle, FfiScopeStack, FfiThreadScopeStackBinding, FfiToolHandle, NemoRelayScopeType, + FfiLlmSanitizeRequestCodec, FfiLlmSanitizeResponseCodec, FfiOpenInferenceSubscriber, + FfiOpenTelemetrySubscriber, FfiPluginActivation, FfiPluginContext, FfiScopeHandle, + FfiScopeStack, FfiThreadScopeStackBinding, FfiToolHandle, NemoRelayScopeType, }; pub use crate::types::{nemo_relay_openinference_subscriber_free, nemo_relay_otel_subscriber_free}; use libc::c_char; diff --git a/crates/ffi/src/api/plugin.rs b/crates/ffi/src/api/plugin.rs index 9130a2f94..09cd151f8 100644 --- a/crates/ffi/src/api/plugin.rs +++ b/crates/ffi/src/api/plugin.rs @@ -4,19 +4,19 @@ use super::{ Arc, CStr, ConfigDiagnostic, DiagnosticLevel, DynamicPluginActivationSpec, FfiPluginActivation, FfiPluginContext, Future, NemoRelayEventSanitizeCb, NemoRelayEventSubscriberCb, - NemoRelayFreeFn, NemoRelayJsonCb, NemoRelayLlmConditionalCb, NemoRelayLlmExecInterceptCb, - NemoRelayLlmRequestCb, NemoRelayLlmRequestInterceptCb, 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_response_fn, wrap_llm_sanitize_request_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, + 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, }; use crate::api::event_registry::Surface; use nemo_relay_pii_redaction::component::register_pii_redaction_component; @@ -731,7 +731,7 @@ pub unsafe extern "C" fn nemo_relay_plugin_context_register_llm_sanitize_request ctx: *mut FfiPluginContext, name: *const c_char, priority: i32, - cb: NemoRelayLlmRequestCb, + cb: NemoRelayLlmSanitizeRequestCb, user_data: *mut libc::c_void, free_fn: NemoRelayFreeFn, ) -> NemoRelayStatus { @@ -763,7 +763,7 @@ pub unsafe extern "C" fn nemo_relay_plugin_context_register_llm_sanitize_respons ctx: *mut FfiPluginContext, name: *const c_char, priority: i32, - cb: NemoRelayJsonCb, + cb: NemoRelayLlmSanitizeResponseCb, user_data: *mut libc::c_void, free_fn: NemoRelayFreeFn, ) -> NemoRelayStatus { @@ -776,7 +776,7 @@ pub unsafe extern "C" fn nemo_relay_plugin_context_register_llm_sanitize_respons Ok(value) => value, Err(status) => return status, }; - let wrapped = wrap_llm_response_fn(cb, user_data, free_fn); + let wrapped = wrap_llm_sanitize_response_fn(cb, user_data, free_fn); match unsafe { &mut *((*ctx).0) } .register_llm_sanitize_response_guardrail(&name, priority, wrapped) { diff --git a/crates/ffi/src/api/scope_registry.rs b/crates/ffi/src/api/scope_registry.rs index 0b2887155..50efd644b 100644 --- a/crates/ffi/src/api/scope_registry.rs +++ b/crates/ffi/src/api/scope_registry.rs @@ -2,15 +2,15 @@ // SPDX-License-Identifier: Apache-2.0 use super::{ - NemoRelayEventSubscriberCb, NemoRelayFreeFn, NemoRelayJsonCb, NemoRelayLlmConditionalCb, - NemoRelayLlmExecInterceptCb, NemoRelayLlmRequestCb, NemoRelayLlmRequestInterceptCb, - 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_response_fn, wrap_llm_sanitize_request_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, + 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, }; // --------------------------------------------------------------------------- @@ -363,7 +363,7 @@ pub unsafe extern "C" fn nemo_relay_scope_register_llm_sanitize_request_guardrai scope_uuid: *const c_char, name: *const c_char, priority: i32, - cb: NemoRelayLlmRequestCb, + cb: NemoRelayLlmSanitizeRequestCb, user_data: *mut libc::c_void, free_fn: NemoRelayFreeFn, ) -> NemoRelayStatus { @@ -426,7 +426,7 @@ pub unsafe extern "C" fn nemo_relay_scope_register_llm_sanitize_response_guardra scope_uuid: *const c_char, name: *const c_char, priority: i32, - cb: NemoRelayJsonCb, + cb: NemoRelayLlmSanitizeResponseCb, user_data: *mut libc::c_void, free_fn: NemoRelayFreeFn, ) -> NemoRelayStatus { @@ -439,7 +439,7 @@ pub unsafe extern "C" fn nemo_relay_scope_register_llm_sanitize_response_guardra Ok(s) => s, Err(status) => return status, }; - let wrapped = wrap_llm_response_fn(cb, user_data, free_fn); + let wrapped = wrap_llm_sanitize_response_fn(cb, user_data, free_fn); match core_registry_api::scope_register_llm_sanitize_response_guardrail( &uuid, &name, priority, wrapped, ) { diff --git a/crates/ffi/src/callable.rs b/crates/ffi/src/callable.rs index 3b968a193..f327e01d2 100644 --- a/crates/ffi/src/callable.rs +++ b/crates/ffi/src/callable.rs @@ -23,9 +23,10 @@ use std::sync::Arc; use libc::c_char; use nemo_relay::api::runtime::{ - EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, LlmJsonStream, - LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, - ToolConditionalFn, ToolExecutionFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, + EventSanitizeFn, EventSubscriberFn, LlmCodecIdentity, LlmConditionalFn, LlmExecutionNextFn, + LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn, + LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, ToolConditionalFn, + ToolExecutionFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; use serde_json::Value as Json; use tokio_stream::StreamExt; @@ -37,7 +38,7 @@ use nemo_relay::codec::request::AnnotatedLlmRequest as AnnotatedLLMRequest; use nemo_relay::codec::traits::LlmCodec; use nemo_relay::error::{FlowError, Result}; -use crate::convert::json_to_c_string; +use crate::convert::{c_str_to_json, json_to_c_string}; use crate::error::{NemoRelayStatus, clear_last_error, last_error_message, set_last_error}; use crate::types::{FfiEvent, FfiLLMRequest, FfiPluginContext}; @@ -97,18 +98,63 @@ pub type NemoRelayToolExecInterceptCb = unsafe extern "C" fn( next_ctx: *mut libc::c_void, ) -> *mut c_char; -/// Generic JSON-to-JSON callback, used for LLM response sanitization and intercepts. -/// The returned string must be allocated with `malloc` or equivalent. -pub type NemoRelayJsonCb = - unsafe extern "C" fn(user_data: *mut libc::c_void, json: *const c_char) -> *mut c_char; +/// Codec identity kind supplied to an LLM sanitizer. +#[repr(u32)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum NemoRelayLlmSanitizeCodecKind { + /// No codec was active. + None = 0, + /// A Relay built-in codec was active. + BuiltIn = 1, + /// A runtime-registered codec was active. + Runtime = 2, + /// A codec was active but has no registered identity. + Opaque = 3, +} + +/// Codec identity supplied to an LLM sanitizer. `codec_id` is null for +/// `None` and `Opaque`, and is valid only for the duration of the callback. +#[repr(C)] +pub struct NemoRelayLlmSanitizeRequestContext { + /// Kind of active codec identity. + pub codec_kind: NemoRelayLlmSanitizeCodecKind, + /// Built-in or runtime codec ID, when applicable. + pub codec_id: *const c_char, + /// Borrowed request codec capability, or null when no codec is active. + pub codec: *const crate::types::FfiLlmSanitizeRequestCodec, +} -/// Callback for LLM request sanitization. Receives an `FfiLLMRequest` and returns -/// a new (possibly modified) `FfiLLMRequest`. Return null to use defaults. -pub type NemoRelayLlmRequestCb = unsafe extern "C" fn( +/// Directional codec context supplied to an LLM response sanitizer. +#[repr(C)] +pub struct NemoRelayLlmSanitizeResponseContext { + /// Kind of active codec identity. + pub codec_kind: NemoRelayLlmSanitizeCodecKind, + /// Built-in or runtime codec ID, when applicable. + pub codec_id: *const c_char, + /// Borrowed response codec capability, or null when no codec is active. + pub codec: *const crate::types::FfiLlmSanitizeResponseCodec, +} + +/// LLM request sanitizer. It receives the request first and its codec context +/// second. Return null to omit the observability payload. The request is +/// borrowed, but returning that same pointer is supported as a pass-through. +/// Any other non-null result transfers ownership to Relay. +pub type NemoRelayLlmSanitizeRequestCb = unsafe extern "C" fn( user_data: *mut libc::c_void, request: *const FfiLLMRequest, + context: NemoRelayLlmSanitizeRequestContext, ) -> *mut FfiLLMRequest; +/// LLM response sanitizer. It receives response JSON first and its codec +/// context second. Return null to omit the observability payload. The response +/// is borrowed, but returning that same pointer is supported as a pass-through. +/// Any other non-null result transfers ownership to Relay. +pub type NemoRelayLlmSanitizeResponseCb = unsafe extern "C" fn( + user_data: *mut libc::c_void, + response_json: *const c_char, + context: NemoRelayLlmSanitizeResponseContext, +) -> *mut c_char; + /// Callback for LLM conditional execution guardrails. /// Returns NULL to allow execution, or an error message string to reject. pub type NemoRelayLlmConditionalCb = unsafe extern "C" fn( @@ -560,23 +606,6 @@ pub fn wrap_llm_stream_exec_intercept_fn( ) } -/// Wrap a generic C JSON callback into a Rust closure. -pub fn wrap_json_fn( - cb: NemoRelayJsonCb, - user_data: *mut libc::c_void, - free_fn: NemoRelayFreeFn, -) -> Box Json + Send + Sync> { - let ud = make_user_data(user_data, free_fn); - Box::new(move |value: Json| { - let c_json = json_to_c_string(&value); - let result_ptr = unsafe { cb(ud.ptr, c_json) }; - unsafe { nemo_relay_string_free_internal(c_json) }; - let result = ptr_to_json(result_ptr); - unsafe { nemo_relay_string_free_internal(result_ptr) }; - result - }) -} - /// Wrap a C LLM request intercept callback (annotated-aware) into a Rust /// `LlmRequestInterceptFn` closure. The callback receives the intercept name, /// the opaque `FfiLLMRequest`, and the annotated JSON (or null). It writes one @@ -650,47 +679,107 @@ pub fn wrap_llm_request_intercept_fn( ) } -/// Wrap a C JSON callback into a `Fn(Json) -> Json` closure for LLM response -/// sanitization. The callback receives the response as a JSON string and -/// returns the (possibly modified) JSON string. -pub fn wrap_llm_response_fn( - cb: NemoRelayJsonCb, +/// Wrap a C LLM request sanitizer into a Rust closure. +pub fn wrap_llm_sanitize_request_fn( + cb: NemoRelayLlmSanitizeRequestCb, user_data: *mut libc::c_void, free_fn: NemoRelayFreeFn, -) -> LlmSanitizeResponseFn { +) -> LlmSanitizeRequestFn { let ud = make_user_data(user_data, free_fn); - Arc::new(move |response: Json| { - let c_json = json_to_c_string(&response); - let result_ptr = unsafe { cb(ud.ptr, c_json) }; - unsafe { nemo_relay_string_free_internal(c_json) }; - let result_json = ptr_to_json(result_ptr); - unsafe { nemo_relay_string_free_internal(result_ptr) }; - result_json - }) + Arc::new( + move |request: LlmRequest, context: LlmSanitizeRequestContext| { + clear_last_error(); + let (codec_kind, codec_id) = match ffi_codec_identity(context.codec()) { + Ok(identity) => identity, + Err(error) => { + set_last_error(&error.to_string()); + return None; + } + }; + let codec = context + .resolve_codec() + .map(crate::types::FfiLlmSanitizeRequestCodec); + let ffi_context = NemoRelayLlmSanitizeRequestContext { + codec_kind, + codec_id: codec_id + .as_ref() + .map_or(std::ptr::null(), |name| name.as_ptr()), + codec: codec.as_ref().map_or(std::ptr::null(), std::ptr::from_ref), + }; + let ffi_req = Box::into_raw(Box::new(FfiLLMRequest(request))); + let result_ptr = unsafe { cb(ud.ptr, ffi_req, ffi_context) }; + if result_ptr.is_null() { + unsafe { drop(Box::from_raw(ffi_req)) }; + return None; + } + if result_ptr == ffi_req { + return Some(unsafe { Box::from_raw(ffi_req) }.0); + } + unsafe { drop(Box::from_raw(ffi_req)) }; + Some(unsafe { Box::from_raw(result_ptr) }.0) + }, + ) } -/// Wrap a C LLM request sanitize callback into a Rust closure. -pub fn wrap_llm_sanitize_request_fn( - cb: NemoRelayLlmRequestCb, +/// Wrap a C LLM response sanitizer into a Rust closure. +pub fn wrap_llm_sanitize_response_fn( + cb: NemoRelayLlmSanitizeResponseCb, user_data: *mut libc::c_void, free_fn: NemoRelayFreeFn, -) -> LlmSanitizeRequestFn { +) -> LlmSanitizeResponseFn { let ud = make_user_data(user_data, free_fn); - Arc::new(move |request: LlmRequest| { - let ffi_req = Box::into_raw(Box::new(FfiLLMRequest(request))); - let result_ptr = unsafe { cb(ud.ptr, ffi_req) }; - // Free the input request - unsafe { drop(Box::from_raw(ffi_req)) }; + Arc::new(move |response: Json, context: LlmSanitizeResponseContext| { + clear_last_error(); + let (codec_kind, codec_id) = match ffi_codec_identity(context.codec()) { + Ok(identity) => identity, + Err(error) => { + set_last_error(&error.to_string()); + return None; + } + }; + let codec = context + .resolve_codec() + .map(crate::types::FfiLlmSanitizeResponseCodec); + let ffi_context = NemoRelayLlmSanitizeResponseContext { + codec_kind, + codec_id: codec_id + .as_ref() + .map_or(std::ptr::null(), |name| name.as_ptr()), + codec: codec.as_ref().map_or(std::ptr::null(), std::ptr::from_ref), + }; + let response_json = json_to_c_string(&response); + let result_ptr = unsafe { cb(ud.ptr, response_json, ffi_context) }; if result_ptr.is_null() { - // If callback returns null, return a default - LlmRequest { - headers: serde_json::Map::new(), - content: Json::Null, + unsafe { nemo_relay_string_free_internal(response_json) }; + return None; + } + let result = c_str_to_json(result_ptr); + unsafe { + nemo_relay_string_free_internal(response_json); + if result_ptr != response_json { + nemo_relay_string_free_internal(result_ptr); } - } else { - let result = unsafe { Box::from_raw(result_ptr) }; - result.0 } + result + }) +} + +fn ffi_codec_identity( + identity: &LlmCodecIdentity, +) -> Result<(NemoRelayLlmSanitizeCodecKind, Option)> { + Ok(match identity { + LlmCodecIdentity::None => (NemoRelayLlmSanitizeCodecKind::None, None), + LlmCodecIdentity::BuiltIn(codec) => ( + NemoRelayLlmSanitizeCodecKind::BuiltIn, + Some(CString::new(codec.id()).expect("built-in codec IDs never contain NUL")), + ), + LlmCodecIdentity::Runtime(id) => ( + NemoRelayLlmSanitizeCodecKind::Runtime, + Some(CString::new(id.as_str()).map_err(|_| { + FlowError::InvalidArgument("runtime codec ID contains an embedded NUL".to_string()) + })?), + ), + LlmCodecIdentity::Opaque => (NemoRelayLlmSanitizeCodecKind::Opaque, None), }) } diff --git a/crates/ffi/src/types/mod.rs b/crates/ffi/src/types/mod.rs index e66f0c4e1..26bf32577 100644 --- a/crates/ffi/src/types/mod.rs +++ b/crates/ffi/src/types/mod.rs @@ -47,6 +47,10 @@ pub struct FfiToolHandle(pub ToolHandle); pub struct FfiLLMHandle(pub LlmHandle); /// Opaque wrapper around an LLM request (headers, content). pub struct FfiLLMRequest(pub LlmRequest); +/// Borrowed, callback-scoped request codec capability supplied to an LLM sanitizer. +pub struct FfiLlmSanitizeRequestCodec(pub std::sync::Arc); +/// Borrowed, callback-scoped response codec capability supplied to an LLM sanitizer. +pub struct FfiLlmSanitizeResponseCodec(pub std::sync::Arc); /// Opaque wrapper around a lifecycle event emitted by the runtime. pub struct FfiEvent(pub Event); /// Opaque handle to an isolated scope stack for per-request/per-task isolation. diff --git a/crates/ffi/tests/integration/api_tests.rs b/crates/ffi/tests/integration/api_tests.rs index fd6a26951..140f089ce 100644 --- a/crates/ffi/tests/integration/api_tests.rs +++ b/crates/ffi/tests/integration/api_tests.rs @@ -14,7 +14,10 @@ use nemo_relay::plugin::PluginRegistrationContext; use serde_json::{Value as Json, json}; use uuid::Uuid; -use nemo_relay_ffi::callable::{NemoRelayLlmExecNextFn, NemoRelayToolExecNextFn}; +use nemo_relay_ffi::callable::{ + NemoRelayLlmExecNextFn, NemoRelayLlmSanitizeCodecKind, NemoRelayLlmSanitizeRequestContext, + NemoRelayLlmSanitizeResponseContext, NemoRelayToolExecNextFn, +}; use nemo_relay_ffi::convert::nemo_relay_string_free; use nemo_relay_ffi::error::{NemoRelayStatus, nemo_relay_last_error, set_last_error}; use nemo_relay_ffi::types::{ @@ -343,6 +346,7 @@ unsafe extern "C" fn tool_exec_intercept_cb( unsafe extern "C" fn llm_request_cb( _user_data: *mut libc::c_void, request: *const FfiLLMRequest, + _context: NemoRelayLlmSanitizeRequestContext, ) -> *mut FfiLLMRequest { let request = unsafe { &*request }; let mut content = request.0.content.clone(); @@ -356,6 +360,7 @@ unsafe extern "C" fn llm_request_cb( unsafe extern "C" fn llm_response_cb( _user_data: *mut libc::c_void, response_json: *const c_char, + _context: NemoRelayLlmSanitizeResponseContext, ) -> *mut c_char { let mut response: Json = serde_json::from_str( unsafe { CStr::from_ptr(response_json) } @@ -388,7 +393,17 @@ unsafe extern "C" fn llm_request_intercept_cb( _annotated_json: *const c_char, out_outcome_json: *mut *mut c_char, ) -> NemoRelayStatus { - let transformed = unsafe { Box::from_raw(llm_request_cb(ptr::null_mut(), request)) }; + let transformed = unsafe { + Box::from_raw(llm_request_cb( + ptr::null_mut(), + request, + NemoRelayLlmSanitizeRequestContext { + codec_kind: NemoRelayLlmSanitizeCodecKind::None, + codec_id: ptr::null(), + codec: ptr::null(), + }, + )) + }; let outcome = json!({ "request": transformed.0, "annotated_request": null, diff --git a/crates/ffi/tests/integration/callable_extra_tests.rs b/crates/ffi/tests/integration/callable_extra_tests.rs index ca7eaf50d..b176f4988 100644 --- a/crates/ffi/tests/integration/callable_extra_tests.rs +++ b/crates/ffi/tests/integration/callable_extra_tests.rs @@ -69,10 +69,52 @@ unsafe extern "C" fn llm_request_intercept_invalid_annotated_cb( unsafe extern "C" fn llm_request_passthrough_cb( _user_data: *mut libc::c_void, request: *const FfiLLMRequest, + _context: NemoRelayLlmSanitizeRequestContext, ) -> *mut FfiLLMRequest { Box::into_raw(Box::new(FfiLLMRequest(unsafe { (&*request).0.clone() }))) } +unsafe extern "C" fn llm_request_codec_round_trip_cb( + _user_data: *mut libc::c_void, + request: *const FfiLLMRequest, + context: NemoRelayLlmSanitizeRequestContext, +) -> *mut FfiLLMRequest { + assert_eq!(context.codec_kind, NemoRelayLlmSanitizeCodecKind::BuiltIn); + assert_eq!( + unsafe { CStr::from_ptr(context.codec_id) } + .to_str() + .unwrap(), + "openai_chat" + ); + let annotated = unsafe { nemo_relay_llm_sanitize_request_codec_decode(context.codec, request) }; + assert!(!annotated.is_null()); + let normalized: Json = + serde_json::from_str(unsafe { CStr::from_ptr(annotated) }.to_str().unwrap()).unwrap(); + assert_eq!(normalized["model"], json!("gpt-test")); + let encoded = + unsafe { nemo_relay_llm_sanitize_request_codec_encode(context.codec, annotated, request) }; + unsafe { nemo_relay_string_free_internal(annotated) }; + encoded +} + +unsafe extern "C" fn llm_response_codec_decode_cb( + _user_data: *mut libc::c_void, + response: *const c_char, + context: NemoRelayLlmSanitizeResponseContext, +) -> *mut c_char { + assert_eq!(context.codec_kind, NemoRelayLlmSanitizeCodecKind::BuiltIn); + let annotated = + unsafe { nemo_relay_llm_sanitize_response_codec_decode(context.codec, response) }; + assert!(!annotated.is_null()); + let normalized: Json = + serde_json::from_str(unsafe { CStr::from_ptr(annotated) }.to_str().unwrap()).unwrap(); + assert_eq!(normalized["model"], json!("gpt-test")); + unsafe { nemo_relay_string_free_internal(annotated) }; + CString::new(unsafe { CStr::from_ptr(response) }.to_bytes()) + .unwrap() + .into_raw() +} + unsafe extern "C" fn llm_conditional_error_cb( _user_data: *mut libc::c_void, _request: *const FfiLLMRequest, @@ -227,7 +269,11 @@ fn test_callable_extra_request_intercept_and_codec_paths() { ); let sanitize = wrap_llm_sanitize_request_fn(llm_request_passthrough_cb, ptr::null_mut(), None); - let sanitized = sanitize(request.clone()); + let sanitized = sanitize( + request.clone(), + nemo_relay::api::runtime::LlmSanitizeRequestContext::default(), + ) + .expect("non-null sanitizer result"); assert_eq!(sanitized.content, request.content); let conditional = wrap_llm_conditional_fn(llm_conditional_error_cb, ptr::null_mut(), None); @@ -298,3 +344,43 @@ fn test_callable_extra_request_intercept_and_codec_paths() { let encode_err = invalid_encode.encode(&annotated, &request).unwrap_err(); assert!(encode_err.to_string().contains("invalid result JSON")); } + +#[test] +fn test_sanitizer_context_resolves_directional_ffi_codecs() { + use nemo_relay::api::runtime::{LlmSanitizeRequestContext, LlmSanitizeResponseContext}; + use nemo_relay::codec::openai_chat::OpenAIChatCodec; + + let codec = Arc::new(OpenAIChatCodec); + let request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-test", + "messages": [{"role": "user", "content": "secret"}], + "preserve": true + }), + }; + let sanitized = + wrap_llm_sanitize_request_fn(llm_request_codec_round_trip_cb, ptr::null_mut(), None)( + request.clone(), + LlmSanitizeRequestContext::for_request_codec(Some(codec.clone())), + ) + .expect("codec round trip returns a request"); + assert_eq!(sanitized.content, request.content); + + let response = json!({ + "id": "chatcmpl-test", + "model": "gpt-test", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "secret"}, + "finish_reason": "stop" + }] + }); + let sanitized = + wrap_llm_sanitize_response_fn(llm_response_codec_decode_cb, ptr::null_mut(), None)( + response.clone(), + LlmSanitizeResponseContext::for_response_codec(Some(codec)), + ) + .expect("codec decode returns a response"); + assert_eq!(sanitized, response); +} diff --git a/crates/ffi/tests/unit/api_tests.rs b/crates/ffi/tests/unit/api_tests.rs index 4c8c19c7b..b7661aae1 100644 --- a/crates/ffi/tests/unit/api_tests.rs +++ b/crates/ffi/tests/unit/api_tests.rs @@ -6,31 +6,37 @@ use super::*; use std::ffi::{CStr, CString}; use std::ptr; -use std::sync::{Mutex, OnceLock}; +use std::sync::{Arc, Mutex, OnceLock}; +use nemo_relay::codec::request::AnnotatedLlmRequest; +use nemo_relay::codec::response::AnnotatedLlmResponse; +use nemo_relay::codec::traits::{LlmCodec, LlmResponseCodec}; use nemo_relay::plugin::PluginRegistrationContext; use serde_json::{Value as Json, json}; use uuid::Uuid; -use crate::callable::{NemoRelayLlmExecNextFn, NemoRelayToolExecNextFn}; +use crate::callable::{ + NemoRelayLlmExecNextFn, NemoRelayLlmSanitizeCodecKind, NemoRelayLlmSanitizeRequestContext, + NemoRelayLlmSanitizeResponseContext, NemoRelayToolExecNextFn, +}; use crate::convert::nemo_relay_string_free; use crate::error::{NemoRelayStatus, nemo_relay_last_error}; use crate::types::{ - FfiAtifExporter, FfiEvent, FfiLLMHandle, FfiLLMRequest, FfiOpenTelemetrySubscriber, - FfiPluginActivation, FfiScopeStack, FfiToolHandle, nemo_relay_atif_exporter_free, - nemo_relay_event_data, nemo_relay_event_input, nemo_relay_event_metadata, - nemo_relay_event_model_name, nemo_relay_event_name, nemo_relay_event_output, - nemo_relay_event_parent_uuid, nemo_relay_event_scope_type, nemo_relay_event_timestamp, - nemo_relay_event_tool_call_id, nemo_relay_event_uuid, nemo_relay_llm_handle_attributes, - nemo_relay_llm_handle_free, nemo_relay_llm_handle_name, nemo_relay_llm_handle_parent_uuid, - nemo_relay_llm_handle_uuid, nemo_relay_llm_request_content, nemo_relay_llm_request_free, - nemo_relay_llm_request_headers, nemo_relay_llm_request_new, nemo_relay_otel_subscriber_free, - nemo_relay_scope_handle_attributes, nemo_relay_scope_handle_data, nemo_relay_scope_handle_free, - nemo_relay_scope_handle_metadata, nemo_relay_scope_handle_name, - nemo_relay_scope_handle_parent_uuid, nemo_relay_scope_handle_scope_type, - nemo_relay_scope_handle_uuid, nemo_relay_scope_stack_free, nemo_relay_tool_handle_attributes, - nemo_relay_tool_handle_free, nemo_relay_tool_handle_name, nemo_relay_tool_handle_parent_uuid, - nemo_relay_tool_handle_uuid, + FfiAtifExporter, FfiEvent, FfiLLMHandle, FfiLLMRequest, FfiLlmSanitizeRequestCodec, + FfiLlmSanitizeResponseCodec, FfiOpenTelemetrySubscriber, FfiPluginActivation, FfiScopeStack, + FfiToolHandle, nemo_relay_atif_exporter_free, nemo_relay_event_data, nemo_relay_event_input, + nemo_relay_event_metadata, nemo_relay_event_model_name, nemo_relay_event_name, + nemo_relay_event_output, nemo_relay_event_parent_uuid, nemo_relay_event_scope_type, + nemo_relay_event_timestamp, nemo_relay_event_tool_call_id, nemo_relay_event_uuid, + nemo_relay_llm_handle_attributes, nemo_relay_llm_handle_free, nemo_relay_llm_handle_name, + nemo_relay_llm_handle_parent_uuid, nemo_relay_llm_handle_uuid, nemo_relay_llm_request_content, + nemo_relay_llm_request_free, nemo_relay_llm_request_headers, nemo_relay_llm_request_new, + nemo_relay_otel_subscriber_free, nemo_relay_scope_handle_attributes, + nemo_relay_scope_handle_data, nemo_relay_scope_handle_free, nemo_relay_scope_handle_metadata, + nemo_relay_scope_handle_name, nemo_relay_scope_handle_parent_uuid, + nemo_relay_scope_handle_scope_type, nemo_relay_scope_handle_uuid, nemo_relay_scope_stack_free, + nemo_relay_tool_handle_attributes, nemo_relay_tool_handle_free, nemo_relay_tool_handle_name, + nemo_relay_tool_handle_parent_uuid, nemo_relay_tool_handle_uuid, }; use crate::{api, callable, types}; @@ -340,6 +346,7 @@ unsafe extern "C" fn tool_exec_intercept_cb( unsafe extern "C" fn llm_request_cb( _user_data: *mut libc::c_void, request: *const FfiLLMRequest, + _context: NemoRelayLlmSanitizeRequestContext, ) -> *mut FfiLLMRequest { let request = unsafe { &*request }; let mut content = request.0.content.clone(); @@ -353,6 +360,7 @@ unsafe extern "C" fn llm_request_cb( unsafe extern "C" fn llm_response_cb( _user_data: *mut libc::c_void, response_json: *const c_char, + _context: NemoRelayLlmSanitizeResponseContext, ) -> *mut c_char { let mut response: Json = serde_json::from_str( unsafe { CStr::from_ptr(response_json) } @@ -385,7 +393,17 @@ unsafe extern "C" fn llm_request_intercept_cb( _annotated_json: *const c_char, out_outcome_json: *mut *mut c_char, ) -> NemoRelayStatus { - let transformed = unsafe { Box::from_raw(llm_request_cb(ptr::null_mut(), request)) }; + let transformed = unsafe { + Box::from_raw(llm_request_cb( + ptr::null_mut(), + request, + NemoRelayLlmSanitizeRequestContext { + codec_kind: NemoRelayLlmSanitizeCodecKind::None, + codec_id: ptr::null(), + codec: ptr::null(), + }, + )) + }; let outcome = json!({ "request": transformed.0, "annotated_request": null, @@ -628,6 +646,86 @@ unsafe extern "C" fn plugin_register_fail_with_last_error( NemoRelayStatus::Internal } +struct PanickingFfiSanitizerCodec; + +impl LlmCodec for PanickingFfiSanitizerCodec { + fn decode( + &self, + _request: &nemo_relay::api::llm::LlmRequest, + ) -> nemo_relay::error::Result { + panic!("request decode panic") + } + + fn encode( + &self, + _annotated: &AnnotatedLlmRequest, + _original: &nemo_relay::api::llm::LlmRequest, + ) -> nemo_relay::error::Result { + panic!("request encode panic") + } +} + +impl LlmResponseCodec for PanickingFfiSanitizerCodec { + fn decode_response(&self, _response: &Json) -> nemo_relay::error::Result { + panic!("response decode panic") + } +} + +#[test] +fn ffi_sanitizer_codec_entrypoints_contain_codec_panics() { + let request_codec = + FfiLlmSanitizeRequestCodec(Arc::new(PanickingFfiSanitizerCodec) as Arc); + let response_codec = FfiLlmSanitizeResponseCodec( + Arc::new(PanickingFfiSanitizerCodec) as Arc + ); + let request = FfiLLMRequest(nemo_relay::api::llm::LlmRequest { + headers: serde_json::Map::new(), + content: json!({"model": "gpt-test", "messages": []}), + }); + let annotated_json = cstring(&serde_json::to_string(&AnnotatedLlmRequest::default()).unwrap()); + let response_json = cstring("{}"); + + let decoded = unsafe { + api::nemo_relay_llm_sanitize_request_codec_decode( + ptr::from_ref(&request_codec), + ptr::from_ref(&request), + ) + }; + assert!(decoded.is_null()); + assert!( + unsafe { read_last_error() } + .unwrap() + .contains("sanitizer request codec decode panicked") + ); + + let encoded = unsafe { + api::nemo_relay_llm_sanitize_request_codec_encode( + ptr::from_ref(&request_codec), + annotated_json.as_ptr(), + ptr::from_ref(&request), + ) + }; + assert!(encoded.is_null()); + assert!( + unsafe { read_last_error() } + .unwrap() + .contains("sanitizer request codec encode panicked") + ); + + let decoded = unsafe { + api::nemo_relay_llm_sanitize_response_codec_decode( + ptr::from_ref(&response_codec), + response_json.as_ptr(), + ) + }; + assert!(decoded.is_null()); + assert!( + unsafe { read_last_error() } + .unwrap() + .contains("sanitizer response codec decode panicked") + ); +} + #[path = "api/core_tests.rs"] mod core_tests; #[path = "api/coverage_sweeps_tests.rs"] diff --git a/crates/ffi/tests/unit/callable_tests.rs b/crates/ffi/tests/unit/callable_tests.rs index c127dd132..c58098b4a 100644 --- a/crates/ffi/tests/unit/callable_tests.rs +++ b/crates/ffi/tests/unit/callable_tests.rs @@ -152,10 +152,19 @@ unsafe extern "C" fn llm_request_intercept_cb( unsafe extern "C" fn llm_request_null_cb( _user_data: *mut libc::c_void, _request: *const FfiLLMRequest, + _context: NemoRelayLlmSanitizeRequestContext, ) -> *mut FfiLLMRequest { std::ptr::null_mut() } +unsafe extern "C" fn llm_request_alias_cb( + _user_data: *mut libc::c_void, + request: *const FfiLLMRequest, + _context: NemoRelayLlmSanitizeRequestContext, +) -> *mut FfiLLMRequest { + request.cast_mut() +} + unsafe extern "C" fn llm_conditional_cb( _user_data: *mut libc::c_void, request: *const FfiLLMRequest, @@ -167,13 +176,41 @@ unsafe extern "C" fn llm_conditional_cb( } } -unsafe extern "C" fn json_cb(_user_data: *mut libc::c_void, json: *const c_char) -> *mut c_char { +unsafe extern "C" fn json_cb( + _user_data: *mut libc::c_void, + json: *const c_char, + _context: NemoRelayLlmSanitizeResponseContext, +) -> *mut c_char { let mut value: Json = serde_json::from_str(unsafe { CStr::from_ptr(json) }.to_str().unwrap()).unwrap(); value["wrapped"] = json!(true); CString::new(value.to_string()).unwrap().into_raw() } +unsafe extern "C" fn json_alias_cb( + _user_data: *mut libc::c_void, + json: *const c_char, + _context: NemoRelayLlmSanitizeResponseContext, +) -> *mut c_char { + json.cast_mut() +} + +unsafe extern "C" fn invalid_json_cb( + _user_data: *mut libc::c_void, + _json: *const c_char, + _context: NemoRelayLlmSanitizeResponseContext, +) -> *mut c_char { + CString::new("not-json").unwrap().into_raw() +} + +unsafe extern "C" fn invalid_utf8_cb( + _user_data: *mut libc::c_void, + _json: *const c_char, + _context: NemoRelayLlmSanitizeResponseContext, +) -> *mut c_char { + CString::new([0xff]).unwrap().into_raw() +} + unsafe extern "C" fn llm_exec_cb( _user_data: *mut libc::c_void, native_json: *const c_char, @@ -379,9 +416,23 @@ fn test_wrap_llm_request_response_and_conditional_callbacks() { let sanitize_request = wrap_llm_sanitize_request_fn(llm_request_null_cb, std::ptr::null_mut(), None); - let sanitized = sanitize_request(make_request()); - assert_eq!(sanitized.headers.len(), 0); - assert_eq!(sanitized.content, Json::Null); + assert_eq!( + sanitize_request( + make_request(), + nemo_relay::api::runtime::LlmSanitizeRequestContext::default(), + ), + None + ); + + let alias_request = + wrap_llm_sanitize_request_fn(llm_request_alias_cb, std::ptr::null_mut(), None); + assert_eq!( + alias_request( + make_request(), + nemo_relay::api::runtime::LlmSanitizeRequestContext::default(), + ), + Some(make_request()) + ); let conditional = wrap_llm_conditional_fn(llm_conditional_cb, std::ptr::null_mut(), None); assert_eq!( @@ -394,14 +445,74 @@ fn test_wrap_llm_request_response_and_conditional_callbacks() { ); assert_eq!(conditional(&make_request()).unwrap(), None); - let wrapped_json = wrap_json_fn(json_cb, std::ptr::null_mut(), None); - assert_eq!(wrapped_json(json!({"value": 1}))["wrapped"], json!(true)); - - let wrapped_response = wrap_llm_response_fn(json_cb, std::ptr::null_mut(), None); + let wrapped_response = wrap_llm_sanitize_response_fn(json_cb, std::ptr::null_mut(), None); assert_eq!( - wrapped_response(json!({"value": 2}))["wrapped"], + wrapped_response( + json!({"value": 2}), + nemo_relay::api::runtime::LlmSanitizeResponseContext::default(), + ) + .unwrap()["wrapped"], json!(true) ); + + let alias_response = wrap_llm_sanitize_response_fn(json_alias_cb, std::ptr::null_mut(), None); + assert_eq!( + alias_response( + json!({"value": 2}), + nemo_relay::api::runtime::LlmSanitizeResponseContext::default(), + ), + Some(json!({"value": 2})) + ); + + for callback in [invalid_json_cb, invalid_utf8_cb] { + let malformed_response = + wrap_llm_sanitize_response_fn(callback, std::ptr::null_mut(), None); + assert_eq!( + malformed_response( + json!({"secret": "must be omitted"}), + nemo_relay::api::runtime::LlmSanitizeResponseContext::default(), + ), + None + ); + } +} + +#[test] +fn test_llm_sanitizers_fail_closed_for_runtime_codec_ids_with_embedded_nul() { + let runtime_identity = + nemo_relay::api::runtime::LlmCodecIdentity::Runtime("runtime\0codec".to_string()); + + let request_sanitizer = + wrap_llm_sanitize_request_fn(llm_request_alias_cb, std::ptr::null_mut(), None); + assert_eq!( + request_sanitizer( + make_request(), + nemo_relay::api::runtime::LlmSanitizeRequestContext::with_identity( + runtime_identity.clone(), + ), + ), + None + ); + assert!( + last_error_message() + .unwrap() + .contains("runtime codec ID contains an embedded NUL") + ); + + let response_sanitizer = + wrap_llm_sanitize_response_fn(json_alias_cb, std::ptr::null_mut(), None); + assert_eq!( + response_sanitizer( + json!({"secret": "must be omitted"}), + nemo_relay::api::runtime::LlmSanitizeResponseContext::with_identity(runtime_identity), + ), + None + ); + assert!( + last_error_message() + .unwrap() + .contains("runtime codec ID contains an embedded NUL") + ); } #[test] diff --git a/crates/node/package.json b/crates/node/package.json index 9cfa7278c..6f2f371a4 100644 --- a/crates/node/package.json +++ b/crates/node/package.json @@ -62,8 +62,8 @@ "triples": {} }, "scripts": { - "build": "napi build --platform --release", - "build-debug": "napi build --platform", + "build": "napi build --platform --release --dts index.d.ts", + "build-debug": "napi build --platform --dts index.d.ts", "check:docstrings": "node ../../scripts/lint/check_public_docstrings.mjs node", "check:docstrings:all": "node ../../scripts/lint/check_public_docstrings.mjs all", "format": "prettier --write \"*.{js,ts,d.ts}\" \"tests/**/*.{mjs,cjs,js,mts,cts,ts,tsx}\"", diff --git a/crates/node/plugin.d.ts b/crates/node/plugin.d.ts index 608289762..158e25b2f 100644 --- a/crates/node/plugin.d.ts +++ b/crates/node/plugin.d.ts @@ -4,6 +4,28 @@ /// import type { EventSanitizeFields, Json } from './index'; +import type { LlmCodec, LlmResponseCodec } from './typed'; + +/** Codec identity available while a managed LLM event is sanitized. */ +export type LlmCodecIdentity = + | { kind: 'none' } + | { kind: 'builtin'; id: 'openai_chat' | 'openai_responses' | 'anthropic_messages' } + | { kind: 'runtime'; id: string } + | { kind: 'opaque' }; + +/** Codec context available while an LLM request is sanitized. */ +export interface LlmSanitizeRequestContext { + codec: LlmCodecIdentity; + /** Resolve the active codec for this callback. Do not retain the result after the callback returns. */ + resolveCodec(): LlmCodec | null; +} + +/** Codec context available while an LLM response is sanitized. */ +export interface LlmSanitizeResponseContext { + codec: LlmCodecIdentity; + /** Resolve the active codec for this callback. Do not retain the result after the callback returns. */ + resolveCodec(): LlmResponseCodec | null; +} /** Policy behavior for unsupported configuration. */ export type UnsupportedBehavior = 'ignore' | 'warn' | 'error'; @@ -205,10 +227,18 @@ export interface PluginContext { priority: number, callback: (name: string, args: Json) => string | null, ): void; - /** Register an LLM sanitize-request guardrail for this component. */ - registerLlmSanitizeRequestGuardrail(name: string, priority: number, callback: (request: Json) => Json): void; - /** Register an LLM sanitize-response guardrail for this component. */ - registerLlmSanitizeResponseGuardrail(name: string, priority: number, callback: (response: Json) => Json): void; + /** Register an LLM sanitize-request guardrail. The callback receives `(request, context)`. */ + registerLlmSanitizeRequestGuardrail( + name: string, + priority: number, + callback: (request: Json, context: LlmSanitizeRequestContext) => Json | null, + ): void; + /** Register an LLM sanitize-response guardrail. The callback receives `(response, context)`. */ + registerLlmSanitizeResponseGuardrail( + name: string, + priority: number, + callback: (response: Json, context: LlmSanitizeResponseContext) => Json | null, + ): void; /** Register an LLM conditional-execution guardrail for this component. */ registerLlmConditionalExecutionGuardrail( name: string, diff --git a/crates/node/src/api/mod.rs b/crates/node/src/api/mod.rs index 12cc22b96..1f8384fae 100644 --- a/crates/node/src/api/mod.rs +++ b/crates/node/src/api/mod.rs @@ -599,6 +599,61 @@ fn middleware_json_callback_tsfn( Ok(tsfn) } +fn middleware_llm_sanitize_request_callback_tsfn( + env: &Env, + func: &JsFunction, +) -> napi::Result< + ThreadsafeFunction<(Json, callable::JsLlmSanitizeRequestContext), ErrorStrategy::Fatal>, +> { + let callback = callable::safe_middleware_callback(env, func)?; + let mut tsfn = callback.create_threadsafe_function( + 0, + |ctx: napi::threadsafe_function::ThreadSafeCallContext<( + Json, + callable::JsLlmSanitizeRequestContext, + )>| { + let first = unsafe { + JsUnknown::from_raw_unchecked( + ctx.env.raw(), + Json::to_napi_value(ctx.env.raw(), ctx.value.0)?, + ) + }; + let context = callable::js_llm_sanitize_request_context_to_napi(&ctx.env, ctx.value.1)?; + Ok(vec![first, context]) + }, + )?; + tsfn.unref(env)?; + Ok(tsfn) +} + +fn middleware_llm_sanitize_response_callback_tsfn( + env: &Env, + func: &JsFunction, +) -> napi::Result< + ThreadsafeFunction<(Json, callable::JsLlmSanitizeResponseContext), ErrorStrategy::Fatal>, +> { + let callback = callable::safe_middleware_callback(env, func)?; + let mut tsfn = callback.create_threadsafe_function( + 0, + |ctx: napi::threadsafe_function::ThreadSafeCallContext<( + Json, + callable::JsLlmSanitizeResponseContext, + )>| { + let first = unsafe { + JsUnknown::from_raw_unchecked( + ctx.env.raw(), + Json::to_napi_value(ctx.env.raw(), ctx.value.0)?, + ) + }; + let context = + callable::js_llm_sanitize_response_context_to_napi(&ctx.env, ctx.value.1)?; + Ok(vec![first, context]) + }, + )?; + tsfn.unref(env)?; + Ok(tsfn) +} + #[allow(clippy::too_many_arguments)] fn add_plugin_event_sanitizer( env: &Env, @@ -854,9 +909,9 @@ fn build_plugin_context( core_registry_api::register_llm_sanitize_request_guardrail( &name, priority, - callable::wrap_js_llm_sanitize_request_fn(middleware_json_callback_tsfn( - ctx.env, &callback, - )?), + callable::wrap_js_llm_sanitize_request_fn( + middleware_llm_sanitize_request_callback_tsfn(ctx.env, &callback)?, + ), ) .map_err(to_napi_err)?; @@ -900,9 +955,9 @@ fn build_plugin_context( core_registry_api::register_llm_sanitize_response_guardrail( &name, priority, - callable::wrap_js_llm_response_fn(middleware_json_callback_tsfn( - ctx.env, &callback, - )?), + callable::wrap_js_llm_sanitize_response_fn( + middleware_llm_sanitize_response_callback_tsfn(ctx.env, &callback)?, + ), ) .map_err(to_napi_err)?; @@ -1204,19 +1259,42 @@ struct NodePluginRegisterCall { /// struct. `reference` must be a valid N-API reference created for a live /// JavaScript function in `env`, and `env` must not be used after the /// corresponding Node.js environment has been torn down. -struct PersistentJsFunction { +pub(crate) struct PersistentJsFunction { env: napi::sys::napi_env, reference: napi::sys::napi_ref, + cleanup: napi::sys::napi_threadsafe_function, } -// SAFETY: `PersistentJsFunction` only stores raw N-API handles. Callers are -// responsible for constructing it from a live environment and function -// reference, and all access goes back through that same environment. +// SAFETY: Direct function access is restricted by callers to the registration +// thread. Releasing the cleanup TSFN is thread-safe, and its event-loop +// finalizer deletes the N-API reference. unsafe impl Send for PersistentJsFunction {} -// SAFETY: The same invariants as `Send` apply. The struct does not provide -// interior mutation beyond the N-API reference lifecycle managed by Node. +// SAFETY: The same invariants as `Send` apply. The stored handles are immutable, +// and reference deletion is serialized by the cleanup TSFN finalizer. unsafe impl Sync for PersistentJsFunction {} +unsafe extern "C" fn delete_persistent_js_function_reference( + env: napi::sys::napi_env, + finalize_data: *mut std::ffi::c_void, + _finalize_hint: *mut std::ffi::c_void, +) { + if !env.is_null() && !finalize_data.is_null() { + // SAFETY: `finalize_data` is the live N-API reference passed to the + // cleanup TSFN at creation. This finalizer runs once on Node's event + // loop after the final TSFN release. + let _ = + unsafe { napi::sys::napi_delete_reference(env, finalize_data as napi::sys::napi_ref) }; + } +} + +unsafe extern "C" fn persistent_js_function_cleanup_call( + _env: napi::sys::napi_env, + _js_callback: napi::sys::napi_value, + _context: *mut std::ffi::c_void, + _data: *mut std::ffi::c_void, +) { +} + impl PersistentJsFunction { fn new(env: &Env, func: &JsFunction) -> napi::Result { let mut reference = ptr::null_mut(); @@ -1225,17 +1303,74 @@ impl PersistentJsFunction { // writable storage for the created reference. let status = unsafe { napi::sys::napi_create_reference(env.raw(), func.raw(), 1, &mut reference) }; - if status == napi::sys::Status::napi_ok { - Ok(Self { - env: env.raw(), - reference, - }) - } else { - Err(napi::Error::from_reason(format!( + if status != napi::sys::Status::napi_ok { + return Err(napi::Error::from_reason(format!( "failed to create JS function reference: {:?}", napi::Status::from(status) - ))) + ))); + } + + let mut resource_name = ptr::null_mut(); + let resource_name_bytes = b"nemo_relay_persistent_js_function\0"; + let status = unsafe { + napi::sys::napi_create_string_utf8( + env.raw(), + resource_name_bytes.as_ptr().cast(), + resource_name_bytes.len() - 1, + &mut resource_name, + ) + }; + if status != napi::sys::Status::napi_ok { + let _ = unsafe { napi::sys::napi_delete_reference(env.raw(), reference) }; + return Err(napi::Error::from_reason(format!( + "failed to create persistent JS function cleanup resource name: {:?}", + napi::Status::from(status) + ))); + } + + let mut cleanup = ptr::null_mut(); + let status = unsafe { + napi::sys::napi_create_threadsafe_function( + env.raw(), + ptr::null_mut(), + ptr::null_mut(), + resource_name, + 0, + 1, + reference.cast(), + Some(delete_persistent_js_function_reference), + ptr::null_mut(), + Some(persistent_js_function_cleanup_call), + &mut cleanup, + ) + }; + if status != napi::sys::Status::napi_ok { + let _ = unsafe { napi::sys::napi_delete_reference(env.raw(), reference) }; + return Err(napi::Error::from_reason(format!( + "failed to create persistent JS function cleanup handle: {:?}", + napi::Status::from(status) + ))); } + + let status = unsafe { napi::sys::napi_unref_threadsafe_function(env.raw(), cleanup) }; + if status != napi::sys::Status::napi_ok { + let _ = unsafe { + napi::sys::napi_release_threadsafe_function( + cleanup, + napi::sys::ThreadsafeFunctionReleaseMode::release, + ) + }; + return Err(napi::Error::from_reason(format!( + "failed to unref persistent JS function cleanup handle: {:?}", + napi::Status::from(status) + ))); + } + + Ok(Self { + env: env.raw(), + reference, + cleanup, + }) } fn call_validate(&self, plugin_config: &Json) -> napi::Result { @@ -1303,6 +1438,29 @@ impl PersistentJsFunction { // SAFETY: `returned` is the live result of invoking `func` in this environment. unsafe { Option::::from_napi_value(self.env, returned.raw()) }.map(callback_json) } + + fn call_json(&self, argument: Json) -> napi::Result { + let mut value = ptr::null_mut(); + // SAFETY: `self.reference` is a live N-API reference created in + // `self.env`, and `value` is writable storage for the borrowed + // function value. + let status = + unsafe { napi::sys::napi_get_reference_value(self.env, self.reference, &mut value) }; + if status != napi::sys::Status::napi_ok { + return Err(napi::Error::from_reason("failed to borrow codec function")); + } + // SAFETY: `value` was resolved from this struct's function reference, + // so it is a live function value in `self.env` for this call. + let func = unsafe { JsFunction::from_raw_unchecked(self.env, value) }; + // SAFETY: `Json::to_napi_value` created this argument in `self.env`, + // so wrapping it as `JsUnknown` is valid for the immediate callback. + let argument = unsafe { + JsUnknown::from_raw_unchecked(self.env, Json::to_napi_value(self.env, argument)?) + }; + let returned = func.call(None, &[argument])?; + // SAFETY: `returned` is the live result of invoking `func` in this environment. + unsafe { Option::::from_napi_value(self.env, returned.raw()) }.map(callback_json) + } } fn core_event_fields( @@ -1390,11 +1548,91 @@ fn node_event_sanitize_fn(env: &Env, func: &JsFunction) -> napi::Result, + Vec>, +); +type NodeLlmResponseCodec = ( + Arc, + Vec>, +); + +fn node_llm_codec( + env: &Env, + decode: &JsFunction, + encode: &JsFunction, +) -> napi::Result { + let direct_decode = Arc::new(PersistentJsFunction::new(env, decode)?); + let direct_encode = Arc::new(PersistentJsFunction::new(env, encode)?); + let references = vec![direct_decode.clone(), direct_encode.clone()]; + let register_thread = std::thread::current().id(); + + let mut decode_tsfn = decode.create_threadsafe_function( + 0, + |ctx: napi::threadsafe_function::ThreadSafeCallContext| Ok(vec![ctx.value]), + )?; + decode_tsfn.unref(env)?; + let mut encode_tsfn = encode.create_threadsafe_function( + 0, + |ctx: napi::threadsafe_function::ThreadSafeCallContext| Ok(vec![ctx.value]), + )?; + encode_tsfn.unref(env)?; + + Ok(( + callable::wrap_js_codec( + decode_tsfn, + encode_tsfn, + register_thread, + Arc::new(move |argument| { + direct_decode.call_json(argument).map_err(|error| { + FlowError::Internal(format!("JS codec decode callback failed: {error}")) + }) + }), + Arc::new(move |argument| { + direct_encode.call_json(argument).map_err(|error| { + FlowError::Internal(format!("JS codec encode callback failed: {error}")) + }) + }), + ), + references, + )) +} + +fn node_llm_response_codec(env: &Env, decode: &JsFunction) -> napi::Result { + let direct_decode = Arc::new(PersistentJsFunction::new(env, decode)?); + let references = vec![direct_decode.clone()]; + let register_thread = std::thread::current().id(); + let mut decode_tsfn = decode.create_threadsafe_function( + 0, + |ctx: napi::threadsafe_function::ThreadSafeCallContext| Ok(vec![ctx.value]), + )?; + decode_tsfn.unref(env)?; + Ok(( + callable::wrap_js_response_codec( + decode_tsfn, + register_thread, + Arc::new(move |argument| { + direct_decode.call_json(argument).map_err(|error| { + FlowError::Internal(format!( + "JS response codec decode callback failed: {error}" + )) + }) + }), + ), + references, + )) +} + impl Drop for PersistentJsFunction { fn drop(&mut self) { - // SAFETY: `self.reference` was created by `napi_create_reference` for - // `self.env` and is deleted exactly once here during drop. - let _ = unsafe { napi::sys::napi_delete_reference(self.env, self.reference) }; + // SAFETY: N-API permits releasing a TSFN from any thread. Its finalizer + // runs on the event loop and deletes `self.reference` exactly once. + let _ = unsafe { + napi::sys::napi_release_threadsafe_function( + self.cleanup, + napi::sys::ThreadsafeFunctionReleaseMode::release, + ) + }; } } @@ -1523,8 +1761,8 @@ pub fn scope_stack_active() -> bool { /// Returns the most recent callback error that could not be surfaced through a direct exception. /// -/// This is primarily used for sanitize callback paths that fail open and cannot -/// surface their errors directly. +/// This is primarily used for sanitize callback paths that omit observability +/// payloads and cannot surface their errors directly. #[napi] pub fn get_last_callback_error() -> Option { get_recorded_callback_error() @@ -1549,30 +1787,32 @@ pub fn test_closed_tool_callback( wrapped(&name, args) } -/// Internal test helper: invoke a closed JS LLM sanitize-request wrapper and return the fallback request. +/// Internal test helper: model a closed JS LLM request sanitizer. #[napi(js_name = "__testClosedLlmSanitizeRequestCallback")] pub fn test_closed_llm_sanitize_request_callback( callback: ThreadsafeFunction, request: Json, -) -> Result { +) -> Result> { clear_recorded_callback_error(); let _ = callback.clone().abort(); let llm_request: LlmRequest = serde_json::from_value(request) .map_err(|e| napi::Error::from_reason(format!("invalid LlmRequest: {e}")))?; - let wrapped = callable::wrap_js_llm_sanitize_request_fn(callback); - Ok(serde_json::to_value(wrapped(llm_request)).unwrap_or(Json::Null)) + drop(llm_request); + record_callback_error("nemo_relay: failed to queue JS LLM sanitize request callback"); + Ok(None) } -/// Internal test helper: invoke a closed JS LLM sanitize-response wrapper and return the fallback response. +/// Internal test helper: model a closed JS LLM response sanitizer. #[napi(js_name = "__testClosedLlmResponseCallback")] pub fn test_closed_llm_response_callback( callback: ThreadsafeFunction, response: Json, -) -> Json { +) -> Option { clear_recorded_callback_error(); let _ = callback.clone().abort(); - let wrapped = callable::wrap_js_llm_response_fn(callback); - wrapped(response) + drop(response); + record_callback_error("nemo_relay: failed to queue JS LLM sanitize response callback"); + None } /// Internal test helper: invoke a closed JS collector wrapper and surface the queue failure. @@ -2123,9 +2363,9 @@ pub fn llm_call_execute( data: Option, metadata: Option, model_name: Option, - codec_decode: Option>, - codec_encode: Option>, - response_codec_decode: Option>, + #[napi(ts_arg_type = "(arg: Json) => any")] codec_decode: Option, + #[napi(ts_arg_type = "(arg: Json) => any")] codec_encode: Option, + #[napi(ts_arg_type = "(arg: Json) => any")] response_codec_decode: Option, ) -> Result { let attrs = LlmAttributes::from_bits_truncate(attributes.unwrap_or(0)); let parent = handle @@ -2136,11 +2376,28 @@ pub fn llm_call_execute( let callback = callable::safe_execution_callback(&env, &func)?; let exec_fn = callable::wrap_js_llm_exec_fn(json_callback_tsfn(&env, &callback)?); let default_fn: LlmExecutionNextFn = std::sync::Arc::new(move |req| exec_fn(req)); - let codec = match (codec_decode, codec_encode) { - (Some(d), Some(e)) => Some(callable::wrap_js_codec(d, e)), - _ => None, + let mut codec_references = Vec::new(); + let codec = match (codec_decode.as_ref(), codec_encode.as_ref()) { + (Some(d), Some(e)) => { + let (codec, references) = node_llm_codec(&env, d, e)?; + codec_references.extend(references); + Some(codec) + } + (None, None) => None, + _ => { + return Err(napi::Error::from_reason( + "codecDecode and codecEncode must be provided together", + )); + } }; - let response_codec = response_codec_decode.map(callable::wrap_js_response_codec); + let response_codec = response_codec_decode + .as_ref() + .map(|decode| node_llm_response_codec(&env, decode)) + .transpose()? + .map(|(codec, references)| { + codec_references.extend(references); + codec + }); let scope_stack = current_scope_stack_handle(); env.execute_tokio_future( @@ -2165,7 +2422,10 @@ pub fn llm_call_execute( }) .await }, - |_env, result| Ok(result), + move |_env, result| { + drop(codec_references); + Ok(result) + }, ) } @@ -2185,9 +2445,9 @@ pub fn llm_call_execute_async( data: Option, metadata: Option, model_name: Option, - codec_decode: Option>, - codec_encode: Option>, - response_codec_decode: Option>, + #[napi(ts_arg_type = "(arg: Json) => any")] codec_decode: Option, + #[napi(ts_arg_type = "(arg: Json) => any")] codec_encode: Option, + #[napi(ts_arg_type = "(arg: Json) => any")] response_codec_decode: Option, ) -> Result { let attrs = LlmAttributes::from_bits_truncate(attributes.unwrap_or(0)); let parent = handle @@ -2209,11 +2469,28 @@ pub fn llm_call_execute_async( Box::pin(async move { pa_fn.call(req_json).await }) }); - let codec = match (codec_decode, codec_encode) { - (Some(d), Some(e)) => Some(callable::wrap_js_codec(d, e)), - _ => None, + let mut codec_references = Vec::new(); + let codec = match (codec_decode.as_ref(), codec_encode.as_ref()) { + (Some(d), Some(e)) => { + let (codec, references) = node_llm_codec(&env, d, e)?; + codec_references.extend(references); + Some(codec) + } + (None, None) => None, + _ => { + return Err(napi::Error::from_reason( + "codecDecode and codecEncode must be provided together", + )); + } }; - let response_codec = response_codec_decode.map(callable::wrap_js_response_codec); + let response_codec = response_codec_decode + .as_ref() + .map(|decode| node_llm_response_codec(&env, decode)) + .transpose()? + .map(|(codec, references)| { + codec_references.extend(references); + codec + }); env.execute_tokio_future( async move { @@ -2237,7 +2514,10 @@ pub fn llm_call_execute_async( }) .await }, - |_env, result| Ok(result), + move |_env, result| { + drop(codec_references); + Ok(result) + }, ) } @@ -2272,9 +2552,9 @@ pub fn llm_stream_call_execute( data: Option, metadata: Option, model_name: Option, - codec_decode: Option>, - codec_encode: Option>, - response_codec_decode: Option>, + #[napi(ts_arg_type = "(arg: Json) => any")] codec_decode: Option, + #[napi(ts_arg_type = "(arg: Json) => any")] codec_encode: Option, + #[napi(ts_arg_type = "(arg: Json) => any")] response_codec_decode: Option, ) -> Result { let attrs = LlmAttributes::from_bits_truncate(attributes.unwrap_or(0)); let parent = handle @@ -2325,11 +2605,29 @@ pub fn llm_stream_call_execute( }) }); - let codec = match (codec_decode, codec_encode) { - (Some(d), Some(e)) => Some(callable::wrap_js_codec(d, e)), - _ => None, + let mut codec_references = Vec::new(); + let codec = match (codec_decode.as_ref(), codec_encode.as_ref()) { + (Some(d), Some(e)) => { + let (codec, references) = node_llm_codec(&env, d, e)?; + codec_references.extend(references); + Some(codec) + } + (None, None) => None, + _ => { + return Err(napi::Error::from_reason( + "codecDecode and codecEncode must be provided together", + )); + } }; - let response_codec = response_codec_decode.map(callable::wrap_js_response_codec); + let response_codec = response_codec_decode + .as_ref() + .map(|decode| node_llm_response_codec(&env, decode)) + .transpose()? + .map(|(codec, references)| { + codec_references.extend(references); + codec + }); + let completion_codec_references = codec_references.clone(); let scope_stack = current_scope_stack_handle(); env.execute_tokio_future( @@ -2368,11 +2666,15 @@ pub fn llm_stream_call_execute( receiver: tokio::sync::Mutex::new(rx), cancel, closed: closed_rx, + codec_references, }) }) .await }, - |_env, result| Ok(result), + move |_env, result| { + drop(completion_codec_references); + Ok(result) + }, ) } @@ -2610,18 +2912,21 @@ pub fn deregister_tool_execution_intercept(name: String) -> Result { /// Register a guardrail that sanitizes LLM request data before execution. /// -/// The `guardrail` callback receives the LLM request as JSON and must return the sanitized request. -/// Higher `priority` values run first. Throws if a guardrail with the same `name` already exists. -/// If the callback throws, Relay preserves the current emitted payload and records the error for -/// `getLastCallbackError()`. +/// The `guardrail` callback receives `(request, context)` and must return the sanitized request, +/// or `null` to omit the observability payload. Lower `priority` values run first. Throws if a +/// guardrail with the same `name` already exists. If the callback throws, Relay omits the payload +/// and records the error for `getLastCallbackError()`. #[napi] pub fn register_llm_sanitize_request_guardrail( env: Env, name: String, priority: i32, + #[napi( + ts_arg_type = "(request: Json, context: import('./plugin').LlmSanitizeRequestContext) => Json | null" + )] guardrail: JsFunction, ) -> Result<()> { - let callback = middleware_json_callback_tsfn(&env, &guardrail)?; + let callback = middleware_llm_sanitize_request_callback_tsfn(&env, &guardrail)?; core_registry_api::register_llm_sanitize_request_guardrail( &name, priority, @@ -2640,23 +2945,25 @@ pub fn deregister_llm_sanitize_request_guardrail(name: String) -> Result { /// Register a guardrail that sanitizes LLM response data after execution. /// -/// The `guardrail` callback receives the LLM response as a JSON value and must return -/// the sanitized response as JSON. Higher `priority` values run first. Throws if a guardrail -/// with the same `name` already exists. -/// If the callback throws, Relay preserves the current emitted payload and records the error for -/// `getLastCallbackError()`. +/// The `guardrail` callback receives `(response, context)` and must return the sanitized response, +/// or `null` to omit the observability payload. Lower `priority` values run first. Throws if a +/// guardrail with the same `name` already exists. If the callback throws, Relay omits the payload +/// and records the error for `getLastCallbackError()`. #[napi] pub fn register_llm_sanitize_response_guardrail( env: Env, name: String, priority: i32, + #[napi( + ts_arg_type = "(response: Json, context: import('./plugin').LlmSanitizeResponseContext) => Json | null" + )] guardrail: JsFunction, ) -> Result<()> { - let callback = middleware_json_callback_tsfn(&env, &guardrail)?; + let callback = middleware_llm_sanitize_response_callback_tsfn(&env, &guardrail)?; core_registry_api::register_llm_sanitize_response_guardrail( &name, priority, - callable::wrap_js_llm_response_fn(callback), + callable::wrap_js_llm_sanitize_response_fn(callback), ) .map_err(to_napi_err) } @@ -3140,17 +3447,19 @@ pub fn scope_deregister_tool_execution_intercept(scope_uuid: String, name: Strin /// Register a scope-local guardrail that sanitizes LLM request data before execution. /// -/// The `guardrail` callback receives the LLM request as JSON and must return the sanitized request. -/// Higher `priority` values run first. Throws if a guardrail with the same `name` already exists -/// on the specified scope. -/// If the callback throws, Relay preserves the current emitted payload and records the error for -/// `getLastCallbackError()`. +/// The `guardrail` callback receives `(request, context)` and must return the sanitized request, +/// or `null` to omit the observability payload. Lower `priority` values run first. Throws if a +/// guardrail with the same `name` already exists on the specified scope. If the callback throws, +/// Relay omits the payload and records the error for `getLastCallbackError()`. #[napi] pub fn scope_register_llm_sanitize_request_guardrail( env: Env, scope_uuid: String, name: String, priority: i32, + #[napi( + ts_arg_type = "(request: Json, context: import('./plugin').LlmSanitizeRequestContext) => Json | null" + )] guardrail: JsFunction, ) -> Result<()> { let uuid = uuid::Uuid::parse_str(&scope_uuid) @@ -3159,7 +3468,9 @@ pub fn scope_register_llm_sanitize_request_guardrail( &uuid, &name, priority, - callable::wrap_js_llm_sanitize_request_fn(middleware_json_callback_tsfn(&env, &guardrail)?), + callable::wrap_js_llm_sanitize_request_fn(middleware_llm_sanitize_request_callback_tsfn( + &env, &guardrail, + )?), ) .map_err(to_napi_err) } @@ -3180,17 +3491,19 @@ pub fn scope_deregister_llm_sanitize_request_guardrail( /// Register a scope-local guardrail that sanitizes LLM response data after execution. /// -/// The `guardrail` callback receives the LLM response as a JSON value and must return -/// the sanitized response as JSON. Higher `priority` values run first. Throws if a guardrail -/// with the same `name` already exists on the specified scope. -/// If the callback throws, Relay preserves the current emitted payload and records the error for -/// `getLastCallbackError()`. +/// The `guardrail` callback receives `(response, context)` and must return the sanitized response, +/// or `null` to omit the observability payload. Lower `priority` values run first. Throws if a +/// guardrail with the same `name` already exists on the specified scope. If the callback throws, +/// Relay omits the payload and records the error for `getLastCallbackError()`. #[napi] pub fn scope_register_llm_sanitize_response_guardrail( env: Env, scope_uuid: String, name: String, priority: i32, + #[napi( + ts_arg_type = "(response: Json, context: import('./plugin').LlmSanitizeResponseContext) => Json | null" + )] guardrail: JsFunction, ) -> Result<()> { let uuid = uuid::Uuid::parse_str(&scope_uuid) @@ -3199,7 +3512,9 @@ pub fn scope_register_llm_sanitize_response_guardrail( &uuid, &name, priority, - callable::wrap_js_llm_response_fn(middleware_json_callback_tsfn(&env, &guardrail)?), + callable::wrap_js_llm_sanitize_response_fn(middleware_llm_sanitize_response_callback_tsfn( + &env, &guardrail, + )?), ) .map_err(to_napi_err) } diff --git a/crates/node/src/callable.rs b/crates/node/src/callable.rs index 6937e8dd5..ba8abafc2 100644 --- a/crates/node/src/callable.rs +++ b/crates/node/src/callable.rs @@ -13,12 +13,15 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use napi::bindgen_prelude::ToNapiValue; use napi::threadsafe_function::{ErrorStrategy, ThreadsafeFunction, ThreadsafeFunctionCallMode}; -use napi::{Env, JsFunction, JsUnknown, NapiRaw, NapiValue}; +use napi::{Env, JsFunction, JsObject, JsUnknown, NapiRaw, NapiValue}; +use napi_derive::napi; use nemo_relay::api::runtime::{ - EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, LlmJsonStream, - LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, - ToolConditionalFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, + EventSanitizeFn, EventSubscriberFn, LlmCodecIdentity, LlmConditionalFn, LlmExecutionNextFn, + LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn, + LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, ToolConditionalFn, + ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; use serde::{Deserialize, Serialize}; use serde_json::Value as Json; @@ -37,10 +40,32 @@ use nemo_relay::codec::traits::{LlmCodec, LlmResponseCodec}; use nemo_relay::error::{FlowError, Result}; use crate::callback_factory; -use crate::convert::{callback_json, record_callback_error}; +use crate::convert::{callback_json, record_callback_error, to_napi_err}; use crate::promise_call::{JsonNextFn, JsonStreamNextFn, PromiseAwareFn}; use crate::types::{EventSanitizeFields, JsEvent, event_sanitize_fields_from_json}; +/// Structured codec identity delivered to JavaScript LLM sanitizers. +#[napi(object)] +#[derive(Clone)] +pub(crate) struct JsLlmCodecIdentity { + pub kind: String, + pub id: Option, +} + +/// Structured per-call request context delivered to JavaScript LLM sanitizers. +#[derive(Clone)] +pub(crate) struct JsLlmSanitizeRequestContext { + pub codec: JsLlmCodecIdentity, + resolved: Option>, +} + +/// Structured per-call response context delivered to JavaScript LLM sanitizers. +#[derive(Clone)] +pub(crate) struct JsLlmSanitizeResponseContext { + pub codec: JsLlmCodecIdentity, + resolved: Option>, +} + /// JavaScript-facing pending mark DTO. #[derive(Debug, Deserialize, Serialize)] #[serde(rename_all = "camelCase", deny_unknown_fields)] @@ -398,74 +423,227 @@ pub fn wrap_js_llm_request_intercept_fn( ) } -/// Wrap a JS function for LLM sanitize request: `(request: LlmRequest) => LlmRequest`. -/// Since ThreadsafeFunction requires serde-serializable args, we serialize the request as JSON. +/// Wrap a JS function for LLM request sanitization. The callback receives +/// `(request, context)`. pub fn wrap_js_llm_sanitize_request_fn( - func: ThreadsafeFunction, + func: ThreadsafeFunction<(Json, JsLlmSanitizeRequestContext), ErrorStrategy::Fatal>, ) -> LlmSanitizeRequestFn { let func = Arc::new(func); - Arc::new(move |request: LlmRequest| { - let func = func.clone(); - let req_json = serde_json::to_value(&request).unwrap_or(Json::Null); - let (tx, rx) = std::sync::mpsc::channel(); - let status = func.call_with_return_value( - req_json, - ThreadsafeFunctionCallMode::Blocking, - move |val: Option| { - let _ = tx.send(callback_json(val)); - Ok(()) + Arc::new( + move |request: LlmRequest, context: LlmSanitizeRequestContext| { + let context = js_llm_sanitize_request_context(&context); + let request = serde_json::to_value(request).unwrap_or(Json::Null); + let (tx, rx) = std::sync::mpsc::channel(); + if func.call_with_return_value( + (request.clone(), context), + ThreadsafeFunctionCallMode::Blocking, + move |value: Option| { + let _ = tx.send(callback_json(value)); + Ok(()) + }, + ) != napi::Status::Ok + { + record_callback_error( + "nemo_relay: failed to queue JS LLM sanitize request callback", + ); + return None; + } + let value = recv_middleware_json_or_value( + rx, + "nemo_relay: JS LLM request sanitizer callback failed", + Json::Null, + ); + if value.is_null() { + return None; + } + serde_json::from_value(value).map_or_else( + |error| { + record_callback_error(format!( + "nemo_relay: JS LLM sanitize request callback failed: failed to deserialize LlmRequest: {error}" + )); + None }, - ); - if status != napi::Status::Ok { - record_callback_error(format!( - "nemo_relay: failed to queue JS LLM sanitize request callback: {status:?}" - )); - return request; - } - // TODO: This closure returns LlmRequest (not Result), so we cannot propagate - // errors through the type system. Log the error so failures are not silent. - let result = recv_middleware_json_or_value( - rx, - "nemo_relay: JS LLM sanitize request callback failed", - serde_json::to_value(&request).unwrap_or(Json::Null), - ); - serde_json::from_value(result).unwrap_or_else(|error| { - record_callback_error(format!( - "nemo_relay: JS LLM sanitize request callback failed: failed to deserialize LlmRequest: {error}" - )); - request - }) - }) + Some, + ) + }, + ) } -/// Wrap a JS function for LLM sanitize response: `(response: Json) => Json`. -pub fn wrap_js_llm_response_fn( - func: ThreadsafeFunction, +/// Wrap a JS function for LLM response sanitization. The callback receives +/// `(response, context)`; returning `null` omits the event payload. +pub fn wrap_js_llm_sanitize_response_fn( + func: ThreadsafeFunction<(Json, JsLlmSanitizeResponseContext), ErrorStrategy::Fatal>, ) -> LlmSanitizeResponseFn { let func = Arc::new(func); - Arc::new(move |response: Json| { - let func = func.clone(); + Arc::new(move |response: Json, context: LlmSanitizeResponseContext| { + let context = js_llm_sanitize_response_context(&context); let (tx, rx) = std::sync::mpsc::channel(); - let status = func.call_with_return_value( - response.clone(), + if func.call_with_return_value( + (response, context), ThreadsafeFunctionCallMode::Blocking, - move |val: Option| { - let _ = tx.send(callback_json(val)); + move |value: Option| { + let _ = tx.send(callback_json(value)); Ok(()) }, - ); - if status != napi::Status::Ok { - record_callback_error(format!( - "nemo_relay: failed to queue JS LLM response callback: {status:?}" - )); - return response; + ) != napi::Status::Ok + { + record_callback_error("nemo_relay: failed to queue JS LLM sanitize response callback"); + return None; } - // TODO: This closure returns Json (not Result), so we cannot propagate - // errors through the type system. Log the error and fall back to original response. - recv_middleware_json_or_value(rx, "nemo_relay: JS LLM response callback failed", response) + let value = recv_middleware_json_or_value( + rx, + "nemo_relay: JS LLM response sanitizer callback failed", + Json::Null, + ); + Some(value).and_then(|value| (!value.is_null()).then_some(value)) }) } +fn js_llm_sanitize_request_context( + context: &LlmSanitizeRequestContext, +) -> JsLlmSanitizeRequestContext { + JsLlmSanitizeRequestContext { + codec: js_codec_identity(context.codec()), + resolved: context.resolve_codec(), + } +} + +fn js_codec_identity(identity: &LlmCodecIdentity) -> JsLlmCodecIdentity { + match identity { + LlmCodecIdentity::None => JsLlmCodecIdentity { + kind: "none".into(), + id: None, + }, + LlmCodecIdentity::BuiltIn(codec) => JsLlmCodecIdentity { + kind: "builtin".into(), + id: Some(codec.id().into()), + }, + LlmCodecIdentity::Runtime(id) => JsLlmCodecIdentity { + kind: "runtime".into(), + id: Some(id.clone()), + }, + LlmCodecIdentity::Opaque => JsLlmCodecIdentity { + kind: "opaque".into(), + id: None, + }, + } +} + +fn js_llm_sanitize_response_context( + context: &LlmSanitizeResponseContext, +) -> JsLlmSanitizeResponseContext { + JsLlmSanitizeResponseContext { + codec: js_codec_identity(context.codec()), + resolved: context.resolve_codec(), + } +} + +fn js_object_to_unknown(env: &Env, object: JsObject) -> JsUnknown { + unsafe { JsUnknown::from_raw_unchecked(env.raw(), object.raw()) } +} + +fn request_codec_object(env: &Env, codec: Arc) -> napi::Result { + let mut object = env.create_object()?; + let decode_codec = codec.clone(); + let decode = env.create_function_from_closure("decode", move |ctx| { + let request = ctx.get::(0)?; + let request = serde_json::from_value(request) + .map_err(|error| napi::Error::from_reason(format!("invalid LlmRequest: {error}")))?; + serde_json::to_value(decode_codec.decode(&request).map_err(to_napi_err)?) + .map_err(|error| napi::Error::from_reason(error.to_string())) + })?; + let encode_codec = codec; + let encode = env.create_function_from_closure("encode", move |ctx| { + let annotated = ctx.get::(0)?; + let original = ctx.get::(1)?; + let annotated = serde_json::from_value(annotated).map_err(|error| { + napi::Error::from_reason(format!("invalid AnnotatedLlmRequest: {error}")) + })?; + let original = serde_json::from_value(original) + .map_err(|error| napi::Error::from_reason(format!("invalid LlmRequest: {error}")))?; + serde_json::to_value( + encode_codec + .encode(&annotated, &original) + .map_err(to_napi_err)?, + ) + .map_err(|error| napi::Error::from_reason(error.to_string())) + })?; + object.set_named_property("decode", decode)?; + object.set_named_property("encode", encode)?; + Ok(object) +} + +fn response_codec_object(env: &Env, codec: Arc) -> napi::Result { + let mut object = env.create_object()?; + let decode = env.create_function_from_closure("decodeResponse", move |ctx| { + let response = ctx.get::(0)?; + serde_json::to_value(codec.decode_response(&response).map_err(to_napi_err)?) + .map_err(|error| napi::Error::from_reason(error.to_string())) + })?; + object.set_named_property("decodeResponse", decode)?; + Ok(object) +} + +/// Convert a request sanitizer context into the JavaScript object passed to a callback. +pub(crate) fn js_llm_sanitize_request_context_to_napi( + env: &Env, + context: JsLlmSanitizeRequestContext, +) -> napi::Result { + let mut object = env.create_object()?; + let codec = unsafe { + JsUnknown::from_raw_unchecked( + env.raw(), + JsLlmCodecIdentity::to_napi_value(env.raw(), context.codec)?, + ) + }; + object.set_named_property("codec", codec)?; + + let resolved = context.resolved; + let resolve_codec = + env.create_function_from_closure("resolveCodec", move |ctx| match resolved.clone() { + Some(codec) => Ok(js_object_to_unknown( + ctx.env, + request_codec_object(ctx.env, codec)?, + )), + None => ctx + .env + .get_null() + .map(|value| unsafe { JsUnknown::from_raw_unchecked(ctx.env.raw(), value.raw()) }), + })?; + object.set_named_property("resolveCodec", resolve_codec)?; + Ok(js_object_to_unknown(env, object)) +} + +/// Convert a response sanitizer context into the JavaScript object passed to a callback. +pub(crate) fn js_llm_sanitize_response_context_to_napi( + env: &Env, + context: JsLlmSanitizeResponseContext, +) -> napi::Result { + let mut object = env.create_object()?; + let codec = unsafe { + JsUnknown::from_raw_unchecked( + env.raw(), + JsLlmCodecIdentity::to_napi_value(env.raw(), context.codec)?, + ) + }; + object.set_named_property("codec", codec)?; + + let resolved = context.resolved; + let resolve_codec = + env.create_function_from_closure("resolveCodec", move |ctx| match resolved.clone() { + Some(codec) => Ok(js_object_to_unknown( + ctx.env, + response_codec_object(ctx.env, codec)?, + )), + None => ctx + .env + .get_null() + .map(|value| unsafe { JsUnknown::from_raw_unchecked(ctx.env.raw(), value.raw()) }), + })?; + object.set_named_property("resolveCodec", resolve_codec)?; + Ok(js_object_to_unknown(env, object)) +} + /// Wrap a JS function for LLM conditional guardrails: `(request: object) => string | null`. pub fn wrap_js_llm_conditional_fn( func: ThreadsafeFunction, @@ -683,11 +861,22 @@ pub fn wrap_js_event_sanitize_fn( struct NapiCodec { decode: Arc>, encode: Arc>, + register_thread: std::thread::ThreadId, + direct_decode: Arc Result + Send + Sync>, + direct_encode: Arc Result + Send + Sync>, } impl LlmCodec for NapiCodec { fn decode(&self, request: &LlmRequest) -> Result { let req_json = serde_json::to_value(request).unwrap_or(Json::Null); + if std::thread::current().id() == self.register_thread { + let result = (self.direct_decode)(req_json)?; + return serde_json::from_value(result).map_err(|e| { + FlowError::Internal(format!( + "JS codec decode callback: failed to deserialize AnnotatedLlmRequest: {e}" + )) + }); + } let (tx, rx) = std::sync::mpsc::channel(); let status = self.decode.call_with_return_value( req_json, @@ -714,6 +903,13 @@ impl LlmCodec for NapiCodec { let annotated_json = serde_json::to_value(annotated).unwrap_or(Json::Null); let original_json = serde_json::to_value(original).unwrap_or(Json::Null); let arg = serde_json::json!({"annotated": annotated_json, "original": original_json}); + if std::thread::current().id() == self.register_thread { + return serde_json::from_value((self.direct_encode)(arg)?).map_err(|e| { + FlowError::Internal(format!( + "JS codec encode callback: failed to deserialize LlmRequest: {e}" + )) + }); + } let (tx, rx) = std::sync::mpsc::channel(); let status = self.encode.call_with_return_value( arg, @@ -737,10 +933,16 @@ impl LlmCodec for NapiCodec { pub fn wrap_js_codec( decode: ThreadsafeFunction, encode: ThreadsafeFunction, + register_thread: std::thread::ThreadId, + direct_decode: Arc Result + Send + Sync>, + direct_encode: Arc Result + Send + Sync>, ) -> Arc { Arc::new(NapiCodec { decode: Arc::new(decode), encode: Arc::new(encode), + register_thread, + direct_decode, + direct_encode, }) } @@ -752,10 +954,20 @@ pub fn wrap_js_codec( /// delegating `decode_response` to a JavaScript function via `ThreadsafeFunction`. struct NapiResponseCodec { decode_response: Arc>, + register_thread: std::thread::ThreadId, + direct_decode_response: Arc Result + Send + Sync>, } impl LlmResponseCodec for NapiResponseCodec { fn decode_response(&self, response: &Json) -> Result { + if std::thread::current().id() == self.register_thread { + let result = (self.direct_decode_response)(response.clone())?; + return serde_json::from_value(result).map_err(|e| { + FlowError::Internal(format!( + "decode_response returned invalid AnnotatedLlmResponse: {e}" + )) + }); + } let (tx, rx) = std::sync::mpsc::channel(); let status = self.decode_response.call_with_return_value( response.clone(), @@ -784,9 +996,13 @@ impl LlmResponseCodec for NapiResponseCodec { /// Wrap a JS decode_response function into an `Arc`. pub fn wrap_js_response_codec( decode_response: ThreadsafeFunction, + register_thread: std::thread::ThreadId, + direct_decode_response: Arc Result + Send + Sync>, ) -> Arc { Arc::new(NapiResponseCodec { decode_response: Arc::new(decode_response), + register_thread, + direct_decode_response, }) } diff --git a/crates/node/src/stream.rs b/crates/node/src/stream.rs index c9308d66a..37b1e43a3 100644 --- a/crates/node/src/stream.rs +++ b/crates/node/src/stream.rs @@ -11,6 +11,9 @@ use napi::bindgen_prelude::*; use napi_derive::napi; use nemo_relay::error::Result as FlowResult; use serde_json::Value as Json; +use std::sync::Arc; + +use crate::api::PersistentJsFunction; /// An async iterator over chunks from a streaming LLM response. /// @@ -21,6 +24,7 @@ pub struct LlmStream { pub(crate) receiver: tokio::sync::Mutex>>, pub(crate) cancel: tokio::sync::watch::Sender, pub(crate) closed: tokio::sync::watch::Receiver>>, + pub(crate) codec_references: Vec>, } #[napi] diff --git a/crates/node/tests/callback_error_tests.mjs b/crates/node/tests/callback_error_tests.mjs index 94130142b..83e684a5f 100644 --- a/crates/node/tests/callback_error_tests.mjs +++ b/crates/node/tests/callback_error_tests.mjs @@ -35,7 +35,7 @@ function makeNative() { describe('callback error helpers', () => { it('getLastCallbackError and clearLastCallbackError expose malformed sanitize-request failures', async () => { clearLastCallbackError(); - registerLlmSanitizeRequestGuardrail('node_llm_san_req_public_error', 10, () => null); + registerLlmSanitizeRequestGuardrail('node_llm_san_req_public_error', 10, () => ({ broken: true })); try { const result = await llmCallExecute( 'san_req_public_error_llm', @@ -83,7 +83,7 @@ describe('callback error helpers', () => { clearLastCallbackError(); }); - it('closed llm sanitize-request callbacks fall back to the original request and record the queue failure', () => { + it('closed llm sanitize-request callbacks omit the payload and record the queue failure', () => { const request = makeNative(); const result = __testClosedLlmSanitizeRequestCallback( () => ({ @@ -91,12 +91,12 @@ describe('callback error helpers', () => { }), request, ); - assert.deepEqual(result, request); + assert.equal(result, null); assert.match(getLastCallbackError() ?? '', /failed to queue JS LLM sanitize request callback/i); clearLastCallbackError(); }); - it('closed llm sanitize-response callbacks fall back to the original response and record the queue failure', () => { + it('closed llm sanitize-response callbacks omit the payload and record the queue failure', () => { const response = { ok: true, }; @@ -106,8 +106,8 @@ describe('callback error helpers', () => { }), response, ); - assert.deepEqual(result, response); - assert.match(getLastCallbackError() ?? '', /failed to queue JS LLM response callback/i); + assert.equal(result, null); + assert.match(getLastCallbackError() ?? '', /failed to queue JS LLM sanitize response callback/i); clearLastCallbackError(); }); diff --git a/crates/node/tests/event_sanitizers_tests.mjs b/crates/node/tests/event_sanitizers_tests.mjs index fdd67520d..812a2a0ef 100644 --- a/crates/node/tests/event_sanitizers_tests.mjs +++ b/crates/node/tests/event_sanitizers_tests.mjs @@ -4,6 +4,9 @@ import assert from 'node:assert/strict'; import { describe, it } from 'node:test'; import { createRequire } from 'node:module'; +import { mkdtempSync, rmSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import path from 'node:path'; const require = createRequire(import.meta.url); const lib = require('../index.js'); @@ -28,6 +31,18 @@ function assertSanitizerFieldsCleared(event) { assert.equal(event.metadata, null); } +async function initializeWithoutDiscoveredPluginConfig(config) { + const previousDirectory = process.cwd(); + const directory = mkdtempSync(path.join(tmpdir(), 'nemo-relay-node-')); + try { + process.chdir(directory); + return await plugin.initialize(config); + } finally { + process.chdir(previousDirectory); + rmSync(directory, { recursive: true, force: true }); + } +} + describe('event sanitizer registries', () => { it('orders mark sanitizers and supports field removal', async () => { const events = capture('node-event-sanitize-order-sub'); @@ -251,7 +266,10 @@ describe('event sanitizer registries', () => { }, }); try { - await plugin.initialize({ version: 1, components: [plugin.ComponentSpec(kind)] }); + await initializeWithoutDiscoveredPluginConfig({ + version: 1, + components: [plugin.ComponentSpec(kind)], + }); lib.event('configured', null, { raw: true }); lib.flushSubscribers(); await waitFor(events, 1); @@ -289,7 +307,10 @@ describe('event sanitizer registries', () => { }); lib.clearLastCallbackError(); try { - await plugin.initialize({ version: 1, components: [plugin.ComponentSpec(kind)] }); + await initializeWithoutDiscoveredPluginConfig({ + version: 1, + components: [plugin.ComponentSpec(kind)], + }); lib.event('plugin-throw', null, { raw: true }, { raw: true }); lib.flushSubscribers(); await waitFor(events, 1); diff --git a/crates/node/tests/llm_tests.mjs b/crates/node/tests/llm_tests.mjs index bb4e9e821..3115b2136 100644 --- a/crates/node/tests/llm_tests.mjs +++ b/crates/node/tests/llm_tests.mjs @@ -3,11 +3,15 @@ import { describe, it } from 'node:test'; import assert from 'node:assert/strict'; +import { execFileSync } from 'node:child_process'; import { createRequire } from 'node:module'; import { readFileSync } from 'node:fs'; +import path from 'node:path'; +import { fileURLToPath } from 'node:url'; const require = createRequire(import.meta.url); const lib = require('../index.js'); +const nodeDir = fileURLToPath(new URL('..', import.meta.url)); const { pushScope, @@ -356,6 +360,275 @@ describe('LLM execute', () => { // =========================================================================== describe('LLM guardrails', () => { + it('contextual sanitizers receive payload first and codec context second', async () => { + const events = []; + let requestContextChecked = false; + let responseContextChecked = false; + registerSubscriber('node_contextual_llm_sanitize_events', (event) => events.push(event)); + registerLlmSanitizeRequestGuardrail('node_contextual_llm_sanitize_request', 10, (request, context) => { + assert.deepEqual(context.codec, { kind: 'none' }); + assert.equal(context.resolveCodec(), null); + requestContextChecked = true; + return { + ...request, + headers: { ...request.headers, 'X-Contextual-Sanitized': 'request' }, + }; + }); + registerLlmSanitizeResponseGuardrail('node_contextual_llm_sanitize_response', 10, (response, context) => { + assert.deepEqual(context.codec, { kind: 'none' }); + assert.equal(context.resolveCodec(), null); + responseContextChecked = true; + return { ...response, contextualSanitized: true }; + }); + + try { + const result = await llmCallExecute('contextual_sanitize_llm', makeNative(), () => ({ ok: true })); + assert.deepEqual(result, { ok: true }); + assert.equal(requestContextChecked, true); + assert.equal(responseContextChecked, true); + await flushSubscriberCallbacks(); + const start = events.find( + (event) => event.name === 'contextual_sanitize_llm' && event.scope_category === 'start', + ); + const end = events.find((event) => event.name === 'contextual_sanitize_llm' && event.scope_category === 'end'); + assert.equal(start.data.headers['X-Contextual-Sanitized'], 'request'); + assert.equal(end.data.contextualSanitized, true); + } finally { + deregisterLlmSanitizeRequestGuardrail('node_contextual_llm_sanitize_request'); + deregisterLlmSanitizeResponseGuardrail('node_contextual_llm_sanitize_response'); + deregisterSubscriber('node_contextual_llm_sanitize_events'); + } + }); + + it('rejects incomplete request codec callback pairs', () => { + const decode = (request) => request; + const encode = ({ annotated }) => annotated; + + assert.throws( + () => llmCallExecute('partial_codec_execute', makeNative(), () => ({}), null, null, null, null, null, decode), + /codecDecode and codecEncode must be provided together/, + ); + assert.throws( + () => + llmCallExecuteAsync( + 'partial_codec_execute_async', + makeNative(), + async () => ({}), + null, + null, + null, + null, + null, + null, + encode, + ), + /codecDecode and codecEncode must be provided together/, + ); + assert.throws( + () => + llmStreamCallExecute( + 'partial_codec_stream', + makeNative(), + () => {}, + null, + null, + null, + null, + null, + null, + null, + decode, + ), + /codecDecode and codecEncode must be provided together/, + ); + }); + + it('resolved sanitizer codecs expose directional operations', async () => { + const codec = new lib.OpenAIChatCodec(); + let requestDecoded = false; + let responseDecoded = false; + registerLlmSanitizeRequestGuardrail('node_resolved_llm_request_codec', 10, (request, context) => { + assert.deepEqual(context.codec, { kind: 'opaque' }); + const resolved = context.resolveCodec(); + assert.notEqual(resolved, null); + const annotated = resolved.decode(request); + requestDecoded = annotated.model === 'test-model'; + return resolved.encode(annotated, request); + }); + registerLlmSanitizeResponseGuardrail('node_resolved_llm_response_codec', 10, (response, context) => { + assert.deepEqual(context.codec, { kind: 'opaque' }); + const resolved = context.resolveCodec(); + assert.notEqual(resolved, null); + const annotated = resolved.decodeResponse(response); + responseDecoded = annotated.model === 'test-model'; + return response; + }); + + try { + const response = { + id: 'chatcmpl-test', + model: 'test-model', + choices: [ + { + index: 0, + message: { role: 'assistant', content: 'ok' }, + finish_reason: 'stop', + }, + ], + }; + const result = await llmCallExecute( + 'resolved_sanitizer_codec_llm', + makeNative(), + () => response, + null, + null, + null, + null, + null, + codec.decode.bind(codec), + ({ annotated, original }) => codec.encode(annotated, original), + codec.decodeResponse.bind(codec), + ); + assert.deepEqual(result, response); + assert.equal(requestDecoded, true); + assert.equal(responseDecoded, true); + } finally { + deregisterLlmSanitizeRequestGuardrail('node_resolved_llm_request_codec'); + deregisterLlmSanitizeResponseGuardrail('node_resolved_llm_response_codec'); + } + }); + + it('streaming sanitizers resolve codecs only during callbacks and can omit the final payload', async () => { + const codec = new lib.OpenAIChatCodec(); + const events = []; + let requestCodecUsed = false; + let responseCodecUsed = false; + registerSubscriber('node_streaming_codec_sanitize_events', (event) => events.push(event)); + registerLlmSanitizeRequestGuardrail('node_streaming_request_codec', 10, (request, context) => { + const resolved = context.resolveCodec(); + assert.notEqual(resolved, null); + const annotated = resolved.decode(request); + requestCodecUsed = true; + return resolved.encode(annotated, request); + }); + registerLlmSanitizeResponseGuardrail('node_streaming_response_codec', 10, (response, context) => { + const resolved = context.resolveCodec(); + assert.notEqual(resolved, null); + resolved.decodeResponse(response); + responseCodecUsed = true; + return null; + }); + + const response = { + id: 'chatcmpl-stream', + model: 'test-model', + choices: [{ index: 0, message: { role: 'assistant', content: 'secret' }, finish_reason: 'stop' }], + }; + try { + const stream = await llmStreamCallExecute( + 'streaming_resolved_sanitizer_codec', + makeNative(), + (wrapper) => { + lib.pushStreamChunk(wrapper.__nemo_relay_stream_id, { delta: 'secret' }); + lib.endStream(wrapper.__nemo_relay_stream_id); + }, + null, + () => response, + null, + null, + null, + null, + null, + codec.decode.bind(codec), + ({ annotated, original }) => codec.encode(annotated, original), + codec.decodeResponse.bind(codec), + ); + assert.deepEqual(await stream.next(), { delta: 'secret' }); + assert.equal(await stream.next(), null); + await flushSubscriberCallbacks(); + assert.equal(requestCodecUsed, true); + assert.equal(responseCodecUsed, true); + const end = events.find( + (event) => event.name === 'streaming_resolved_sanitizer_codec' && event.scope_category === 'end', + ); + assert.equal(end.data, null); + assert.equal(end.category_profile.annotated_response, undefined); + } finally { + deregisterLlmSanitizeRequestGuardrail('node_streaming_request_codec'); + deregisterLlmSanitizeResponseGuardrail('node_streaming_response_codec'); + deregisterSubscriber('node_streaming_codec_sanitize_events'); + } + }); + + it('releases custom stream codec references safely after early garbage collection', () => { + const modulePath = path.join(nodeDir, 'index.js'); + const script = ` + import { createRequire } from 'node:module'; + const require = createRequire(${JSON.stringify(path.join(nodeDir, 'package.json'))}); + const lib = require(${JSON.stringify(modulePath)}); + const codec = new lib.OpenAIChatCodec(); + let ended = false; + lib.registerSubscriber('early_drop_custom_codec_stream_events', (event) => { + if ( + event.name === 'early_drop_custom_codec_stream' && + event.scope_category === 'end' + ) { + ended = true; + } + }); + let stream = await lib.llmStreamCallExecute( + 'early_drop_custom_codec_stream', + { + headers: {}, + content: { model: 'test-model', messages: [] }, + }, + (wrapper) => { + setTimeout(() => { + lib.endStream(wrapper.__nemo_relay_stream_id); + }, 25); + }, + null, + () => ({ model: 'test-model', choices: [] }), + null, + null, + null, + null, + null, + codec.decode.bind(codec), + ({ annotated, original }) => codec.encode(annotated, original), + codec.decodeResponse.bind(codec), + ); + const weak = new WeakRef(stream); + stream = null; + let collected = false; + for (let index = 0; index < 100; index += 1) { + global.gc(); + await new Promise((resolve) => setImmediate(resolve)); + if (weak.deref() === undefined) { + collected = true; + break; + } + await new Promise((resolve) => setImmediate(resolve)); + } + if (!collected) { + throw new Error('unfinished custom-codec stream was not garbage collected'); + } + for (let index = 0; index < 100 && !ended; index += 1) { + global.gc(); + await new Promise((resolve) => setImmediate(resolve)); + } + if (!ended) { + throw new Error('early-dropped custom-codec stream did not finish cleanup'); + } + lib.deregisterSubscriber('early_drop_custom_codec_stream_events'); + await new Promise((resolve) => setImmediate(resolve)); + `; + execFileSync(process.execPath, ['--expose-gc', '--input-type=module', '--eval', script], { + stdio: 'inherit', + timeout: 30_000, + }); + }); + it('sanitize request guardrail', () => { registerLlmSanitizeRequestGuardrail('node_llm_san_req', 10, (request) => { request.extra = 'sanitized'; @@ -420,7 +693,7 @@ describe('LLM guardrails', () => { } }); - it('sanitize request guardrail falls back on malformed return', async () => { + it('sanitize request guardrail can omit the observability payload', async () => { registerLlmSanitizeRequestGuardrail('node_llm_san_req_bad', 10, () => null); try { const result = await llmCallExecute( @@ -445,7 +718,7 @@ describe('LLM guardrails', () => { } }); - it('sanitize request guardrail failures preserve the payload and remain usable', async () => { + it('sanitize request guardrail failures omit the payload and remain usable', async () => { const events = []; clearLastCallbackError(); registerSubscriber('node_llm_san_req_throw_sub', (event) => events.push(event)); @@ -463,7 +736,7 @@ describe('LLM guardrails', () => { event.category === 'llm' && event.scope_category === 'start', ); - assert.deepEqual(start.data, request); + assert.equal(start.data, null); assert.match(getLastCallbackError() ?? '', /JavaScript callback threw/i); deregisterLlmSanitizeRequestGuardrail('node_llm_san_req_throw'); @@ -567,7 +840,7 @@ describe('LLM guardrails', () => { } }); - it('sanitize response guardrail failures preserve the payload and remain usable', async () => { + it('sanitize response guardrail failures omit the payload and remain usable', async () => { const events = []; clearLastCallbackError(); registerSubscriber('node_llm_san_resp_throw_sub', (event) => events.push(event)); @@ -585,7 +858,7 @@ describe('LLM guardrails', () => { event.category === 'llm' && event.scope_category === 'end', ); - assert.deepEqual(end.data, response); + assert.equal(end.data, null); assert.match(getLastCallbackError() ?? '', /response sanitizer boom/i); deregisterLlmSanitizeResponseGuardrail('node_llm_san_resp_throw'); @@ -1074,6 +1347,14 @@ describe('LLM intercepts', () => { assert.equal(declarations.split(openKind).length - 1, 3); }); + it('generated LLM sanitizer declarations expose directional codec contexts', () => { + const declarations = readFileSync(new URL('../index.d.ts', import.meta.url), 'utf8'); + + assert.equal(declarations.split("context: import('./plugin').LlmSanitizeRequestContext").length - 1, 2); + assert.equal(declarations.split("context: import('./plugin').LlmSanitizeResponseContext").length - 1, 2); + assert.doesNotMatch(declarations, /registerLlmSanitizeRequestGuardrail\([^\n]*\.\.\.args: any\[\]/); + }); + it('standalone conditional execution helper throws on rejection', async () => { registerLlmConditionalExecutionGuardrail('node_llm_cond_helper', 10, () => 'llm blocked by helper'); try { diff --git a/crates/pii-redaction/README.md b/crates/pii-redaction/README.md index 1b0f820da..4c8528582 100644 --- a/crates/pii-redaction/README.md +++ b/crates/pii-redaction/README.md @@ -85,7 +85,6 @@ sanitization surface: kind = "pii_redaction" [components.config] -codec = "openai_chat" [[components.config.profiles]] mode = "builtin" @@ -124,7 +123,6 @@ enabled = true [components.config] version = 1 -codec = "anthropic_messages" [[components.config.profiles]] mode = "builtin" diff --git a/crates/pii-redaction/src/builtin.rs b/crates/pii-redaction/src/builtin.rs index 43cc9cdb8..6e6d1dc5c 100644 --- a/crates/pii-redaction/src/builtin.rs +++ b/crates/pii-redaction/src/builtin.rs @@ -6,16 +6,19 @@ use std::sync::Arc; use regex::Regex; use serde::Serialize; use serde::de::DeserializeOwned; -use serde_json::Value as Json; +use serde_json::{Map, Value as Json}; use sha2::{Digest, Sha256}; use nemo_relay::api::event::{CategoryProfile, Event}; use nemo_relay::api::llm::LlmRequest; use nemo_relay::api::runtime::{ - EventSanitizeFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, ToolSanitizeFn, + BuiltinLlmCodec, EventSanitizeFn, LlmCodecIdentity, LlmSanitizeRequestFn, + LlmSanitizeResponseFn, ToolSanitizeFn, }; +use nemo_relay::codec::request::AnnotatedLlmRequest; use nemo_relay::codec::resolve::{ - ProviderSurface, request_codec as build_request_codec, response_codec as build_response_codec, + ProviderSurface, detect_response_surface, request_codec as build_request_codec, + response_codec as build_response_codec, }; use nemo_relay::codec::traits::{LlmCodec, LlmResponseCodec}; use nemo_relay::plugin::{PluginError, Result as PluginResult}; @@ -29,9 +32,7 @@ use super::trajectory::{CustomMarkPayloadPolicy, TrajectorySanitizer}; pub(super) struct CompiledBuiltinBackend { action: BuiltinAction, target_paths: Arc>, - request_codec: Option>, - response_codec: Option>, - codec_name: Option, + legacy_surface: Option, trajectory: Option, } @@ -168,9 +169,7 @@ impl CompiledBuiltinBackend { Ok(Self { action, target_paths: Arc::new(config.target_paths), - request_codec: surface.map(build_request_codec), - response_codec: surface.map(build_response_codec), - codec_name: surface.map(BuiltinCodecName::from_provider_surface), + legacy_surface: surface, trajectory, }) } @@ -286,19 +285,171 @@ impl CompiledBuiltinBackend { } } - fn sanitize_request_with_codec(&self, request: &LlmRequest) -> Option { - let codec = self.request_codec.as_ref()?; + fn selected_surface(&self, codec: &LlmCodecIdentity) -> Option { + match codec { + LlmCodecIdentity::None => self.legacy_surface, + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) => { + Some(ProviderSurface::OpenAIChat) + } + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiResponses) => { + Some(ProviderSurface::OpenAIResponses) + } + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) => { + Some(ProviderSurface::AnthropicMessages) + } + LlmCodecIdentity::Runtime(_) | LlmCodecIdentity::Opaque => None, + } + } + + fn uses_compatible_legacy_response_codec(&self, payload: &Json) -> bool { + self.legacy_surface + .is_some_and(|surface| detect_response_surface(payload) == Some(surface)) + } + + fn sanitize_request_with_codec( + &self, + codec: &dyn LlmCodec, + request: &LlmRequest, + ) -> Option { let annotated = codec.decode(request).ok()?; let sanitized_annotated = sanitize_serializable_with_backend(self, annotated).ok()?; - codec.encode(&sanitized_annotated, request).ok() + codec + .encode(&sanitized_annotated, request) + .ok() + .or_else(|| { + self.sanitize_request_target_paths_incrementally( + codec, + request, + sanitized_annotated, + ) + }) + } + + fn sanitize_request_target_paths_incrementally( + &self, + codec: &dyn LlmCodec, + request: &LlmRequest, + sanitized_annotated: AnnotatedLlmRequest, + ) -> Option { + let sanitized = serde_json::to_value(sanitized_annotated).ok()?; + let mut sanitized_request = request.clone(); + + for target_path in self.target_paths.iter() { + let target_segments = json_pointer_segments(target_path)?; + let current_annotated = codec.decode(&sanitized_request).ok()?; + let mut current = serde_json::to_value(¤t_annotated).ok()?; + match ( + sanitized_json_pointer_value(&sanitized, &target_segments), + sanitized_json_pointer_value(¤t, &target_segments), + ) { + (None, None) => continue, + (Some(target_value), Some(current_value)) if current_value == target_value => { + continue; + } + (Some(target_value), Some(_)) => { + replace_sanitized_json_pointer_value( + &mut current, + &target_segments, + target_value.clone(), + )?; + } + (None, Some(_)) if matches!(self.action, BuiltinAction::Remove) => { + remove_sanitized_json_pointer_value(&mut current, &target_segments)?; + } + _ => return None, + } + let updated = serde_json::from_value(current).ok()?; + sanitized_request = codec.encode(&updated, &sanitized_request).ok()?; + } + + Some(sanitized_request) + } + + fn sanitize_request_headers(&self, headers: Map) -> Map { + let sanitized = self.sanitize_json_preorder_dfs(Json::Object( + [("headers".to_string(), Json::Object(headers))] + .into_iter() + .collect(), + )); + sanitized + .get("headers") + .and_then(Json::as_object) + .cloned() + .unwrap_or_default() } - fn sanitize_response_with_codec(&self, payload: Json) -> Option { - let codec = self.response_codec.as_ref()?; - let codec_name = self.codec_name?; + fn sanitize_response_with_codec( + &self, + codec: &dyn LlmResponseCodec, + surface: ProviderSurface, + payload: Json, + ) -> Option { + if surface == ProviderSurface::OpenAIChat + && payload + .get("choices") + .and_then(Json::as_array) + .is_some_and(|choices| choices.len() > 1) + && self.targets_normalized_openai_chat_choice() + { + return None; + } + let codec_name = BuiltinCodecName::from_provider_surface(surface); let annotated = codec.decode_response(&payload).ok()?; + let annotated_json = serde_json::to_value(&annotated).ok()?; let sanitized_annotated = sanitize_serializable_with_backend(self, annotated).ok()?; - Some(codec_name.overlay_response_payload(payload, &sanitized_annotated)) + let sanitized_json = serde_json::to_value(&sanitized_annotated).ok()?; + let payload = codec_name.overlay_response_payload(payload, &sanitized_annotated); + let payload = self.sanitize_json_preorder_dfs(payload); + let target_segments = self + .target_paths + .iter() + .map(|target_path| json_pointer_segments(target_path)) + .collect::>>()?; + let has_normalized_target = target_segments.iter().any(|target_segments| { + sanitized_json_pointer_value(&annotated_json, target_segments).is_some() + || sanitized_json_pointer_value(&sanitized_json, target_segments).is_some() + }); + if !has_normalized_target { + return Some(payload); + } + let projected = codec.decode_response(&payload).ok()?; + let projected = serde_json::to_value(projected).ok()?; + Self::normalized_response_targets_match( + &target_segments, + &annotated_json, + &sanitized_json, + &projected, + ) + .then_some(payload) + } + + fn normalized_response_targets_match( + target_paths: &[Vec], + annotated: &Json, + sanitized: &Json, + projected: &Json, + ) -> bool { + target_paths.iter().all(|target_segments| { + let original = sanitized_json_pointer_value(annotated, target_segments); + let expected = sanitized_json_pointer_value(sanitized, target_segments); + if original.is_none() && expected.is_none() { + return true; + } + sanitized_json_pointer_value(projected, target_segments) == expected + }) + } + + fn targets_normalized_openai_chat_choice(&self) -> bool { + self.target_paths.iter().any(|path| { + json_pointer_segments(path) + .and_then(|segments| segments.into_iter().next()) + .is_some_and(|segment| { + matches!( + segment.as_str(), + "message" | "tool_calls" | "finish_reason" | "api_specific" + ) + }) + }) } } @@ -368,37 +519,183 @@ fn event_sanitize_callback_with_scope_categories( pub(super) fn llm_sanitize_request_callback( backend: CompiledBuiltinBackend, ) -> LlmSanitizeRequestFn { - Arc::new(move |mut request: LlmRequest| { + Arc::new(move |mut request: LlmRequest, context| { if let Some(trajectory) = backend.trajectory.as_ref() { + request.headers = trajectory + .sanitize_tool_payload(Json::Object(request.headers)) + .as_object() + .cloned() + .unwrap_or_default(); request.content = trajectory.sanitize_provider_payload(request.content); - return request; + return Some(request); } - if let Some(encoded) = backend.sanitize_request_with_codec(&request) { - return encoded; + request.headers = backend.sanitize_request_headers(request.headers); + if backend.target_paths.is_empty() { + request.content = backend.sanitize_json_preorder_dfs(request.content); + return Some(request); + } + let resolved = context.resolve_codec(); + let fallback = if resolved.is_none() { + backend + .selected_surface(context.codec()) + .map(build_request_codec) + } else { + None + }; + let Some(codec) = resolved.as_deref().or(fallback.as_deref()) else { + log_llm_payload_omitted("request", context.codec(), "no usable request codec"); + return None; + }; + let sanitized = backend.sanitize_request_with_codec(codec, &request); + if sanitized.is_none() { + log_llm_payload_omitted( + "request", + context.codec(), + "codec decode, sanitize, or encode failure", + ); } - request.content = backend.sanitize_json_preorder_dfs(request.content); - request + sanitized }) } pub(super) fn llm_sanitize_response_callback( backend: CompiledBuiltinBackend, ) -> LlmSanitizeResponseFn { - Arc::new(move |payload: Json| { + Arc::new(move |payload: Json, context| { if let Some(trajectory) = backend.trajectory.as_ref() { - return trajectory.sanitize_provider_payload(payload); + return Some(trajectory.sanitize_provider_payload(payload)); } if backend.target_paths.is_empty() { - return backend.sanitize_json_preorder_dfs(payload); + return Some(backend.sanitize_json_preorder_dfs(payload)); } - - let payload = backend - .sanitize_response_with_codec(payload.clone()) - .unwrap_or(payload); - backend.sanitize_json_preorder_dfs(payload) + if matches!(context.codec(), LlmCodecIdentity::None) + && !backend.uses_compatible_legacy_response_codec(&payload) + { + log_llm_payload_omitted( + "response", + context.codec(), + "no active response codec or compatible legacy codec", + ); + return None; + } + let Some(surface) = backend.selected_surface(context.codec()) else { + log_llm_payload_omitted( + "response", + context.codec(), + "no recognized response codec surface", + ); + return None; + }; + let resolved = context.resolve_codec(); + let fallback = if resolved.is_none() { + Some(build_response_codec(surface)) + } else { + None + }; + let Some(codec) = resolved.as_deref().or(fallback.as_deref()) else { + log_llm_payload_omitted("response", context.codec(), "no usable response codec"); + return None; + }; + let sanitized = backend.sanitize_response_with_codec(codec, surface, payload); + if sanitized.is_none() { + log_llm_payload_omitted( + "response", + context.codec(), + "codec decode, sanitize, or encode failure", + ); + } + sanitized }) } +fn log_llm_payload_omitted(direction: &str, codec: &LlmCodecIdentity, reason: &str) { + let codec_kind = match codec { + LlmCodecIdentity::None => "none", + LlmCodecIdentity::BuiltIn(_) => "builtin", + LlmCodecIdentity::Runtime(_) => "runtime", + LlmCodecIdentity::Opaque => "opaque", + }; + log::warn!( + target: "nemo_relay.plugin", + event = "pii_llm_payload_omitted", + codec_kind, + reason; + "PII redaction omitted an LLM {direction} payload" + ); +} + +fn json_pointer_segments(pointer: &str) -> Option> { + pointer + .strip_prefix('/') + .map(|path| path.split('/').map(unescape_json_pointer_segment).collect()) +} + +fn unescape_json_pointer_segment(segment: &str) -> String { + segment.replace("~1", "/").replace("~0", "~") +} + +fn sanitized_json_pointer_value<'a>(value: &'a Json, segments: &[String]) -> Option<&'a Json> { + segments + .iter() + .try_fold(value, |value, segment| match value { + Json::Object(values) => values.get(segment), + Json::Array(values) => segment + .parse::() + .ok() + .and_then(|index| values.get(index)), + _ => None, + }) +} + +fn replace_sanitized_json_pointer_value( + value: &mut Json, + segments: &[String], + replacement: Json, +) -> Option<()> { + let (last, parents) = segments.split_last()?; + let parent = parents + .iter() + .try_fold(value, |value, segment| match value { + Json::Object(values) => values.get_mut(segment), + Json::Array(values) => segment + .parse::() + .ok() + .and_then(|index| values.get_mut(index)), + _ => None, + })?; + match parent { + Json::Object(values) => { + values.insert(last.clone(), replacement); + Some(()) + } + Json::Array(values) => { + let index = last.parse::().ok()?; + let value = values.get_mut(index)?; + *value = replacement; + Some(()) + } + _ => None, + } +} + +fn remove_sanitized_json_pointer_value(value: &mut Json, segments: &[String]) -> Option<()> { + let (last, parents) = segments.split_last()?; + let parent = parents + .iter() + .try_fold(value, |value, segment| match value { + Json::Object(values) => values.get_mut(segment), + Json::Array(values) => segment + .parse::() + .ok() + .and_then(|index| values.get_mut(index)), + _ => None, + })?; + match parent { + Json::Object(values) => values.remove(last).map(|_| ()), + Json::Array(_) | Json::Null | Json::Bool(_) | Json::Number(_) | Json::String(_) => None, + } +} + fn render_json_pointer_path(path_segments: &[String]) -> String { if path_segments.is_empty() { return String::new(); diff --git a/crates/pii-redaction/src/component.rs b/crates/pii-redaction/src/component.rs index 0c4c3b415..2f6b9bfd3 100644 --- a/crates/pii-redaction/src/component.rs +++ b/crates/pii-redaction/src/component.rs @@ -97,7 +97,8 @@ pub struct PiiRedactionConfig { skip_serializing_if = "is_default_priority" )] pub priority: i32, - /// Provider request/response codec for LLM-managed surfaces. + /// Compatibility fallback codec for LLM-managed surfaces without an active + /// per-call codec. #[serde(default, skip_serializing_if = "Option::is_none")] #[cfg_attr(feature = "schema", schemars(schema_with = "codec_schema"))] pub codec: Option, @@ -1113,14 +1114,6 @@ fn validate_codec_requirements( } let Some(codec) = config.codec.as_deref() else { - push_policy_diag( - diagnostics, - policy.unsupported_value, - "pii_redaction.unsupported_value", - Some(PII_REDACTION_PLUGIN_KIND.to_string()), - Some("codec".to_string()), - "codec is required when any LLM surface is enabled".to_string(), - ); return; }; diff --git a/crates/pii-redaction/src/overlay.rs b/crates/pii-redaction/src/overlay.rs index 2cab7c35f..e74da6e26 100644 --- a/crates/pii-redaction/src/overlay.rs +++ b/crates/pii-redaction/src/overlay.rs @@ -88,8 +88,12 @@ fn overlay_openai_responses_response(mut payload: Json, annotated: &AnnotatedLlm .map(openai_responses_status), ); + let message_text = annotated_message_text(annotated.message.as_ref()); + if root.contains_key("output_text") { + set_optional_string_field(root, "output_text", message_text.as_deref()); + } if let Some(items) = root.get_mut("output").and_then(Json::as_array_mut) { - overlay_output_text_blocks(items, annotated_message_text(annotated.message.as_ref())); + overlay_output_text_blocks(items, message_text); overlay_openai_responses_tool_calls(items, annotated.tool_calls.as_deref()); } payload diff --git a/crates/pii-redaction/tests/unit/component_tests.rs b/crates/pii-redaction/tests/unit/component_tests.rs index 3ac54253e..cd7ad4026 100644 --- a/crates/pii-redaction/tests/unit/component_tests.rs +++ b/crates/pii-redaction/tests/unit/component_tests.rs @@ -14,8 +14,9 @@ use crate::api::llm::{ llm_call_execute, llm_stream_call_execute, }; use crate::api::runtime::{ - LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, NemoRelayContextState, - create_scope_stack, global_context, set_thread_scope_stack, + BuiltinLlmCodec, LlmCodecIdentity, LlmExecutionNextFn, LlmJsonStream, + LlmSanitizeRequestContext, LlmSanitizeResponseContext, LlmStreamExecutionNextFn, + NemoRelayContextState, create_scope_stack, global_context, set_thread_scope_stack, }; use crate::api::scope::{ EmitMarkEventParams, PopScopeParams, PushScopeParams, ScopeType, event, pop_scope, push_scope, @@ -24,12 +25,13 @@ use crate::api::subscriber::{deregister_subscriber, register_subscriber}; use crate::api::tool::{ToolCallEndParams, ToolCallParams, tool_call, tool_call_end}; use crate::codec::openai_chat::OpenAIChatCodec; use crate::codec::openai_responses::OpenAIResponsesCodec; +use crate::codec::request::AnnotatedLlmRequest; use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::plugin::{ ConfigPolicy, DiagnosticLevel, PluginComponentSpec, PluginConfig, PluginError, PluginRegistrationContext, UnsupportedBehavior, clear_plugin_configuration, - ensure_builtin_plugins_registered, initialize_plugins, list_plugin_kinds, - rollback_registrations, validate_plugin_config, + ensure_builtin_plugins_registered, initialize_plugins_exact as initialize_plugins, + list_plugin_kinds, rollback_registrations, validate_plugin_config, }; use futures::StreamExt; use nemo_relay::observability::atif::{AtifAgentInfo, AtifExporter}; @@ -263,6 +265,253 @@ fn trajectory_backend(codec: Option<&str>, policy: &str) -> crate::builtin::Comp .unwrap() } +fn no_codec_context() -> LlmSanitizeResponseContext { + LlmSanitizeResponseContext::default() +} + +fn no_codec_request_context() -> LlmSanitizeRequestContext { + LlmSanitizeRequestContext::default() +} + +struct IdentifiedRequestCodec { + identity: LlmCodecIdentity, + inner: OpenAIResponsesCodec, +} + +impl LlmCodec for IdentifiedRequestCodec { + fn codec_identity(&self) -> LlmCodecIdentity { + self.identity.clone() + } + + fn decode(&self, request: &LlmRequest) -> nemo_relay::error::Result { + self.inner.decode(request) + } + + fn encode( + &self, + annotated: &AnnotatedLlmRequest, + original: &LlmRequest, + ) -> nemo_relay::error::Result { + self.inner.encode(annotated, original) + } +} + +#[test] +fn normalized_llm_paths_use_the_active_codec_and_fail_closed_for_unknown_codecs() { + let backend = crate::builtin::CompiledBuiltinBackend::new( + BuiltinBackendConfig { + action: "regex_replace".to_string(), + pattern: Some("sk-[A-Za-z0-9_-]+".to_string()), + replacement: Some("[REDACTED]".to_string()), + target_paths: vec![ + "/messages/0/content/0/text".to_string(), + "/message".to_string(), + ], + ..BuiltinBackendConfig::default() + }, + Some("openai_chat".to_string()), + ) + .unwrap(); + let sanitize_request = crate::builtin::llm_sanitize_request_callback(backend.clone()); + let sanitize_response = crate::builtin::llm_sanitize_response_callback(backend); + let active_request = sanitize_request( + LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4.1-mini", + "input": [{"role": "user", "content": [{"type": "input_text", "text": "sk-request-secret"}]}] + }), + }, + LlmSanitizeRequestContext::for_request_codec(Some(Arc::new(OpenAIResponsesCodec))), + ) + .expect("the active OpenAI Responses codec must override the legacy fallback"); + assert_eq!( + active_request.content["input"][0]["content"][0]["text"], + json!("[REDACTED]") + ); + + for identity in [ + LlmCodecIdentity::Runtime("com.example.responses.v1".to_owned()), + LlmCodecIdentity::Opaque, + ] { + let active_request = sanitize_request( + LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4.1-mini", + "input": [{ + "role": "user", + "content": [{"type": "input_text", "text": "sk-request-secret"}] + }] + }), + }, + LlmSanitizeRequestContext::for_request_codec(Some(Arc::new(IdentifiedRequestCodec { + identity, + inner: OpenAIResponsesCodec, + }))), + ) + .expect("an active runtime or opaque request codec must remain usable"); + assert_eq!( + active_request.content["input"][0]["content"][0]["text"], + json!("[REDACTED]") + ); + } + + let responses_payload = json!({ + "id": "resp_123", + "model": "gpt-4.1-mini", + "status": "completed", + "output": [{ + "type": "message", + "content": [{"type": "output_text", "text": "sk-responses-secret"}] + }] + }); + + let active_responses = sanitize_response( + responses_payload.clone(), + LlmSanitizeResponseContext::with_identity(LlmCodecIdentity::BuiltIn( + BuiltinLlmCodec::OpenAiResponses, + )), + ) + .expect("the active OpenAI Responses codec must override the legacy fallback"); + assert_eq!( + active_responses["output"][0]["content"][0]["text"], + json!("[REDACTED]") + ); + + assert!( + sanitize_response(responses_payload.clone(), no_codec_context()).is_none(), + "an incompatible configured fallback codec must omit a normalized payload" + ); + + assert!( + sanitize_response( + responses_payload, + LlmSanitizeResponseContext::with_identity(LlmCodecIdentity::Opaque), + ) + .is_none(), + "a normalized-path policy must omit an unknown active provider payload" + ); + + assert!( + sanitize_response( + json!({ + "id": "resp_123", + "output": [{"content": [{"type": "output_text", "text": "sk-runtime-secret"}]}] + }), + LlmSanitizeResponseContext::with_identity(LlmCodecIdentity::Runtime( + "com.example.chat.v1".to_owned(), + )), + ) + .is_none(), + "a normalized-path policy must omit a runtime codec until it has a compatible projection" + ); +} + +#[test] +fn normalized_llm_paths_omit_payloads_when_legacy_codec_decode_fails() { + let backend = crate::builtin::CompiledBuiltinBackend::new( + BuiltinBackendConfig { + action: "regex_replace".to_string(), + pattern: Some("sk-[A-Za-z0-9_-]+".to_string()), + replacement: Some("[REDACTED]".to_string()), + target_paths: vec!["/messages/0/content".to_string()], + ..BuiltinBackendConfig::default() + }, + Some("openai_chat".to_string()), + ) + .unwrap(); + let sanitize_request = crate::builtin::llm_sanitize_request_callback(backend.clone()); + let sanitize_response = crate::builtin::llm_sanitize_response_callback(backend); + + assert!( + sanitize_request( + LlmRequest { + headers: serde_json::Map::new(), + content: json!({"messages": "sk-request-secret"}), + }, + no_codec_request_context(), + ) + .is_none(), + "a shallow legacy surface match must not enable a raw-payload fallback" + ); + assert!( + sanitize_response(json!({"choices": "sk-response-secret"}), no_codec_context()).is_none(), + "a legacy response codec failure must omit the payload instead of emitting raw content" + ); +} + +#[test] +fn normalized_openai_chat_api_specific_policy_omits_multiple_choices() { + let backend = crate::builtin::CompiledBuiltinBackend::new( + BuiltinBackendConfig { + action: "remove".to_string(), + target_paths: vec!["/api_specific".to_string()], + ..BuiltinBackendConfig::default() + }, + None, + ) + .unwrap(); + let sanitize_response = crate::builtin::llm_sanitize_response_callback(backend); + + assert!( + sanitize_response( + json!({ + "id": "chatcmpl-multi", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "first"}, + "logprobs": {"content": [{"token": "SECRET-FIRST"}]} + }, + { + "index": 1, + "message": {"role": "assistant", "content": "second"}, + "logprobs": {"content": [{"token": "SECRET-SECOND"}]} + } + ] + }), + LlmSanitizeResponseContext::with_identity(LlmCodecIdentity::BuiltIn( + BuiltinLlmCodec::OpenAiChat, + )), + ) + .is_none() + ); +} + +#[test] +fn normalized_llm_paths_use_configured_anthropic_codec_without_a_system_message() { + let backend = crate::builtin::CompiledBuiltinBackend::new( + BuiltinBackendConfig { + action: "regex_replace".to_string(), + pattern: Some("sk-[A-Za-z0-9_-]+".to_string()), + replacement: Some("[REDACTED]".to_string()), + target_paths: vec!["/messages/0/content".to_string()], + ..BuiltinBackendConfig::default() + }, + Some("anthropic_messages".to_string()), + ) + .unwrap(); + let sanitize_request = crate::builtin::llm_sanitize_request_callback(backend); + + let sanitized = sanitize_request( + LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "claude-sonnet-4-6", + "messages": [{"role": "user", "content": "sk-anthropic-secret"}], + }), + }, + no_codec_request_context(), + ) + .expect("the configured Anthropic codec must sanitize a valid message-only request"); + + assert_eq!( + sanitized.content["messages"][0]["content"], + json!("[REDACTED]") + ); +} + #[test] fn trajectory_preset_redacts_chat_content_without_erasing_request_structure() { let callback = crate::builtin::llm_sanitize_request_callback(trajectory_backend( @@ -298,7 +547,8 @@ fn trajectory_preset_redacts_chat_content_without_erasing_request_structure() { "participant": {"name": "Alice Example", "username": "alice"}, "person_name": "Alice Example" }), - }); + }, no_codec_request_context()) + .unwrap(); assert_eq!(request.content["model"], "claude-sonnet-4-6"); assert_eq!(request.content["temperature"], 0.2); @@ -351,19 +601,23 @@ fn trajectory_preset_preserves_response_analytics_and_redacts_response_content() Some("openai_chat"), "preserve", )); - let sanitized = callback(json!({ - "id": "chatcmpl_1", - "model": "claude-opus-4-6", - "choices": [{"index": 0, "finish_reason": "tool_calls", "message": { - "role": "assistant", - "content": "private answer", - "tool_calls": [{"id": "call_1", "type": "function", "function": { - "name": "terminal", "arguments": "{\"command\":\"cat secret.txt\"}" - }}] - }, "logprobs": {"content": [{"token": "secret", "logprob": -0.5}]}}], - "usage": {"prompt_tokens": 20, "completion_tokens": 5, "total_tokens": 25}, - "cost": {"total": 1.25} - })); + let sanitized = callback( + json!({ + "id": "chatcmpl_1", + "model": "claude-opus-4-6", + "choices": [{"index": 0, "finish_reason": "tool_calls", "message": { + "role": "assistant", + "content": "private answer", + "tool_calls": [{"id": "call_1", "type": "function", "function": { + "name": "terminal", "arguments": "{\"command\":\"cat secret.txt\"}" + }}] + }, "logprobs": {"content": [{"token": "secret", "logprob": -0.5}]}}], + "usage": {"prompt_tokens": 20, "completion_tokens": 5, "total_tokens": 25}, + "cost": {"total": 1.25} + }), + no_codec_context(), + ) + .unwrap(); assert_eq!(sanitized["id"], "chatcmpl_1"); assert_eq!(sanitized["model"], "claude-opus-4-6"); @@ -402,7 +656,8 @@ fn trajectory_preset_covers_responses_and_anthropic_provider_shapes() { "reasoning": {"effort": "high", "summary": "private reasoning"}, "max_output_tokens": 100 }), - }); + }, no_codec_request_context()) + .unwrap(); assert_eq!(responses_request.content["model"], "gpt-5"); assert_eq!(responses_request.content["input"][0]["role"], "user"); assert_eq!( @@ -414,15 +669,19 @@ fn trajectory_preset_covers_responses_and_anthropic_provider_shapes() { let responses_response = crate::builtin::llm_sanitize_response_callback(trajectory_backend( Some("openai_responses"), "preserve", - ))(json!({ - "id": "resp_1", - "model": "gpt-5", - "status": "completed", - "output": [{"id": "msg_1", "type": "message", "role": "assistant", "content": [ - {"type": "output_text", "text": "private output"} - ]}], - "usage": {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14} - })); + ))( + json!({ + "id": "resp_1", + "model": "gpt-5", + "status": "completed", + "output": [{"id": "msg_1", "type": "message", "role": "assistant", "content": [ + {"type": "output_text", "text": "private output"} + ]}], + "usage": {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14} + }), + no_codec_context(), + ) + .unwrap(); assert_eq!(responses_response["id"], "resp_1"); assert_eq!(responses_response["status"], "completed"); assert_eq!(responses_response["output"][0]["id"], "msg_1"); @@ -446,7 +705,8 @@ fn trajectory_preset_covers_responses_and_anthropic_provider_shapes() { ]}], "max_tokens": 128 }), - }); + }, no_codec_request_context()) + .unwrap(); assert_eq!(anthropic_request.content["model"], "claude-sonnet-4-6"); assert_eq!(anthropic_request.content["system"], "[REDACTED]"); assert_eq!( @@ -474,7 +734,8 @@ fn trajectory_preset_covers_responses_and_anthropic_provider_shapes() { ], "stop_reason": "end_turn", "usage": {"input_tokens": 12, "output_tokens": 6, "cache_read_input_tokens": 8} - })); + }), no_codec_context()) + .unwrap(); assert_eq!(anthropic_response["id"], "msg_1"); assert_eq!(anthropic_response["role"], "assistant"); assert_eq!(anthropic_response["content"][0]["type"], "thinking"); @@ -1451,7 +1712,7 @@ fn validate_rejects_builtin_mode_without_builtin_section() { } #[test] -fn validate_rejects_llm_surfaces_without_codec() { +fn validate_allows_llm_surfaces_without_codec() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); reset_runtime(); @@ -1464,12 +1725,7 @@ fn validate_rejects_llm_surfaces_without_codec() { "output": false, }))); - assert!(report.diagnostics.iter().any(|diag| { - diag.field.as_deref() == Some("codec") - && diag - .message - .contains("codec is required when any LLM surface is enabled") - })); + assert!(report.diagnostics.is_empty(), "{report:?}"); } #[test] @@ -2029,6 +2285,93 @@ async fn trajectory_preset_sanitizes_stream_finalization_without_changing_client clear_plugin_configuration().unwrap(); } +#[tokio::test] +async fn normalized_paths_use_the_active_codec_for_stream_finalization() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "openai_responses", + "input": false, + "output": true, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "regex_replace", + "pattern": "sk-[A-Za-z0-9_-]+", + "replacement": "[REDACTED]", + "target_paths": ["/message"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-active-codec-stream-finalization"); + let provider: LlmStreamExecutionNextFn = Arc::new(move |_request| { + Box::pin(async move { + Ok(LlmJsonStream::new(futures::stream::iter(vec![Ok(json!({ + "id": "chatcmpl-stream", + "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {"content": "sk-client-visible"}}] + }))]))) + }) + }); + let request_codec: Arc = Arc::new(OpenAIChatCodec); + let response_codec: Arc = Arc::new(OpenAIChatCodec); + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("openai") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hello"}]}), + }) + .func(provider) + .collector(Box::new(|_| Ok(()))) + .finalizer(Box::new(|| { + json!({ + "id": "chatcmpl-stream", + "model": "gpt-4o-mini", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "sk-stream-secret"}, + "finish_reason": "stop" + }] + }) + })) + .codec(request_codec) + .response_codec(response_codec) + .build(), + ) + .await + .unwrap(); + + assert_eq!( + stream.next().await.unwrap().unwrap()["choices"][0]["delta"]["content"], + json!("sk-client-visible") + ); + assert!(stream.next().await.is_none()); + + let captured = captured_events_snapshot(&events); + let end = captured + .iter() + .find(|event| event.scope_category() == Some(ScopeCategory::End)) + .unwrap(); + assert_eq!( + end.output().unwrap()["choices"][0]["message"]["content"], + json!("[REDACTED]") + ); + assert_eq!( + end.annotated_response() + .and_then(|response| response.response_text()), + Some("[REDACTED]") + ); + + deregister_subscriber("pii-active-codec-stream-finalization").unwrap(); + clear_plugin_configuration().unwrap(); +} + #[test] fn builtin_backend_sanitizes_tool_start_and_end_payloads_with_preorder_targets() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); @@ -3743,7 +4086,7 @@ fn builtin_backend_sanitizes_llm_start_payload_via_codec_and_reencodes_provider_ } #[tokio::test] -async fn builtin_backend_sanitizes_llm_end_payload_and_response_codec_decodes_sanitized_output() { +async fn builtin_backend_removes_targeted_message_names_and_ignores_missing_normalized_paths() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); reset_runtime(); setup_isolated_thread(); @@ -3751,50 +4094,286 @@ async fn builtin_backend_sanitizes_llm_end_payload_and_response_codec_decodes_sa initialize_plugins(plugin_config(json!({ "mode": "builtin", "codec": "openai_chat", - "input": false, - "output": true, + "input": true, + "output": false, "tool_input": false, "tool_output": false, "builtin": { - "action": "regex_replace", - "pattern": "sk-[A-Za-z0-9_-]+", - "replacement": "[REDACTED]", - "target_paths": ["/choices/0/message/content", "/audit_owner"] + "action": "remove", + "target_paths": [ + "/messages/0/missing", + "/messages/0/name", + "/messages/1/name" + ] } }))) .await .unwrap(); - let events = capture_events("pii-redaction-llm-end-events"); + let events = capture_events("pii-redaction-remove-message-names"); let request = LlmRequest { headers: serde_json::Map::new(), content: json!({ "model": "gpt-4o-mini", "messages": [ - {"role": "user", "content": "hello"} + {"role": "system", "name": "SECRET-SYSTEM", "content": "instructions"}, + {"role": "user", "name": "SECRET-USER", "content": "hello"} ] }), }; - let response = json!({ - "id": "chatcmpl-123", - "model": "gpt-4o-mini", - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": "sk-response-secret" - }, - "finish_reason": "stop" - } - ], - "usage": { - "prompt_tokens": 3, - "completion_tokens": 2, - "total_tokens": 5 - } - }); - let response_codec: Arc = Arc::new(OpenAIChatCodec); + + llm_call( + LlmCallParams::builder() + .name("openai") + .request(&request) + .build(), + ) + .unwrap(); + + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 1); + assert_eq!( + captured_events[0].input(), + Some(&json!({ + "headers": {}, + "content": { + "model": "gpt-4o-mini", + "messages": [ + {"role": "system", "content": "instructions"}, + {"role": "user", "content": "hello"} + ] + } + })) + ); + assert!( + !serde_json::to_string(&captured_events[0]) + .unwrap() + .contains("SECRET-") + ); + + deregister_subscriber("pii-redaction-remove-message-names").unwrap(); + clear_plugin_configuration().unwrap(); +} + +#[tokio::test] +async fn builtin_backend_omits_request_and_annotation_for_unsafe_normalized_array_removal() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "openai_responses", + "input": true, + "output": false, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "remove", + "target_paths": ["/messages/0/content/0"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-unsafe-normalized-array-removal"); + let request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4.1-mini", + "input": [{ + "role": "user", + "content": [{"type": "input_text", "text": "SECRET-REQUEST"}] + }] + }), + }; + let annotated_request = Arc::new(OpenAIResponsesCodec.decode(&request).unwrap()); + + llm_call( + LlmCallParams::builder() + .name("openai") + .request(&request) + .annotated_request(annotated_request) + .build(), + ) + .unwrap(); + + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 1); + assert!(captured_events[0].input().is_none()); + assert!(captured_events[0].annotated_request().is_none()); + assert!( + !serde_json::to_string(&captured_events[0]) + .unwrap() + .contains("SECRET-REQUEST") + ); + + deregister_subscriber("pii-redaction-unsafe-normalized-array-removal").unwrap(); + clear_plugin_configuration().unwrap(); +} + +#[tokio::test] +async fn builtin_backend_sanitizes_observable_llm_request_headers() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "input": true, + "output": false, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "redact", + "detector": "email" + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-llm-request-headers"); + let request = LlmRequest { + headers: [("x-user-email".to_string(), json!("alice@example.com"))] + .into_iter() + .collect(), + content: json!({ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}] + }), + }; + + llm_call( + LlmCallParams::builder() + .name("openai") + .request(&request) + .build(), + ) + .unwrap(); + + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 1); + assert_eq!( + captured_events[0].input().unwrap()["headers"]["x-user-email"], + json!("[REDACTED]") + ); + assert!( + !serde_json::to_string(&captured_events[0]) + .unwrap() + .contains("alice@example.com") + ); + + deregister_subscriber("pii-redaction-llm-request-headers").unwrap(); + clear_plugin_configuration().unwrap(); +} + +#[tokio::test] +async fn trajectory_preset_redacts_opaque_llm_request_header_values() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "input": true, + "output": false, + "tool_input": false, + "tool_output": false, + "builtin": { + "preset": "trajectory_context" + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-trajectory-llm-request-headers"); + let request = LlmRequest { + headers: [("model".to_string(), json!("SECRET-HEADER"))] + .into_iter() + .collect(), + content: json!({ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}] + }), + }; + + llm_call( + LlmCallParams::builder() + .name("openai") + .request(&request) + .build(), + ) + .unwrap(); + + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 1); + assert_eq!( + captured_events[0].input().unwrap()["headers"]["model"], + json!("[REDACTED]") + ); + assert!( + !serde_json::to_string(&captured_events[0]) + .unwrap() + .contains("SECRET-HEADER") + ); + + deregister_subscriber("pii-trajectory-llm-request-headers").unwrap(); + clear_plugin_configuration().unwrap(); +} + +#[tokio::test] +async fn builtin_backend_sanitizes_llm_end_payload_and_response_codec_decodes_sanitized_output() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "openai_chat", + "input": false, + "output": true, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "regex_replace", + "pattern": "sk-[A-Za-z0-9_-]+", + "replacement": "[REDACTED]", + "target_paths": ["/choices/0/message/content", "/audit_owner"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-llm-end-events"); + let request = LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4o-mini", + "messages": [ + {"role": "user", "content": "hello"} + ] + }), + }; + let response = json!({ + "id": "chatcmpl-123", + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "sk-response-secret" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 3, + "completion_tokens": 2, + "total_tokens": 5 + } + }); + let response_codec: Arc = Arc::new(OpenAIChatCodec); let result = llm_call_execute( LlmCallExecuteParams::builder() @@ -3915,6 +4494,142 @@ async fn builtin_backend_sanitizes_openai_chat_response_from_normalized_message_ clear_plugin_configuration().unwrap(); } +#[tokio::test] +async fn builtin_backend_omits_unprojectable_openai_chat_api_specific_response() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "openai_chat", + "input": false, + "output": true, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "remove", + "target_paths": ["/api_specific/system_fingerprint"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-openai-chat-api-specific-response"); + let response = json!({ + "id": "chatcmpl-api-specific", + "model": "gpt-4o-mini", + "system_fingerprint": "SECRET-FINGERPRINT", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop" + }] + }); + + let result = llm_call_execute( + LlmCallExecuteParams::builder() + .name("openai") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}] + }), + }) + .func(noop_openai_chat_exec_fn(response.clone())) + .response_codec(Arc::new(OpenAIChatCodec)) + .build(), + ) + .await + .unwrap(); + + assert_eq!(result, response); + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 2); + assert!(captured_events[1].output().is_none()); + assert!(captured_events[1].annotated_response().is_none()); + assert!( + !serde_json::to_string(&captured_events[1]) + .unwrap() + .contains("SECRET-FINGERPRINT") + ); + + deregister_subscriber("pii-redaction-openai-chat-api-specific-response").unwrap(); + clear_plugin_configuration().unwrap(); +} + +#[tokio::test] +async fn builtin_backend_omits_multi_choice_openai_chat_normalized_response() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "openai_chat", + "input": false, + "output": true, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "regex_replace", + "pattern": "sk-[A-Za-z0-9_-]+", + "replacement": "[REDACTED]", + "target_paths": ["/message"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-openai-chat-multi-choice-response"); + let response = json!({ + "id": "chatcmpl-multi", + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "sk-first-secret"}, + "finish_reason": "stop" + }, + { + "index": 1, + "message": {"role": "assistant", "content": "sk-second-secret"}, + "finish_reason": "stop" + } + ] + }); + + let result = llm_call_execute( + LlmCallExecuteParams::builder() + .name("openai") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}] + }), + }) + .func(noop_openai_chat_exec_fn(response.clone())) + .response_codec(Arc::new(OpenAIChatCodec)) + .build(), + ) + .await + .unwrap(); + + assert_eq!(result, response); + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 2); + assert!(captured_events[1].output().is_none()); + assert!(captured_events[1].annotated_response().is_none()); + let serialized_end = serde_json::to_string(&captured_events[1]).unwrap(); + assert!(!serialized_end.contains("sk-first-secret")); + assert!(!serialized_end.contains("sk-second-secret")); + + deregister_subscriber("pii-redaction-openai-chat-multi-choice-response").unwrap(); + clear_plugin_configuration().unwrap(); +} + #[tokio::test] async fn builtin_redact_sanitizes_openai_chat_response_from_detector_path() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); @@ -4044,6 +4759,65 @@ async fn builtin_backend_sanitizes_anthropic_response_from_normalized_message_pa clear_plugin_configuration().unwrap(); } +#[tokio::test] +async fn builtin_backend_omits_unprojectable_anthropic_normalized_usage_response() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "codec": "anthropic_messages", + "input": false, + "output": true, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "remove", + "target_paths": ["/usage/prompt_tokens"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-anthropic-normalized-usage-response"); + let response = json!({ + "id": "msg_usage", + "model": "claude-sonnet-4-20250514", + "role": "assistant", + "type": "message", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 123, "output_tokens": 4} + }); + + let result = llm_call_execute( + LlmCallExecuteParams::builder() + .name("anthropic") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "claude-sonnet-4-20250514", + "messages": [{"role": "user", "content": "hello"}] + }), + }) + .func(noop_openai_chat_exec_fn(response.clone())) + .response_codec(Arc::new(crate::codec::anthropic::AnthropicMessagesCodec)) + .build(), + ) + .await + .unwrap(); + + assert_eq!(result, response); + let captured_events = captured_events_snapshot(&events); + assert_eq!(captured_events.len(), 2); + assert!(captured_events[1].output().is_none()); + assert!(captured_events[1].annotated_response().is_none()); + + deregister_subscriber("pii-redaction-anthropic-normalized-usage-response").unwrap(); + clear_plugin_configuration().unwrap(); +} + #[tokio::test] async fn builtin_backend_sanitizes_openai_responses_response_from_normalized_message_path() { let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); @@ -4052,7 +4826,7 @@ async fn builtin_backend_sanitizes_openai_responses_response_from_normalized_mes initialize_plugins(plugin_config(json!({ "mode": "builtin", - "codec": "openai_responses", + "codec": "openai_chat", "input": false, "output": true, "tool_input": false, @@ -4081,6 +4855,7 @@ async fn builtin_backend_sanitizes_openai_responses_response_from_normalized_mes "id": "resp_123", "model": "gpt-4.1-mini", "status": "completed", + "output_text": "sk-responses-secret", "output": [ { "type": "message", @@ -4101,13 +4876,117 @@ async fn builtin_backend_sanitizes_openai_responses_response_from_normalized_mes captured_events[1].output().unwrap()["output"][0]["content"][0]["text"], json!("[REDACTED]") ); + assert_eq!( + captured_events[1].output().unwrap()["output_text"], + json!("[REDACTED]") + ); assert_eq!( captured_events[1] .annotated_response() .and_then(|response| response.response_text()), Some("[REDACTED]") ); + assert!( + !captured_events[1] + .to_json_string() + .unwrap() + .contains("sk-responses-secret") + ); deregister_subscriber("pii-redaction-openai-responses-normalized-response").unwrap(); clear_plugin_configuration().unwrap(); } + +#[tokio::test] +async fn builtin_backend_sanitizes_openai_responses_output_text_alias_on_stream_finalization() { + let _guard = crate::plugins::pii_redaction::test_mutex().lock().unwrap(); + reset_runtime(); + setup_isolated_thread(); + + initialize_plugins(plugin_config(json!({ + "mode": "builtin", + "input": false, + "output": true, + "tool_input": false, + "tool_output": false, + "builtin": { + "action": "regex_replace", + "pattern": "sk-[A-Za-z0-9_-]+", + "replacement": "[REDACTED]", + "target_paths": ["/message"] + } + }))) + .await + .unwrap(); + + let events = capture_events("pii-redaction-openai-responses-output-text-stream"); + let provider: LlmStreamExecutionNextFn = Arc::new(move |_request| { + Box::pin(async move { + Ok(LlmJsonStream::new(futures::stream::iter(vec![Ok(json!({ + "type": "response.output_text.delta", + "delta": "sk-client-visible" + }))]))) + }) + }); + let request_codec: Arc = Arc::new(OpenAIResponsesCodec); + let response_codec: Arc = Arc::new(OpenAIResponsesCodec); + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("openai") + .request(LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "gpt-4.1-mini", + "input": "hello" + }), + }) + .func(provider) + .collector(Box::new(|_| Ok(()))) + .finalizer(Box::new(|| { + json!({ + "id": "resp_stream", + "model": "gpt-4.1-mini", + "status": "completed", + "output_text": "sk-stream-secret", + "output": [{ + "type": "message", + "content": [{ + "type": "output_text", + "text": "sk-stream-secret" + }] + }] + }) + })) + .codec(request_codec) + .response_codec(response_codec) + .build(), + ) + .await + .unwrap(); + + assert_eq!( + stream.next().await.unwrap().unwrap()["delta"], + json!("sk-client-visible") + ); + assert!(stream.next().await.is_none()); + + let captured_events = captured_events_snapshot(&events); + let end = captured_events + .iter() + .find(|event| event.scope_category() == Some(ScopeCategory::End)) + .unwrap(); + assert_eq!( + end.output().unwrap()["output"][0]["content"][0]["text"], + json!("[REDACTED]") + ); + assert_eq!(end.output().unwrap()["output_text"], json!("[REDACTED]")); + assert_eq!( + end.annotated_response() + .and_then(|response| response.response_text()), + Some("[REDACTED]") + ); + assert!(!end.to_json_string().unwrap().contains("sk-stream-secret")); + + deregister_subscriber("pii-redaction-openai-responses-output-text-stream").unwrap(); + clear_plugin_configuration().unwrap(); +} diff --git a/crates/plugin/src/lib.rs b/crates/plugin/src/lib.rs index 056661db0..727e09aa2 100644 --- a/crates/plugin/src/lib.rs +++ b/crates/plugin/src/lib.rs @@ -36,7 +36,52 @@ use serde::{Serialize, de::DeserializeOwned}; use serde_json::Map; /// Native plugin ABI version supported by this crate. -pub const NEMO_RELAY_NATIVE_ABI_VERSION: u32 = 1; +pub const NEMO_RELAY_NATIVE_ABI_VERSION: u32 = 2; + +/// Built-in LLM codec identities available to native plugins. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BuiltinLlmCodec { + /// OpenAI Chat Completions. + #[serde(rename = "openai_chat")] + OpenAiChat, + /// OpenAI Responses. + #[serde(rename = "openai_responses")] + OpenAiResponses, + /// Anthropic Messages. + #[serde(rename = "anthropic_messages")] + AnthropicMessages, +} + +/// Per-call LLM codec identity delivered to native plugins. +#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, serde::Deserialize)] +#[serde(tag = "kind", content = "id", rename_all = "snake_case")] +pub enum LlmCodecIdentity { + /// No codec was active. + #[default] + None, + /// A Relay built-in codec was active. + #[serde(rename = "builtin")] + BuiltIn(BuiltinLlmCodec), + /// A runtime-registered codec was active, identified by its stable ID. + Runtime(String), + /// A codec was active but has no registered identity. + Opaque, +} + +/// Per-call request codec context delivered to an LLM sanitizer. +pub struct LlmSanitizeRequestContext<'a> { + /// Identity of the active codec. + pub codec: LlmCodecIdentity, + resolved: Option>, +} + +/// Per-call response codec context delivered to an LLM sanitizer. +pub struct LlmSanitizeResponseContext<'a> { + /// Identity of the active codec. + pub codec: LlmCodecIdentity, + resolved: Option>, +} /// Status codes returned by stable native ABI functions. #[repr(i32)] @@ -73,6 +118,144 @@ pub struct NemoRelayNativeString { _marker: PhantomData<(*mut u8, PhantomPinned)>, } +/// Opaque callback-scoped request codec capability owned by the host. +#[repr(C)] +pub struct NemoRelayNativeLlmRequestCodec { + _private: [u8; 0], + _marker: PhantomData<(*mut u8, PhantomPinned)>, +} + +/// Opaque callback-scoped response codec capability owned by the host. +#[repr(C)] +pub struct NemoRelayNativeLlmResponseCodec { + _private: [u8; 0], + _marker: PhantomData<(*mut u8, PhantomPinned)>, +} + +/// Discriminator for the codec supplied to an LLM sanitizer over the native ABI. +#[repr(u32)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum NemoRelayNativeLlmCodecKind { + /// No codec was active for this call. + None = 0, + /// A Relay built-in codec was active. + BuiltIn = 1, + /// A runtime-registered codec was active. + Runtime = 2, + /// A codec was active but has no registered identity. + Opaque = 3, +} + +/// Per-call LLM sanitizer context passed over the native ABI. +/// +/// `codec_id` is borrowed for the duration of the callback. It is null for +/// [`NemoRelayNativeLlmCodecKind::None`] and +/// [`NemoRelayNativeLlmCodecKind::Opaque`]. For `BuiltIn`, it is one of the +/// stable built-in codec IDs; for `Runtime`, it is the registered codec ID. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NemoRelayNativeLlmSanitizeRequestContext { + /// Discriminator for the active codec. + pub codec_kind: NemoRelayNativeLlmCodecKind, + /// Optional borrowed codec identifier. + pub codec_id: *const NemoRelayNativeString, + /// Borrowed request codec capability, or null when no codec is active. + pub codec: *const NemoRelayNativeLlmRequestCodec, +} + +/// Per-call response sanitizer context passed over the native ABI. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NemoRelayNativeLlmSanitizeResponseContext { + /// Discriminator for the active codec. + pub codec_kind: NemoRelayNativeLlmCodecKind, + /// Optional borrowed codec identifier. + pub codec_id: *const NemoRelayNativeString, + /// Borrowed response codec capability, or null when no codec is active. + pub codec: *const NemoRelayNativeLlmResponseCodec, +} + +/// Safe callback-scoped request codec facade for typed native plugins. +pub struct LlmSanitizeRequestCodec<'a> { + host: NemoRelayNativeHostApiV1, + handle: *const NemoRelayNativeLlmRequestCodec, + _lifetime: PhantomData<&'a NemoRelayNativeLlmRequestCodec>, +} + +impl LlmSanitizeRequestCodec<'_> { + /// Decode an opaque request into Relay's normalized request model. + pub fn decode(&self, request: &LlmRequest) -> Result { + native_codec_call(&self.host, |out| unsafe { + let request = HostString::from_json(&self.host, request) + .ok_or_else(|| "failed to serialize LLM request".to_string())?; + codec_status( + &self.host, + (self.host.llm_request_codec_decode)(self.handle, request.as_ptr(), out), + ) + }) + } + + /// Encode normalized changes onto the original opaque request. + pub fn encode( + &self, + annotated: &AnnotatedLlmRequest, + original: &LlmRequest, + ) -> Result { + native_codec_call(&self.host, |out| unsafe { + let annotated = HostString::from_json(&self.host, annotated) + .ok_or_else(|| "failed to serialize annotated request".to_string())?; + let original = HostString::from_json(&self.host, original) + .ok_or_else(|| "failed to serialize original request".to_string())?; + codec_status( + &self.host, + (self.host.llm_request_codec_encode)( + self.handle, + annotated.as_ptr(), + original.as_ptr(), + out, + ), + ) + }) + } +} + +/// Safe callback-scoped response codec facade for typed native plugins. +pub struct LlmSanitizeResponseCodec<'a> { + host: NemoRelayNativeHostApiV1, + handle: *const NemoRelayNativeLlmResponseCodec, + _lifetime: PhantomData<&'a NemoRelayNativeLlmResponseCodec>, +} + +impl LlmSanitizeResponseCodec<'_> { + /// Decode an opaque response into Relay's normalized response model. + pub fn decode(&self, response: &Json) -> Result { + native_codec_call(&self.host, |out| unsafe { + let response = HostString::from_json(&self.host, response) + .ok_or_else(|| "failed to serialize LLM response".to_string())?; + codec_status( + &self.host, + (self.host.llm_response_codec_decode)(self.handle, response.as_ptr(), out), + ) + }) + } +} + +impl<'a> LlmSanitizeRequestContext<'a> { + /// Resolve the active request codec capability. + #[must_use] + pub fn resolve_codec(&self) -> Option<&LlmSanitizeRequestCodec<'a>> { + self.resolved.as_ref() + } +} + +impl<'a> LlmSanitizeResponseContext<'a> { + /// Resolve the active response codec capability. + #[must_use] + pub fn resolve_codec(&self) -> Option<&LlmSanitizeResponseCodec<'a>> { + self.resolved.as_ref() + } +} + /// Opaque plugin registration context borrowed from the host during registration. #[repr(C)] pub struct NemoRelayNativePluginContext { @@ -241,17 +424,27 @@ pub type NemoRelayNativeToolExecutionCb = unsafe extern "C" fn( out_outcome_json: *mut *mut NemoRelayNativeString, ) -> NemoRelayStatus; -/// Native LLM request transform callback for request sanitizers. -pub type NemoRelayNativeLlmRequestCb = unsafe extern "C" fn( +/// 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 +/// host releases an aliased input/output once. Any other non-null output must +/// be host-allocated and transfers ownership to the host. +pub type NemoRelayNativeLlmSanitizeRequestCb = unsafe extern "C" fn( user_data: *mut c_void, request_json: *const NemoRelayNativeString, + context: NemoRelayNativeLlmSanitizeRequestContext, out_request_json: *mut *mut NemoRelayNativeString, ) -> NemoRelayStatus; -/// Native JSON transform callback for LLM response sanitizers. -pub type NemoRelayNativeJsonCb = unsafe extern "C" fn( +/// Native LLM response sanitizer callback. Return a successful null output to +/// omit the observability payload and annotation. `payload_json` is borrowed, +/// but may be written directly to `out_json` as a pass-through; the host +/// releases an aliased input/output once. Any other non-null output must be +/// host-allocated and transfers ownership to the host. +pub type NemoRelayNativeLlmSanitizeResponseCb = unsafe extern "C" fn( user_data: *mut c_void, payload_json: *const NemoRelayNativeString, + context: NemoRelayNativeLlmSanitizeResponseContext, out_json: *mut *mut NemoRelayNativeString, ) -> NemoRelayStatus; @@ -334,6 +527,25 @@ pub struct NemoRelayNativeHostApiV1 { pub last_error_clear: unsafe extern "C" fn(), /// Sets the host thread-local native ABI error message. pub last_error_set: unsafe extern "C" fn(message: *const NemoRelayNativeString), + /// Decodes an LLM request through a callback-scoped codec capability. + pub llm_request_codec_decode: unsafe extern "C" fn( + codec: *const NemoRelayNativeLlmRequestCodec, + request_json: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, + ) -> NemoRelayStatus, + /// Encodes normalized request changes through a callback-scoped codec capability. + pub llm_request_codec_encode: unsafe extern "C" fn( + codec: *const NemoRelayNativeLlmRequestCodec, + annotated_json: *const NemoRelayNativeString, + original_json: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, + ) -> NemoRelayStatus, + /// Decodes an LLM response through a callback-scoped codec capability. + pub llm_response_codec_decode: unsafe extern "C" fn( + codec: *const NemoRelayNativeLlmResponseCodec, + response_json: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, + ) -> NemoRelayStatus, /// Registers an event subscriber through the plugin context. pub plugin_context_register_subscriber: unsafe extern "C" fn( ctx: *mut NemoRelayNativePluginContext, @@ -399,7 +611,7 @@ pub struct NemoRelayNativeHostApiV1 { ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, priority: i32, - cb: NemoRelayNativeLlmRequestCb, + cb: NemoRelayNativeLlmSanitizeRequestCb, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus, @@ -409,7 +621,7 @@ pub struct NemoRelayNativeHostApiV1 { ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, priority: i32, - cb: NemoRelayNativeJsonCb, + cb: NemoRelayNativeLlmSanitizeResponseCb, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus, @@ -1526,7 +1738,10 @@ impl<'a> PluginContext<'a> { callback: F, ) -> Result<()> where - F: Fn(LlmRequest) -> LlmRequest + Send + Sync + 'static, + F: for<'ctx> Fn(LlmRequest, LlmSanitizeRequestContext<'ctx>) -> Option + + Send + + Sync + + 'static, { let user_data = typed_callback_user_data(self.host, callback); let status = unsafe { @@ -1554,7 +1769,10 @@ impl<'a> PluginContext<'a> { callback: F, ) -> Result<()> where - F: Fn(Json) -> Json + Send + Sync + 'static, + F: for<'ctx> Fn(Json, LlmSanitizeResponseContext<'ctx>) -> Option + + Send + + Sync + + 'static, { let user_data = typed_callback_user_data(self.host, callback); let status = unsafe { @@ -1655,7 +1873,7 @@ impl<'a> PluginContext<'a> { /// Registers a typed LLM stream execution intercept. /// - /// Native ABI v1 represents stream execution as one JSON result. The host + /// Native ABI v2 represents stream execution as one JSON result. The host /// wraps that result as a one-chunk stream. pub fn register_llm_stream_execution_intercept( &mut self, @@ -1890,7 +2108,7 @@ impl<'a> PluginContext<'a> { &mut self, name: &str, priority: i32, - cb: NemoRelayNativeLlmRequestCb, + cb: NemoRelayNativeLlmSanitizeRequestCb, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus { @@ -1911,7 +2129,7 @@ impl<'a> PluginContext<'a> { &mut self, name: &str, priority: i32, - cb: NemoRelayNativeJsonCb, + cb: NemoRelayNativeLlmSanitizeResponseCb, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus { @@ -2274,10 +2492,14 @@ where unsafe extern "C" fn typed_llm_sanitize_request_trampoline( user_data: *mut c_void, request_json: *const NemoRelayNativeString, + context: NemoRelayNativeLlmSanitizeRequestContext, out_request_json: *mut *mut NemoRelayNativeString, ) -> NemoRelayStatus where - F: Fn(LlmRequest) -> LlmRequest + Send + Sync + 'static, + F: for<'a> Fn(LlmRequest, LlmSanitizeRequestContext<'a>) -> Option + + Send + + Sync + + 'static, { if user_data.is_null() || out_request_json.is_null() { return NemoRelayStatus::NullPointer; @@ -2285,9 +2507,14 @@ where unsafe { *out_request_json = ptr::null_mut() }; let state = unsafe { &*(user_data as *const TypedCallback) }; let result = catch_unwind(AssertUnwindSafe(|| { + let context = llm_sanitize_request_context_from_native(&state.host, context)?; let request: LlmRequest = read_json_value(&state.host, request_json, "LLM request")?; - let output = (state.callback)(request); - Ok::<_, NemoRelayStatus>(write_json(&state.host, &output, out_request_json)) + match (state.callback)(request, context) { + Some(output) => { + Ok::<_, NemoRelayStatus>(write_json(&state.host, &output, out_request_json)) + } + None => Ok(NemoRelayStatus::Ok), + } })); match result { Ok(Ok(status)) => status, @@ -2299,10 +2526,11 @@ where unsafe extern "C" fn typed_llm_sanitize_response_trampoline( user_data: *mut c_void, payload_json: *const NemoRelayNativeString, + context: NemoRelayNativeLlmSanitizeResponseContext, out_json: *mut *mut NemoRelayNativeString, ) -> NemoRelayStatus where - F: Fn(Json) -> Json + Send + Sync + 'static, + F: for<'a> Fn(Json, LlmSanitizeResponseContext<'a>) -> Option + Send + Sync + 'static, { if user_data.is_null() || out_json.is_null() { return NemoRelayStatus::NullPointer; @@ -2310,9 +2538,12 @@ where unsafe { *out_json = ptr::null_mut() }; let state = unsafe { &*(user_data as *const TypedCallback) }; let result = catch_unwind(AssertUnwindSafe(|| { + let context = llm_sanitize_response_context_from_native(&state.host, context)?; let payload: Json = read_json_value(&state.host, payload_json, "LLM response")?; - let output = (state.callback)(payload); - Ok::<_, NemoRelayStatus>(write_json(&state.host, &output, out_json)) + match (state.callback)(payload, context) { + Some(output) => Ok::<_, NemoRelayStatus>(write_json(&state.host, &output, out_json)), + None => Ok(NemoRelayStatus::Ok), + } })); match result { Ok(Ok(status)) => status, @@ -2625,6 +2856,29 @@ impl Drop for HostString<'_> { } } +fn codec_status(host: &NemoRelayNativeHostApiV1, status: NemoRelayStatus) -> Result<()> { + if status == NemoRelayStatus::Ok { + Ok(()) + } else { + Err(status_error(host, status, "LLM codec operation")) + } +} + +fn native_codec_call( + host: &NemoRelayNativeHostApiV1, + call: impl FnOnce(*mut *mut NemoRelayNativeString) -> Result<()>, +) -> Result { + let mut out = ptr::null_mut(); + call(&mut out)?; + if out.is_null() { + return Err("LLM codec operation returned null".into()); + } + let out = HostString { host, ptr: out }; + let text = read_host_string(host, out.as_ptr()) + .map_err(|_| "LLM codec operation returned invalid UTF-8".to_string())?; + serde_json::from_str(&text).map_err(|error| format!("invalid LLM codec result: {error}")) +} + struct OptionalHostJson<'a>(Option>); impl<'a> OptionalHostJson<'a> { @@ -2765,6 +3019,60 @@ fn read_optional_json_value( } } +fn llm_codec_identity_from_native( + host: &NemoRelayNativeHostApiV1, + codec_kind: NemoRelayNativeLlmCodecKind, + codec_id: *const NemoRelayNativeString, +) -> std::result::Result { + let codec = match codec_kind { + NemoRelayNativeLlmCodecKind::None => LlmCodecIdentity::None, + NemoRelayNativeLlmCodecKind::Opaque => LlmCodecIdentity::Opaque, + NemoRelayNativeLlmCodecKind::BuiltIn => { + let id = read_required_host_string(host, codec_id, "LLM built-in codec ID")?; + let builtin = match id.as_str() { + "openai_chat" => BuiltinLlmCodec::OpenAiChat, + "openai_responses" => BuiltinLlmCodec::OpenAiResponses, + "anthropic_messages" => BuiltinLlmCodec::AnthropicMessages, + _ => { + set_last_error(host, &format!("unknown built-in LLM codec ID: {id}")); + return Err(NemoRelayStatus::InvalidArg); + } + }; + LlmCodecIdentity::BuiltIn(builtin) + } + NemoRelayNativeLlmCodecKind::Runtime => LlmCodecIdentity::Runtime( + read_required_host_string(host, codec_id, "LLM runtime codec ID")?, + ), + }; + Ok(codec) +} + +fn llm_sanitize_request_context_from_native<'a>( + host: &NemoRelayNativeHostApiV1, + context: NemoRelayNativeLlmSanitizeRequestContext, +) -> std::result::Result, NemoRelayStatus> { + let codec = llm_codec_identity_from_native(host, context.codec_kind, context.codec_id)?; + let resolved = (!context.codec.is_null()).then_some(LlmSanitizeRequestCodec { + host: *host, + handle: context.codec, + _lifetime: PhantomData, + }); + Ok(LlmSanitizeRequestContext { codec, resolved }) +} + +fn llm_sanitize_response_context_from_native<'a>( + host: &NemoRelayNativeHostApiV1, + context: NemoRelayNativeLlmSanitizeResponseContext, +) -> std::result::Result, NemoRelayStatus> { + let codec = llm_codec_identity_from_native(host, context.codec_kind, context.codec_id)?; + let resolved = (!context.codec.is_null()).then_some(LlmSanitizeResponseCodec { + host: *host, + handle: context.codec, + _lifetime: PhantomData, + }); + Ok(LlmSanitizeResponseContext { codec, resolved }) +} + enum HostStringReadError { Null, InvalidUtf8, diff --git a/crates/plugin/tests/typed_callbacks.rs b/crates/plugin/tests/typed_callbacks.rs index 4a1b20041..fb3562bd3 100644 --- a/crates/plugin/tests/typed_callbacks.rs +++ b/crates/plugin/tests/typed_callbacks.rs @@ -13,13 +13,16 @@ use std::sync::{ }; use nemo_relay_plugin::{ - AnnotatedLlmRequest, CategoryProfile, ConfigDiagnostic, DiagnosticLevel, Event, EventCategory, - EventSanitizeFields, Json, LlmJsonStream, LlmNext, LlmRequest, LlmRequestInterceptOutcome, - LlmStream, LlmStreamNext, NEMO_RELAY_NATIVE_ABI_VERSION, NativePlugin, - NemoRelayNativeEventSanitizeCb, NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, - NemoRelayNativeHostApiV1, NemoRelayNativeJsonCb, NemoRelayNativeLlmConditionalCb, - NemoRelayNativeLlmExecutionCb, NemoRelayNativeLlmRequestCb, - NemoRelayNativeLlmRequestInterceptCb, NemoRelayNativeLlmStreamExecutionCb, + AnnotatedLlmRequest, BuiltinLlmCodec, CategoryProfile, ConfigDiagnostic, DiagnosticLevel, + Event, EventCategory, EventSanitizeFields, Json, LlmCodecIdentity, LlmJsonStream, LlmNext, + LlmRequest, LlmRequestInterceptOutcome, LlmStream, LlmStreamNext, + NEMO_RELAY_NATIVE_ABI_VERSION, NativePlugin, NemoRelayNativeEventSanitizeCb, + NemoRelayNativeEventSubscriberCb, NemoRelayNativeFreeFn, NemoRelayNativeHostApiV1, + NemoRelayNativeLlmCodecKind, NemoRelayNativeLlmConditionalCb, NemoRelayNativeLlmExecutionCb, + NemoRelayNativeLlmRequestCodec, NemoRelayNativeLlmRequestInterceptCb, + NemoRelayNativeLlmResponseCodec, NemoRelayNativeLlmSanitizeRequestCb, + NemoRelayNativeLlmSanitizeRequestContext, NemoRelayNativeLlmSanitizeResponseCb, + NemoRelayNativeLlmSanitizeResponseContext, NemoRelayNativeLlmStreamExecutionCb, NemoRelayNativeLlmStreamV1, NemoRelayNativePluginContext, NemoRelayNativePluginV1, NemoRelayNativeScopeHandle, NemoRelayNativeScopeStack, NemoRelayNativeScopeStackBinding, NemoRelayNativeScopeType, NemoRelayNativeString, NemoRelayNativeToolConditionalCb, @@ -114,7 +117,7 @@ impl RegisteredToolExecution { struct RegisteredLlmRequest { name: String, priority: i32, - cb: NemoRelayNativeLlmRequestCb, + cb: NemoRelayNativeLlmSanitizeRequestCb, user_data: usize, free_fn: NemoRelayNativeFreeFn, } @@ -130,7 +133,7 @@ impl RegisteredLlmRequest { struct RegisteredLlmJson { name: String, priority: i32, - cb: NemoRelayNativeJsonCb, + cb: NemoRelayNativeLlmSanitizeResponseCb, user_data: usize, free_fn: NemoRelayNativeFreeFn, } @@ -297,8 +300,8 @@ static LLM_REQUEST_INTERCEPT_REGISTRATION: Mutex(), test_host().struct_size @@ -316,13 +319,13 @@ fn native_abi_v1_struct_sizes_are_self_describing() { #[cfg(target_pointer_width = "64")] { assert_eq!(align_of::(), 8); - assert_eq!(size_of::(), 296); + assert_eq!(size_of::(), 320); assert_eq!( host_api_offsets(), [ 0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 264, 272, - 280, 288, + 280, 288, 296, 304, 312, ] ); assert_eq!(align_of::(), 8); @@ -336,12 +339,13 @@ fn native_abi_v1_struct_sizes_are_self_describing() { #[cfg(target_pointer_width = "32")] { assert_eq!(align_of::(), 4); - assert_eq!(size_of::(), 148); + assert_eq!(size_of::(), 160); assert_eq!( host_api_offsets(), [ 0, 4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80, - 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, + 84, 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, 148, + 152, 156, ] ); assert_eq!(align_of::(), 4); @@ -353,7 +357,7 @@ fn native_abi_v1_struct_sizes_are_self_describing() { } } -fn host_api_offsets() -> [usize; 37] { +fn host_api_offsets() -> [usize; 40] { [ offset_of!(NemoRelayNativeHostApiV1, abi_version), offset_of!(NemoRelayNativeHostApiV1, struct_size), @@ -364,6 +368,9 @@ fn host_api_offsets() -> [usize; 37] { offset_of!(NemoRelayNativeHostApiV1, string_free), offset_of!(NemoRelayNativeHostApiV1, last_error_clear), offset_of!(NemoRelayNativeHostApiV1, last_error_set), + offset_of!(NemoRelayNativeHostApiV1, llm_request_codec_decode), + offset_of!(NemoRelayNativeHostApiV1, llm_request_codec_encode), + offset_of!(NemoRelayNativeHostApiV1, llm_response_codec_decode), offset_of!(NemoRelayNativeHostApiV1, plugin_context_register_subscriber), offset_of!( NemoRelayNativeHostApiV1, @@ -688,7 +695,7 @@ unsafe extern "C" fn capture_llm_request( _ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, priority: i32, - cb: NemoRelayNativeLlmRequestCb, + cb: NemoRelayNativeLlmSanitizeRequestCb, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus { @@ -717,7 +724,7 @@ unsafe extern "C" fn capture_llm_json( _ctx: *mut NemoRelayNativePluginContext, name: *const NemoRelayNativeString, priority: i32, - cb: NemoRelayNativeJsonCb, + cb: NemoRelayNativeLlmSanitizeResponseCb, user_data: *mut c_void, free_fn: NemoRelayNativeFreeFn, ) -> NemoRelayStatus { @@ -1126,6 +1133,57 @@ unsafe extern "C" fn true_scope_stack_active() -> bool { true } +unsafe extern "C" fn unavailable_request_codec_decode( + _codec: *const NemoRelayNativeLlmRequestCodec, + _request: *const NemoRelayNativeString, + _out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + NemoRelayStatus::Internal +} + +unsafe extern "C" fn unavailable_request_codec_encode( + _codec: *const NemoRelayNativeLlmRequestCodec, + _annotated: *const NemoRelayNativeString, + _original: *const NemoRelayNativeString, + _out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + NemoRelayStatus::Internal +} + +unsafe extern "C" fn unavailable_response_codec_decode( + _codec: *const NemoRelayNativeLlmResponseCodec, + _response: *const NemoRelayNativeString, + _out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + NemoRelayStatus::Internal +} + +unsafe extern "C" fn successful_request_codec_decode( + _codec: *const NemoRelayNativeLlmRequestCodec, + _request: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { test_string_new(c"{}".as_ptr().cast(), 2, out) } +} + +unsafe extern "C" fn successful_request_codec_encode( + _codec: *const NemoRelayNativeLlmRequestCodec, + _annotated: *const NemoRelayNativeString, + original: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + let bytes = unsafe { &*(original.cast::()) }.0.as_slice(); + unsafe { test_string_new(bytes.as_ptr(), bytes.len(), out) } +} + +unsafe extern "C" fn successful_response_codec_decode( + _codec: *const NemoRelayNativeLlmResponseCodec, + _response: *const NemoRelayNativeString, + out: *mut *mut NemoRelayNativeString, +) -> NemoRelayStatus { + unsafe { test_string_new(c"{}".as_ptr().cast(), 2, out) } +} + fn test_host() -> NemoRelayNativeHostApiV1 { NemoRelayNativeHostApiV1 { abi_version: NEMO_RELAY_NATIVE_ABI_VERSION, @@ -1137,6 +1195,9 @@ fn test_host() -> NemoRelayNativeHostApiV1 { string_free: test_string_free, last_error_clear: test_last_error_clear, last_error_set: test_last_error_set, + llm_request_codec_decode: unavailable_request_codec_decode, + llm_request_codec_encode: unavailable_request_codec_encode, + llm_response_codec_decode: unavailable_response_codec_decode, plugin_context_register_subscriber: capture_register_subscriber, plugin_context_register_tool_sanitize_request_guardrail: capture_tool_json, plugin_context_register_tool_sanitize_response_guardrail: capture_tool_json, @@ -1284,6 +1345,22 @@ fn json_host_string(host: &NemoRelayNativeHostApiV1, value: Json) -> *mut NemoRe host_string(host, &serde_json::to_string(&value).unwrap()) } +fn native_no_codec_context() -> NemoRelayNativeLlmSanitizeRequestContext { + NemoRelayNativeLlmSanitizeRequestContext { + codec_kind: NemoRelayNativeLlmCodecKind::None, + codec_id: ptr::null(), + codec: ptr::null(), + } +} + +fn native_no_response_codec_context() -> NemoRelayNativeLlmSanitizeResponseContext { + NemoRelayNativeLlmSanitizeResponseContext { + codec_kind: NemoRelayNativeLlmCodecKind::None, + codec_id: ptr::null(), + codec: ptr::null(), + } +} + fn read_json_and_free(host: &NemoRelayNativeHostApiV1, value: *mut NemoRelayNativeString) -> Json { let result: Json = serde_json::from_str(&read_host_string(host, value).unwrap()).unwrap(); unsafe { (host.string_free)(value) }; @@ -2802,12 +2879,21 @@ fn typed_callbacks_reject_null_abi_pointers_before_decoding_inputs() { } let mut ctx = test_context(&host); - ctx.register_llm_sanitize_request_guardrail("llm-request", 0, |request| request) - .unwrap(); + ctx.register_llm_sanitize_request_guardrail("llm-request", 0, |request, _context| { + Some(request) + }) + .unwrap(); let registration = take_llm_request_registration(); let request = json_host_string(&host, serde_json::to_value(test_llm_request()).unwrap()); assert_eq!( - unsafe { (registration.cb)(ptr::null_mut(), request, &mut out) }, + unsafe { + (registration.cb)( + ptr::null_mut(), + request, + native_no_codec_context(), + &mut out, + ) + }, NemoRelayStatus::NullPointer ); assert_eq!( @@ -2815,6 +2901,7 @@ fn typed_callbacks_reject_null_abi_pointers_before_decoding_inputs() { (registration.cb)( registration.user_data as *mut c_void, request, + native_no_codec_context(), ptr::null_mut(), ) }, @@ -2826,12 +2913,19 @@ fn typed_callbacks_reject_null_abi_pointers_before_decoding_inputs() { } let mut ctx = test_context(&host); - ctx.register_llm_sanitize_response_guardrail("llm-response", 0, |value| value) + ctx.register_llm_sanitize_response_guardrail("llm-response", 0, |value, _context| Some(value)) .unwrap(); let registration = take_llm_json_registration(); let response = json_host_string(&host, json!({})); assert_eq!( - unsafe { (registration.cb)(ptr::null_mut(), response, &mut out) }, + unsafe { + (registration.cb)( + ptr::null_mut(), + response, + native_no_response_codec_context(), + &mut out, + ) + }, NemoRelayStatus::NullPointer ); assert_eq!( @@ -2839,6 +2933,7 @@ fn typed_callbacks_reject_null_abi_pointers_before_decoding_inputs() { (registration.cb)( registration.user_data as *mut c_void, response, + native_no_response_codec_context(), ptr::null_mut(), ) }, @@ -3166,14 +3261,24 @@ fn typed_callbacks_report_invalid_json_for_each_decoder_family() { } let mut ctx = test_context(&host); - ctx.register_llm_sanitize_request_guardrail("llm-request", 0, |request| request) - .unwrap(); + ctx.register_llm_sanitize_request_guardrail("llm-request", 0, |request, _context| { + Some(request) + }) + .unwrap(); let registration = take_llm_request_registration(); let request = host_string(&host, "{not json"); + let context = native_no_codec_context(); let stale_out = host_string(&host, r#"{"stale":true}"#); let mut out = stale_out; assert_eq!( - unsafe { (registration.cb)(registration.user_data as *mut c_void, request, &mut out) }, + unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + request, + context, + &mut out, + ) + }, NemoRelayStatus::InvalidJson ); assert!(out.is_null()); @@ -3184,14 +3289,22 @@ fn typed_callbacks_report_invalid_json_for_each_decoder_family() { } let mut ctx = test_context(&host); - ctx.register_llm_sanitize_response_guardrail("llm-response", 0, |value| value) + ctx.register_llm_sanitize_response_guardrail("llm-response", 0, |value, _context| Some(value)) .unwrap(); let registration = take_llm_json_registration(); let response = host_string(&host, "{not json"); + let context = native_no_response_codec_context(); let stale_out = host_string(&host, r#"{"stale":true}"#); let mut out = stale_out; assert_eq!( - unsafe { (registration.cb)(registration.user_data as *mut c_void, response, &mut out) }, + unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + response, + context, + &mut out, + ) + }, NemoRelayStatus::InvalidJson ); assert!(out.is_null()); @@ -4073,20 +4186,31 @@ fn typed_llm_sanitize_guardrails_transform_request_and_response() { let _guard = begin_test(); let host = test_host(); let mut ctx = test_context(&host); - ctx.register_llm_sanitize_request_guardrail("llm-sanitize-request", 12, |mut request| { - request.headers.insert("x-policy".into(), json!("sdk")); - request.content["sanitized"] = json!(true); - request - }) + ctx.register_llm_sanitize_request_guardrail( + "llm-sanitize-request", + 12, + |mut request, _context| { + request.headers.insert("x-policy".into(), json!("sdk")); + request.content["sanitized"] = json!(true); + Some(request) + }, + ) .unwrap(); let registration = take_llm_request_registration(); assert_eq!(registration.name, "llm-sanitize-request"); assert_eq!(registration.priority, 12); let request = json_host_string(&host, serde_json::to_value(test_llm_request()).unwrap()); + let context = native_no_codec_context(); let mut out = ptr::null_mut(); - let status = - unsafe { (registration.cb)(registration.user_data as *mut c_void, request, &mut out) }; + let status = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + request, + context, + &mut out, + ) + }; assert_eq!(status, NemoRelayStatus::Ok); let output = read_json_and_free(&host, out); assert_eq!(output["headers"]["x-policy"], json!("sdk")); @@ -4097,19 +4221,30 @@ fn typed_llm_sanitize_guardrails_transform_request_and_response() { } let mut ctx = test_context(&host); - ctx.register_llm_sanitize_response_guardrail("llm-sanitize-response", 13, |mut payload| { - payload["sanitized"] = json!(true); - payload - }) + ctx.register_llm_sanitize_response_guardrail( + "llm-sanitize-response", + 13, + |mut payload, _context| { + payload["sanitized"] = json!(true); + Some(payload) + }, + ) .unwrap(); let registration = take_llm_json_registration(); assert_eq!(registration.name, "llm-sanitize-response"); assert_eq!(registration.priority, 13); let response = json_host_string(&host, json!({ "output": true })); + let context = native_no_response_codec_context(); let mut out = ptr::null_mut(); - let status = - unsafe { (registration.cb)(registration.user_data as *mut c_void, response, &mut out) }; + let status = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + response, + context, + &mut out, + ) + }; assert_eq!(status, NemoRelayStatus::Ok); assert_eq!(read_json_and_free(&host, out)["sanitized"], json!(true)); unsafe { @@ -4118,6 +4253,176 @@ fn typed_llm_sanitize_guardrails_transform_request_and_response() { } } +#[test] +fn typed_contextual_llm_sanitize_guardrails_receive_payload_before_context() { + let _guard = begin_test(); + let mut host = test_host(); + host.llm_request_codec_decode = successful_request_codec_decode; + host.llm_request_codec_encode = successful_request_codec_encode; + host.llm_response_codec_decode = successful_response_codec_decode; + let mut ctx = test_context(&host); + ctx.register_llm_sanitize_request_guardrail( + "contextual-request", + 14, + |mut request, callback_context| { + assert_eq!( + callback_context.codec, + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) + ); + let codec = callback_context + .resolve_codec() + .expect("active request codec must resolve"); + let annotated = codec.decode(&request).expect("request decode succeeds"); + request = codec + .encode(&annotated, &request) + .expect("request encode succeeds"); + request.headers.insert("x-contextual".into(), json!(true)); + Some(request) + }, + ) + .unwrap(); + + let registration = take_llm_request_registration(); + assert_eq!(registration.name, "contextual-request"); + assert_eq!(registration.priority, 14); + let request = json_host_string(&host, serde_json::to_value(test_llm_request()).unwrap()); + let context_id = host_string(&host, "openai_chat"); + let request_codec_placeholder = Box::new(0_usize); + let native_context = NemoRelayNativeLlmSanitizeRequestContext { + codec_kind: NemoRelayNativeLlmCodecKind::BuiltIn, + codec_id: context_id, + codec: std::ptr::from_ref(request_codec_placeholder.as_ref()) + .cast::(), + }; + let mut out = ptr::null_mut(); + let status = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + request, + native_context, + &mut out, + ) + }; + assert_eq!(status, NemoRelayStatus::Ok); + assert_eq!( + read_json_and_free(&host, out)["headers"]["x-contextual"], + json!(true) + ); + unsafe { + (host.string_free)(request); + (host.string_free)(context_id); + registration.free(); + } + + let mut ctx = test_context(&host); + ctx.register_llm_sanitize_response_guardrail( + "contextual-response", + 15, + |mut payload, callback_context| { + assert_eq!( + callback_context.codec, + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) + ); + callback_context + .resolve_codec() + .expect("active response codec must resolve") + .decode(&payload) + .expect("response decode succeeds"); + payload["contextual"] = json!(true); + Some(payload) + }, + ) + .unwrap(); + + let registration = take_llm_json_registration(); + assert_eq!(registration.name, "contextual-response"); + assert_eq!(registration.priority, 15); + let response = json_host_string(&host, json!({ "output": true })); + let context_id = host_string(&host, "openai_chat"); + let response_codec_placeholder = Box::new(0_usize); + let native_context = NemoRelayNativeLlmSanitizeResponseContext { + codec_kind: NemoRelayNativeLlmCodecKind::BuiltIn, + codec_id: context_id, + codec: std::ptr::from_ref(response_codec_placeholder.as_ref()) + .cast::(), + }; + let mut out = ptr::null_mut(); + let status = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + response, + native_context, + &mut out, + ) + }; + assert_eq!(status, NemoRelayStatus::Ok); + assert_eq!(read_json_and_free(&host, out)["contextual"], json!(true)); + unsafe { + (host.string_free)(response); + (host.string_free)(context_id); + registration.free(); + } +} + +#[test] +fn typed_contextual_llm_sanitizer_uses_null_output_to_omit_payload() { + let _guard = begin_test(); + let host = test_host(); + let mut ctx = test_context(&host); + ctx.register_llm_sanitize_request_guardrail( + "contextual-omit-request", + 16, + |_request, _context| None, + ) + .unwrap(); + + let registration = take_llm_request_registration(); + let request = json_host_string(&host, serde_json::to_value(test_llm_request()).unwrap()); + let context = native_no_codec_context(); + let mut out = ptr::null_mut(); + let status = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + request, + context, + &mut out, + ) + }; + assert_eq!(status, NemoRelayStatus::Ok); + assert!(out.is_null(), "null native output must represent omission"); + unsafe { + (host.string_free)(request); + registration.free(); + } + + let mut ctx = test_context(&host); + ctx.register_llm_sanitize_response_guardrail("contextual-omit", 16, |_payload, _context| None) + .unwrap(); + + let registration = take_llm_json_registration(); + let response = json_host_string(&host, json!({"secret": "value"})); + let context = NemoRelayNativeLlmSanitizeResponseContext { + codec_kind: NemoRelayNativeLlmCodecKind::None, + codec_id: ptr::null(), + codec: ptr::null(), + }; + let mut out = ptr::null_mut(); + let status = unsafe { + (registration.cb)( + registration.user_data as *mut c_void, + response, + context, + &mut out, + ) + }; + assert_eq!(status, NemoRelayStatus::Ok); + assert!(out.is_null(), "null native output must represent omission"); + unsafe { + (host.string_free)(response); + registration.free(); + } +} + #[test] fn typed_llm_conditional_guardrail_returns_optional_reason() { let _guard = begin_test(); diff --git a/crates/python/src/py_api/mod.rs b/crates/python/src/py_api/mod.rs index bc026d76e..64e2341e3 100644 --- a/crates/python/src/py_api/mod.rs +++ b/crates/python/src/py_api/mod.rs @@ -73,6 +73,26 @@ fn py_llm_response_codec( }) } +fn py_llm_codec(codec: Option<&Bound<'_, PyAny>>) -> Option> { + codec.and_then(|codec| -> Option> { + if codec.is_none() { + return None; + } + if let Ok(builtin) = codec.extract::>() { + return Some(builtin.inner_codec.clone()); + } + if let Ok(builtin) = codec.extract::>() { + return Some(builtin.inner_codec.clone()); + } + if let Ok(builtin) = codec.extract::>() { + return Some(builtin.inner_codec.clone()); + } + Some(Arc::new(py_callable::PyLlmCodecWrapper { + py_codec: codec.clone().unbind(), + })) + }) +} + fn py_annotated_llm_response( annotated_response: Option<&Bound<'_, PyAny>>, ) -> PyResult>> { @@ -750,11 +770,7 @@ fn llm_call_execute<'py>( let exec_fn = py_callable::wrap_py_llm_exec_fn(func); let default_fn: LlmExecutionNextFn = Arc::new(move |req| exec_fn(req)); let parent_handle = handle.map(|h| h.inner).unwrap_or_else(task_scope_top); - let codec_arc: Option> = codec.map(|c| { - Arc::new(py_callable::PyLlmCodecWrapper { - py_codec: c.clone().unbind(), - }) as Arc - }); + let codec_arc = py_llm_codec(codec); let response_codec_arc = py_llm_response_codec(response_codec); let scope_stack = current_scope_stack_handle(); @@ -855,11 +871,7 @@ fn llm_stream_call_execute<'py>( let collector_fn = py_callable::wrap_py_collector_fn(collector); let finalizer_fn = py_callable::wrap_py_finalizer_fn(finalizer); let parent_handle = handle.map(|h| h.inner).unwrap_or_else(task_scope_top); - let codec_arc: Option> = codec.map(|c| { - Arc::new(py_callable::PyLlmCodecWrapper { - py_codec: c.clone().unbind(), - }) as Arc - }); + let codec_arc = py_llm_codec(codec); let response_codec_arc = py_llm_response_codec(response_codec); let scope_stack = current_scope_stack_handle(); @@ -1081,7 +1093,8 @@ fn deregister_tool_execution_intercept(name: &str) -> PyResult { /// Register an LLM sanitize-request guardrail. /// -/// Callback: ``(request: LlmRequest) -> LlmRequest`` — returns a sanitized request. +/// Callback: ``(request: LlmRequest, context: LlmSanitizeRequestContext) -> +/// Optional[LlmRequest]``. Return ``None`` to omit the observability payload and annotation. #[pyfunction] fn register_llm_sanitize_request_guardrail( name: &str, @@ -1091,7 +1104,7 @@ fn register_llm_sanitize_request_guardrail( core_registry_api::register_llm_sanitize_request_guardrail( name, priority, - py_callable::wrap_py_llm_sanitize_request_fn(guardrail), + py_callable::wrap_py_llm_sanitize_request_fn(guardrail)?, ) .map_err(to_py_err) } @@ -1104,7 +1117,8 @@ fn deregister_llm_sanitize_request_guardrail(name: &str) -> PyResult { /// Register an LLM sanitize-response guardrail. /// -/// Callback: ``(response: dict) -> dict`` — returns a sanitized response. +/// Callback: ``(response: Json, context: LlmSanitizeResponseContext) -> Optional[Json]``. +/// Return ``None`` to omit the observability payload and annotation. #[pyfunction] fn register_llm_sanitize_response_guardrail( name: &str, @@ -1114,7 +1128,7 @@ fn register_llm_sanitize_response_guardrail( core_registry_api::register_llm_sanitize_response_guardrail( name, priority, - py_callable::wrap_py_llm_sanitize_response_fn(guardrail), + py_callable::wrap_py_llm_sanitize_response_fn(guardrail)?, ) .map_err(to_py_err) } @@ -1542,7 +1556,7 @@ fn scope_register_llm_sanitize_request_guardrail( &uuid, name, priority, - py_callable::wrap_py_llm_sanitize_request_fn(guardrail), + py_callable::wrap_py_llm_sanitize_request_fn(guardrail)?, ) .map_err(to_py_err) } @@ -1568,7 +1582,7 @@ fn scope_register_llm_sanitize_response_guardrail( &uuid, name, priority, - py_callable::wrap_py_llm_sanitize_response_fn(guardrail), + py_callable::wrap_py_llm_sanitize_response_fn(guardrail)?, ) .map_err(to_py_err) } diff --git a/crates/python/src/py_callable.rs b/crates/python/src/py_callable.rs index b08fb2f2d..57fffd5d5 100644 --- a/crates/python/src/py_callable.rs +++ b/crates/python/src/py_callable.rs @@ -27,8 +27,9 @@ use std::task::{Context, Poll}; use nemo_relay::api::runtime::{ EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, LlmJsonStream, - LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, - LlmStreamInner, ToolConditionalFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, + LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn, + LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, LlmStreamInner, + ToolConditionalFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn, }; use nemo_relay::error::{FlowError, Result as FlowResult}; use pyo3::prelude::*; @@ -47,11 +48,31 @@ use nemo_relay::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::convert::{json_to_py, py_to_json}; use crate::py_types::{ PyAnnotatedLLMRequest, PyAnnotatedLLMResponse, PyLLMRequest, PyLLMRequestInterceptOutcome, - PyToolExecutionInterceptOutcome, + PyLlmSanitizeRequestContext, PyLlmSanitizeResponseContext, PyToolExecutionInterceptOutcome, }; type PyValueFuture = Pin>> + Send>>; +fn validate_python_llm_sanitizer_signature(py_fn: &Py) -> PyResult<()> { + Python::attach(|py| { + let inspect = py.import("inspect")?; + let signature = inspect.call_method1("signature", (py_fn.bind(py),)).map_err(|_| { + pyo3::exceptions::PyTypeError::new_err( + "LLM sanitizer callback signature cannot be inspected; use a callable that accepts `(payload, context)`", + ) + })?; + if signature + .call_method1("bind", (py.None(), py.None())) + .is_ok() + { + return Ok(()); + } + Err(pyo3::exceptions::PyTypeError::new_err( + "LLM sanitizer callback must accept `(payload, context)`", + )) + }) +} + fn split_json_or_future( py: Python<'_>, result: Py, @@ -792,33 +813,35 @@ pub fn wrap_py_llm_stream_exec_intercept_fn( ) } -/// Wrap a Python callable `(LlmRequest) -> LlmRequest` for LLM sanitize request guardrails. -pub fn wrap_py_llm_sanitize_request_fn(py_fn: Py) -> LlmSanitizeRequestFn { - Arc::new(move |request: LlmRequest| { - Python::attach(|py| { - let py_req = PyLLMRequest { - inner: request.clone(), - }; - let result = match py_fn.call1(py, (py_req,)) { - Ok(v) => v, - Err(e) => { - eprintln!("nemo_relay: LLM sanitize request guardrail callable failed: {e}"); - return request; +/// Wrap a Python callable `(LlmRequest, LlmSanitizeRequestContext) -> Optional`. +fn wrap_py_llm_sanitize_request_callback(py_fn: Py) -> LlmSanitizeRequestFn { + Arc::new( + move |request: LlmRequest, context: LlmSanitizeRequestContext| { + Python::attach(|py| { + let py_context = PyLlmSanitizeRequestContext { inner: context }; + let py_request = PyLLMRequest { inner: request }; + let result = match py_fn.call1(py, (py_request, py_context)) { + Ok(value) => value, + Err(error) => { + eprintln!("nemo_relay: LLM sanitize request callable failed: {error}"); + return None; + } + }; + if result.is_none(py) { + return None; } - }; - let extracted = result.extract::(py); - match extracted { - Ok(r) => r.inner, - Err(e) => { - eprintln!( - "nemo_relay: LLM sanitize request guardrail returned unexpected type \ - (expected LlmRequest): {e}" - ); - request + match result.extract::(py) { + Ok(request) => Some(request.inner), + Err(error) => { + eprintln!( + "nemo_relay: LLM sanitize request callable returned unexpected type: {error}" + ); + None + } } - } - }) - }) + }) + }, + ) } /// Wrap a Python callable `(LlmRequest) -> Optional[str]` for LLM conditional guardrails. @@ -989,34 +1012,51 @@ pub fn wrap_py_finalizer_fn(py_fn: Py) -> Box Json + Send }) } -/// Wrap a Python callable `(dict) -> dict` for LLM sanitize response guardrails. -pub fn wrap_py_llm_sanitize_response_fn(py_fn: Py) -> LlmSanitizeResponseFn { - Arc::new(move |response: Json| { +/// Wrap a Python callable `(Json, LlmSanitizeResponseContext) -> Optional[Json]`. +fn wrap_py_llm_sanitize_response_callback(py_fn: Py) -> LlmSanitizeResponseFn { + Arc::new(move |response: Json, context: LlmSanitizeResponseContext| { Python::attach(|py| { - let py_resp = match json_to_py(py, &response) { - Ok(v) => v, - Err(e) => { - eprintln!( - "nemo_relay: json_to_py failed in LLM sanitize response guardrail: {e}" - ); - return response.clone(); + let py_context = PyLlmSanitizeResponseContext { inner: context }; + let py_response = match json_to_py(py, &response) { + Ok(value) => value, + Err(error) => { + eprintln!("nemo_relay: json_to_py failed in LLM sanitize response: {error}"); + return None; } }; - let result = match py_fn.call1(py, (py_resp,)) { - Ok(v) => v, - Err(e) => { - eprintln!("nemo_relay: LLM sanitize response guardrail callable failed: {e}"); - return response.clone(); + let result = match py_fn.call1(py, (py_response, py_context)) { + Ok(value) => value, + Err(error) => { + eprintln!("nemo_relay: LLM sanitize response callable failed: {error}"); + return None; } }; - py_to_json(result.bind(py)).unwrap_or_else(|e| { - eprintln!("nemo_relay: py_to_json failed in LLM sanitize response guardrail: {e}"); - response.clone() - }) + if result.is_none(py) { + return None; + } + match py_to_json(result.bind(py)) { + Ok(response) => Some(response), + Err(error) => { + eprintln!("nemo_relay: py_to_json failed in LLM sanitize response: {error}"); + None + } + } }) }) } +/// Wrap a Python LLM sanitize-request callback. +pub fn wrap_py_llm_sanitize_request_fn(py_fn: Py) -> PyResult { + validate_python_llm_sanitizer_signature(&py_fn)?; + Ok(wrap_py_llm_sanitize_request_callback(py_fn)) +} + +/// Wrap a Python LLM sanitize-response callback. +pub fn wrap_py_llm_sanitize_response_fn(py_fn: Py) -> PyResult { + validate_python_llm_sanitizer_signature(&py_fn)?; + Ok(wrap_py_llm_sanitize_response_callback(py_fn)) +} + /// Wrap a Python callable `(Event) -> None` for event subscribers. pub fn wrap_py_event_subscriber(py_fn: Py) -> EventSubscriberFn { Arc::new(move |event: &Event| { diff --git a/crates/python/src/py_plugin.rs b/crates/python/src/py_plugin.rs index e709fc769..2cf76a4a8 100644 --- a/crates/python/src/py_plugin.rs +++ b/crates/python/src/py_plugin.rs @@ -388,14 +388,11 @@ impl PyPluginContext { priority: i32, callback: Py, ) -> PyResult<()> { + let callback = wrap_py_llm_sanitize_request_fn(callback)?; self.register_callback( name, |qualified_name| { - register_llm_sanitize_request_guardrail( - qualified_name, - priority, - wrap_py_llm_sanitize_request_fn(callback), - ) + register_llm_sanitize_request_guardrail(qualified_name, priority, callback) }, deregister_llm_sanitize_request_guardrail, "llm sanitize request guardrail", @@ -409,14 +406,11 @@ impl PyPluginContext { priority: i32, callback: Py, ) -> PyResult<()> { + let callback = wrap_py_llm_sanitize_response_fn(callback)?; self.register_callback( name, |qualified_name| { - register_llm_sanitize_response_guardrail( - qualified_name, - priority, - wrap_py_llm_sanitize_response_fn(callback), - ) + register_llm_sanitize_response_guardrail(qualified_name, priority, callback) }, deregister_llm_sanitize_response_guardrail, "llm sanitize response guardrail", diff --git a/crates/python/src/py_types/codecs.rs b/crates/python/src/py_types/codecs.rs index 73efdf5a8..1b3a22875 100644 --- a/crates/python/src/py_types/codecs.rs +++ b/crates/python/src/py_types/codecs.rs @@ -26,6 +26,53 @@ use super::{ }; use nemo_relay::codec::response::FinishReason; +/// A resolved request codec available while an LLM request sanitizer runs. +#[pyclass(name = "LlmSanitizeRequestCodec")] +pub struct PyLlmSanitizeRequestCodec { + pub(crate) inner: Arc, +} + +#[pymethods] +impl PyLlmSanitizeRequestCodec { + /// Parse an opaque request into its normalized representation. + fn decode(&self, request: &PyLLMRequest) -> PyResult { + self.inner + .decode(&request.inner) + .map(|inner| PyAnnotatedLLMRequest { inner }) + .map_err(|error| pyo3::exceptions::PyRuntimeError::new_err(error.to_string())) + } + + /// Merge a normalized request back into its provider representation. + fn encode( + &self, + annotated: &PyAnnotatedLLMRequest, + original: &PyLLMRequest, + ) -> PyResult { + self.inner + .encode(&annotated.inner, &original.inner) + .map(|inner| PyLLMRequest { inner }) + .map_err(|error| pyo3::exceptions::PyRuntimeError::new_err(error.to_string())) + } +} + +/// A resolved response codec available while an LLM response sanitizer runs. +#[pyclass(name = "LlmSanitizeResponseCodec")] +pub struct PyLlmSanitizeResponseCodec { + pub(crate) inner: Arc, +} + +#[pymethods] +impl PyLlmSanitizeResponseCodec { + /// Parse an opaque response into its normalized representation. + fn decode_response(&self, response: &Bound<'_, PyAny>) -> PyResult { + let response = py_to_json(response)?; + self.inner + .decode_response(&response) + .map(|inner| PyAnnotatedLLMResponse { inner }) + .map_err(|error| pyo3::exceptions::PyRuntimeError::new_err(error.to_string())) + } +} + // --------------------------------------------------------------------------- // AnnotatedLLMRequest // --------------------------------------------------------------------------- diff --git a/crates/python/src/py_types/core.rs b/crates/python/src/py_types/core.rs index aeca45305..d0f122311 100644 --- a/crates/python/src/py_types/core.rs +++ b/crates/python/src/py_types/core.rs @@ -5,15 +5,96 @@ use std::sync::Arc; use pyo3::prelude::*; +use super::codecs::{PyLlmSanitizeRequestCodec, PyLlmSanitizeResponseCodec}; use super::{ - AnnotatedLLMRequest, Bound, CoreScopeType, FlowResult, LlmAttributes, LlmHandle, LlmRequest, - PyAnnotatedLLMRequest, PyAny, PyErr, PyRef, PyResult, Python, ScopeAttributes, ScopeHandle, - ScopeStackHandle, ToolAttributes, ToolHandle, json_to_py, opt_json_to_py, py_to_json, + AnnotatedLLMRequest, Bound, CoreScopeType, FlowResult, LlmAttributes, LlmCodecIdentity, + LlmHandle, LlmRequest, PyAnnotatedLLMRequest, PyAny, PyErr, PyRef, PyResult, Python, + ScopeAttributes, ScopeHandle, ScopeStackHandle, ToolAttributes, ToolHandle, json_to_py, + opt_json_to_py, py_to_json, }; use nemo_relay::api::event::{CategoryProfile, EventCategory, PendingMarkSpec}; use nemo_relay::api::llm::LlmRequestInterceptOutcome; +use nemo_relay::api::runtime::{LlmSanitizeRequestContext, LlmSanitizeResponseContext}; use nemo_relay::api::tool::ToolExecutionInterceptOutcome; +/// Structured identity of the codec active during LLM sanitization. +#[pyclass(name = "LlmCodecIdentity", frozen)] +pub struct PyLlmCodecIdentity { + pub(crate) inner: LlmCodecIdentity, +} + +#[pymethods] +impl PyLlmCodecIdentity { + /// Identity variant: ``none``, ``builtin``, ``runtime``, or ``opaque``. + #[getter] + fn kind(&self) -> &'static str { + match self.inner { + LlmCodecIdentity::None => "none", + LlmCodecIdentity::BuiltIn(_) => "builtin", + LlmCodecIdentity::Runtime(_) => "runtime", + LlmCodecIdentity::Opaque => "opaque", + } + } + + /// Stable built-in or runtime codec ID, when this identity has one. + #[getter] + fn id(&self) -> Option { + match &self.inner { + LlmCodecIdentity::BuiltIn(codec) => Some(codec.id().to_owned()), + LlmCodecIdentity::Runtime(id) => Some(id.clone()), + LlmCodecIdentity::None | LlmCodecIdentity::Opaque => None, + } + } +} + +/// Structured per-call context delivered to LLM request sanitizer callbacks. +#[pyclass(name = "LlmSanitizeRequestContext", frozen)] +pub struct PyLlmSanitizeRequestContext { + pub(crate) inner: LlmSanitizeRequestContext, +} + +#[pymethods] +impl PyLlmSanitizeRequestContext { + /// The active codec identity for this request or response payload. + #[getter] + fn codec(&self) -> PyLlmCodecIdentity { + PyLlmCodecIdentity { + inner: self.inner.codec().clone(), + } + } + + /// Resolve the active request codec. + fn resolve_codec(&self) -> Option { + self.inner + .resolve_codec() + .map(|inner| PyLlmSanitizeRequestCodec { inner }) + } +} + +/// Structured per-call context delivered to LLM response sanitizer callbacks. +#[pyclass(name = "LlmSanitizeResponseContext", frozen)] +pub struct PyLlmSanitizeResponseContext { + pub(crate) inner: LlmSanitizeResponseContext, +} + +#[pymethods] +impl PyLlmSanitizeResponseContext { + /// The active codec identity for this request or response payload. + #[getter] + fn codec(&self) -> PyLlmCodecIdentity { + PyLlmCodecIdentity { + inner: self.inner.codec().clone(), + } + } + + /// Resolve the active response codec. + fn resolve_codec(&self) -> Option { + self.inner + .resolve_codec() + .map(|inner| PyLlmSanitizeResponseCodec { inner }) + } +} + // --------------------------------------------------------------------------- // LlmStream (async iterator) // --------------------------------------------------------------------------- diff --git a/crates/python/src/py_types/mod.rs b/crates/python/src/py_types/mod.rs index 046b580ff..eb60374d9 100644 --- a/crates/python/src/py_types/mod.rs +++ b/crates/python/src/py_types/mod.rs @@ -15,7 +15,7 @@ use std::time::Duration; use nemo_relay::api::event::{MarkEvent, ScopeEvent}; use nemo_relay::api::llm::{LlmAttributes, LlmHandle, LlmRequest}; -use nemo_relay::api::runtime::ScopeStackHandle; +use nemo_relay::api::runtime::{LlmCodecIdentity, ScopeStackHandle}; use nemo_relay::api::scope::{ScopeAttributes, ScopeHandle, ScopeType as CoreScopeType}; use nemo_relay::api::tool::{ToolAttributes, ToolHandle}; use nemo_relay::codec::request::{ @@ -128,6 +128,14 @@ pub use events::*; pub use observability::*; pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { + register_runtime_types(m)?; + register_llm_types(m)?; + register_event_types(m)?; + register_observability_types(m)?; + register_codec_types(m) +} + +fn register_runtime_types(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; @@ -138,13 +146,30 @@ pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + Ok(()) +} + +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::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; + Ok(()) +} + +fn register_event_types(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; + Ok(()) +} + +fn register_observability_types(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; @@ -154,6 +179,10 @@ pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + Ok(()) +} + +fn register_codec_types(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/crates/python/tests/coverage/coverage_tests.rs b/crates/python/tests/coverage/coverage_tests.rs index 765c3850a..67cebd3ba 100644 --- a/crates/python/tests/coverage/coverage_tests.rs +++ b/crates/python/tests/coverage/coverage_tests.rs @@ -4,6 +4,7 @@ //! Coverage tests for coverage in the NeMo Relay Python crate. use std::ffi::CString; +use std::path::PathBuf; use std::sync::Arc; use pyo3::prelude::*; @@ -51,6 +52,29 @@ fn make_request() -> LlmRequest { } } +/// Restores the process working directory after a plugin-layering test. +/// +/// Plugin discovery intentionally walks from the working directory, so this +/// guard keeps coverage tests independent of a developer's parent project +/// configuration. The Python test guard serializes these process-global test +/// changes. +struct CurrentDirectoryGuard(PathBuf); + +impl CurrentDirectoryGuard { + fn move_to_temporary_directory() -> Self { + let original = std::env::current_dir().expect("test process has a working directory"); + std::env::set_current_dir(std::env::temp_dir()) + .expect("temporary directory is available for plugin config discovery"); + Self(original) + } +} + +impl Drop for CurrentDirectoryGuard { + fn drop(&mut self) { + std::env::set_current_dir(&self.0).expect("restore test working directory"); + } +} + fn with_event_loop(py: Python<'_>, f: impl FnOnce(Bound<'_, PyAny>) -> T) -> T { let asyncio = py.import("asyncio").unwrap(); #[cfg(windows)] @@ -346,6 +370,7 @@ fn test_py_adaptive_binding_rejects_zero_sensitivity() { fn test_plugin_bindings_validate_configure_and_clear() { let _python = crate::test_support::init_python_test(); let _plugin_test_state = crate::py_plugin::lock_plugin_test_state_for_tests(); + let _working_directory = CurrentDirectoryGuard::move_to_temporary_directory(); Python::attach(|py| { nemo_relay_adaptive::plugin_component::register_adaptive_component().unwrap(); @@ -405,10 +430,10 @@ def tool_passthrough(name, value): def tool_conditional(name, value): return None -def llm_sanitize_request(request): +def llm_sanitize_request(request, context): return request -def llm_sanitize_response(response): +def llm_sanitize_response(response, context): return response def llm_conditional(request): @@ -610,7 +635,7 @@ def tool_fail(name, args): def tool_cond_bad(name, args): return 123 -def llm_sanitize_bad(request): +def llm_sanitize_bad(request, context): return {"bad": True} def llm_cond_bad(request): @@ -622,7 +647,7 @@ def llm_cond_none(request): def llm_req_bad(name, request): return {"bad": True} -def llm_resp_fail(response): +def llm_resp_fail(response, context): raise RuntimeError("resp boom") def collector_fail(chunk): @@ -656,8 +681,15 @@ def event_fail(event): let request = make_request(); let llm_sanitize = - wrap_py_llm_sanitize_request_fn(module.getattr("llm_sanitize_bad").unwrap().unbind()); - assert_eq!(llm_sanitize(request.clone()).content, request.content); + wrap_py_llm_sanitize_request_fn(module.getattr("llm_sanitize_bad").unwrap().unbind()) + .unwrap(); + assert_eq!( + llm_sanitize( + request.clone(), + nemo_relay::api::runtime::LlmSanitizeRequestContext::default(), + ), + None + ); let llm_cond = wrap_py_llm_conditional_fn(module.getattr("llm_cond_bad").unwrap().unbind()); assert!( @@ -689,8 +721,15 @@ def event_fail(event): ); let llm_resp = - wrap_py_llm_sanitize_response_fn(module.getattr("llm_resp_fail").unwrap().unbind()); - assert_eq!(llm_resp(json!({"ok": true})), json!({"ok": true})); + wrap_py_llm_sanitize_response_fn(module.getattr("llm_resp_fail").unwrap().unbind()) + .unwrap(); + assert_eq!( + llm_resp( + json!({"ok": true}), + nemo_relay::api::runtime::LlmSanitizeResponseContext::default(), + ), + None + ); let mut collector = wrap_py_collector_fn(module.getattr("collector_fail").unwrap().unbind()); diff --git a/crates/python/tests/coverage/py_api_coverage_tests.rs b/crates/python/tests/coverage/py_api_coverage_tests.rs index 0dee9e556..ec4c74fec 100644 --- a/crates/python/tests/coverage/py_api_coverage_tests.rs +++ b/crates/python/tests/coverage/py_api_coverage_tests.rs @@ -179,10 +179,10 @@ async def tool_exec_intercept(name, args, next): result["tool_intercepted"] = True return ToolExecutionInterceptOutcome(result) -def llm_sanitize_request(request): +def llm_sanitize_request(request, context): return request -def llm_sanitize_response(response): +def llm_sanitize_response(response, context): updated = dict(response) updated["llm_sanitized_response"] = True return updated diff --git a/crates/python/tests/coverage/py_callable_coverage_tests.rs b/crates/python/tests/coverage/py_callable_coverage_tests.rs index e6d8e381a..fe815cc1b 100644 --- a/crates/python/tests/coverage/py_callable_coverage_tests.rs +++ b/crates/python/tests/coverage/py_callable_coverage_tests.rs @@ -73,7 +73,7 @@ def collector_ok(chunk): def finalizer_bad_json(): return object() -def llm_resp_bad_json(response): +def llm_resp_bad_json(response, context): return object() class BadCodec: @@ -187,8 +187,15 @@ class RaisingResponseCodec: assert_eq!(finalizer(), serde_json::Value::Null); let llm_response = - wrap_py_llm_sanitize_response_fn(module.getattr("llm_resp_bad_json").unwrap().unbind()); - assert_eq!(llm_response(json!({"ok": true})), json!({"ok": true})); + wrap_py_llm_sanitize_response_fn(module.getattr("llm_resp_bad_json").unwrap().unbind()) + .unwrap(); + assert_eq!( + llm_response( + json!({"ok": true}), + nemo_relay::api::runtime::LlmSanitizeResponseContext::default() + ), + None + ); let bad_codec = PyLlmCodecWrapper { py_codec: module diff --git a/crates/python/tests/coverage/py_plugin_coverage_tests.rs b/crates/python/tests/coverage/py_plugin_coverage_tests.rs index c17b7e469..83bc2c9b5 100644 --- a/crates/python/tests/coverage/py_plugin_coverage_tests.rs +++ b/crates/python/tests/coverage/py_plugin_coverage_tests.rs @@ -67,6 +67,36 @@ fn plugin_context_helpers_and_error_conversion_work() { assert!(err.to_string().contains("boom")); } +#[test] +fn plugin_context_rejects_legacy_and_uninspectable_llm_sanitizers() { + let _python = crate::test_support::init_python_test(); + let context = PyPluginContext { + registrations: Arc::new(Mutex::new(vec![])), + namespace_prefix: "invalid.".to_string(), + }; + + Python::attach(|py| { + let helpers = load_module( + py, + r#" +def one_argument(payload): + return payload +"#, + ); + for callback in [helpers.getattr("one_argument").unwrap().unbind(), py.None()] { + let request_error = context + .register_llm_sanitize_request_guardrail("request", 1, callback.clone_ref(py)) + .unwrap_err(); + assert!(request_error.to_string().contains("payload, context")); + + let response_error = context + .register_llm_sanitize_response_guardrail("response", 1, callback) + .unwrap_err(); + assert!(response_error.to_string().contains("payload, context")); + } + }); +} + #[test] fn register_adds_plugin_management_bindings() { let _python = crate::test_support::init_python_test(); @@ -148,10 +178,10 @@ def tool_fn(name, value): def tool_conditional(name, value): return None -def llm_sanitize_request(request): +def llm_sanitize_request(request, context): return request -def llm_sanitize_response(response): +def llm_sanitize_response(response, context): return response def llm_conditional(request): @@ -628,10 +658,10 @@ def tool_fn(name, value): def tool_conditional(name, value): return None -def llm_sanitize_request(request): +def llm_sanitize_request(request, context): return request -def llm_sanitize_response(response): +def llm_sanitize_response(response, context): return response def llm_conditional(request): @@ -912,10 +942,10 @@ def tool_fn(name, value): def tool_conditional(name, value): return None -def llm_sanitize_request(request): +def llm_sanitize_request(request, context): return request -def llm_sanitize_response(response): +def llm_sanitize_response(response, context): return response def llm_conditional(request): 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 0d17d8b6b..75307e56d 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 @@ -25,6 +25,9 @@ service RelayHostRuntime { rpc ToolNext(ToolNextRequest) returns (JsonResult); rpc LlmNext(LlmNextRequest) returns (JsonResult); rpc LlmStreamNext(LlmStreamNextRequest) returns (stream StreamChunk); + rpc DecodeLlmCodecRequest(LlmCodecDecodeRequest) returns (JsonResult); + rpc EncodeLlmCodecRequest(LlmCodecEncodeRequest) returns (JsonResult); + rpc DecodeLlmCodecResponse(LlmCodecDecodeResponse) returns (JsonResult); } message JsonEnvelope { @@ -51,6 +54,13 @@ enum RegistrationSurface { SCOPE_SANITIZE_END_GUARDRAIL = 32; } +enum LlmCodecKind { + LLM_CODEC_KIND_UNSPECIFIED = 0; + LLM_CODEC_KIND_BUILTIN = 1; + LLM_CODEC_KIND_RUNTIME = 2; + LLM_CODEC_KIND_OPAQUE = 3; +} + enum ScopeType { SCOPE_TYPE_UNSPECIFIED = 0; AGENT = 1; @@ -132,6 +142,7 @@ message Registration { RegistrationSurface surface = 2; int32 priority = 3; bool break_chain = 4; + reserved 5; } message InvokeRequest { @@ -160,6 +171,28 @@ message LlmInvocation { JsonEnvelope request = 2; JsonEnvelope annotated_request = 3; JsonEnvelope response = 4; + reserved 5, 6, 7, 8; + oneof sanitize_context { + LlmSanitizeRequestContext request_sanitize_context = 9; + LlmSanitizeResponseContext response_sanitize_context = 10; + } +} + +message LlmCodecIdentity { + LlmCodecKind kind = 1; + optional string id = 2; +} + +message LlmSanitizeRequestContext { + LlmCodecIdentity codec = 1; + // Opaque, invocation-scoped host capability. SDKs must not expose this value. + optional string codec_capability_id = 2; +} + +message LlmSanitizeResponseContext { + LlmCodecIdentity codec = 1; + // Opaque, invocation-scoped host capability. SDKs must not expose this value. + optional string codec_capability_id = 2; } message InvokeResponse { @@ -302,3 +335,28 @@ message LlmStreamNextRequest { string continuation_id = 3; JsonEnvelope request = 4; } + +message LlmCodecDecodeRequest { + string activation_id = 1; + string auth_token = 2; + string codec_capability_id = 3; + JsonEnvelope request = 4; + string invocation_id = 5; +} + +message LlmCodecEncodeRequest { + string activation_id = 1; + string auth_token = 2; + string codec_capability_id = 3; + JsonEnvelope annotated_request = 4; + JsonEnvelope original_request = 5; + string invocation_id = 6; +} + +message LlmCodecDecodeResponse { + string activation_id = 1; + string auth_token = 2; + string codec_capability_id = 3; + JsonEnvelope response = 4; + string invocation_id = 5; +} diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 982501b64..f2e523f4a 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -52,7 +52,8 @@ use nemo_relay_worker_proto::v1::relay_host_runtime_client::RelayHostRuntimeClie use nemo_relay_worker_proto::v1::{ CancelInvocationRequest, CreateScopeStackRequest, DropScopeStackRequest, EmitMarkRequest, EmptyResult, GuardrailResult, HandshakeRequest, HandshakeResponse, HealthRequest, - HealthResponse, InvokeRequest, InvokeResponse, JsonEnvelope, JsonResult, LlmNextRequest, + HealthResponse, InvokeRequest, InvokeResponse, JsonEnvelope, JsonResult, LlmCodecDecodeRequest, + LlmCodecDecodeResponse, LlmCodecEncodeRequest, LlmCodecKind, LlmNextRequest, LlmRequestInterceptResult, LlmStreamNextRequest, PopScopeRequest, PushScopeRequest, RegisterRequest, RegisterResponse, Registration, RegistrationSurface, ScopeContext, ShutdownRequest, StreamChunk, ToolExecutionInterceptResult, ToolNextRequest, ValidateRequest, @@ -80,6 +81,9 @@ pub type BoxFutureResult = Pin> + Send>>; /// Boxed JSON stream returned by streaming worker callbacks. pub type JsonStream = Pin> + Send>>; +const JSON_SCHEMA: &str = "nemo.relay.Json@1"; +const LLM_REQUEST_SCHEMA: &str = "nemo.relay.LlmRequest@1"; + tokio::task_local! { static TASK_SCOPE_CONTEXT: Option; } @@ -126,15 +130,137 @@ pub trait WorkerPlugin: Send + Sync + 'static { type SubscriberFn = Arc; type EventSanitizeFn = - Arc EventSanitizeFields + Send + Sync>; -type ToolSanitizeFn = Arc Json + Send + Sync>; + Arc BoxFutureResult + Send + Sync>; +type ToolSanitizeFn = Arc BoxFutureResult + Send + Sync>; type ToolConditionalFn = Arc Result> + Send + Sync>; type ToolRequestFn = Arc Result + Send + Sync>; type ToolExecutionFn = Arc< dyn Fn(&str, Json, ToolNext) -> BoxFutureResult + Send + Sync, >; -type LlmSanitizeRequestFn = Arc LlmRequest + Send + Sync>; -type LlmSanitizeResponseFn = Arc Json + Send + Sync>; +type LlmSanitizeRequestFn = Arc< + dyn Fn(LlmRequest, LlmSanitizeRequestContext) -> BoxFutureResult> + + Send + + Sync, +>; +type LlmSanitizeResponseFn = + Arc BoxFutureResult> + Send + Sync>; + +/// Relay built-in codec identities supplied to worker sanitizers. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BuiltinLlmCodec { + /// OpenAI Chat Completions. + OpenAiChat, + /// OpenAI Responses. + OpenAiResponses, + /// Anthropic Messages. + AnthropicMessages, +} + +/// Per-call LLM codec identity supplied to worker sanitizers. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum LlmCodecIdentity { + /// No codec was active. + None, + /// A Relay built-in codec was active. + BuiltIn(BuiltinLlmCodec), + /// A runtime-registered codec was active, identified by its stable ID. + Runtime(String), + /// A codec was active but has no registered identity. + Opaque, +} + +/// Active codec context supplied to an LLM request sanitizer. +#[derive(Clone)] +pub struct LlmSanitizeRequestContext { + /// Identity of the active codec. + pub codec: LlmCodecIdentity, + runtime: Option, + codec_capability_id: Option, + invocation_id: Option, +} + +/// Active codec context supplied to an LLM response sanitizer. +#[derive(Clone)] +pub struct LlmSanitizeResponseContext { + /// Identity of the active codec. + pub codec: LlmCodecIdentity, + runtime: Option, + codec_capability_id: Option, + invocation_id: Option, +} + +impl LlmSanitizeRequestContext { + /// Resolves the active request codec for this callback. + #[must_use] + pub fn resolve_codec(&self) -> Option { + Some(WorkerRequestCodec { + runtime: self.runtime.clone()?, + capability_id: self.codec_capability_id.clone()?, + invocation_id: self.invocation_id.clone()?, + }) + } +} + +impl LlmSanitizeResponseContext { + /// Resolves the active response codec for this callback. + #[must_use] + pub fn resolve_codec(&self) -> Option { + Some(WorkerResponseCodec { + runtime: self.runtime.clone()?, + capability_id: self.codec_capability_id.clone()?, + invocation_id: self.invocation_id.clone()?, + }) + } +} + +/// Invocation-scoped proxy for the active LLM request codec. +#[derive(Clone)] +pub struct WorkerRequestCodec { + runtime: PluginRuntime, + capability_id: String, + invocation_id: String, +} + +impl WorkerRequestCodec { + /// Decodes an opaque request into its normalized representation. + pub async fn decode(&self, request: &LlmRequest) -> Result { + self.runtime + .decode_llm_codec_request(&self.capability_id, &self.invocation_id, request) + .await + } + /// Encodes normalized request changes onto the original opaque request. + pub async fn encode( + &self, + annotated: &AnnotatedLlmRequest, + original: &LlmRequest, + ) -> Result { + self.runtime + .encode_llm_codec_request( + &self.capability_id, + &self.invocation_id, + annotated, + original, + ) + .await + } +} + +/// Invocation-scoped proxy for the active LLM response codec. +#[derive(Clone)] +pub struct WorkerResponseCodec { + runtime: PluginRuntime, + capability_id: String, + invocation_id: String, +} + +impl WorkerResponseCodec { + /// Decodes an opaque response into its normalized representation. + pub async fn decode(&self, response: &Json) -> Result { + self.runtime + .decode_llm_codec_response(&self.capability_id, &self.invocation_id, response) + .await + } +} type LlmConditionalFn = Arc Result> + Send + Sync>; type LlmRequestFn = Arc< dyn Fn(&str, LlmRequest, Option) -> Result @@ -204,14 +330,15 @@ impl PluginContext { .insert(name.into(), Arc::new(callback)); } - fn register_event_sanitizer( + fn register_event_sanitizer( &mut self, name: &str, priority: i32, surface: RegistrationSurface, callback: F, ) where - F: Fn(&Event, EventSanitizeFields) -> EventSanitizeFields + Send + Sync + 'static, + F: Fn(&Event, EventSanitizeFields) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, { self.push_registration(name, surface, priority, false); let sanitizers = match surface { @@ -224,13 +351,21 @@ impl PluginContext { } _ => unreachable!("event sanitizer registration requires an event sanitizer surface"), }; - sanitizers.insert(name.into(), Arc::new(callback)); + sanitizers.insert( + name.into(), + Arc::new(move |event, fields| Box::pin(callback(event, fields))), + ); } /// Registers a mark event sanitizer. - pub fn register_mark_sanitize_guardrail(&mut self, name: &str, priority: i32, callback: F) - where - F: Fn(&Event, EventSanitizeFields) -> EventSanitizeFields + Send + Sync + 'static, + pub fn register_mark_sanitize_guardrail( + &mut self, + name: &str, + priority: i32, + callback: F, + ) where + F: Fn(&Event, EventSanitizeFields) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, { self.register_event_sanitizer( name, @@ -241,13 +376,14 @@ impl PluginContext { } /// Registers a scope-start event sanitizer. - pub fn register_scope_sanitize_start_guardrail( + pub fn register_scope_sanitize_start_guardrail( &mut self, name: &str, priority: i32, callback: F, ) where - F: Fn(&Event, EventSanitizeFields) -> EventSanitizeFields + Send + Sync + 'static, + F: Fn(&Event, EventSanitizeFields) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, { self.register_event_sanitizer( name, @@ -258,13 +394,14 @@ impl PluginContext { } /// Registers a scope-end event sanitizer. - pub fn register_scope_sanitize_end_guardrail( + pub fn register_scope_sanitize_end_guardrail( &mut self, name: &str, priority: i32, callback: F, ) where - F: Fn(&Event, EventSanitizeFields) -> EventSanitizeFields + Send + Sync + 'static, + F: Fn(&Event, EventSanitizeFields) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, { self.register_event_sanitizer( name, @@ -275,13 +412,14 @@ impl PluginContext { } /// Registers a tool sanitize-request guardrail. - pub fn register_tool_sanitize_request_guardrail( + pub fn register_tool_sanitize_request_guardrail( &mut self, name: &str, priority: i32, callback: F, ) where - F: Fn(&str, Json) -> Json + Send + Sync + 'static, + F: Fn(&str, Json) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, { self.push_registration( name, @@ -289,19 +427,21 @@ impl PluginContext { priority, false, ); - self.handlers - .tool_sanitize_requests - .insert(name.into(), Arc::new(callback)); + self.handlers.tool_sanitize_requests.insert( + name.into(), + Arc::new(move |tool_name, value| Box::pin(callback(tool_name, value))), + ); } /// Registers a tool sanitize-response guardrail. - pub fn register_tool_sanitize_response_guardrail( + pub fn register_tool_sanitize_response_guardrail( &mut self, name: &str, priority: i32, callback: F, ) where - F: Fn(&str, Json) -> Json + Send + Sync + 'static, + F: Fn(&str, Json) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, { self.push_registration( name, @@ -309,9 +449,10 @@ impl PluginContext { priority, false, ); - self.handlers - .tool_sanitize_responses - .insert(name.into(), Arc::new(callback)); + self.handlers.tool_sanitize_responses.insert( + name.into(), + Arc::new(move |tool_name, value| Box::pin(callback(tool_name, value))), + ); } /// Registers a tool conditional-execution guardrail. @@ -382,13 +523,14 @@ impl PluginContext { } /// Registers an LLM sanitize-request guardrail. - pub fn register_llm_sanitize_request_guardrail( + pub fn register_llm_sanitize_request_guardrail( &mut self, name: &str, priority: i32, callback: F, ) where - F: Fn(LlmRequest) -> LlmRequest + Send + Sync + 'static, + F: Fn(LlmRequest, LlmSanitizeRequestContext) -> Fut + Send + Sync + 'static, + Fut: Future>> + Send + 'static, { self.push_registration( name, @@ -396,19 +538,21 @@ impl PluginContext { priority, false, ); - self.handlers - .llm_sanitize_requests - .insert(name.into(), Arc::new(callback)); + self.handlers.llm_sanitize_requests.insert( + name.into(), + Arc::new(move |request, context| Box::pin(callback(request, context))), + ); } /// Registers an LLM sanitize-response guardrail. - pub fn register_llm_sanitize_response_guardrail( + pub fn register_llm_sanitize_response_guardrail( &mut self, name: &str, priority: i32, callback: F, ) where - F: Fn(Json) -> Json + Send + Sync + 'static, + F: Fn(Json, LlmSanitizeResponseContext) -> Fut + Send + Sync + 'static, + Fut: Future>> + Send + 'static, { self.push_registration( name, @@ -416,9 +560,10 @@ impl PluginContext { priority, false, ); - self.handlers - .llm_sanitize_responses - .insert(name.into(), Arc::new(callback)); + self.handlers.llm_sanitize_responses.insert( + name.into(), + Arc::new(move |response, context| Box::pin(callback(response, context))), + ); } /// Registers an LLM conditional-execution guardrail. @@ -541,6 +686,70 @@ pub struct PluginRuntime { } impl PluginRuntime { + async fn decode_llm_codec_request( + &self, + capability_id: &str, + invocation_id: &str, + request: &LlmRequest, + ) -> Result { + let mut client = self.host_client().await?; + let response = client + .decode_llm_codec_request(Request::new(LlmCodecDecodeRequest { + activation_id: self.activation_id.clone(), + auth_token: self.auth_token.clone(), + codec_capability_id: capability_id.into(), + invocation_id: invocation_id.into(), + request: Some(json_envelope(LLM_REQUEST_SCHEMA, request)?), + })) + .await + .map_err(|err| WorkerSdkError::Transport(err.to_string()))? + .into_inner(); + decode_typed_json_result(response, ANNOTATED_LLM_REQUEST_SCHEMA) + } + + async fn encode_llm_codec_request( + &self, + capability_id: &str, + invocation_id: &str, + annotated: &AnnotatedLlmRequest, + original: &LlmRequest, + ) -> Result { + let mut client = self.host_client().await?; + let response = client + .encode_llm_codec_request(Request::new(LlmCodecEncodeRequest { + activation_id: self.activation_id.clone(), + auth_token: self.auth_token.clone(), + codec_capability_id: capability_id.into(), + invocation_id: invocation_id.into(), + annotated_request: Some(json_envelope(ANNOTATED_LLM_REQUEST_SCHEMA, annotated)?), + original_request: Some(json_envelope(LLM_REQUEST_SCHEMA, original)?), + })) + .await + .map_err(|err| WorkerSdkError::Transport(err.to_string()))? + .into_inner(); + decode_typed_json_result(response, LLM_REQUEST_SCHEMA) + } + + async fn decode_llm_codec_response( + &self, + capability_id: &str, + invocation_id: &str, + response: &Json, + ) -> Result { + let mut client = self.host_client().await?; + let response = client + .decode_llm_codec_response(Request::new(LlmCodecDecodeResponse { + activation_id: self.activation_id.clone(), + auth_token: self.auth_token.clone(), + codec_capability_id: capability_id.into(), + invocation_id: invocation_id.into(), + response: Some(json_envelope(JSON_SCHEMA, response)?), + })) + .await + .map_err(|err| WorkerSdkError::Transport(err.to_string()))? + .into_inner(); + decode_typed_json_result(response, JSON_SCHEMA) + } /// Emits a mark event through the host runtime. pub async fn emit_mark( &self, @@ -713,7 +922,7 @@ impl ToolNext { activation_id: self.runtime.activation_id.clone(), auth_token: self.runtime.auth_token.clone(), continuation_id: self.continuation_id.clone(), - value: Some(json_envelope("nemo.relay.Json@1", &value)?), + value: Some(json_envelope(JSON_SCHEMA, &value)?), })) .await .map_err(|err| WorkerSdkError::Transport(err.to_string()))? @@ -738,7 +947,7 @@ impl LlmNext { activation_id: self.runtime.activation_id.clone(), auth_token: self.runtime.auth_token.clone(), continuation_id: self.continuation_id.clone(), - request: Some(json_envelope("nemo.relay.LlmRequest@1", &request)?), + request: Some(json_envelope(LLM_REQUEST_SCHEMA, &request)?), })) .await .map_err(|err| WorkerSdkError::Transport(err.to_string()))? @@ -764,7 +973,7 @@ impl LlmStreamNext { activation_id: self.runtime.activation_id.clone(), auth_token: self.runtime.auth_token.clone(), continuation_id: self.continuation_id.clone(), - request: Some(json_envelope("nemo.relay.LlmRequest@1", &request)?), + request: Some(json_envelope(LLM_REQUEST_SCHEMA, &request)?), })) .await .map_err(|err| WorkerSdkError::Transport(err.to_string()))?; @@ -1222,7 +1431,7 @@ impl PluginWorker for WorkerService { let chunk = match item { Ok(value) => StreamChunk { item: Some(nemo_relay_worker_proto::v1::stream_chunk::Item::Value( - match json_envelope("nemo.relay.Json@1", &value) { + match json_envelope(JSON_SCHEMA, &value) { Ok(value) => value, Err(err) => { let _ = @@ -1431,127 +1640,267 @@ impl WorkerService { let surface = RegistrationSurface::try_from(request.surface) .map_err(|_| WorkerSdkError::InvalidInput("unknown registration surface".into()))?; match surface { - RegistrationSurface::Subscriber => { - let event = event_payload(request.payload)?; - let handler = self.subscriber(&request.registration_name)?; - with_thread_scope(&scope, || handler(&event)); - Ok(empty_response()) - } + RegistrationSurface::Subscriber => self.invoke_subscriber_response(request, &scope), RegistrationSurface::MarkSanitizeGuardrail | RegistrationSurface::ScopeSanitizeStartGuardrail | RegistrationSurface::ScopeSanitizeEndGuardrail => { - let event = event_payload(request.payload)?; - let fields = event.sanitize_fields(); - let handler = self.event_sanitizer(surface, &request.registration_name)?; - let fields = with_thread_scope(&scope, || handler(&event, fields)); - Ok(json_response( - serde_json::to_value(fields) - .expect("event sanitize fields are JSON serializable"), + self.invoke_event_sanitize_response(request, &scope, surface) + .await + } + RegistrationSurface::ToolSanitizeRequestGuardrail + | RegistrationSurface::ToolSanitizeResponseGuardrail + | RegistrationSurface::ToolConditionalExecutionGuardrail + | RegistrationSurface::ToolRequestIntercept + | RegistrationSurface::ToolExecutionIntercept => { + self.invoke_tool_response(request, &scope, surface).await + } + RegistrationSurface::LlmSanitizeRequestGuardrail + | RegistrationSurface::LlmSanitizeResponseGuardrail + | RegistrationSurface::LlmConditionalExecutionGuardrail + | RegistrationSurface::LlmRequestIntercept + | RegistrationSurface::LlmExecutionIntercept => { + self.invoke_llm_response(request, &scope, surface).await + } + RegistrationSurface::LlmStreamExecutionIntercept | RegistrationSurface::Unspecified => { + Err(WorkerSdkError::InvalidInput( + "surface must use InvokeStream or is unspecified".into(), )) } + } + } + + fn invoke_subscriber_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let event = event_payload(request.payload)?; + let handler = self.subscriber(&request.registration_name)?; + with_thread_scope(scope, || handler(&event)); + Ok(empty_response()) + } + + async fn invoke_event_sanitize_response( + &self, + request: InvokeRequest, + scope: &Option, + surface: RegistrationSurface, + ) -> Result { + let event = event_payload(request.payload)?; + let fields = event.sanitize_fields(); + let handler = self.event_sanitizer(surface, &request.registration_name)?; + let fields = with_thread_scope(scope, || handler(&event, fields)).await?; + Ok(json_response( + serde_json::to_value(fields).expect("event sanitize fields are JSON serializable"), + )) + } + + async fn invoke_tool_response( + &self, + request: InvokeRequest, + scope: &Option, + surface: RegistrationSurface, + ) -> Result { + match surface { RegistrationSurface::ToolSanitizeRequestGuardrail => { - let payload = tool_payload(request.payload)?; - let handler = self.tool_sanitize_request(&request.registration_name)?; - Ok(json_response(with_thread_scope(&scope, || { - handler(&payload.tool_name, payload.value) - }))) + self.invoke_tool_sanitize_request_response(request, scope) + .await } RegistrationSurface::ToolSanitizeResponseGuardrail => { - let payload = tool_payload(request.payload)?; - let handler = self.tool_sanitize_response(&request.registration_name)?; - Ok(json_response(with_thread_scope(&scope, || { - handler(&payload.tool_name, payload.value) - }))) + self.invoke_tool_sanitize_response_response(request, scope) + .await } RegistrationSurface::ToolConditionalExecutionGuardrail => { - let payload = tool_payload(request.payload)?; - let handler = self.tool_conditional(&request.registration_name)?; - Ok(guardrail_response(with_thread_scope(&scope, || { - handler(&payload.tool_name, &payload.value) - })?)) + self.invoke_tool_conditional_response(request, scope) } RegistrationSurface::ToolRequestIntercept => { - let payload = tool_payload(request.payload)?; - let handler = self.tool_request(&request.registration_name)?; - Ok(json_response(with_thread_scope(&scope, || { - handler(&payload.tool_name, payload.value) - })?)) + self.invoke_tool_request_response(request, scope) } RegistrationSurface::ToolExecutionIntercept => { - let payload = tool_payload(request.payload)?; - let handler = self.tool_execution(&request.registration_name)?; - let next = ToolNext { - runtime: self.runtime.clone(), - continuation_id: request.continuation_id, - }; - let future = - with_thread_scope(&scope, || handler(&payload.tool_name, payload.value, next)); - Ok(tool_execution_response(future.await?)?) + self.invoke_tool_execution_response(request, scope).await } + _ => unreachable!("tool surface was pre-filtered"), + } + } + + async fn invoke_tool_sanitize_request_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = tool_payload(request.payload)?; + let handler = self.tool_sanitize_request(&request.registration_name)?; + Ok(json_response( + with_thread_scope(scope, || handler(&payload.tool_name, payload.value)).await?, + )) + } + + async fn invoke_tool_sanitize_response_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = tool_payload(request.payload)?; + let handler = self.tool_sanitize_response(&request.registration_name)?; + Ok(json_response( + with_thread_scope(scope, || handler(&payload.tool_name, payload.value)).await?, + )) + } + + fn invoke_tool_conditional_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = tool_payload(request.payload)?; + let handler = self.tool_conditional(&request.registration_name)?; + Ok(guardrail_response(with_thread_scope(scope, || { + handler(&payload.tool_name, &payload.value) + })?)) + } + + fn invoke_tool_request_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = tool_payload(request.payload)?; + let handler = self.tool_request(&request.registration_name)?; + Ok(json_response(with_thread_scope(scope, || { + handler(&payload.tool_name, payload.value) + })?)) + } + + async fn invoke_tool_execution_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = tool_payload(request.payload)?; + let handler = self.tool_execution(&request.registration_name)?; + let next = ToolNext { + runtime: self.runtime.clone(), + continuation_id: request.continuation_id, + }; + let future = with_thread_scope(scope, || handler(&payload.tool_name, payload.value, next)); + tool_execution_response(future.await?) + } + + async fn invoke_llm_response( + &self, + request: InvokeRequest, + scope: &Option, + surface: RegistrationSurface, + ) -> Result { + match surface { RegistrationSurface::LlmSanitizeRequestGuardrail => { - let payload = llm_payload(request.payload)?; - let request_value = required_json::(payload.request, "llm request")?; - let handler = self.llm_sanitize_request(&request.registration_name)?; - let request = with_thread_scope(&scope, || handler(request_value)); - Ok(json_response( - serde_json::to_value(request).expect("LLM request is JSON serializable"), - )) + self.invoke_llm_sanitize_request_response(request, scope) + .await } RegistrationSurface::LlmSanitizeResponseGuardrail => { - let payload = llm_payload(request.payload)?; - let response = required_json::(payload.response, "llm response")?; - let handler = self.llm_sanitize_response(&request.registration_name)?; - Ok(json_response(with_thread_scope(&scope, || { - handler(response) - }))) + self.invoke_llm_sanitize_response_response(request, scope) + .await } RegistrationSurface::LlmConditionalExecutionGuardrail => { - let payload = llm_payload(request.payload)?; - let request_value = required_json::(payload.request, "llm request")?; - let handler = self.llm_conditional(&request.registration_name)?; - Ok(guardrail_response(with_thread_scope(&scope, || { - handler(&request_value) - })?)) + self.invoke_llm_conditional_response(request, scope) } RegistrationSurface::LlmRequestIntercept => { - let payload = llm_payload(request.payload)?; - let request_value = required_json::(payload.request, "llm request")?; - let annotated = payload - .annotated_request - .map(|value| { - decode_expected_json_envelope::( - &value, - "annotated llm request", - ANNOTATED_LLM_REQUEST_SCHEMA, - ) - }) - .transpose()?; - let handler = self.llm_request(&request.registration_name)?; - let outcome = with_thread_scope(&scope, || { - handler(&payload.model_name, request_value, annotated) - })?; - Ok(llm_request_response(outcome)?) + self.invoke_llm_request_response(request, scope) } RegistrationSurface::LlmExecutionIntercept => { - let payload = llm_payload(request.payload)?; - let request_value = required_json::(payload.request, "llm request")?; - let handler = self.llm_execution(&request.registration_name)?; - let next = LlmNext { - runtime: self.runtime.clone(), - continuation_id: request.continuation_id, - }; - let future = - with_thread_scope(&scope, || handler(&payload.model_name, request_value, next)); - Ok(json_response(future.await?)) - } - RegistrationSurface::LlmStreamExecutionIntercept | RegistrationSurface::Unspecified => { - Err(WorkerSdkError::InvalidInput( - "surface must use InvokeStream or is unspecified".into(), - )) + self.invoke_llm_execution_response(request, scope).await } + _ => unreachable!("LLM surface was pre-filtered"), } } + async fn invoke_llm_sanitize_request_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = llm_payload(request.payload)?; + let mut context = payload.sanitize_request_context(&request.invocation_id); + context.runtime = Some(self.runtime.clone()); + let request_value = required_json::(payload.request, "llm request")?; + let handler = self.llm_sanitize_request(&request.registration_name)?; + match with_thread_scope(scope, || handler(request_value, context)).await? { + Some(request) => Ok(json_response( + serde_json::to_value(request).expect("LLM request is JSON serializable"), + )), + None => Ok(empty_response()), + } + } + + async fn invoke_llm_sanitize_response_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = llm_payload(request.payload)?; + let mut context = payload.sanitize_response_context(&request.invocation_id); + context.runtime = Some(self.runtime.clone()); + let response = required_json::(payload.response, "llm response")?; + let handler = self.llm_sanitize_response(&request.registration_name)?; + match with_thread_scope(scope, || handler(response, context)).await? { + Some(response) => Ok(json_response(response)), + None => Ok(empty_response()), + } + } + + fn invoke_llm_conditional_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = llm_payload(request.payload)?; + let request_value = required_json::(payload.request, "llm request")?; + let handler = self.llm_conditional(&request.registration_name)?; + Ok(guardrail_response(with_thread_scope(scope, || { + handler(&request_value) + })?)) + } + + fn invoke_llm_request_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = llm_payload(request.payload)?; + let request_value = required_json::(payload.request, "llm request")?; + let annotated = payload + .annotated_request + .map(|value| { + decode_expected_json_envelope::( + &value, + "annotated llm request", + ANNOTATED_LLM_REQUEST_SCHEMA, + ) + }) + .transpose()?; + let handler = self.llm_request(&request.registration_name)?; + let outcome = with_thread_scope(scope, || { + handler(&payload.model_name, request_value, annotated) + })?; + llm_request_response(outcome) + } + + async fn invoke_llm_execution_response( + &self, + request: InvokeRequest, + scope: &Option, + ) -> Result { + let payload = llm_payload(request.payload)?; + let request_value = required_json::(payload.request, "llm request")?; + let handler = self.llm_execution(&request.registration_name)?; + let next = LlmNext { + runtime: self.runtime.clone(), + continuation_id: request.continuation_id, + }; + let future = with_thread_scope(scope, || handler(&payload.model_name, request_value, next)); + Ok(json_response(future.await?)) + } + fn subscriber(&self, name: &str) -> Result { self.handlers .lock() @@ -1719,6 +2068,65 @@ struct LlmPayload { request: Option, annotated_request: Option, response: Option, + sanitize_context: Option, +} + +impl LlmPayload { + fn sanitize_request_context(&self, invocation_id: &str) -> LlmSanitizeRequestContext { + let codec = match self.sanitize_context.as_ref() { + Some(nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::RequestSanitizeContext(context)) => context.codec.as_ref(), + _ => None, + }; + LlmSanitizeRequestContext { + codec: codec_identity_from_proto(codec), + runtime: None, + codec_capability_id: match self.sanitize_context.as_ref() { + Some(nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::RequestSanitizeContext(context)) => context.codec_capability_id.clone(), + _ => None, + }, + invocation_id: Some(invocation_id.to_owned()), + } + } + + fn sanitize_response_context(&self, invocation_id: &str) -> LlmSanitizeResponseContext { + let codec = match self.sanitize_context.as_ref() { + Some(nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::ResponseSanitizeContext(context)) => context.codec.as_ref(), + _ => None, + }; + LlmSanitizeResponseContext { + codec: codec_identity_from_proto(codec), + runtime: None, + codec_capability_id: match self.sanitize_context.as_ref() { + Some(nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::ResponseSanitizeContext(context)) => context.codec_capability_id.clone(), + _ => None, + }, + invocation_id: Some(invocation_id.to_owned()), + } + } +} + +fn codec_identity_from_proto( + codec: Option<&nemo_relay_worker_proto::v1::LlmCodecIdentity>, +) -> LlmCodecIdentity { + let codec_kind = codec + .map(|codec| codec.kind) + .unwrap_or(LlmCodecKind::Unspecified as i32); + let codec_id = codec.and_then(|codec| codec.id.clone()); + match LlmCodecKind::try_from(codec_kind).ok() { + Some(LlmCodecKind::Unspecified) => LlmCodecIdentity::None, + Some(LlmCodecKind::Builtin) => match codec_id.as_deref() { + Some("openai_chat") => LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat), + Some("openai_responses") => LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiResponses), + Some("anthropic_messages") => { + LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::AnthropicMessages) + } + _ => LlmCodecIdentity::Opaque, + }, + Some(LlmCodecKind::Runtime) => codec_id + .filter(|id| !id.is_empty()) + .map_or(LlmCodecIdentity::Opaque, LlmCodecIdentity::Runtime), + Some(LlmCodecKind::Opaque) | None => LlmCodecIdentity::Opaque, + } } fn event_payload( @@ -1758,6 +2166,7 @@ fn llm_payload( request: value.request, annotated_request: value.annotated_request, response: value.response, + sanitize_context: value.sanitize_context, }), _ => Err(WorkerSdkError::InvalidInput("expected llm payload".into())), } @@ -1797,7 +2206,7 @@ fn json_response(value: Json) -> InvokeResponse { InvokeResponse { result: Some(nemo_relay_worker_proto::v1::invoke_response::Result::Json( JsonResult { - value: Some(infallible_json_envelope("nemo.relay.Json@1", &value)), + value: Some(infallible_json_envelope(JSON_SCHEMA, &value)), error: None, }, )), @@ -1863,10 +2272,23 @@ fn json_result_to_sdk(result: JsonResult) -> Result { required_json(result.value, "json result") } +fn decode_typed_json_result( + result: JsonResult, + expected_schema: &str, +) -> Result { + if let Some(error) = result.error { + return Err(worker_error_to_sdk(error)); + } + let value = result + .value + .ok_or_else(|| WorkerSdkError::InvalidInput("json result is missing".into()))?; + decode_expected_json_envelope(&value, "json result", expected_schema) +} + fn optional_json_envelope(value: Option) -> Result> { value .as_ref() - .map(|value| json_envelope("nemo.relay.Json@1", value).map_err(WorkerSdkError::from)) + .map(|value| json_envelope(JSON_SCHEMA, value).map_err(WorkerSdkError::from)) .transpose() } diff --git a/crates/worker/tests/worker_sdk_tests.rs b/crates/worker/tests/worker_sdk_tests.rs index bedf2fb71..80653155f 100644 --- a/crates/worker/tests/worker_sdk_tests.rs +++ b/crates/worker/tests/worker_sdk_tests.rs @@ -200,7 +200,22 @@ 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(), 19); + assert_eq!(registrations.len(), 21); + for local_name in [ + "llm-sanitize-request", + "llm-sanitize-response", + "llm-sanitize-omit-request", + "llm-sanitize-omit-response", + ] { + assert_eq!( + registrations + .iter() + .filter(|registration| registration.local_name == local_name) + .count(), + 1, + "expected exactly one {local_name} registration" + ); + } assert_eq!( registrations .iter() @@ -484,7 +499,20 @@ async fn worker_service_invokes_every_registration_surface() { let plugin = Arc::new(SurfacePlugin::default()); let events = plugin.events.clone(); let (worker_handle, mut client) = spawn_worker(plugin, tcp_endpoint(&host_endpoint)).await; - register_plugin(&mut client).await; + let registrations = register_plugin(&mut client).await; + for name in [ + "llm-sanitize-request", + "llm-sanitize-response", + "llm-sanitize-omit-request", + "llm-sanitize-omit-response", + ] { + assert!( + registrations + .iter() + .any(|registration| registration.local_name == name), + "{name} must be registered" + ); + } let subscriber_response = client .invoke(Request::new(event_invoke("subscriber"))) @@ -632,6 +660,47 @@ async fn worker_service_invokes_every_registration_surface() { "phase", "llm_sanitize_response", ); + let omitted_request = client + .invoke(Request::new(llm_invoke( + "llm-sanitize-omit-request", + RegistrationSurface::LlmSanitizeRequestGuardrail, + llm_request(), + None, + None, + ))) + .await + .expect("request omission invoke") + .into_inner(); + assert_empty_response(omitted_request); + let omitted_response = client + .invoke(Request::new(llm_invoke( + "llm-sanitize-omit-response", + RegistrationSurface::LlmSanitizeResponseGuardrail, + llm_request(), + None, + Some(json!({})), + ))) + .await + .expect("response omission invoke") + .into_inner(); + assert_empty_response(omitted_response); + let codec_request_result = invoke_json(&mut client, llm_request_codec_invoke()).await; + assert_eq!( + codec_request_result, + json!({ + "headers": {}, + "content": {"phase": "llm_sanitize_request"}, + }) + ); + + let codec_response_result = invoke_json(&mut client, llm_response_codec_invoke()).await; + assert_eq!( + codec_response_result, + json!({ + "value": "secret", + "phase": "llm_sanitize_response", + }) + ); assert_eq!( invoke_guardrail( &mut client, @@ -718,6 +787,9 @@ async fn worker_service_invokes_every_registration_surface() { assert!(calls.contains(&"tool_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())); + assert!(calls.contains(&"codec_request_encode:request-capability".into())); + assert!(calls.contains(&"codec_response_decode:response-capability".into())); assert!(calls.contains(&"mark:stream-poll:stack-1:parent-1".into())); assert!(calls.contains(&"push:scope-agent:explicit-stack:".into())); assert!(calls.contains(&"push:scope-unknown:explicit-stack:".into())); @@ -1439,6 +1511,43 @@ async fn worker_service_propagates_host_runtime_errors() { "llm next failed", ); + for (failures, request, expected) in [ + ( + MockHostFailures { + codec_request_decode: true, + ..Default::default() + }, + llm_request_codec_invoke(), + "codec request decode failed", + ), + ( + MockHostFailures { + codec_request_encode: true, + ..Default::default() + }, + llm_request_codec_invoke(), + "codec request encode failed", + ), + ( + MockHostFailures { + codec_response_decode: true, + ..Default::default() + }, + llm_response_codec_invoke(), + "codec response decode failed", + ), + ] { + host.set_failures(failures); + assert_worker_error( + client + .invoke(Request::new(request)) + .await + .expect("codec failure returns structured error") + .into_inner(), + expected, + ); + } + for (mode, expected) in [ (MockStreamMode::WorkerError, "stream worker failed"), (MockStreamMode::EmptyChunk, "empty stream chunk"), @@ -1489,8 +1598,8 @@ impl WorkerPlugin for DuplicateEventSanitizerPlugin { } fn register(&self, ctx: &mut PluginContext, _config: &Json) -> Result<()> { - ctx.register_mark_sanitize_guardrail("duplicate", 0, |_, fields| fields); - ctx.register_mark_sanitize_guardrail("duplicate", 1, |_, fields| fields); + ctx.register_mark_sanitize_guardrail("duplicate", 0, |_, fields| async move { Ok(fields) }); + ctx.register_mark_sanitize_guardrail("duplicate", 1, |_, fields| async move { Ok(fields) }); Ok(()) } } @@ -1634,22 +1743,33 @@ impl WorkerPlugin for SurfacePlugin { .push(event.name().into()); }); ctx.register_mark_sanitize_guardrail("event-sanitize", 1, |event, mut fields| { - fields.data = Some(json!({"name": event.name(), "phase": "mark"})); - fields - }); - ctx.register_scope_sanitize_start_guardrail("event-sanitize", 1, |_, mut fields| { - fields.metadata = Some(json!({"phase": "scope_start"})); - fields - }); - ctx.register_scope_sanitize_end_guardrail("event-sanitize", 1, |_, mut fields| { - fields.metadata = Some(json!({"phase": "scope_end"})); - fields + let event_name = event.name().to_owned(); + async move { + fields.data = Some(json!({"name": event_name, "phase": "mark"})); + Ok(fields) + } }); - ctx.register_tool_sanitize_request_guardrail("tool-sanitize", 1, |_, value| { - set_json_field(value, "phase", "tool_sanitize_request") + ctx.register_scope_sanitize_start_guardrail( + "event-sanitize", + 1, + |_, mut fields| async move { + fields.metadata = Some(json!({"phase": "scope_start"})); + Ok(fields) + }, + ); + ctx.register_scope_sanitize_end_guardrail( + "event-sanitize", + 1, + |_, mut fields| async move { + fields.metadata = Some(json!({"phase": "scope_end"})); + Ok(fields) + }, + ); + ctx.register_tool_sanitize_request_guardrail("tool-sanitize", 1, |_, value| async move { + Ok(set_json_field(value, "phase", "tool_sanitize_request")) }); - ctx.register_tool_sanitize_response_guardrail("tool-sanitize", 1, |_, value| { - set_json_field(value, "phase", "tool_sanitize_response") + ctx.register_tool_sanitize_response_guardrail("tool-sanitize", 1, |_, value| async move { + Ok(set_json_field(value, "phase", "tool_sanitize_response")) }); ctx.register_tool_conditional_execution_guardrail("tool-conditional", 1, |_, value| { Ok(value @@ -1730,12 +1850,41 @@ impl WorkerPlugin for SurfacePlugin { } }); - ctx.register_llm_sanitize_request_guardrail("llm-sanitize-request", 1, |request| { - set_llm_phase(request, "llm_sanitize_request") - }); - ctx.register_llm_sanitize_response_guardrail("llm-sanitize-response", 1, |value| { - set_json_field(value, "phase", "llm_sanitize_response") - }); + ctx.register_llm_sanitize_request_guardrail( + "llm-sanitize-request", + 1, + |mut request, context| async move { + if let Some(codec) = context.resolve_codec() { + let annotated = codec.decode(&request).await?; + request = codec.encode(&annotated, &request).await?; + } + Ok(Some(set_llm_phase(request, "llm_sanitize_request"))) + }, + ); + ctx.register_llm_sanitize_response_guardrail( + "llm-sanitize-response", + 1, + |value, context| async move { + if let Some(codec) = context.resolve_codec() { + codec.decode(&value).await?; + } + Ok(Some(set_json_field( + value, + "phase", + "llm_sanitize_response", + ))) + }, + ); + ctx.register_llm_sanitize_request_guardrail( + "llm-sanitize-omit-request", + 1, + |_request, _context| async { Ok(None) }, + ); + ctx.register_llm_sanitize_response_guardrail( + "llm-sanitize-omit-response", + 1, + |_response, _context| async { Ok(None) }, + ); ctx.register_llm_conditional_execution_guardrail("llm-conditional", 1, |request| { Ok(request .content @@ -1854,6 +2003,9 @@ struct MockHostFailures { tool_next: bool, llm_next: bool, llm_stream_mode: MockStreamMode, + codec_request_decode: bool, + codec_request_encode: bool, + codec_response_decode: bool, } #[derive(Clone, Default)] @@ -2062,6 +2214,81 @@ impl RelayHostRuntime for MockHost { let stream = tokio_stream::iter(chunks); Ok(Response::new(Box::pin(stream))) } + + async fn decode_llm_codec_request( + &self, + request: Request, + ) -> std::result::Result, Status> { + let request = request.into_inner(); + authorize_host(&request.activation_id, &request.auth_token)?; + self.record(format!( + "codec_request_decode:{}", + request.codec_capability_id + )); + if self.failures().codec_request_decode { + return Ok(Response::new(JsonResult { + value: None, + error: Some(worker_error("codec request decode failed")), + })); + } + Ok(Response::new(JsonResult { + value: Some( + json_envelope("nemo.relay.AnnotatedLlmRequest@2", &json!({})) + .expect("encode annotated request"), + ), + error: None, + })) + } + + async fn encode_llm_codec_request( + &self, + request: Request, + ) -> std::result::Result, Status> { + let request = request.into_inner(); + authorize_host(&request.activation_id, &request.auth_token)?; + self.record(format!( + "codec_request_encode:{}", + request.codec_capability_id + )); + if self.failures().codec_request_encode { + return Ok(Response::new(JsonResult { + value: None, + error: Some(worker_error("codec request encode failed")), + })); + } + Ok(Response::new(JsonResult { + value: Some( + json_envelope( + "nemo.relay.LlmRequest@1", + &json!({"headers": {}, "content": {}}), + ) + .expect("encode LLM request"), + ), + error: None, + })) + } + + async fn decode_llm_codec_response( + &self, + request: Request, + ) -> std::result::Result, Status> { + let request = request.into_inner(); + authorize_host(&request.activation_id, &request.auth_token)?; + self.record(format!( + "codec_response_decode:{}", + request.codec_capability_id + )); + if self.failures().codec_response_decode { + return Ok(Response::new(JsonResult { + value: None, + error: Some(worker_error("codec response decode failed")), + })); + } + Ok(Response::new(JsonResult { + value: Some(json_env(json!({}))), + error: None, + })) + } } async fn spawn_worker( @@ -2317,6 +2544,7 @@ fn llm_invoke( ), annotated_request: annotated_request.map(json_env), response: response.map(json_env), + sanitize_context: None, }, )), } @@ -2340,6 +2568,7 @@ fn llm_invoke_without_request( request: None, annotated_request: None, response: None, + sanitize_context: None, }, )), } @@ -2404,6 +2633,58 @@ fn llm_request_with_block() -> LlmRequest { } } +fn llm_request_codec_invoke() -> InvokeRequest { + let mut request = llm_invoke( + "llm-sanitize-request", + RegistrationSurface::LlmSanitizeRequestGuardrail, + llm_request(), + None, + None, + ); + if let Some(nemo_relay_worker_proto::v1::invoke_request::Payload::Llm(invocation)) = + request.payload.as_mut() + { + invocation.sanitize_context = Some( + nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::RequestSanitizeContext( + nemo_relay_worker_proto::v1::LlmSanitizeRequestContext { + codec: Some(nemo_relay_worker_proto::v1::LlmCodecIdentity { + kind: nemo_relay_worker_proto::v1::LlmCodecKind::Builtin as i32, + id: Some("openai_chat".into()), + }), + codec_capability_id: Some("request-capability".into()), + }, + ), + ); + } + request +} + +fn llm_response_codec_invoke() -> InvokeRequest { + let mut request = llm_invoke( + "llm-sanitize-response", + RegistrationSurface::LlmSanitizeResponseGuardrail, + llm_request(), + None, + Some(json!({"value": "secret"})), + ); + if let Some(nemo_relay_worker_proto::v1::invoke_request::Payload::Llm(invocation)) = + request.payload.as_mut() + { + invocation.sanitize_context = Some( + nemo_relay_worker_proto::v1::llm_invocation::SanitizeContext::ResponseSanitizeContext( + nemo_relay_worker_proto::v1::LlmSanitizeResponseContext { + codec: Some(nemo_relay_worker_proto::v1::LlmCodecIdentity { + kind: nemo_relay_worker_proto::v1::LlmCodecKind::Opaque as i32, + id: None, + }), + codec_capability_id: Some("response-capability".into()), + }, + ), + ); + } + request +} + fn set_json_field(mut value: Json, key: &str, field_value: &str) -> Json { value .as_object_mut() diff --git a/docs/about-nemo-relay/concepts/middleware.mdx b/docs/about-nemo-relay/concepts/middleware.mdx index 8b4011b63..8f2b9aca1 100644 --- a/docs/about-nemo-relay/concepts/middleware.mdx +++ b/docs/about-nemo-relay/concepts/middleware.mdx @@ -191,6 +191,124 @@ guardrails important: - If you need to change the real execution path, use an intercept - If you need to change only the emitted payload, use a sanitize guardrail +## Codec-Aware LLM Sanitizers + +Every LLM sanitize guardrail receives the payload first and a required +directional, per-call context second: + +```text +request: (LlmRequest, LlmSanitizeRequestContext) -> Option +response: (Json, LlmSanitizeResponseContext) -> Option +``` + +`context.codec` is a binding-native codec identity, not a JSON payload. +Python and Node.js expose `kind` and, when applicable, `id` properties. +In-process Rust and the typed native Rust SDK expose enum variants. The raw +native ABI exposes the same information through `codec_kind` and `codec_id`. + +`codec.kind` is `none` for a call with no codec, `builtin` for +Relay's built-in `openai_chat`, `openai_responses`, and `anthropic_messages` +codecs, `runtime` for a named runtime-registered codec, and `opaque` for an +active codec without a registered identity. `codec.id` is present only for +`builtin` and `runtime`. Do not infer a provider from an opaque request shape. + +Return a payload to continue the sanitizer chain. Return `None` (or `null` in +JavaScript) only when observing the payload would be unsafe: Relay omits the +LLM event payload and its annotation, while leaving the client-visible request +and response unchanged. Omission short-circuits later LLM sanitizers. + +All LLM sanitizer callbacks must implement this two-parameter contract. For an +in-process sanitizer, use `resolve_codec()` to access the active codec +implementation. A resolver returns no codec for manual calls without one. +Worker-plugin contexts resolve to an invocation-scoped asynchronous proxy with +the same directional operations: request codecs provide `decode(request)` and +`encode(annotated, original)`, while response codecs provide `decode(response)`. +Python in-process response codecs expose that operation as `decode_response`, +and Node.js exposes it as `decodeResponse`. An active runtime or opaque codec +resolves just like a built-in codec; its identity does not limit the available +operations. Node.js codecs supplied through decode and encode callbacks have +an `opaque` identity but remain resolvable. The worker proxy is valid only +while its sanitizer callback is running; do not retain it after the callback +returns. + + + + +```python +import nemo_relay +from nemo_relay import LLMRequest +from nemo_relay import guardrails + +def redact_request( + request: LLMRequest, + context: nemo_relay.LlmSanitizeRequestContext, +) -> LLMRequest | None: + codec = context.resolve_codec() + if ( + context.codec.kind == "builtin" + and context.codec.id == "openai_chat" + and codec is not None + ): + annotated = codec.decode(request) + # Apply a policy to the normalized request, then preserve the wire shape. + annotated.messages = [] + return codec.encode(annotated, request) + return request + +guardrails.register_llm_sanitize_request("redact-openai-chat", 10, redact_request) +``` + + + + +```js +const relay = require("nemo-relay-node"); + +relay.registerLlmSanitizeRequestGuardrail( + "redact-openai-chat", + 10, + (request, context) => { + const codec = context.resolveCodec(); + if (codec) { + const annotated = codec.decode(request); + annotated.messages = []; + return codec.encode(annotated, request); + } + return request; + }, +); +``` + + + + +```rust +use std::sync::Arc; +use nemo_relay::api::runtime::{BuiltinLlmCodec, LlmCodecIdentity}; +use nemo_relay::api::registry::register_llm_sanitize_request_guardrail; + +register_llm_sanitize_request_guardrail( + "redact-openai-chat", + 10, + Arc::new(|mut request, context| { + if context.codec() == &LlmCodecIdentity::BuiltIn(BuiltinLlmCodec::OpenAiChat) + && let Some(codec) = context.resolve_codec() + && let Ok(mut annotated) = codec.decode(&request) + { + annotated.messages.clear(); + request = codec.encode(&annotated, &request).ok()?; + } + Some(request) + }), +)?; +``` + + + + +The same registration names support scope-local and plugin-context +registrations. Priority and name tie-break ordering are unchanged. + ## Detailed Execution Flow diff --git a/docs/about-nemo-relay/release-notes/index.mdx b/docs/about-nemo-relay/release-notes/index.mdx index 8d9550fa2..a87eb2176 100644 --- a/docs/about-nemo-relay/release-notes/index.mdx +++ b/docs/about-nemo-relay/release-notes/index.mdx @@ -27,25 +27,50 @@ Relay](/about-nemo-relay/overview). ## Release 0.7 -NVIDIA NeMo Relay 0.7 release notes are in preparation. +NVIDIA NeMo Relay 0.7 release notes are in preparation. The following +compatibility information applies to the current 0.7 prerelease. ### Highlights -- _Highlights will be added for the 0.7 release._ +- LLM observability sanitizers now receive the active request or response codec + for each managed call. Sanitizers can normalize built-in, + runtime-registered, and opaque codec payloads without changing the + client-visible request or response. +- The PII redaction plugin now selects the active codec per call. One component + can safely sanitize mixed OpenAI Chat, OpenAI Responses, and Anthropic + Messages traffic without a fixed provider codec. +- Native and worker plugins can resolve invocation-scoped codec capabilities. + Request sanitizers can decode and encode requests, and response sanitizers + can decode responses. ### Support Matrix and Compatibility Updates + +LLM sanitizer callbacks require migration in 0.7. Callbacks now receive the +payload first and a directional context second, and they return an optional +payload. There is no one-argument compatibility adapter. Rust worker mark, +scope, tool, and LLM sanitizer callbacks are also asynchronous in 0.7. Native +plugins and workers must be rebuilt with the matching 0.7 SDK and protocol +definitions. + + The [Support Matrix](/reference/support-matrix) is the canonical reference for supported platforms and architectures, worker runtimes, coding agents, and integrations. It also records current limitations, including platform-specific worker requirements. -Migration guidance for upgrading from 0.6 to 0.7 will be added to the -[Migration Guides](/reference/migration-guides) before the release. +Before upgrading from 0.6, follow the +[0.7 Migration Guide](/reference/migration-guides#upgrade-to-nemo-relay-07). +For the new callback contract and codec operations, refer to +[Codec-Aware LLM Sanitizers](/about-nemo-relay/concepts/middleware#codec-aware-llm-sanitizers). ### Fixed Known Issues in 0.7 -- _Fixed issues will be added for the 0.7 release._ +- LLM payload redaction now follows the codec active for each call instead of a + codec captured from plugin configuration. Codec-dependent policies omit the + observability payload and annotation when Relay cannot safely normalize the + payload. This change resolves the mixed-provider PII sanitization issue + [#526](https://github.com/NVIDIA/NeMo-Relay/issues/526). ## Known Issues in 0.7 diff --git a/docs/build-plugins/dynamic-plugins/about.mdx b/docs/build-plugins/dynamic-plugins/about.mdx index 68f2331ba..772a0b827 100644 --- a/docs/build-plugins/dynamic-plugins/about.mdx +++ b/docs/build-plugins/dynamic-plugins/about.mdx @@ -13,7 +13,7 @@ two execution lanes: | Lane | Use when | Stable boundary | | --- | --- | --- | -| `rust_dynamic` | Behavior must run in the Relay process. | Native ABI v1 | +| `rust_dynamic` | Behavior must run in the Relay process. | Native ABI v2 | | `worker` | Behavior should run in a separate local process. | `grpc-v1` | The manifest describes compatibility, capabilities, artifact integrity, and the diff --git a/docs/build-plugins/dynamic-plugins/grpc-worker/grpc-worker-protocol.mdx b/docs/build-plugins/dynamic-plugins/grpc-worker/grpc-worker-protocol.mdx index 872a6bde9..b8ca9bfbb 100644 --- a/docs/build-plugins/dynamic-plugins/grpc-worker/grpc-worker-protocol.mdx +++ b/docs/build-plugins/dynamic-plugins/grpc-worker/grpc-worker-protocol.mdx @@ -74,6 +74,61 @@ invocation, or LLM invocation. `InvokeResponse` returns an empty result, JSON result, guardrail result, LLM request-intercept result, tool-execution result, or `WorkerError`. `InvokeStream` emits JSON chunks or `WorkerError` chunks. +Every LLM sanitizer invocation includes a directional context with tagged codec +identity: `none`, `builtin(id)`, `runtime(id)`, or `opaque`. Worker SDKs expose +`context.resolve_codec()` for active codecs. The resulting invocation-scoped +proxy supports request decode/encode or response decode and calls Relay through +the host runtime; the capability identifier is protocol-internal and expires +when the callback completes. `resolve_codec()` returns no proxy only when no +codec is active. Runtime and opaque codecs remain resolvable. + +Request and response sanitizer handlers always receive `(payload, context)`. +Return the sanitized payload to continue the chain, or return no payload to +omit the observability payload and annotation without changing the +client-visible value. Rust worker sanitizer callbacks are async for every +sanitizer surface: mark, scope, tool, and LLM. Python worker sanitizers can +return either an immediate value or an awaitable. + + + + +```rust +ctx.register_llm_sanitize_request_guardrail( + "normalize-request", + 10, + |request, context| async move { + let Some(codec) = context.resolve_codec() else { + return Ok(Some(request)); + }; + let mut annotated = codec.decode(&request).await?; + annotated.messages.clear(); + Ok(Some(codec.encode(&annotated, &request).await?)) + }, +); +``` + + + + +```python +async def normalize_request(request, context): + codec = context.resolve_codec() + if codec is None: + return request + annotated = await codec.decode(request) + annotated["messages"] = [] + return await codec.encode(annotated, request) + +ctx.register_llm_sanitize_request_guardrail( + "normalize-request", + normalize_request, + priority=10, +) +``` + + + + ## RelayHostRuntime RPCs Relay implements these RPCs for worker callbacks: @@ -83,10 +138,15 @@ Relay implements these RPCs for worker callbacks: | Marks and scopes | `EmitMark`, `PushScope`, `PopScope` | | Isolated scope stacks | `CreateScopeStack`, `DropScopeStack` | | Execution continuations | `ToolNext`, `LlmNext`, `LlmStreamNext` | +| Codec capabilities | `DecodeLlmCodecRequest`, `EncodeLlmCodecRequest`, `DecodeLlmCodecResponse` | Every host-runtime request also carries the activation ID and authentication token. Scope operations include a `ScopeContext`; continuation calls include -the continuation ID that Relay supplied for the active intercept. +the continuation ID that Relay supplied for the active intercept. Codec +operations additionally require the unforgeable capability ID supplied for the +current sanitizer invocation. Relay rejects missing, forged, expired, +wrong-direction, or activation-mismatched capabilities. Codec transformation +failures are non-retryable worker errors. ## Authentication and Endpoints diff --git a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx index db11f2acc..a4c162368 100644 --- a/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx +++ b/docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx @@ -1,6 +1,6 @@ --- title: "Native Dynamic Plugins (Rust)" -description: "Build in-process Rust shared-library plugins against the NeMo Relay Native ABI v1." +description: "Build in-process Rust shared-library plugins against the NeMo Relay Native ABI v2." position: 10 --- {/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. @@ -116,7 +116,7 @@ path, then replace `` with that library's SHA-256 digest. Use Native Plugin](/build-plugins/dynamic-plugins/native-dynamic/rust-native-plugin-example) for a complete example with validation, middleware, scopes, and configuration schema support. -## Native ABI v1 +## Native ABI v2 The host passes a `NemoRelayNativeHostApiV1` table to the entry symbol. The plugin returns a `NemoRelayNativePluginV1` descriptor: @@ -143,6 +143,35 @@ ABI callbacks can register these runtime surfaces: Relay keeps the library alive while those registrations exist and deregisters them before unloading it. +LLM sanitize callbacks receive their request or response JSON first, followed +by `NemoRelayNativeLlmSanitizeRequestContext` or +`NemoRelayNativeLlmSanitizeResponseContext`. Each context contains structured +codec identity and a borrowed, callback-lifetime codec handle. `codec_kind` is +`None`, `BuiltIn`, `Runtime`, or `Opaque`. `codec_id` is present for `BuiltIn` +(one of `openai_chat`, `openai_responses`, or `anthropic_messages`) and +`Runtime`, and null for `None` and `Opaque`. + +The request handle supports host operations to decode an `LlmRequest` into an +`AnnotatedLlmRequest` and encode normalized changes onto the original request. +The response handle supports decoding response JSON into an +`AnnotatedLlmResponse`. A null handle means no codec is active. Runtime and +opaque codecs still have non-null handles and support the same operations as a +built-in codec. Do not retain a handle or resolved SDK facade after the +sanitizer callback returns. + +The Rust SDK converts the raw structures into +`LlmSanitizeRequestContext` and `LlmSanitizeResponseContext`. Call +`resolve_codec()` to obtain the safe directional facade. Raw ABI consumers use +the `llm_request_codec_decode`, `llm_request_codec_encode`, and +`llm_response_codec_decode` host-table functions. The host owns all returned +JSON string handles, which callers release with the ordinary host string +release operation. + +A successful null sanitizer output omits the LLM observability payload and its +annotation. Returning an error also fails closed rather than exposing the +unsanitized payload. Neither case changes the client-visible request or +response. + Use the `nemo-relay-plugin` crate rather than the host `nemo-relay` runtime crate. Refer to [Build a Rust Native Plugin](/build-plugins/dynamic-plugins/native-dynamic/rust-native-plugin-example) for the SDK-backed example. diff --git a/docs/build-plugins/language-binding/about.mdx b/docs/build-plugins/language-binding/about.mdx index 73359fd44..ec63654b3 100644 --- a/docs/build-plugins/language-binding/about.mdx +++ b/docs/build-plugins/language-binding/about.mdx @@ -73,6 +73,14 @@ A plugin can install one or more of these runtime surfaces: conditional-execution guardrails, request intercepts, execution intercepts, and stream execution intercepts. +LLM sanitize guardrails always receive the request or response first, followed +by a structured context whose `codec` is `none`, `builtin(id)`, `runtime(id)`, +or `opaque`. Every callback declares `(payload, context)` and may omit the LLM +event payload by returning no value. In-process sanitizer contexts can resolve +the active codec for normalized processing. Worker sanitizer contexts resolve +an invocation-scoped asynchronous codec proxy, so workers can perform the same +directional normalization without receiving a host-process object. + Start with one surface. Add a bundle only when one configuration document clearly controls related behavior, such as a subscriber plus the request intercepts needed to add correlation metadata. ## Registration Lifecycle diff --git a/docs/configure-plugins/pii-redaction/about.mdx b/docs/configure-plugins/pii-redaction/about.mdx index 982308d77..a37cd4812 100644 --- a/docs/configure-plugins/pii-redaction/about.mdx +++ b/docs/configure-plugins/pii-redaction/about.mdx @@ -91,9 +91,11 @@ The current backend boundary is intentional: - Managed tool surfaces are sanitized as JSON payloads with exact JSON-pointer targeting. -- Managed LLM surfaces use the selected built-in codec so redaction can target - normalized Relay request and response shapes such as `/messages/0/content` - and `/message`. +- Managed LLM requests use the active resolved codec for each call, including + built-in, runtime, and opaque codecs. Normalized response projection requires + a recognized built-in codec because response codec capabilities are + decode-only. This lets redaction target normalized Relay shapes such as + `/messages/0/content` and `/message`. - `mark` sanitizes `data`, `category_profile`, and `metadata` independently on every mark event. It defaults to `true`; set `mark = false` to opt out. - `input`, `output`, `tool_input`, and `tool_output` sanitize scope metadata on diff --git a/docs/configure-plugins/pii-redaction/configuration.mdx b/docs/configure-plugins/pii-redaction/configuration.mdx index 329f78892..d10f84813 100644 --- a/docs/configure-plugins/pii-redaction/configuration.mdx +++ b/docs/configure-plugins/pii-redaction/configuration.mdx @@ -72,7 +72,7 @@ The top-level PII redaction object contains: | `tool_input` | Enables sanitization of emitted tool-request observability payloads. Defaults to `true`. | | `tool_output` | Enables sanitization of emitted tool-response observability payloads. Defaults to `true`. | | `priority` | Guardrail priority. Lower values run earlier. Defaults to `100`. | -| `codec` | Managed LLM provider codec. Required when `input` or `output` is enabled. | +| `codec` | Optional compatibility fallback for legacy/manual managed LLM calls that do not expose an active codec. | | `profiles` | Ordered redaction profiles. Each profile selects `builtin` or `local_model`; enabled profiles cover every supported sanitization surface. | | `builtin` | Built-in backend settings used when `mode = "builtin"`. | | `local` | Local-backend settings used when `mode = "local_model"`. | @@ -90,6 +90,37 @@ profiles do not require IDs. Disabled profiles are validated but do not register callbacks. The complete array is replaced as one value during config layering. +### Mixed-Provider Gateway Example + +For a gateway that routes requests to more than one supported provider, omit +`codec`. Relay selects the active codec for each managed call, so the same +policy safely applies to OpenAI Chat, OpenAI Responses, and Anthropic Messages +traffic: + +```toml +[[components]] +kind = "pii_redaction" +enabled = true + +[components.config] + +[[components.config.profiles]] +mode = "builtin" +priority = 80 + +[components.config.profiles.builtin] +action = "redact" +detector = "email" +target_paths = ["/messages/0/content", "/message"] +``` + +Use `codec = "openai_chat"` only as a compatibility fallback for a legacy +manual call where Relay cannot identify an active codec. It is ignored when a +managed call supplies any active codec. Request policies use that codec's +resolved decode/encode capability, including runtime and opaque codecs. +Normalized response projection currently requires a recognized built-in codec +because response codec capabilities are decode-only. + ## Backend Support The following table compares the available PII redaction backends: @@ -120,7 +151,7 @@ paths, and supported codecs. To use `mode = "builtin"`: - `builtin` settings are required. -- `codec` is required when `input` or `output` is enabled. +- Set `codec` only when legacy/manual managed LLM calls need a fallback codec. - `builtin.action` must be `remove`, `redact`, `regex_replace`, `hash`, or `mask`. - `builtin.pattern` or `builtin.detector` is required when `builtin.action = "regex_replace"` or `builtin.action = "redact"`. @@ -140,7 +171,6 @@ enabled = true [components.config] version = 1 -codec = "openai_chat" [[components.config.profiles]] mode = "builtin" @@ -220,7 +250,9 @@ The `builtin` section contains: When a target matches: - Object fields are removed. -- Array elements become `null`. +- Array elements become `null` on JSON-native surfaces. If a normalized LLM + array removal cannot be safely re-encoded, Relay omits the event payload and + its annotation instead. - Targeted scalar or root values become `null`. ### `regex_replace` @@ -290,7 +322,10 @@ The plugin uses different payload boundaries for tools and LLMs: - Tools use JSON-native payloads. Paths point into the emitted tool args or tool result shape directly. -- LLMs use the selected built-in codec. Prefer normalized Relay paths such as: +- LLM requests use the active resolved codec for that managed call. LLM + responses use the active built-in codec for normalized projection. Prefer + normalized Relay paths such as: + - `/headers/authorization` for an exact request header field - `/messages/0/content` for request message content - `/message` for the normalized assistant response text - Marks and non-tool, non-LLM scopes sanitize `data`, `category_profile`, and @@ -299,8 +334,14 @@ The plugin uses different payload boundaries for tools and LLMs: `/metadata`. The current implementation also preserves provider-shaped response-path -compatibility for the supported codecs, but normalized LLM paths are the -recommended contract for new configuration. +compatibility for the supported built-in codecs, but normalized LLM paths are +the recommended contract for new configuration. If normalized paths are +configured and Relay cannot resolve a request codec or cannot project a +response through a recognized built-in or compatible fallback codec, the LLM +event payload and its annotation are omitted rather than emitting unsanitized +content. In particular, a manual LLM call with normalized `target_paths`, no +active codec, and no configured `codec` fallback omits both the payload and its +annotation. ## Decision Rule @@ -355,7 +396,7 @@ That means: - The emitted NeMo Relay start, end, or mark event observability fields are sanitized. - `annotated_response` is regenerated from the sanitized end-event payload when - the configured response codec can decode that payload. + the active response codec can decode that payload. - Subscribers and exporters consume only the sanitized canonical event stream. Each exporter can omit or transform fields as part of its normal projection, but raw values removed by the plugin do not bypass the sanitizer. diff --git a/docs/reference/llm-request-intercept-outcomes.mdx b/docs/reference/llm-request-intercept-outcomes.mdx index 5e968a4b8..734a6327d 100644 --- a/docs/reference/llm-request-intercept-outcomes.mdx +++ b/docs/reference/llm-request-intercept-outcomes.mdx @@ -91,12 +91,11 @@ or object shape: - Python callbacks return `LLMRequestInterceptOutcome`. - Rust callbacks return `LlmRequestInterceptOutcome`. -- Go callbacks return `LLMRequestInterceptOutcome`. - Node.js callbacks return `{ request, annotated?, pendingMarks? }`. JavaScript pending-mark DTOs use `categoryProfile`; canonical JSON retains `pending_marks` and `category_profile`. - Public C callbacks return one owned canonical outcome JSON string, and native - ABI v1 callbacks return one host-owned outcome JSON string. + ABI v2 callbacks return one host-owned outcome JSON string. - Rust and Python `grpc-v1` worker SDKs return their canonical outcome in a `JsonEnvelope` with schema `nemo.relay.LlmRequestInterceptOutcome@2`. @@ -121,10 +120,10 @@ codec input, sanitizer input, or start payload. ## Migration -This finalizes unpublished native ABI v1 and `grpc-v1` contracts. Rebuild all +This finalizes unpublished native ABI v2 and `grpc-v1` contracts. Rebuild all development native plugins and workers against the same NeMo Relay release that -hosts them. Replace tuple results, split C/Go -outputs, metadata envelopes, and parallel mark-aware registrations with the +hosts them. Replace tuple results, split outputs, metadata envelopes, and +parallel mark-aware registrations with the canonical outcome and the existing `register_llm_request_intercept` registration name. diff --git a/docs/reference/migration-guides.mdx b/docs/reference/migration-guides.mdx index debc2cda7..edfeada6f 100644 --- a/docs/reference/migration-guides.mdx +++ b/docs/reference/migration-guides.mdx @@ -1,6 +1,6 @@ --- title: "Migration Guides" -description: "Upgrade NeMo Relay integrations, plugins, exporters, and public API consumers." +description: "Upgrade NeMo Relay integrations and migrate LLM sanitizer callbacks, plugins, workers, and PII policy." position: 6 --- {/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. @@ -13,7 +13,292 @@ intervening release in sequence. ## Upgrade to NeMo Relay 0.7 -Migration guidance will be added before the 0.7 release. +NeMo Relay 0.7 changes the LLM observability sanitizer contract across +in-process bindings, native plugins, raw C FFI consumers, and worker plugins. +Complete the following migrations before you run an existing sanitizer with a +0.7 host. + + +Do not deploy a 0.6 sanitizer plugin or worker against a 0.7 host. The LLM +callback signature, native ABI layout, and worker invocation schema changed. +NeMo Relay does not adapt one-argument LLM sanitizer callbacks. + + +### Update LLM Sanitizer Callbacks + +The registration names remain unchanged for global, plugin-context, and +scope-local sanitizers. Update the callback itself as follows: + +| Surface | 0.6 Contract | 0.7 Contract | +| --- | --- | --- | +| Request | `(LlmRequest) -> LlmRequest` | `(LlmRequest, LlmSanitizeRequestContext) -> Option` | +| Response | `(Json) -> Json` | `(Json, LlmSanitizeResponseContext) -> Option` | + +The payload is always the first argument. The directional context is always the +second argument. Update request sanitizers with the following pattern: + + + + +```rust +use std::sync::Arc; + +use nemo_relay::api::registry::register_llm_sanitize_request_guardrail; + +register_llm_sanitize_request_guardrail( + "redact-request", + 10, + Arc::new(|request, context| { + let _active_codec = context.resolve_codec(); + // Apply policy, using _active_codec when normalized access is required. + Some(request) + }), +)?; +``` + + + + +```python +import nemo_relay +from nemo_relay import LLMRequest +from nemo_relay import guardrails + +def redact_request( + request: LLMRequest, + context: nemo_relay.LlmSanitizeRequestContext, +) -> LLMRequest | None: + _active_codec = context.resolve_codec() + # Apply policy, using _active_codec when normalized access is required. + return request + +guardrails.register_llm_sanitize_request( + "redact-request", + 10, + redact_request, +) +``` + + + + +```js +const relay = require("nemo-relay-node"); + +relay.registerLlmSanitizeRequestGuardrail( + "redact-request", + 10, + (request, context) => { + const _activeCodec = context.resolveCodec(); + // Apply policy, using _activeCodec when normalized access is required. + return request; + }, +); +``` + + + + +Apply the same change to response sanitizers with +`LlmSanitizeResponseContext`. Return the response payload to keep it in the +observability event. + +### Handle Payload Omission + +In 0.7, `None` or `null` has a specific fail-closed meaning for an LLM +sanitizer: Relay omits the observability payload and its annotation. Omission +short-circuits the remaining LLM sanitizer chain but does not change the +request or response returned to the client. + +Check every callback for an implicit empty return. In particular: + +- A Python function that reaches the end without `return` omits the payload. +- A JavaScript callback that returns `null` or `undefined` omits the payload. +- A Rust callback must return `Some(payload)` to retain the payload. +- A sanitizer error reported through a plugin or binding boundary also omits + the payload and annotation. + +Use omission only when recording the payload would be unsafe. + +### Resolve the Active Codec + +The request and response contexts expose `codec`, which has one of the +following identities: + +| Identity | Meaning | `resolve_codec()` | +| --- | --- | --- | +| `none` | No codec is active for this payload direction. | Returns no codec. | +| `builtin(id)` | A built-in codec is active. | Returns the active codec. | +| `runtime(id)` | A named runtime-registered codec is active. | Returns the active codec. | +| `opaque` | A codec without a public ID is active. | Returns the active codec. | + +Treat identity as descriptive and codec resolution as authoritative. Do not +infer a provider from the request shape. A request codec supports +`decode(request)` and `encode(annotated, original)`. A response codec supports +`decode_response(response)` in in-process Rust and Python, +`decode(response).await` in Rust workers, `await decode(response)` in Python +workers, and `decodeResponse(response)` in Node.js; response encoding is not +available. + +Worker codec proxies are scoped to the sanitizer callback and must not be +stored for later work. Worker codec operations are asynchronous because they +call the host over `grpc-v1`. In-process resolvers return the active codec +facade directly. + +For complete in-process examples, refer to +[Codec-Aware LLM Sanitizers](/about-nemo-relay/concepts/middleware#codec-aware-llm-sanitizers). + +### Migrate Worker Sanitizers + +All Rust worker sanitizer registrations now require callbacks that return +futures. This change applies to mark, scope-start, scope-end, tool-request, +tool-response, LLM-request, and LLM-response sanitizers. Conditional guardrails +and request intercepts keep their existing synchronous contracts. + +Update Rust worker callbacks to use `async move` and return `Result` from the +future. Python worker sanitizers can return either an immediate value or an +awaitable, but Python LLM sanitizers must still accept both the payload and +directional context. + + + + +```rust +ctx.register_llm_sanitize_request_guardrail( + "redact-request", + 10, + |request, context| async move { + if let Some(codec) = context.resolve_codec() { + let annotated = codec.decode(&request).await?; + let request = codec.encode(&annotated, &request).await?; + return Ok(Some(request)); + } + Ok(Some(request)) + }, +); +``` + + + + +```python +from nemo_relay_plugin import LlmSanitizeRequestContext + +async def redact_request(request, context: LlmSanitizeRequestContext): + codec = context.resolve_codec() + if codec is None: + return request + annotated = await codec.decode(request) + return await codec.encode(annotated, request) + +ctx.register_llm_sanitize_request_guardrail( + "redact-request", + redact_request, + priority=10, +) +``` + + + + +Regenerate worker bindings from the 0.7 +`nemo.relay.worker.v1` definition and deploy the 0.7 host and worker SDK +together. The protocol identifier remains `grpc-v1`, but the wire contract now +includes: + +- `LlmCodecIdentity` +- `LlmSanitizeRequestContext` +- `LlmSanitizeResponseContext` +- Host RPCs for request decode, request encode, and response decode + +The previous codec fields on `LlmInvocation` are reserved. Read codec identity +and capability data from the directional sanitizer context instead. + +The codec capability ID in the protocol is SDK-internal. Do not expose or +persist it in plugin code. The host rejects forged, expired, unauthorized, and +wrong-direction capability IDs. + +### Rebuild Native and Raw FFI Plugins + +NeMo Relay 0.7 uses native ABI v2. Recompile native plugins against the 0.7 +`nemo-relay-plugin` crate and rebuild raw FFI consumers against the generated +0.7 header. + +If you already built a plugin against an earlier 0.7 ABI v2 prerelease, rebuild +it again. The ABI version remains 2, but the prerelease LLM sanitizer callback +slots and context layouts changed before release. + +The plugin manifest value remains `compat.native_api = "1"`. This manifest +contract version is separate from the host ABI version; do not change it to +`"2"`. + +Request and response callbacks now receive distinct context structures. Each +structure contains structured codec identity and a borrowed directional codec +handle. The request handle supports decode and encode host operations. The +response handle supports decode. A successful null sanitizer output omits the +observability payload and annotation. + +Do not retain the codec handle, SDK facade, input pointers, or borrowed codec ID +after the callback returns. Release host-owned output strings with the standard +host string release operation. + +For the complete ABI contract, refer to +[Native ABI v2](/build-plugins/dynamic-plugins/native-dynamic/about#native-abi-v2). + +### Update PII Redaction Configuration + +The PII `codec` field is now an optional compatibility fallback. Relay uses it +only for a legacy or manual managed LLM call with no active codec. Any active +codec takes precedence, including runtime and opaque codecs. + +For a mixed-provider gateway, remove the fixed `codec` field: + + + + +```toml +[[components]] +kind = "pii_redaction" +enabled = true + +[components.config] + +[[components.config.profiles]] +mode = "builtin" +priority = 80 + +[components.config.profiles.builtin] +action = "redact" +detector = "email" +target_paths = ["/messages/0/content", "/message"] +``` + + + + +Provider-agnostic all-leaf and `trajectory_context` policies do not require +`codec`. A normalized `target_paths` policy fails closed when Relay has neither +a recognized active codec nor a compatible fallback: Relay omits the LLM +payload and annotation instead of recording unsanitized data. Normalized +`target_paths` response policies for runtime and opaque codecs also omit the +payload because response codec capabilities are decode-only. + +For the complete policy behavior, refer to +[PII Redaction Configuration](/configure-plugins/pii-redaction/configuration). + +### Verify the Upgrade + +Before deployment: + +1. Search for every LLM sanitize-request and sanitize-response registration. +2. Add the required directional context argument and optional payload result. +3. Check that no callback implicitly returns `None`, `null`, or `undefined`. +4. Convert every Rust worker sanitizer callback to an asynchronous callback. +5. Rebuild native and raw FFI plugins against 0.7. +6. Regenerate and redeploy workers with the matching 0.7 protocol and SDK. +7. Remove fixed PII `codec` values from mixed-provider configurations. +8. Test buffered and streaming calls for every provider and custom codec that + your deployment uses. ## Related Release Information diff --git a/examples/rust-native-plugin/src/lib.rs b/examples/rust-native-plugin/src/lib.rs index af19ee8f7..84f51410c 100644 --- a/examples/rust-native-plugin/src/lib.rs +++ b/examples/rust-native-plugin/src/lib.rs @@ -234,11 +234,19 @@ impl NativePlugin for ExampleNativePlugin { ctx.register_llm_sanitize_request_guardrail("example_llm_sanitize_request", 10, { let tag = config.tag.clone(); - move |request| tag_llm_request(request, "native_llm_sanitize_request", &tag) + move |request, _context| { + Some(tag_llm_request( + request, + "native_llm_sanitize_request", + &tag, + )) + } })?; ctx.register_llm_sanitize_response_guardrail("example_llm_sanitize_response", 10, { let tag = config.tag.clone(); - move |response| tag_json(response, "native_llm_sanitize_response", &tag) + move |response, _context| { + Some(tag_json(response, "native_llm_sanitize_response", &tag)) + } })?; ctx.register_llm_conditional_execution_guardrail("example_llm_conditional", 10, { let block_llms = config.block_llms; diff --git a/go/nemo_relay/adaptive_plugin_test.go b/go/nemo_relay/adaptive_plugin_test.go index 58da74050..9931d0187 100644 --- a/go/nemo_relay/adaptive_plugin_test.go +++ b/go/nemo_relay/adaptive_plugin_test.go @@ -90,8 +90,8 @@ func registerLifecycleGuardrails(ctx *PluginContext) error { if err := ctx.RegisterLlmSanitizeRequestGuardrail( "llm_sanitize_request", 7, - func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { - return headers, content + func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { + return request, false }, ); err != nil { return err @@ -99,7 +99,9 @@ func registerLifecycleGuardrails(ctx *PluginContext) error { if err := ctx.RegisterLlmSanitizeResponseGuardrail( "llm_sanitize_response", 7, - func(responseJSON json.RawMessage) json.RawMessage { return responseJSON }, + func(responseJSON json.RawMessage, _ LLMSanitizeResponseContext) (json.RawMessage, bool) { + return responseJSON, false + }, ); err != nil { return err } @@ -532,10 +534,12 @@ func TestPluginFuncsAndClosedContextBranches(t *testing.T) { return closed.RegisterToolConditionalExecutionGuardrail("tool_conditional", 1, func(name string, args json.RawMessage) *string { return nil }) }}, {"llm sanitize request", func() error { - return closed.RegisterLlmSanitizeRequestGuardrail("llm_sanitize_request", 1, func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { return headers, content }) + return closed.RegisterLlmSanitizeRequestGuardrail("llm_sanitize_request", 1, func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { return request, false }) }}, {"llm sanitize response", func() error { - return closed.RegisterLlmSanitizeResponseGuardrail("llm_sanitize_response", 1, func(response json.RawMessage) json.RawMessage { return response }) + return closed.RegisterLlmSanitizeResponseGuardrail("llm_sanitize_response", 1, func(response json.RawMessage, _ LLMSanitizeResponseContext) (json.RawMessage, bool) { + return response, false + }) }}, {"llm conditional", func() error { return closed.RegisterLlmConditionalExecutionGuardrail("llm_conditional", 1, func(headers, content json.RawMessage) *string { return nil }) diff --git a/go/nemo_relay/callbacks.go b/go/nemo_relay/callbacks.go index f8069eda0..8096c7447 100644 --- a/go/nemo_relay/callbacks.go +++ b/go/nemo_relay/callbacks.go @@ -26,15 +26,27 @@ typedef struct FfiToolHandle FfiToolHandle; typedef struct FfiLLMHandle FfiLLMHandle; typedef struct FfiLLMRequest FfiLLMRequest; typedef struct FfiEvent FfiEvent; +typedef struct FfiLlmSanitizeRequestCodec FfiLlmSanitizeRequestCodec; +typedef struct FfiLlmSanitizeResponseCodec FfiLlmSanitizeResponseCodec; +typedef struct NemoRelayLlmSanitizeRequestContext { + uint32_t codec_kind; + const char* codec_id; + const FfiLlmSanitizeRequestCodec* codec; +} NemoRelayLlmSanitizeRequestContext; +typedef struct NemoRelayLlmSanitizeResponseContext { + uint32_t codec_kind; + const char* codec_id; + const FfiLlmSanitizeResponseCodec* codec; +} NemoRelayLlmSanitizeResponseContext; typedef void (*NemoRelayFreeFn)(void* user_data); typedef char* (*NemoRelayToolSanitizeFn)(void* user_data, const char* name, const char* args_json); typedef char* (*NemoRelayToolConditionalFn)(void* user_data, const char* name, const char* args_json); typedef char* (*NemoRelayToolExecFn)(void* user_data, const char* args_json); -typedef FfiLLMRequest* (*NemoRelayLlmRequestCb)(void* user_data, const FfiLLMRequest* request); +typedef FfiLLMRequest* (*NemoRelayLlmSanitizeRequestCb)(void* user_data, const FfiLLMRequest* request, NemoRelayLlmSanitizeRequestContext context); typedef char* (*NemoRelayLlmConditionalCb)(void* user_data, const FfiLLMRequest* request); typedef char* (*NemoRelayLlmExecFn)(void* user_data, const char* native_json); -typedef char* (*NemoRelayLlmResponseFn)(void* user_data, const char* response_json); +typedef char* (*NemoRelayLlmSanitizeResponseCb)(void* user_data, const char* response_json, NemoRelayLlmSanitizeResponseContext context); typedef void (*NemoRelayEventSubscriberFn)(void* user_data, const FfiEvent* event); typedef char* (*NemoRelayEventSanitizeFn)(void* user_data, const FfiEvent* event, const char* fields_json); typedef struct FfiPluginContext FfiPluginContext; @@ -57,10 +69,14 @@ static inline char* callLlmExecNext(NemoRelayLlmExecNextFn next_fn, const char* // LLMRequest accessors (also declared in types.go, needed here for trampolines) extern FfiLLMRequest* nemo_relay_llm_request_new(const char* headers_json, const char* content_json); +extern void nemo_relay_llm_request_free(FfiLLMRequest* request); extern char* nemo_relay_llm_request_headers(const FfiLLMRequest* ptr); extern char* nemo_relay_llm_request_content(const FfiLLMRequest* ptr); extern void nemo_relay_string_free(char* ptr); extern void nemo_relay_set_last_error_message(const char* msg); +extern char* nemo_relay_llm_sanitize_request_codec_decode(const FfiLlmSanitizeRequestCodec*, const FfiLLMRequest*); +extern FfiLLMRequest* nemo_relay_llm_sanitize_request_codec_encode(const FfiLlmSanitizeRequestCodec*, const char*, const FfiLLMRequest*); +extern char* nemo_relay_llm_sanitize_response_codec_decode(const FfiLlmSanitizeResponseCodec*, const char*); // Codec callback typedefs (kept for trampoline use at execute time) typedef char* (*NemoRelayCodecDecodeCb)(void* user_data, const FfiLLMRequest* request); @@ -72,6 +88,7 @@ import "C" import ( "encoding/json" + "errors" "sync" "sync/atomic" "unsafe" @@ -162,16 +179,110 @@ type ToolExecutionFunc func(args json.RawMessage) (json.RawMessage, error) // the canonical outcome containing the tool result and any pending marks. type ToolExecutionInterceptFunc func(args json.RawMessage, next func(json.RawMessage) (json.RawMessage, error)) (ToolExecutionInterceptOutcome, error) -// LLMResponseFunc is a callback that transforms an LLM response. It receives -// the response as plain JSON and must return the (possibly modified) response -// JSON. -type LLMResponseFunc func(responseJSON json.RawMessage) json.RawMessage +// LLMCodecKind identifies the active codec state supplied to a sanitizer. +type LLMCodecKind string + +const ( + // LLMCodecNone means no codec was active. + LLMCodecNone LLMCodecKind = "none" + // LLMCodecBuiltin means a Relay built-in codec was active. + LLMCodecBuiltin LLMCodecKind = "builtin" + // LLMCodecRuntime means a runtime-registered codec was active. + LLMCodecRuntime LLMCodecKind = "runtime" + // LLMCodecOpaque means an active codec has no registered identity. + LLMCodecOpaque LLMCodecKind = "opaque" +) + +// LLMCodec identifies the codec active for one managed LLM sanitizer callback. +// ID is present for built-in and runtime codec identities. +type LLMCodec struct { + CodecKind LLMCodecKind + CodecID *string +} + +// LLMSanitizeRequestContext provides request codec context for one sanitizer call. +type LLMSanitizeRequestContext struct { + Codec LLMCodec + resolved *LLMRequestSanitizeCodec +} + +// ResolveCodec returns the active callback-scoped request codec, if any. +func (context LLMSanitizeRequestContext) ResolveCodec() *LLMRequestSanitizeCodec { + return context.resolved +} + +// LLMSanitizeResponseContext provides response codec context for one sanitizer call. +type LLMSanitizeResponseContext struct { + Codec LLMCodec + resolved *LLMResponseSanitizeCodec +} + +// ResolveCodec returns the active callback-scoped response codec, if any. +func (context LLMSanitizeResponseContext) ResolveCodec() *LLMResponseSanitizeCodec { + return context.resolved +} + +// ErrLLMSanitizeCodecExpired is returned when a callback-scoped codec +// capability is used after its sanitizer callback has returned. +var ErrLLMSanitizeCodecExpired = errors.New("LLM sanitizer codec capability is no longer active") + +type llmSanitizeCodecInvocation struct { + mu sync.RWMutex + active bool +} + +func newLLMSanitizeCodecInvocation() *llmSanitizeCodecInvocation { + return &llmSanitizeCodecInvocation{active: true} +} + +func (invocation *llmSanitizeCodecInvocation) acquire() (func(), error) { + invocation.mu.RLock() + if !invocation.active { + invocation.mu.RUnlock() + return nil, ErrLLMSanitizeCodecExpired + } + return invocation.mu.RUnlock, nil +} + +func (invocation *llmSanitizeCodecInvocation) invalidate() { + invocation.mu.Lock() + invocation.active = false + invocation.mu.Unlock() +} + +// LLMRequestSanitizeCodec is a callback-scoped request codec capability. +type LLMRequestSanitizeCodec struct { + ptr unsafe.Pointer + invocation *llmSanitizeCodecInvocation +} + +// LLMResponseSanitizeCodec is a callback-scoped response codec capability. +type LLMResponseSanitizeCodec struct { + ptr unsafe.Pointer + invocation *llmSanitizeCodecInvocation +} + +func acquireLLMSanitizeCodec( + ptr unsafe.Pointer, + invocation *llmSanitizeCodecInvocation, +) (unsafe.Pointer, func(), error) { + if ptr == nil || invocation == nil { + return nil, nil, ErrLLMSanitizeCodecExpired + } + release, err := invocation.acquire() + if err != nil { + return nil, nil, err + } + return ptr, release, nil +} -// LLMRequestFunc is a callback that transforms an LLM request. It receives -// the headers JSON and content JSON from the FfiLLMRequest, and returns the -// (possibly modified) versions of each. The Go binding uses JSON -// serialization rather than opaque C pointers for ergonomics. -type LLMRequestFunc func(headers, content json.RawMessage) (headers2, content2 json.RawMessage) +// LLMRequestFunc sanitizes an emitted LLM request. It receives the request +// first and codec context second; returning omit true removes observability. +type LLMRequestFunc func(request LLMRequestDTO, context LLMSanitizeRequestContext) (sanitized LLMRequestDTO, omit bool) + +// LLMResponseFunc sanitizes an emitted LLM response. It receives the response +// first and codec context second; returning omit true removes observability. +type LLMResponseFunc func(response json.RawMessage, context LLMSanitizeResponseContext) (sanitized json.RawMessage, omit bool) // LLMConditionalFunc is a callback that decides whether an LLM call should // proceed. It returns nil to allow execution, or a non-nil pointer to an error @@ -236,6 +347,86 @@ type LLMRequestDTO struct { Content json.RawMessage `json:"content"` } +// Decode normalizes an opaque request through the active codec. +func (codec *LLMRequestSanitizeCodec) Decode(request LLMRequestDTO) (json.RawMessage, error) { + if codec == nil { + return nil, ErrLLMSanitizeCodecExpired + } + ptr, release, err := acquireLLMSanitizeCodec(codec.ptr, codec.invocation) + if err != nil { + return nil, err + } + defer release() + cHeaders := C.CString(string(request.Headers)) + defer C.free(unsafe.Pointer(cHeaders)) + cContent := C.CString(string(request.Content)) + defer C.free(unsafe.Pointer(cContent)) + cRequest := C.nemo_relay_llm_request_new(cHeaders, cContent) + if cRequest == nil { + return nil, lastError() + } + defer C.nemo_relay_llm_request_free(cRequest) + out := C.nemo_relay_llm_sanitize_request_codec_decode((*C.FfiLlmSanitizeRequestCodec)(ptr), cRequest) + if out == nil { + return nil, lastError() + } + defer C.nemo_relay_string_free(out) + return json.RawMessage(C.GoString(out)), nil +} + +// Encode merges normalized changes onto the original opaque request. +func (codec *LLMRequestSanitizeCodec) Encode(annotated json.RawMessage, original LLMRequestDTO) (LLMRequestDTO, error) { + if codec == nil { + return LLMRequestDTO{}, ErrLLMSanitizeCodecExpired + } + ptr, release, err := acquireLLMSanitizeCodec(codec.ptr, codec.invocation) + if err != nil { + return LLMRequestDTO{}, err + } + defer release() + cHeaders := C.CString(string(original.Headers)) + defer C.free(unsafe.Pointer(cHeaders)) + cContent := C.CString(string(original.Content)) + defer C.free(unsafe.Pointer(cContent)) + cOriginal := C.nemo_relay_llm_request_new(cHeaders, cContent) + if cOriginal == nil { + return LLMRequestDTO{}, lastError() + } + defer C.nemo_relay_llm_request_free(cOriginal) + cAnnotated := C.CString(string(annotated)) + defer C.free(unsafe.Pointer(cAnnotated)) + out := C.nemo_relay_llm_sanitize_request_codec_encode((*C.FfiLlmSanitizeRequestCodec)(ptr), cAnnotated, cOriginal) + if out == nil { + return LLMRequestDTO{}, lastError() + } + defer C.nemo_relay_llm_request_free(out) + headers := C.nemo_relay_llm_request_headers(out) + defer C.nemo_relay_string_free(headers) + content := C.nemo_relay_llm_request_content(out) + defer C.nemo_relay_string_free(content) + return LLMRequestDTO{Headers: json.RawMessage(C.GoString(headers)), Content: json.RawMessage(C.GoString(content))}, nil +} + +// Decode normalizes an opaque response through the active codec. +func (codec *LLMResponseSanitizeCodec) Decode(response json.RawMessage) (json.RawMessage, error) { + if codec == nil { + return nil, ErrLLMSanitizeCodecExpired + } + ptr, release, err := acquireLLMSanitizeCodec(codec.ptr, codec.invocation) + if err != nil { + return nil, err + } + defer release() + cResponse := C.CString(string(response)) + defer C.free(unsafe.Pointer(cResponse)) + out := C.nemo_relay_llm_sanitize_response_codec_decode((*C.FfiLlmSanitizeResponseCodec)(ptr), cResponse) + if out == nil { + return nil, lastError() + } + defer C.nemo_relay_string_free(out) + return json.RawMessage(C.GoString(out)), nil +} + // PendingMarkSpec describes a mark Relay materializes under a managed lifecycle. type PendingMarkSpec struct { Name string `json:"name"` @@ -452,34 +643,97 @@ func goFreeTrampoline(userData unsafe.Pointer) { } //export goLlmRequestTrampoline -func goLlmRequestTrampoline(userData unsafe.Pointer, request *C.FfiLLMRequest) *C.FfiLLMRequest { +func goLlmRequestTrampoline(userData unsafe.Pointer, request *C.FfiLLMRequest, context C.NemoRelayLlmSanitizeRequestContext) *C.FfiLLMRequest { fn := lookupClosure(userData).(LLMRequestFunc) - - // Extract headers and content from the incoming FfiLLMRequest cHeaders := C.nemo_relay_llm_request_headers(request) cContent := C.nemo_relay_llm_request_content(request) - goHeaders := json.RawMessage(C.GoString(cHeaders)) - goContent := json.RawMessage(C.GoString(cContent)) - C.nemo_relay_string_free(cHeaders) - C.nemo_relay_string_free(cContent) - - // Call the Go callback - newHeaders, newContent := fn(goHeaders, goContent) + defer C.nemo_relay_string_free(cHeaders) + defer C.nemo_relay_string_free(cContent) + invocation := newLLMSanitizeCodecInvocation() + defer invocation.invalidate() + sanitized, omit := fn(LLMRequestDTO{ + Headers: json.RawMessage(C.GoString(cHeaders)), + Content: json.RawMessage(C.GoString(cContent)), + }, llmSanitizeRequestContextFromC(context, invocation)) + if omit { + return nil + } + cHeaders = C.CString(string(sanitized.Headers)) + cContent = C.CString(string(sanitized.Content)) + defer C.free(unsafe.Pointer(cHeaders)) + defer C.free(unsafe.Pointer(cContent)) + return C.nemo_relay_llm_request_new(cHeaders, cContent) +} + +func llmSanitizeRequestContextFromC( + context C.NemoRelayLlmSanitizeRequestContext, + invocation *llmSanitizeCodecInvocation, +) LLMSanitizeRequestContext { + var codecID *string + if context.codec_id != nil { + id := C.GoString(context.codec_id) + codecID = &id + } + result := LLMSanitizeRequestContext{Codec: llmCodecIdentity(uint32(context.codec_kind), codecID)} + if context.codec != nil { + result.resolved = &LLMRequestSanitizeCodec{ + ptr: unsafe.Pointer(context.codec), + invocation: invocation, + } + } + return result +} - // Create a new FfiLLMRequest from the result - cNewHeaders := C.CString(string(newHeaders)) - cNewContent := C.CString(string(newContent)) - defer C.free(unsafe.Pointer(cNewHeaders)) - defer C.free(unsafe.Pointer(cNewContent)) - return C.nemo_relay_llm_request_new(cNewHeaders, cNewContent) +func llmCodecIdentity(codecKind uint32, codecID *string) LLMCodec { + kind := LLMCodecOpaque + switch codecKind { + case 0: + kind = LLMCodecNone + case 1: + kind = LLMCodecBuiltin + case 2: + kind = LLMCodecRuntime + case 3: + kind = LLMCodecOpaque + } + return LLMCodec{CodecKind: kind, CodecID: codecID} } //export goLlmResponseTrampoline -func goLlmResponseTrampoline(userData unsafe.Pointer, responseJSON *C.char) *C.char { +func goLlmResponseTrampoline(userData unsafe.Pointer, responseJSON *C.char, context C.NemoRelayLlmSanitizeResponseContext) *C.char { fn := lookupClosure(userData).(LLMResponseFunc) - goJSON := json.RawMessage(C.GoString(responseJSON)) - result := fn(goJSON) - return C.CString(string(result)) + invocation := newLLMSanitizeCodecInvocation() + defer invocation.invalidate() + sanitized, omit := fn( + json.RawMessage(C.GoString(responseJSON)), + llmSanitizeResponseContextFromC(context, invocation), + ) + if omit { + return nil + } + return C.CString(string(sanitized)) +} + +func llmSanitizeResponseContextFromC( + context C.NemoRelayLlmSanitizeResponseContext, + invocation *llmSanitizeCodecInvocation, +) LLMSanitizeResponseContext { + var codecID *string + if context.codec_id != nil { + id := C.GoString(context.codec_id) + codecID = &id + } + resolved := (*LLMResponseSanitizeCodec)(nil) + if context.codec != nil { + resolved = &LLMResponseSanitizeCodec{ + ptr: unsafe.Pointer(context.codec), + invocation: invocation, + } + } + return LLMSanitizeResponseContext{ + Codec: llmCodecIdentity(uint32(context.codec_kind), codecID), + resolved: resolved, + } } //export goLlmConditionalTrampoline diff --git a/go/nemo_relay/callbacks_test.go b/go/nemo_relay/callbacks_test.go index 2655608d0..4be9b5637 100644 --- a/go/nemo_relay/callbacks_test.go +++ b/go/nemo_relay/callbacks_test.go @@ -5,9 +5,40 @@ package nemo_relay import ( "encoding/json" + "errors" "testing" + "time" ) +func TestLLMSanitizeCodecInvocationInvalidationWaitsForInflightCall(t *testing.T) { + invocation := newLLMSanitizeCodecInvocation() + release, err := invocation.acquire() + if err != nil { + t.Fatalf("acquire active invocation: %v", err) + } + + invalidated := make(chan struct{}) + go func() { + invocation.invalidate() + close(invalidated) + }() + + select { + case <-invalidated: + t.Fatal("invalidation returned while a codec operation was in flight") + case <-time.After(20 * time.Millisecond): + } + release() + select { + case <-invalidated: + case <-time.After(time.Second): + t.Fatal("invalidation did not finish after the codec operation completed") + } + if _, err := invocation.acquire(); !errors.Is(err, ErrLLMSanitizeCodecExpired) { + t.Fatalf("expired invocation acquire returned %v", err) + } +} + func toolExecutionOutcome(result json.RawMessage, err error) (ToolExecutionInterceptOutcome, error) { return ToolExecutionInterceptOutcome{Result: result}, err } @@ -36,3 +67,43 @@ func TestRegisterAndUnregisterClosure(t *testing.T) { t.Fatal("closure registry still contains callback after unregister") } } + +func TestLlmSanitizeDirectionalContextsPreserveEveryCodecIdentity(t *testing.T) { + openAIChat := "openai_chat" + openAIResponses := "openai_responses" + anthropicMessages := "anthropic_messages" + runtimeCodec := "com.example.chat.v1" + + cases := []struct { + name string + kind uint32 + id *string + want LLMCodecKind + }{ + {"none", 0, nil, LLMCodecNone}, + {"openai chat", 1, &openAIChat, LLMCodecBuiltin}, + {"openai responses", 1, &openAIResponses, LLMCodecBuiltin}, + {"anthropic messages", 1, &anthropicMessages, LLMCodecBuiltin}, + {"runtime", 2, &runtimeCodec, LLMCodecRuntime}, + {"opaque", 3, nil, LLMCodecOpaque}, + {"unknown", 99, nil, LLMCodecOpaque}, + } + + for _, test := range cases { + t.Run(test.name, func(t *testing.T) { + codec := llmCodecIdentity(test.kind, test.id) + if codec.CodecKind != test.want { + t.Fatalf("codec kind = %q, want %q", codec.CodecKind, test.want) + } + if codec.CodecID == nil && test.id != nil { + t.Fatal("codec ID was lost") + } + if codec.CodecID != nil && test.id == nil { + t.Fatalf("unexpected codec ID %q", *codec.CodecID) + } + if codec.CodecID != nil && test.id != nil && *codec.CodecID != *test.id { + t.Fatalf("codec ID = %q, want %q", *codec.CodecID, *test.id) + } + }) + } +} diff --git a/go/nemo_relay/deregister_test.go b/go/nemo_relay/deregister_test.go index ab6446e2e..e243bed6b 100644 --- a/go/nemo_relay/deregister_test.go +++ b/go/nemo_relay/deregister_test.go @@ -159,8 +159,8 @@ func TestRegisterDeregisterReregisterToolExecutionIntercept(t *testing.T) { func TestRegisterDeregisterReregisterLlmSanitizeRequestGuardrail(t *testing.T) { name := "go_reregister_llm_san_req" - fn := func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { - return headers, content + fn := func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { + return request, false } err := RegisterLlmSanitizeRequestGuardrail(name, 1, fn) diff --git a/go/nemo_relay/error_test.go b/go/nemo_relay/error_test.go index 1d32b1087..197edef89 100644 --- a/go/nemo_relay/error_test.go +++ b/go/nemo_relay/error_test.go @@ -126,7 +126,7 @@ func TestAlreadyExistsErrorOnDuplicateSubscriber(t *testing.T) { func TestAlreadyExistsErrorOnDuplicateLlmGuardrails(t *testing.T) { t.Run("LlmSanitizeRequest", func(t *testing.T) { name := "go_err_dup_llm_san_req" - fn := func(h, c json.RawMessage) (json.RawMessage, json.RawMessage) { return h, c } + fn := func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { return request, false } err := RegisterLlmSanitizeRequestGuardrail(name, 1, fn) if err != nil { @@ -142,7 +142,7 @@ func TestAlreadyExistsErrorOnDuplicateLlmGuardrails(t *testing.T) { t.Run("LlmSanitizeResponse", func(t *testing.T) { name := "go_err_dup_llm_san_resp" - fn := func(r json.RawMessage) json.RawMessage { return r } + fn := func(r json.RawMessage, _ LLMSanitizeResponseContext) (json.RawMessage, bool) { return r, false } err := RegisterLlmSanitizeResponseGuardrail(name, 1, fn) if err != nil { diff --git a/go/nemo_relay/guardrails/guardrails_test.go b/go/nemo_relay/guardrails/guardrails_test.go index 002906950..c7bc5470d 100644 --- a/go/nemo_relay/guardrails/guardrails_test.go +++ b/go/nemo_relay/guardrails/guardrails_test.go @@ -120,12 +120,13 @@ func runGlobalLLMGuardrailShorthandChecks(t *testing.T, output func() json.RawMe t.Helper() if err := guardrails.RegisterLlmSanitizeRequest("guardrails_llm_req", 1, - func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { + func(request nemo_relay.LLMRequestDTO, _ nemo_relay.LLMSanitizeRequestContext) (nemo_relay.LLMRequestDTO, bool) { var payload map[string]interface{} - _ = json.Unmarshal(content, &payload) + _ = json.Unmarshal(request.Content, &payload) payload["request_sanitized"] = true out, _ := json.Marshal(payload) - return headers, out + request.Content = out + return request, false }, ); err != nil { t.Fatalf("RegisterLlmSanitizeRequest failed: %v", err) @@ -135,12 +136,12 @@ func runGlobalLLMGuardrailShorthandChecks(t *testing.T, output func() json.RawMe }) if err := guardrails.RegisterLlmSanitizeResponse("guardrails_llm_resp", 1, - func(response json.RawMessage) json.RawMessage { + func(response json.RawMessage, _ nemo_relay.LLMSanitizeResponseContext) (json.RawMessage, bool) { var payload map[string]interface{} _ = json.Unmarshal(response, &payload) payload["guarded"] = true out, _ := json.Marshal(payload) - return out + return out, false }, ); err != nil { t.Fatalf("RegisterLlmSanitizeResponse failed: %v", err) @@ -212,14 +213,16 @@ func runScopeLocalLLMGuardrailShorthandChecks(t *testing.T, scopeUUID string) { t.Helper() if err := guardrails.ScopeRegisterLlmSanitizeRequest(scopeUUID, "guardrails_scope_llm_req", 1, - func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { - return headers, content + func(request nemo_relay.LLMRequestDTO, _ nemo_relay.LLMSanitizeRequestContext) (nemo_relay.LLMRequestDTO, bool) { + return request, false }, ); err != nil { t.Fatalf("ScopeRegisterLlmSanitizeRequest failed: %v", err) } if err := guardrails.ScopeRegisterLlmSanitizeResponse(scopeUUID, "guardrails_scope_llm_resp", 1, - func(response json.RawMessage) json.RawMessage { return response }, + func(response json.RawMessage, _ nemo_relay.LLMSanitizeResponseContext) (json.RawMessage, bool) { + return response, false + }, ); err != nil { t.Fatalf("ScopeRegisterLlmSanitizeResponse failed: %v", err) } diff --git a/go/nemo_relay/llm_test.go b/go/nemo_relay/llm_test.go index 9bb402b33..bf2a60fe4 100644 --- a/go/nemo_relay/llm_test.go +++ b/go/nemo_relay/llm_test.go @@ -282,6 +282,141 @@ func TestLlmCallExecuteWithRequestAndResponseCodecs(t *testing.T) { } } +func TestLlmSanitizersResolveDirectionalCodecs(t *testing.T) { + const requestGuard = "go_llm_resolved_request_codec" + const responseGuard = "go_llm_resolved_response_codec" + _ = DeregisterLlmSanitizeRequestGuardrail(requestGuard) + _ = DeregisterLlmSanitizeResponseGuardrail(responseGuard) + defer DeregisterLlmSanitizeRequestGuardrail(requestGuard) + defer DeregisterLlmSanitizeResponseGuardrail(responseGuard) + + var callbackState struct { + sync.Mutex + requestResolved bool + responseResolved bool + retainedRequestCodec *LLMRequestSanitizeCodec + retainedResponseCodec *LLMResponseSanitizeCodec + retainedRequest LLMRequestDTO + errors []string + } + recordCallbackError := func(format string, args ...any) { + callbackState.Lock() + defer callbackState.Unlock() + callbackState.errors = append(callbackState.errors, fmt.Sprintf(format, args...)) + } + if err := RegisterLlmSanitizeRequestGuardrail( + requestGuard, + 0, + func(request LLMRequestDTO, context LLMSanitizeRequestContext) (LLMRequestDTO, bool) { + if context.Codec.CodecKind != LLMCodecOpaque || + context.Codec.CodecID != nil { + recordCallbackError("unexpected request codec identity: %#v", context.Codec) + } + codec := context.ResolveCodec() + if codec == nil { + recordCallbackError("active request codec did not resolve") + return request, false + } + callbackState.Lock() + callbackState.retainedRequestCodec = codec + callbackState.retainedRequest = request + callbackState.Unlock() + annotated, err := codec.Decode(request) + if err != nil { + recordCallbackError("request codec decode failed: %v", err) + return request, false + } + encoded, err := codec.Encode(annotated, request) + if err != nil { + recordCallbackError("request codec encode failed: %v", err) + return request, false + } + callbackState.Lock() + callbackState.requestResolved = true + callbackState.Unlock() + return encoded, false + }, + ); err != nil { + t.Fatalf("request sanitizer registration failed: %v", err) + } + if err := RegisterLlmSanitizeResponseGuardrail( + responseGuard, + 0, + func(response json.RawMessage, context LLMSanitizeResponseContext) (json.RawMessage, bool) { + if context.Codec.CodecKind != LLMCodecBuiltin || + context.Codec.CodecID == nil || + *context.Codec.CodecID != "openai_chat" { + recordCallbackError("unexpected response codec identity: %#v", context.Codec) + } + codec := context.ResolveCodec() + if codec == nil { + recordCallbackError("active response codec did not resolve") + return response, false + } + callbackState.Lock() + callbackState.retainedResponseCodec = codec + callbackState.Unlock() + if _, err := codec.Decode(response); err != nil { + recordCallbackError("response codec decode failed: %v", err) + return response, false + } + callbackState.Lock() + callbackState.responseResolved = true + callbackState.Unlock() + return response, false + }, + ); err != nil { + t.Fatalf("response sanitizer registration failed: %v", err) + } + + response := json.RawMessage(`{ + "id":"chatcmpl-test", + "model":"test-model", + "choices":[{ + "index":0, + "message":{"role":"assistant","content":"ok"}, + "finish_reason":"stop" + }] + }`) + _, err := LlmCallExecute( + "resolved_codec_llm", + makeRequest(), + func(json.RawMessage) (json.RawMessage, error) { return response, nil }, + WithLLMCodec(llmRequestResponseCodec()), + WithLLMResponseCodec(NewOpenAIChatCodec()), + ) + if err != nil { + t.Fatalf(llmCallExecuteFailed, err) + } + callbackState.Lock() + sanitizerErrors := append([]string(nil), callbackState.errors...) + requestResolved := callbackState.requestResolved + responseResolved := callbackState.responseResolved + retainedRequestCodec := callbackState.retainedRequestCodec + retainedResponseCodec := callbackState.retainedResponseCodec + retainedRequest := callbackState.retainedRequest + callbackState.Unlock() + if len(sanitizerErrors) != 0 { + t.Fatalf("sanitizer callbacks failed: %v", sanitizerErrors) + } + if !requestResolved || !responseResolved { + t.Fatalf( + "expected both codec capabilities to resolve, request=%t response=%t", + requestResolved, + responseResolved, + ) + } + if _, err := retainedRequestCodec.Decode(retainedRequest); !errors.Is(err, ErrLLMSanitizeCodecExpired) { + t.Fatalf("retained request codec must expire after callback, got %v", err) + } + if _, err := retainedRequestCodec.Encode(json.RawMessage(`{}`), retainedRequest); !errors.Is(err, ErrLLMSanitizeCodecExpired) { + t.Fatalf("retained request codec encode must expire after callback, got %v", err) + } + if _, err := retainedResponseCodec.Decode(response); !errors.Is(err, ErrLLMSanitizeCodecExpired) { + t.Fatalf("retained response codec must expire after callback, got %v", err) + } +} + func llmRequestResponseCodec() CodecFunc { return CodecFunc{ Decode: func(headersJSON, contentJSON json.RawMessage) (json.RawMessage, error) { @@ -365,8 +500,8 @@ func requireLlmScopeEvents(t *testing.T, events []Event) (*ScopeEvent, *ScopeEve func TestLlmSanitizeRequestGuardrail(t *testing.T) { err := RegisterLlmSanitizeRequestGuardrail("go_llm_san_req", 1, - func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { - return headers, content + func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { + return request, false }, ) if err != nil { @@ -377,7 +512,9 @@ func TestLlmSanitizeRequestGuardrail(t *testing.T) { func TestLlmSanitizeResponseGuardrail(t *testing.T) { err := RegisterLlmSanitizeResponseGuardrail("go_llm_san_resp", 1, - func(responseJSON json.RawMessage) json.RawMessage { return responseJSON }, + func(responseJSON json.RawMessage, _ LLMSanitizeResponseContext) (json.RawMessage, bool) { + return responseJSON, false + }, ) if err != nil { t.Fatalf(llmRegisterFailed, err) @@ -385,6 +522,81 @@ func TestLlmSanitizeResponseGuardrail(t *testing.T) { DeregisterLlmSanitizeResponseGuardrail("go_llm_san_resp") } +func TestLlmSanitizeGuardrailsReceiveContext(t *testing.T) { + var capturedInput, capturedOutput json.RawMessage + var mu sync.Mutex + var callbackErrorsMu sync.Mutex + var callbackErrors []string + recordCallbackError := func(message string) { + callbackErrorsMu.Lock() + defer callbackErrorsMu.Unlock() + callbackErrors = append(callbackErrors, message) + } + if err := RegisterSubscriber("go_contextual_llm_sanitize_events", func(event Event) { + mu.Lock() + defer mu.Unlock() + if event.Kind() == "scope" && event.Category() == "llm" && event.ScopeCategory() == "start" { + capturedInput = append(json.RawMessage(nil), event.Input()...) + } + if event.Kind() == "scope" && event.Category() == "llm" && event.ScopeCategory() == "end" { + capturedOutput = append(json.RawMessage(nil), event.Output()...) + } + }); err != nil { + t.Fatalf("RegisterSubscriber failed: %v", err) + } + defer DeregisterSubscriber("go_contextual_llm_sanitize_events") + + if err := RegisterLlmSanitizeRequestGuardrail("go_contextual_llm_request", 1, + func(request LLMRequestDTO, context LLMSanitizeRequestContext) (LLMRequestDTO, bool) { + if context.Codec.CodecKind != LLMCodecNone { + recordCallbackError("manual registration received an active codec identity") + } + return request, true + }, + ); err != nil { + t.Fatalf(llmRegisterFailed, err) + } + defer DeregisterLlmSanitizeRequestGuardrail("go_contextual_llm_request") + + if err := RegisterLlmSanitizeResponseGuardrail("go_contextual_llm_response", 1, + func(response json.RawMessage, context LLMSanitizeResponseContext) (json.RawMessage, bool) { + if context.Codec.CodecID != nil { + recordCallbackError("manual registration received a codec ID") + } + return response, true + }, + ); err != nil { + t.Fatalf(llmRegisterFailed, err) + } + defer DeregisterLlmSanitizeResponseGuardrail("go_contextual_llm_response") + + result, err := LlmCallExecute("go_contextual_llm_sanitize", makeRequest(), + func(nativeJSON json.RawMessage) (json.RawMessage, error) { + return json.RawMessage(`{"response":"client-visible"}`), nil + }, + ) + if err != nil { + t.Fatalf(llmCallExecuteFailed, err) + } + callbackErrorsMu.Lock() + sanitizerErrors := append([]string(nil), callbackErrors...) + callbackErrorsMu.Unlock() + if len(sanitizerErrors) != 0 { + t.Fatalf("sanitizer callbacks failed: %v", sanitizerErrors) + } + if string(result) != `{"response":"client-visible"}` { + t.Fatalf("contextual sanitizers must not change the client result: %s", result) + } + if err := FlushSubscribers(); err != nil { + t.Fatalf(llmFlushSubscribersFailed, err) + } + mu.Lock() + defer mu.Unlock() + if capturedInput != nil || capturedOutput != nil { + t.Fatalf("contextual omission must remove observability payloads, got input=%s output=%s", capturedInput, capturedOutput) + } +} + func TestLlmConditionalExecutionGuardrail(t *testing.T) { err := RegisterLlmConditionalExecutionGuardrail("go_llm_cond", 1, func(headers, content json.RawMessage) *string { @@ -399,13 +611,13 @@ func TestLlmConditionalExecutionGuardrail(t *testing.T) { func TestLlmDuplicateGuardrailFails(t *testing.T) { RegisterLlmSanitizeRequestGuardrail("go_llm_dup", 1, - func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { - return headers, content + func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { + return request, false }, ) err := RegisterLlmSanitizeRequestGuardrail("go_llm_dup", 1, - func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { - return headers, content + func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { + return request, false }, ) if err == nil { @@ -660,12 +872,13 @@ func TestLlmSanitizeRequestGuardrailModifiesEventInput(t *testing.T) { defer DeregisterSubscriber("go_llm_san_evt_sub") RegisterLlmSanitizeRequestGuardrail("go_llm_content_mod", 1, - func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { + func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { var m map[string]interface{} - json.Unmarshal(content, &m) + json.Unmarshal(request.Content, &m) m["system_prompt_injected"] = true out, _ := json.Marshal(m) - return headers, out + request.Content = out + return request, false }, ) defer DeregisterLlmSanitizeRequestGuardrail("go_llm_content_mod") diff --git a/go/nemo_relay/nemo_relay.go b/go/nemo_relay/nemo_relay.go index 727f8380c..937427e5d 100644 --- a/go/nemo_relay/nemo_relay.go +++ b/go/nemo_relay/nemo_relay.go @@ -38,6 +38,10 @@ typedef struct FfiLLMRequest FfiLLMRequest; typedef struct FfiEvent FfiEvent; typedef struct FfiStream FfiStream; typedef struct FfiCodecHandle FfiCodecHandle; +typedef struct FfiLlmSanitizeRequestCodec FfiLlmSanitizeRequestCodec; +typedef struct FfiLlmSanitizeResponseCodec FfiLlmSanitizeResponseCodec; +typedef struct NemoRelayLlmSanitizeRequestContext { uint32_t codec_kind; const char* codec_id; const FfiLlmSanitizeRequestCodec* codec; } NemoRelayLlmSanitizeRequestContext; +typedef struct NemoRelayLlmSanitizeResponseContext { uint32_t codec_kind; const char* codec_id; const FfiLlmSanitizeResponseCodec* codec; } NemoRelayLlmSanitizeResponseContext; typedef void (*NemoRelayFreeFn)(void* user_data); @@ -135,12 +139,12 @@ extern int32_t nemo_relay_register_tool_execution_intercept(const char* name, in extern int32_t nemo_relay_deregister_tool_execution_intercept(const char* name); // LLM guardrails -typedef FfiLLMRequest* (*NemoRelayLlmRequestCb)(void* user_data, const FfiLLMRequest* request); -extern int32_t nemo_relay_register_llm_sanitize_request_guardrail(const char* name, int32_t priority, NemoRelayLlmRequestCb cb, void* user_data, NemoRelayFreeFn free_fn); +typedef FfiLLMRequest* (*NemoRelayLlmSanitizeRequestCb)(void* user_data, const FfiLLMRequest* request, NemoRelayLlmSanitizeRequestContext context); +extern int32_t nemo_relay_register_llm_sanitize_request_guardrail(const char* name, int32_t priority, NemoRelayLlmSanitizeRequestCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_deregister_llm_sanitize_request_guardrail(const char* name); -typedef char* (*NemoRelayLlmResponseFn)(void* user_data, const char* response_json); -extern int32_t nemo_relay_register_llm_sanitize_response_guardrail(const char* name, int32_t priority, NemoRelayLlmResponseFn cb, void* user_data, NemoRelayFreeFn free_fn); +typedef char* (*NemoRelayLlmSanitizeResponseCb)(void* user_data, const char* response_json, NemoRelayLlmSanitizeResponseContext context); +extern int32_t nemo_relay_register_llm_sanitize_response_guardrail(const char* name, int32_t priority, NemoRelayLlmSanitizeResponseCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_deregister_llm_sanitize_response_guardrail(const char* name); typedef char* (*NemoRelayLlmConditionalCb)(void* user_data, const FfiLLMRequest* request); @@ -193,9 +197,9 @@ extern int32_t nemo_relay_scope_register_tool_execution_intercept(const char* sc extern int32_t nemo_relay_scope_deregister_tool_execution_intercept(const char* scope_uuid, const char* name); // Scope-local LLM guardrails -extern int32_t nemo_relay_scope_register_llm_sanitize_request_guardrail(const char* scope_uuid, const char* name, int32_t priority, NemoRelayLlmRequestCb cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_scope_register_llm_sanitize_request_guardrail(const char* scope_uuid, const char* name, int32_t priority, NemoRelayLlmSanitizeRequestCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_scope_deregister_llm_sanitize_request_guardrail(const char* scope_uuid, const char* name); -extern int32_t nemo_relay_scope_register_llm_sanitize_response_guardrail(const char* scope_uuid, const char* name, int32_t priority, NemoRelayLlmResponseFn cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_scope_register_llm_sanitize_response_guardrail(const char* scope_uuid, const char* name, int32_t priority, NemoRelayLlmSanitizeResponseCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_scope_deregister_llm_sanitize_response_guardrail(const char* scope_uuid, const char* name); extern int32_t nemo_relay_scope_register_llm_conditional_execution_guardrail(const char* scope_uuid, const char* name, int32_t priority, NemoRelayLlmConditionalCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_scope_deregister_llm_conditional_execution_guardrail(const char* scope_uuid, const char* name); @@ -274,8 +278,8 @@ extern char* goToolConditionalTrampoline(void*, const char*, const char*); extern char* goToolExecTrampoline(void*, const char*); extern void goEventSubscriberTrampoline(void*, const FfiEvent*); extern void goFreeTrampoline(void*); -extern FfiLLMRequest* goLlmRequestTrampoline(void*, const FfiLLMRequest*); -extern char* goLlmResponseTrampoline(void*, const char*); +extern FfiLLMRequest* goLlmRequestTrampoline(void*, const FfiLLMRequest*, NemoRelayLlmSanitizeRequestContext); +extern char* goLlmResponseTrampoline(void*, const char*, NemoRelayLlmSanitizeResponseContext); extern char* goLlmConditionalTrampoline(void*, const FfiLLMRequest*); extern char* goLlmExecTrampoline(void*, const char*); extern char* goToolExecInterceptTrampoline(void*, const char*, NemoRelayToolExecNextFn, void*); @@ -1367,17 +1371,15 @@ func DeregisterToolExecutionIntercept(name string) error { // Guardrail/Intercept registration (LLM) // --------------------------------------------------------------------------- -// RegisterLlmSanitizeRequestGuardrail registers a guardrail that sanitizes LLM -// request data before the call is made. The callback receives the request -// headers and content JSON and must return the (possibly modified) versions. -// Guardrails are invoked in priority order (lower values run first). +// RegisterLlmSanitizeRequestGuardrail registers a codec-aware LLM request +// sanitizer. Returning omit true removes only the emitted payload. func RegisterLlmSanitizeRequestGuardrail(name string, priority int32, fn LLMRequestFunc) error { id := registerClosure(fn) cName := C.CString(name) defer C.free(unsafe.Pointer(cName)) return checkStatus(C.nemo_relay_register_llm_sanitize_request_guardrail( cName, C.int32_t(priority), - C.NemoRelayLlmRequestCb(C.goLlmRequestTrampoline), + C.NemoRelayLlmSanitizeRequestCb(C.goLlmRequestTrampoline), id, C.NemoRelayFreeFn(C.goFreeTrampoline), )) @@ -1391,17 +1393,15 @@ func DeregisterLlmSanitizeRequestGuardrail(name string) error { return checkStatus(C.nemo_relay_deregister_llm_sanitize_request_guardrail(cName)) } -// RegisterLlmSanitizeResponseGuardrail registers a guardrail that sanitizes -// LLM response data before it is returned to the caller. The callback receives -// the response as plain JSON and must return the (possibly modified) response -// JSON. Guardrails are invoked in priority order (lower values run first). +// RegisterLlmSanitizeResponseGuardrail registers a codec-aware LLM response +// sanitizer. Returning omit true removes only the emitted payload. func RegisterLlmSanitizeResponseGuardrail(name string, priority int32, fn LLMResponseFunc) error { id := registerClosure(fn) cName := C.CString(name) defer C.free(unsafe.Pointer(cName)) return checkStatus(C.nemo_relay_register_llm_sanitize_response_guardrail( cName, C.int32_t(priority), - C.NemoRelayLlmResponseFn(C.goLlmResponseTrampoline), + C.NemoRelayLlmSanitizeResponseCb(C.goLlmResponseTrampoline), id, C.NemoRelayFreeFn(C.goFreeTrampoline), )) @@ -2426,7 +2426,7 @@ func ScopeRegisterLlmSanitizeRequestGuardrail(scopeUUID, name string, priority i defer C.free(unsafe.Pointer(cName)) return checkStatus(C.nemo_relay_scope_register_llm_sanitize_request_guardrail( cScopeUUID, cName, C.int32_t(priority), - C.NemoRelayLlmRequestCb(C.goLlmRequestTrampoline), + C.NemoRelayLlmSanitizeRequestCb(C.goLlmRequestTrampoline), id, C.NemoRelayFreeFn(C.goFreeTrampoline), )) @@ -2452,7 +2452,7 @@ func ScopeRegisterLlmSanitizeResponseGuardrail(scopeUUID, name string, priority defer C.free(unsafe.Pointer(cName)) return checkStatus(C.nemo_relay_scope_register_llm_sanitize_response_guardrail( cScopeUUID, cName, C.int32_t(priority), - C.NemoRelayLlmResponseFn(C.goLlmResponseTrampoline), + C.NemoRelayLlmSanitizeResponseCb(C.goLlmResponseTrampoline), id, C.NemoRelayFreeFn(C.goFreeTrampoline), )) diff --git a/go/nemo_relay/plugin.go b/go/nemo_relay/plugin.go index 7da42b475..8a7f1942e 100644 --- a/go/nemo_relay/plugin.go +++ b/go/nemo_relay/plugin.go @@ -5,10 +5,15 @@ package nemo_relay /* #include +#include #include typedef struct FfiPluginContext FfiPluginContext; typedef struct FfiPluginActivation FfiPluginActivation; +typedef struct FfiLlmSanitizeRequestCodec FfiLlmSanitizeRequestCodec; +typedef struct FfiLlmSanitizeResponseCodec FfiLlmSanitizeResponseCodec; +typedef struct NemoRelayLlmSanitizeRequestContext { uint32_t codec_kind; const char* codec_id; const FfiLlmSanitizeRequestCodec* codec; } NemoRelayLlmSanitizeRequestContext; +typedef struct NemoRelayLlmSanitizeResponseContext { uint32_t codec_kind; const char* codec_id; const FfiLlmSanitizeResponseCodec* codec; } NemoRelayLlmSanitizeResponseContext; typedef void (*NemoRelayFreeFn)(void* user_data); typedef char* (*NemoRelayPluginValidateCb)(void* user_data, const char* plugin_config_json); @@ -17,8 +22,8 @@ typedef void (*NemoRelayEventSubscriberFn)(void* user_data, const void* event); typedef char* (*NemoRelayEventSanitizeFn)(void* user_data, const void* event, const char* fields_json); typedef char* (*NemoRelayToolSanitizeFn)(void* user_data, const char* name, const char* args_json); typedef char* (*NemoRelayToolConditionalFn)(void* user_data, const char* name, const char* args_json); -typedef void* (*NemoRelayLlmRequestCb)(void* user_data, const void* request); -typedef char* (*NemoRelayLlmResponseFn)(void* user_data, const char* response_json); +typedef void* (*NemoRelayLlmSanitizeRequestCb)(void* user_data, const void* request, NemoRelayLlmSanitizeRequestContext context); +typedef char* (*NemoRelayLlmSanitizeResponseCb)(void* user_data, const char* response_json, NemoRelayLlmSanitizeResponseContext context); typedef char* (*NemoRelayLlmConditionalCb)(void* user_data, const void* request); typedef int32_t (*NemoRelayLlmRequestInterceptCb)(void* user_data, const char* name, const void* request, const char* annotated_json, char** out_outcome_json); typedef char* (*NemoRelayLlmExecNextFn)(const char* native_json, void* next_ctx); @@ -45,8 +50,8 @@ extern int32_t nemo_relay_plugin_context_register_scope_sanitize_end_guardrail(F extern int32_t nemo_relay_plugin_context_register_tool_sanitize_request_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayToolSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_plugin_context_register_tool_sanitize_response_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayToolSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_plugin_context_register_tool_conditional_execution_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayToolConditionalFn cb, void* user_data, NemoRelayFreeFn free_fn); -extern int32_t nemo_relay_plugin_context_register_llm_sanitize_request_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayLlmRequestCb cb, void* user_data, NemoRelayFreeFn free_fn); -extern int32_t nemo_relay_plugin_context_register_llm_sanitize_response_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayLlmResponseFn cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_plugin_context_register_llm_sanitize_request_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayLlmSanitizeRequestCb cb, void* user_data, NemoRelayFreeFn free_fn); +extern int32_t nemo_relay_plugin_context_register_llm_sanitize_response_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayLlmSanitizeResponseCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_plugin_context_register_llm_conditional_execution_guardrail(FfiPluginContext* ctx, const char* name, int32_t priority, NemoRelayLlmConditionalCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_plugin_context_register_llm_request_intercept(FfiPluginContext* ctx, const char* name, int32_t priority, _Bool break_chain, NemoRelayLlmRequestInterceptCb cb, void* user_data, NemoRelayFreeFn free_fn); extern int32_t nemo_relay_plugin_context_register_tool_request_intercept(FfiPluginContext* ctx, const char* name, int32_t priority, _Bool break_chain, NemoRelayToolSanitizeFn cb, void* user_data, NemoRelayFreeFn free_fn); @@ -61,8 +66,8 @@ extern char* goEventSanitizeTrampoline(void*, const void*, const char*); extern void goFreeTrampoline(void*); extern char* goToolSanitizeTrampoline(void*, const char*, const char*); extern char* goToolConditionalTrampoline(void*, const char*, const char*); -extern void* goLlmRequestTrampoline(void*, const void*); -extern char* goLlmResponseTrampoline(void*, const char*); +extern void* goLlmRequestTrampoline(void*, const void*, NemoRelayLlmSanitizeRequestContext); +extern char* goLlmResponseTrampoline(void*, const char*, NemoRelayLlmSanitizeResponseContext); 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**); @@ -668,7 +673,7 @@ func (ctx *PluginContext) RegisterLlmSanitizeRequestGuardrail(name string, prior ctx.ptr, cName, C.int32_t(priority), - (C.NemoRelayLlmRequestCb)(C.goLlmRequestTrampoline), + (C.NemoRelayLlmSanitizeRequestCb)(C.goLlmRequestTrampoline), userData, (C.NemoRelayFreeFn)(C.goFreeTrampoline), )) @@ -686,7 +691,7 @@ func (ctx *PluginContext) RegisterLlmSanitizeResponseGuardrail(name string, prio ctx.ptr, cName, C.int32_t(priority), - (C.NemoRelayLlmResponseFn)(C.goLlmResponseTrampoline), + (C.NemoRelayLlmSanitizeResponseCb)(C.goLlmResponseTrampoline), userData, (C.NemoRelayFreeFn)(C.goFreeTrampoline), )) diff --git a/go/nemo_relay/scope_local_test.go b/go/nemo_relay/scope_local_test.go index 3e553d458..565d0552e 100644 --- a/go/nemo_relay/scope_local_test.go +++ b/go/nemo_relay/scope_local_test.go @@ -923,12 +923,13 @@ func TestScopeLocalLlmSanitizeRequestGuardrailAffectsEvent(t *testing.T) { handle, _ := PushScope("llm_scope_guard", ScopeTypeAgent) defer PopScope(handle) err := ScopeRegisterLlmSanitizeRequestGuardrail(handle.UUID(), "scope_llm_san_req", 1, - func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { + func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { var m map[string]interface{} - json.Unmarshal(content, &m) + json.Unmarshal(request.Content, &m) m["scope_llm_sanitized"] = true out, _ := json.Marshal(m) - return headers, out + request.Content = out + return request, false }) if err != nil { t.Fatalf("ScopeRegisterLlmSanitizeRequestGuardrail failed: %v", err) @@ -1080,9 +1081,9 @@ func assertScopeLocalLLMWrappersDeregister(t *testing.T, scopeUUID string, reque &sanitizeRequestCalls, func() error { return ScopeRegisterLlmSanitizeRequestGuardrail(scopeUUID, "llm_scope_san_req", 1, - func(headers, content json.RawMessage) (json.RawMessage, json.RawMessage) { + func(request LLMRequestDTO, _ LLMSanitizeRequestContext) (LLMRequestDTO, bool) { sanitizeRequestCalls++ - return headers, content + return request, false }, ) }, @@ -1098,9 +1099,9 @@ func assertScopeLocalLLMWrappersDeregister(t *testing.T, scopeUUID string, reque &sanitizeResponseCalls, func() error { return ScopeRegisterLlmSanitizeResponseGuardrail(scopeUUID, "llm_scope_san_resp", 1, - func(responseJSON json.RawMessage) json.RawMessage { + func(responseJSON json.RawMessage, _ LLMSanitizeResponseContext) (json.RawMessage, bool) { sanitizeResponseCalls++ - return responseJSON + return responseJSON, false }, ) }, diff --git a/python/nemo_relay/__init__.py b/python/nemo_relay/__init__.py index b54404246..fb8c89152 100644 --- a/python/nemo_relay/__init__.py +++ b/python/nemo_relay/__init__.py @@ -96,9 +96,14 @@ async def main(): AtofExporterMode, AtofStreamSinkConfig, LLMAttributes, + LlmCodecIdentity, LLMHandle, LLMRequest, LLMRequestInterceptOutcome, + LlmSanitizeRequestCodec, + LlmSanitizeRequestContext, + LlmSanitizeResponseCodec, + LlmSanitizeResponseContext, MarkEvent, OpenInferenceConfig, OpenInferenceSubscriber, @@ -157,12 +162,13 @@ class EventSanitizeFields(TypedDict): #: message. Returning ``None`` allows execution to continue. ToolConditionalExecutionGuardrail: TypeAlias = Callable[[str, Json], Optional[str]] #: Guardrail callback that sanitizes an ``LLMRequest`` used for emitted events. -#: The returned request is recorded for observability and does not replace the -#: caller-visible request value unless the managed LLM API documents otherwise. -LlmSanitizeRequestGuardrail: TypeAlias = Callable[[LLMRequest], LLMRequest] -#: Guardrail callback that sanitizes an emitted JSON LLM response payload. The -#: returned object is recorded on the event; callback exceptions propagate. -LlmSanitizeResponseGuardrail: TypeAlias = Callable[[JsonObject], JsonObject] +#: Callbacks receive ``(request, context)``. Returning ``None`` omits the LLM observability +#: payload and annotation without changing the caller-visible request. +LlmSanitizeRequestGuardrail: TypeAlias = Callable[[LLMRequest, "LlmSanitizeRequestContext"], Optional[LLMRequest]] +#: Guardrail callback that sanitizes an emitted JSON LLM response payload. +#: Callbacks receive ``(response, context)`` and can return ``None`` to omit +#: observability payload and annotation without changing the caller response. +LlmSanitizeResponseGuardrail: TypeAlias = Callable[[Json, "LlmSanitizeResponseContext"], Optional[Json]] #: Guardrail callback that can block an LLM call by returning a rejection #: message. Returning ``None`` allows execution to continue. LlmConditionalExecutionGuardrail: TypeAlias = Callable[[LLMRequest], Optional[str]] @@ -497,6 +503,11 @@ def worker() -> None: "ToolConditionalExecutionGuardrail", "LlmSanitizeRequestGuardrail", "LlmSanitizeResponseGuardrail", + "LlmCodecIdentity", + "LlmSanitizeRequestContext", + "LlmSanitizeResponseContext", + "LlmSanitizeRequestCodec", + "LlmSanitizeResponseCodec", "LlmConditionalExecutionGuardrail", "ToolRequestIntercept", "ToolExecutionIntercept", diff --git a/python/nemo_relay/__init__.pyi b/python/nemo_relay/__init__.pyi index 04cb3429d..7c2372a57 100644 --- a/python/nemo_relay/__init__.pyi +++ b/python/nemo_relay/__init__.pyi @@ -63,6 +63,9 @@ from nemo_relay._native import ( from nemo_relay._native import ( LLMAttributes as LLMAttributes, ) +from nemo_relay._native import ( + LlmCodecIdentity as LlmCodecIdentity, +) from nemo_relay._native import ( LLMHandle as LLMHandle, ) @@ -72,6 +75,18 @@ from nemo_relay._native import ( from nemo_relay._native import ( LLMRequestInterceptOutcome as LLMRequestInterceptOutcome, ) +from nemo_relay._native import ( + LlmSanitizeRequestCodec as LlmSanitizeRequestCodec, +) +from nemo_relay._native import ( + LlmSanitizeRequestContext as LlmSanitizeRequestContext, +) +from nemo_relay._native import ( + LlmSanitizeResponseCodec as LlmSanitizeResponseCodec, +) +from nemo_relay._native import ( + LlmSanitizeResponseContext as LlmSanitizeResponseContext, +) from nemo_relay._native import ( MarkEvent as MarkEvent, ) @@ -172,23 +187,31 @@ Arguments: Return: ``None`` to allow execution, or a rejection message to block it. """ -LlmSanitizeRequestGuardrail: TypeAlias = Callable[[LLMRequest], LLMRequest] +LlmSanitizeRequestGuardrail: TypeAlias = Callable[[LLMRequest, "LlmSanitizeRequestContext"], Optional[LLMRequest]] """Guardrail callback that sanitizes an ``LLMRequest`` used for emitted events. Arguments: - The current LLM request. + The current LLM request and a context object containing a tagged ``codec`` + identity. Its ``kind`` is ``none``, ``builtin``, ``runtime``, or ``opaque``; + ``builtin`` and ``runtime`` identities include ``id``. Use + ``context.resolve_codec()`` to access the active in-process codec. Return: - Request object recorded on the emitted lifecycle event. + Request object recorded on the emitted lifecycle event, or ``None`` to omit + the LLM observability payload and annotation. """ -LlmSanitizeResponseGuardrail: TypeAlias = Callable[[JsonObject], JsonObject] +LlmSanitizeResponseGuardrail: TypeAlias = Callable[[Json, "LlmSanitizeResponseContext"], Optional[Json]] """Guardrail callback that sanitizes an emitted JSON LLM response payload. Arguments: - The response object to sanitize for observability. + The response object and a context object containing a tagged ``codec`` + identity. Its ``kind`` is ``none``, ``builtin``, ``runtime``, or ``opaque``; + ``builtin`` and ``runtime`` identities include ``id``. Use + ``context.resolve_codec()`` to access the active in-process codec. Return: - Response object recorded on the emitted lifecycle event. + Response object recorded on the emitted lifecycle event, or ``None`` to + omit the LLM observability payload and annotation. """ LlmConditionalExecutionGuardrail: TypeAlias = Callable[[LLMRequest], Optional[str]] """Guardrail callback that can block an LLM call. diff --git a/python/nemo_relay/_native.pyi b/python/nemo_relay/_native.pyi index aaec3027a..a76ee5240 100644 --- a/python/nemo_relay/_native.pyi +++ b/python/nemo_relay/_native.pyi @@ -39,8 +39,8 @@ class _EventSanitizeFields(TypedDict): _ToolSanitizeGuardrail: TypeAlias = Callable[[str, _Json], _Json] _ToolConditionalExecutionGuardrail: TypeAlias = Callable[[str, _Json], Optional[str]] -_LlmSanitizeRequestGuardrail: TypeAlias = Callable[["LLMRequest"], "LLMRequest"] -_LlmSanitizeResponseGuardrail: TypeAlias = Callable[[_JsonObject], _JsonObject] +_LlmSanitizeRequestGuardrail: TypeAlias = Callable[["LLMRequest", "LlmSanitizeRequestContext"], Optional["LLMRequest"]] +_LlmSanitizeResponseGuardrail: TypeAlias = Callable[[_Json, "LlmSanitizeResponseContext"], Optional[_Json]] _EventSanitizeGuardrail: TypeAlias = Callable[[ScopeEvent | MarkEvent, _EventSanitizeFields], _EventSanitizeFields] _LlmConditionalExecutionGuardrail: TypeAlias = Callable[["LLMRequest"], Optional[str]] _ToolRequestIntercept: TypeAlias = Callable[[str, _Json], _Json] @@ -61,6 +61,35 @@ _LlmStreamExecutionIntercept: TypeAlias = Callable[ AsyncIterator[_Json] | Awaitable[AsyncIterator[_Json]], ] +class LlmCodecIdentity: + """Structured identity of the active managed LLM codec.""" + + @property + def kind(self) -> Literal["none", "builtin", "runtime", "opaque"]: ... + @property + def id(self) -> str | None: ... + +class LlmSanitizeRequestContext: + """Per-call context passed to an LLM request sanitizer callback.""" + + @property + def codec(self) -> LlmCodecIdentity: ... + def resolve_codec(self) -> LlmSanitizeRequestCodec | None: ... + +class LlmSanitizeResponseContext: + """Per-call context passed to an LLM response sanitizer callback.""" + + @property + def codec(self) -> LlmCodecIdentity: ... + def resolve_codec(self) -> LlmSanitizeResponseCodec | None: ... + +class LlmSanitizeRequestCodec: + def decode(self, request: LLMRequest) -> AnnotatedLLMRequest: ... + def encode(self, annotated: AnnotatedLLMRequest, original: LLMRequest) -> LLMRequest: ... + +class LlmSanitizeResponseCodec: + def decode_response(self, response: _Json) -> AnnotatedLLMResponse: ... + class ScopeAttributes: """Bitflags describing scope properties. diff --git a/python/nemo_relay/guardrails.py b/python/nemo_relay/guardrails.py index 92bc8d9b5..94cb2fd63 100644 --- a/python/nemo_relay/guardrails.py +++ b/python/nemo_relay/guardrails.py @@ -248,8 +248,14 @@ def register_llm_sanitize_request(name: str, priority: int, guardrail: LlmSaniti Args: name: Unique guardrail name used for later replacement or removal. priority: Execution order for the guardrail. Lower values run first. - guardrail: Callable invoked as ``guardrail(request)`` that must return - the sanitized request recorded on the emitted start event. + guardrail: Callable invoked as ``guardrail(request, context)`` that + returns the sanitized request, or ``None`` to omit the LLM + observability payload and its annotation. ``context`` contains + a ``LlmSanitizeRequestContext`` whose ``codec`` is a structured identity + with ``kind`` of ``none``, ``builtin``, ``runtime``, or ``opaque``. + ``builtin`` and ``runtime`` identities include ``id``. Use + ``context.resolve_codec()`` to access the active codec for normalized + processing. Returns: None: This function returns after the guardrail is registered. @@ -263,7 +269,7 @@ def register_llm_sanitize_request(name: str, priority: int, guardrail: LlmSaniti import nemo_relay - def strip_auth(request): + def strip_auth(request, context): headers = {k: v for k, v in request.headers.items() if k.lower() != "authorization"} return nemo_relay.LLMRequest(headers, request.content) @@ -295,8 +301,14 @@ def register_llm_sanitize_response(name: str, priority: int, guardrail: LlmSanit Args: name: Unique guardrail name used for later replacement or removal. priority: Execution order for the guardrail. Lower values run first. - guardrail: Callable invoked as ``guardrail(response)`` that must return - the sanitized payload recorded on the emitted end event. + guardrail: Callable invoked as ``guardrail(response, context)`` that + returns the sanitized payload, or ``None`` to omit the LLM + observability payload and its annotation. ``context`` contains + a ``LlmSanitizeResponseContext`` whose ``codec`` is a structured identity + with ``kind`` of ``none``, ``builtin``, ``runtime``, or ``opaque``. + ``builtin`` and ``runtime`` identities include ``id``. Use + ``context.resolve_codec()`` to access the active codec for normalized + processing. Returns: None: This function returns after the guardrail is registered. diff --git a/python/plugin/src/nemo_relay_plugin/__init__.py b/python/plugin/src/nemo_relay_plugin/__init__.py index f3052e42a..54466a5b5 100644 --- a/python/plugin/src/nemo_relay_plugin/__init__.py +++ b/python/plugin/src/nemo_relay_plugin/__init__.py @@ -20,6 +20,11 @@ Event: A Relay event represented as a JSON object. EventSanitizeFields: Mutable event observability fields. LlmRequest: A Relay LLM request represented as a JSON object. + LlmCodecIdentity: Typed discriminator for the active LLM codec. + LlmSanitizeRequestContext: Per-call context supplied to an LLM request sanitizer. + LlmSanitizeResponseContext: Per-call context supplied to an LLM response sanitizer. + WorkerRequestCodec: Invocation-scoped async proxy for an active request codec. + WorkerResponseCodec: Invocation-scoped async proxy for an active response codec. AnnotatedLlmRequest: An annotated Relay LLM request represented as a JSON object. PendingMarkSpec: A mark Relay emits under its managed lifecycle scope. @@ -71,6 +76,7 @@ EventSanitizeCallback, EventSanitizeFields, Json, + LlmCodecIdentity, LlmConditionalCallback, LlmExecutionCallback, LlmNext, @@ -85,7 +91,9 @@ LlmRequestCallback, LlmRequestInterceptOutcome, LlmSanitizeRequestCallback, + LlmSanitizeRequestContext, LlmSanitizeResponseCallback, + LlmSanitizeResponseContext, LlmStreamExecutionCallback, LlmStreamNext, PendingMarkSpec, @@ -100,6 +108,8 @@ ToolRequestCallback, ToolSanitizeCallback, WorkerPlugin, + WorkerRequestCodec, + WorkerResponseCodec, WorkerSdkError, serve_plugin, ) @@ -113,6 +123,7 @@ "EventSanitizeFields", "Json", "LlmConditionalCallback", + "LlmCodecIdentity", "LlmExecutionCallback", "LlmOptimizationContribution", "LlmOptimizationDataSchema", @@ -123,6 +134,8 @@ "LlmOptimizationTokens", "LlmNext", "LlmRequest", + "LlmSanitizeRequestContext", + "LlmSanitizeResponseContext", "LlmRequestCallback", "LlmRequestInterceptOutcome", "LlmSanitizeRequestCallback", @@ -140,6 +153,8 @@ "ToolNext", "ToolRequestCallback", "ToolSanitizeCallback", + "WorkerRequestCodec", + "WorkerResponseCodec", "WorkerPlugin", "WorkerSdkError", "serve_plugin", diff --git a/python/plugin/src/nemo_relay_plugin/_api.py b/python/plugin/src/nemo_relay_plugin/_api.py index b8facbd76..7f8b6ccb8 100644 --- a/python/plugin/src/nemo_relay_plugin/_api.py +++ b/python/plugin/src/nemo_relay_plugin/_api.py @@ -102,6 +102,105 @@ class EventSanitizeFields(TypedDict): #: An annotated Relay LLM request represented as a JSON object. AnnotatedLlmRequest: TypeAlias = dict[str, Any] + +@dataclass(frozen=True) +class LlmCodecIdentity: + """Structured identity of the codec active for one worker invocation.""" + + kind: str + id: str | None = None + + +@dataclass(frozen=True) +class LlmSanitizeRequestContext: + """Structured per-call context provided to an LLM request sanitizer.""" + + codec: LlmCodecIdentity + _runtime: "PluginRuntime | None" = field(default=None, repr=False, compare=False) + _capability_id: str | None = field(default=None, repr=False, compare=False) + _invocation_id: str | None = field(default=None, repr=False, compare=False) + + def resolve_codec(self) -> "WorkerRequestCodec | None": + """Return the active request-codec proxy for this invocation.""" + if self._runtime is None or self._capability_id is None: + return None + if self._invocation_id is None: + return None + return WorkerRequestCodec(self._runtime, self._capability_id, self._invocation_id) + + +@dataclass(frozen=True) +class LlmSanitizeResponseContext: + """Structured per-call context provided to an LLM response sanitizer.""" + + codec: LlmCodecIdentity + _runtime: "PluginRuntime | None" = field(default=None, repr=False, compare=False) + _capability_id: str | None = field(default=None, repr=False, compare=False) + _invocation_id: str | None = field(default=None, repr=False, compare=False) + + def resolve_codec(self) -> "WorkerResponseCodec | None": + """Return the active response-codec proxy for this invocation.""" + if self._runtime is None or self._capability_id is None: + return None + if self._invocation_id is None: + return None + return WorkerResponseCodec(self._runtime, self._capability_id, self._invocation_id) + + +@dataclass(frozen=True) +class WorkerRequestCodec: + """Invocation-scoped async proxy for an active request codec.""" + + _runtime: "PluginRuntime" + _capability_id: str + _invocation_id: str + + async def decode(self, request: LlmRequest) -> AnnotatedLlmRequest: + """Decode a wire request with the active host codec.""" + return await self._runtime._decode_llm_codec_request(self._capability_id, self._invocation_id, request) + + async def encode(self, annotated: AnnotatedLlmRequest, original: LlmRequest) -> LlmRequest: + """Encode an annotated request with the active host codec.""" + return await self._runtime._encode_llm_codec_request( + self._capability_id, self._invocation_id, annotated, original + ) + + +@dataclass(frozen=True) +class WorkerResponseCodec: + """Invocation-scoped async proxy for an active response codec.""" + + _runtime: "PluginRuntime" + _capability_id: str + _invocation_id: str + + async def decode(self, response: Json) -> Json: + """Decode a wire response with the active host codec.""" + return await self._runtime._decode_llm_codec_response(self._capability_id, self._invocation_id, response) + + +def _llm_codec_identity(invocation: pb.LlmInvocation) -> LlmCodecIdentity: + """Return the codec identity from a worker invocation.""" + context = getattr(invocation, invocation.WhichOneof("sanitize_context") or "", None) + proto_codec = context.codec if context is not None and context.HasField("codec") else None + codec_id = proto_codec.id if proto_codec is not None and proto_codec.HasField("id") else None + codec_kind = proto_codec.kind if proto_codec is not None else pb.LLM_CODEC_KIND_UNSPECIFIED + if codec_kind == pb.LLM_CODEC_KIND_UNSPECIFIED: + identity = LlmCodecIdentity("none") + elif codec_kind == pb.LLM_CODEC_KIND_BUILTIN and codec_id: + identity = LlmCodecIdentity("builtin", codec_id) + elif codec_kind == pb.LLM_CODEC_KIND_RUNTIME and codec_id: + identity = LlmCodecIdentity("runtime", codec_id) + else: + identity = LlmCodecIdentity("opaque") + return identity + + +def _llm_codec_capability(invocation: pb.LlmInvocation) -> str | None: + context = getattr(invocation, invocation.WhichOneof("sanitize_context") or "", None) + return context.codec_capability_id if context is not None and context.HasField("codec_capability_id") else None + + WORKER_PROTOCOL = "grpc-v1" JSON_SCHEMA = "nemo.relay.Json@1" EVENT_SCHEMA = "nemo.relay.Event@1" @@ -730,8 +829,12 @@ def register(self, ctx: PluginContext, config: Json) -> None | Awaitable[None]: [str, Json, "ToolNext"], ToolExecutionInterceptOutcome | Awaitable[ToolExecutionInterceptOutcome], ] -LlmSanitizeRequestCallback: TypeAlias = Callable[[LlmRequest], LlmRequest | Awaitable[LlmRequest]] -LlmSanitizeResponseCallback: TypeAlias = Callable[[Json], Json | Awaitable[Json]] +LlmSanitizeRequestCallback: TypeAlias = Callable[ + [LlmRequest, LlmSanitizeRequestContext], LlmRequest | None | Awaitable[LlmRequest | None] +] +LlmSanitizeResponseCallback: TypeAlias = Callable[ + [Json, LlmSanitizeResponseContext], Json | None | Awaitable[Json | None] +] LlmConditionalCallback: TypeAlias = Callable[[LlmRequest], str | None | Awaitable[str | None]] LlmRequestCallback: TypeAlias = Callable[ [str, LlmRequest, AnnotatedLlmRequest | None], @@ -808,6 +911,15 @@ class PluginContext: pb.SCOPE_SANITIZE_START_GUARDRAIL: "scope_start_sanitizers", pb.SCOPE_SANITIZE_END_GUARDRAIL: "scope_end_sanitizers", } + _LLM_HANDLER_ATTRIBUTES: ClassVar[frozenset[int]] = frozenset( + { + pb.LLM_SANITIZE_REQUEST_GUARDRAIL, + pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, + pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL, + pb.LLM_REQUEST_INTERCEPT, + pb.LLM_EXECUTION_INTERCEPT, + } + ) def __init__(self, runtime: PluginRuntime | None = None) -> None: self._runtime = runtime @@ -1030,10 +1142,9 @@ def register_llm_sanitize_request_guardrail( Args: name: Component-local registration name. - callback: Function receiving an :data:`LlmRequest` and returning + callback: Function receiving ``(request, context)`` and returning the request recorded on the LLM start event, directly or - through an awaitable. It does not change the request sent to - the model. + through an awaitable. Return ``None`` to omit observability. priority: Execution order. Lower values run first. """ self._push_registration(name, pb.LLM_SANITIZE_REQUEST_GUARDRAIL, priority, False) @@ -1050,9 +1161,9 @@ def register_llm_sanitize_response_guardrail( Args: name: Component-local registration name. - callback: Function receiving response JSON and returning the value - recorded on the LLM end event, directly or through an - awaitable. It does not change the real model response. + callback: Function receiving ``(response, context)`` and returning + the value recorded on the LLM end event, directly or through an + awaitable. Return ``None`` to omit observability. priority: Execution order. Lower values run first. """ self._push_registration(name, pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, priority, False) @@ -1190,6 +1301,47 @@ def __init__(self, *, activation_id: str, auth_token: str, host_stub: Any) -> No self._auth_token = auth_token self._host_stub = host_stub + async def _decode_llm_codec_request( + self, capability_id: str, invocation_id: str, request: LlmRequest + ) -> AnnotatedLlmRequest: + result = await self._host_stub.DecodeLlmCodecRequest( + pb.LlmCodecDecodeRequest( + activation_id=self._activation_id, + auth_token=self._auth_token, + codec_capability_id=capability_id, + invocation_id=invocation_id, + request=_json_envelope(LLM_REQUEST_SCHEMA, request), + ) + ) + return _json_result_to_value(result, ANNOTATED_LLM_REQUEST_SCHEMA) + + async def _encode_llm_codec_request( + self, capability_id: str, invocation_id: str, annotated: AnnotatedLlmRequest, original: LlmRequest + ) -> LlmRequest: + result = await self._host_stub.EncodeLlmCodecRequest( + pb.LlmCodecEncodeRequest( + activation_id=self._activation_id, + auth_token=self._auth_token, + codec_capability_id=capability_id, + invocation_id=invocation_id, + annotated_request=_json_envelope(ANNOTATED_LLM_REQUEST_SCHEMA, annotated), + original_request=_json_envelope(LLM_REQUEST_SCHEMA, original), + ) + ) + return _json_result_to_value(result, LLM_REQUEST_SCHEMA) + + async def _decode_llm_codec_response(self, capability_id: str, invocation_id: str, response: Json) -> Json: + result = await self._host_stub.DecodeLlmCodecResponse( + pb.LlmCodecDecodeResponse( + activation_id=self._activation_id, + auth_token=self._auth_token, + codec_capability_id=capability_id, + invocation_id=invocation_id, + response=_json_envelope(JSON_SCHEMA, response), + ) + ) + return _json_result_to_value(result) + async def emit_mark( self, name: str, @@ -1837,6 +1989,8 @@ async def Shutdown(self, request: Any, context: Any) -> Any: async def _invoke_result(self, request: Any) -> Any: with _bind_invocation_scope(request): + if request.surface in PluginContext._LLM_HANDLER_ATTRIBUTES: + return await self._invoke_llm_result(request) if request.surface == pb.SUBSCRIBER: event = _decode_required_envelope(request.event, "event", EVENT_SCHEMA) await _maybe_await(self._handler(self._handlers.subscribers, request.registration_name)(event)) @@ -1906,69 +2060,80 @@ async def _invoke_result(self, request: Any) -> Any: ), ) ) - if request.surface == pb.LLM_SANITIZE_REQUEST_GUARDRAIL: - return _json_response( - await _maybe_await( - self._handler(self._handlers.llm_sanitize_requests, request.registration_name)( - _decode_required_envelope(request.llm.request, "llm request", LLM_REQUEST_SCHEMA) - ) - ) - ) - if request.surface == pb.LLM_SANITIZE_RESPONSE_GUARDRAIL: - return _json_response( - await _maybe_await( - self._handler(self._handlers.llm_sanitize_responses, request.registration_name)( - _decode_required_envelope(request.llm.response, "llm response") - ) - ) - ) - if request.surface == pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL: - result = await _maybe_await( - self._handler(self._handlers.llm_conditionals, request.registration_name)( - _decode_required_envelope(request.llm.request, "llm request", LLM_REQUEST_SCHEMA) - ) + raise WorkerSdkError(f"unsupported registration surface {request.surface}") + + async def _invoke_llm_result(self, request: Any) -> Any: + if request.surface == pb.LLM_SANITIZE_REQUEST_GUARDRAIL: + callback = self._handler(self._handlers.llm_sanitize_requests, request.registration_name) + payload = _decode_required_envelope(request.llm.request, "llm request", LLM_REQUEST_SCHEMA) + return await self._invoke_llm_sanitizer(callback, payload, request, LlmSanitizeRequestContext) + if request.surface == pb.LLM_SANITIZE_RESPONSE_GUARDRAIL: + callback = self._handler(self._handlers.llm_sanitize_responses, request.registration_name) + payload = _decode_required_envelope(request.llm.response, "llm response") + return await self._invoke_llm_sanitizer(callback, payload, request, LlmSanitizeResponseContext) + if request.surface == pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL: + result = await _maybe_await( + self._handler(self._handlers.llm_conditionals, request.registration_name)( + _decode_required_envelope(request.llm.request, "llm request", LLM_REQUEST_SCHEMA) ) - return pb.InvokeResponse(guardrail=pb.GuardrailResult(block_reason=result or "")) - if request.surface == pb.LLM_REQUEST_INTERCEPT: - payload = request.llm - llm_request = _decode_required_envelope(payload.request, "llm request", LLM_REQUEST_SCHEMA) - annotated = ( - _decode_required_envelope( - payload.annotated_request, - "annotated llm request", - ANNOTATED_LLM_REQUEST_SCHEMA, - ) - if payload.HasField("annotated_request") - else None + ) + return pb.InvokeResponse(guardrail=pb.GuardrailResult(block_reason=result or "")) + if request.surface == pb.LLM_REQUEST_INTERCEPT: + payload = request.llm + llm_request = _decode_required_envelope(payload.request, "llm request", LLM_REQUEST_SCHEMA) + annotated = ( + _decode_required_envelope( + payload.annotated_request, + "annotated llm request", + ANNOTATED_LLM_REQUEST_SCHEMA, ) - result = await _maybe_await( - self._handler(self._handlers.llm_requests, request.registration_name)( - payload.model_name, - llm_request, - annotated, - ) + if payload.HasField("annotated_request") + else None + ) + result = await _maybe_await( + self._handler(self._handlers.llm_requests, request.registration_name)( + payload.model_name, + llm_request, + annotated, ) - if not isinstance(result, LlmRequestInterceptOutcome): - raise WorkerSdkError("LLM request intercept must return LlmRequestInterceptOutcome") - return pb.InvokeResponse( - llm_request=pb.LlmRequestInterceptResult( - outcome=_json_envelope( - LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA, - result.to_json(), - ), - ) + ) + if not isinstance(result, LlmRequestInterceptOutcome): + raise WorkerSdkError("LLM request intercept must return LlmRequestInterceptOutcome") + return pb.InvokeResponse( + llm_request=pb.LlmRequestInterceptResult( + outcome=_json_envelope( + LLM_REQUEST_INTERCEPT_OUTCOME_SCHEMA, + result.to_json(), + ), ) - if request.surface == pb.LLM_EXECUTION_INTERCEPT: - payload = request.llm - result = await _maybe_await( - self._handler(self._handlers.llm_executions, request.registration_name)( - payload.model_name, - _decode_required_envelope(payload.request, "llm request", LLM_REQUEST_SCHEMA), - LlmNext(self._runtime, request.continuation_id), - ) + ) + if request.surface == pb.LLM_EXECUTION_INTERCEPT: + payload = request.llm + result = await _maybe_await( + self._handler(self._handlers.llm_executions, request.registration_name)( + payload.model_name, + _decode_required_envelope(payload.request, "llm request", LLM_REQUEST_SCHEMA), + LlmNext(self._runtime, request.continuation_id), ) - return _json_response(result) - raise WorkerSdkError(f"unsupported registration surface {request.surface}") + ) + return _json_response(result) + raise WorkerSdkError(f"unsupported LLM registration surface {request.surface}") + + async def _invoke_llm_sanitizer( + self, + callback: Any, + payload: Any, + request: Any, + context_type: type[LlmSanitizeRequestContext] | type[LlmSanitizeResponseContext], + ) -> Any: + context = context_type( + _llm_codec_identity(request.llm), + self._runtime, + _llm_codec_capability(request.llm), + request.invocation_id, + ) + result = await _maybe_await(callback(payload, context)) + return pb.InvokeResponse(empty=pb.EmptyResult()) if result is None else _json_response(result) def _start_invocation( self, @@ -2116,10 +2281,10 @@ def _json_response(value: Json) -> Any: return pb.InvokeResponse(json=pb.JsonResult(value=_json_envelope(JSON_SCHEMA, value))) -def _json_result_to_value(result: Any) -> Json: +def _json_result_to_value(result: Any, expected_schema: str = JSON_SCHEMA) -> Json: if result.HasField("error"): raise _worker_error_to_sdk(result.error) - return _decode_required_envelope(result.value, "json result") + return _decode_required_envelope(result.value, "json result", expected_schema) def _stream_chunk_to_value(chunk: Any) -> Json: diff --git a/python/tests/plugin/test_worker_sdk.py b/python/tests/plugin/test_worker_sdk.py index 5a521fb1e..7ab918d20 100644 --- a/python/tests/plugin/test_worker_sdk.py +++ b/python/tests/plugin/test_worker_sdk.py @@ -29,6 +29,8 @@ Json, LlmOptimizationContribution, LlmRequestInterceptOutcome, + LlmSanitizeRequestContext, + LlmSanitizeResponseContext, PendingMarkSpec, PluginContext, PluginRuntime, @@ -246,6 +248,37 @@ async def stream() -> AsyncIterator[Any]: return stream() + async def DecodeLlmCodecRequest(self, request: Any) -> Any: + self.requests.append(request) + if self.failures.get("DecodeLlmCodecRequest") == "error": + return pb.JsonResult(error=_worker_error("DecodeLlmCodecRequest failed")) + value = _envelope_value(request.request) + model = value.get("content", {}).get("model") + return pb.JsonResult( + value=_json_envelope( + ANNOTATED_LLM_REQUEST_SCHEMA, + {"messages": [], "model": model}, + ) + ) + + async def EncodeLlmCodecRequest(self, request: Any) -> Any: + self.requests.append(request) + if self.failures.get("EncodeLlmCodecRequest") == "error": + return pb.JsonResult(error=_worker_error("EncodeLlmCodecRequest failed")) + return pb.JsonResult(value=request.original_request) + + async def DecodeLlmCodecResponse(self, request: Any) -> Any: + self.requests.append(request) + if self.failures.get("DecodeLlmCodecResponse") == "error": + return pb.JsonResult(error=_worker_error("DecodeLlmCodecResponse failed")) + value = _envelope_value(request.response) + return pb.JsonResult( + value=_json_envelope( + JSON_SCHEMA, + {"model": value.get("model"), "message": value.get("message")}, + ) + ) + def _host_ack(self, method: str) -> Any: failure = self.failures.get(method) if failure == "empty": @@ -301,10 +334,12 @@ async def tool_execution(name: str, value: Json, next_call: ToolNext) -> ToolExe pending_marks=[PendingMarkSpec("worker.tool.execution")], ) - def llm_sanitize_request(request: Json) -> Json: + def llm_sanitize_request(request: Json, context: LlmSanitizeRequestContext) -> Json: + del context return _tag_llm_request(request, "llm_sanitize_request") - async def llm_sanitize_response(response: Json) -> Json: + async def llm_sanitize_response(response: Json, context: LlmSanitizeResponseContext) -> Json: + del context return _tag(response, "llm_sanitize_response") def llm_block(request: Json) -> str | None: @@ -699,6 +734,266 @@ def callback(tool_name: str, value: Json) -> Json: assert ("shared", pb.TOOL_SANITIZE_RESPONSE_GUARDRAIL) in registrations +def test_plugin_context_registers_llm_sanitizers_under_standard_names(): + context = PluginContext() + + context.register_llm_sanitize_request_guardrail( + "request", + lambda request, codec_context: request if codec_context.codec.kind != "none" else None, + ) + context.register_llm_sanitize_response_guardrail( + "response", + lambda response, codec_context: response if codec_context.codec.kind == "builtin" else None, + ) + + assert [(registration.local_name, registration.surface) for registration in context._handlers.registrations] == [ + ("request", pb.LLM_SANITIZE_REQUEST_GUARDRAIL), + ("response", pb.LLM_SANITIZE_RESPONSE_GUARDRAIL), + ] + + +async def test_llm_sanitizers_receive_codec_context_and_can_omit_payloads(): + seen: list[tuple[str, LlmSanitizeRequestContext | LlmSanitizeResponseContext]] = [] + + class ContextualSanitizerPlugin(WorkerPlugin): + plugin_id = "tests.llm_sanitizer" + + def register(self, ctx: PluginContext, config: Json) -> None: + del config + + def request_sanitizer(request: Json, codec_context: LlmSanitizeRequestContext) -> None: + del request + seen.append(("request", codec_context)) + + def response_sanitizer(response: Json, codec_context: LlmSanitizeResponseContext) -> None: + del response + seen.append(("response", codec_context)) + + ctx.register_llm_sanitize_request_guardrail("request", request_sanitizer) + ctx.register_llm_sanitize_response_guardrail("response", response_sanitizer) + + service = _service(ContextualSanitizerPlugin(), RecordingHostStub()) + await _register(service) + for codec_kind, codec_id in [ + (pb.LLM_CODEC_KIND_UNSPECIFIED, None), + (pb.LLM_CODEC_KIND_BUILTIN, "openai_chat"), + (pb.LLM_CODEC_KIND_BUILTIN, "openai_responses"), + (pb.LLM_CODEC_KIND_BUILTIN, "anthropic_messages"), + (pb.LLM_CODEC_KIND_RUNTIME, "com.example.chat.v1"), + (pb.LLM_CODEC_KIND_OPAQUE, None), + ]: + codec = pb.LlmCodecIdentity(kind=codec_kind) + if codec_id is not None: + codec.id = codec_id + request_payload = _llm_payload( + request={"content": {"prompt": "secret"}}, + response={"secret": "value"}, + ) + request_payload.request_sanitize_context.CopyFrom(pb.LlmSanitizeRequestContext(codec=codec)) + response_payload = _llm_payload( + request={"content": {"prompt": "secret"}}, + response={"secret": "value"}, + ) + response_payload.response_sanitize_context.CopyFrom(pb.LlmSanitizeResponseContext(codec=codec)) + + request_response = await service.Invoke( + _invoke_request( + "request", + pb.LLM_SANITIZE_REQUEST_GUARDRAIL, + llm=request_payload, + ), + AbortContext(), + ) + response_response = await service.Invoke( + _invoke_request( + "response", + pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, + llm=response_payload, + ), + AbortContext(), + ) + assert request_response.WhichOneof("result") == "empty" + assert response_response.WhichOneof("result") == "empty" + + assert seen == [ + ("request", LlmSanitizeRequestContext(plugin_api.LlmCodecIdentity("none"))), + ("response", LlmSanitizeResponseContext(plugin_api.LlmCodecIdentity("none"))), + ("request", LlmSanitizeRequestContext(plugin_api.LlmCodecIdentity("builtin", "openai_chat"))), + ("response", LlmSanitizeResponseContext(plugin_api.LlmCodecIdentity("builtin", "openai_chat"))), + ("request", LlmSanitizeRequestContext(plugin_api.LlmCodecIdentity("builtin", "openai_responses"))), + ("response", LlmSanitizeResponseContext(plugin_api.LlmCodecIdentity("builtin", "openai_responses"))), + ( + "request", + LlmSanitizeRequestContext(plugin_api.LlmCodecIdentity("builtin", "anthropic_messages")), + ), + ( + "response", + LlmSanitizeResponseContext(plugin_api.LlmCodecIdentity("builtin", "anthropic_messages")), + ), + ("request", LlmSanitizeRequestContext(plugin_api.LlmCodecIdentity("runtime", "com.example.chat.v1"))), + ("response", LlmSanitizeResponseContext(plugin_api.LlmCodecIdentity("runtime", "com.example.chat.v1"))), + ("request", LlmSanitizeRequestContext(plugin_api.LlmCodecIdentity("opaque"))), + ("response", LlmSanitizeResponseContext(plugin_api.LlmCodecIdentity("opaque"))), + ] + + +async def test_llm_sanitizers_resolve_directional_codec_proxies(): + calls: list[tuple[str, Json]] = [] + + class CodecSanitizerPlugin(WorkerPlugin): + plugin_id = "tests.llm_codec_proxy" + + def register(self, ctx: PluginContext, config: Json) -> None: + del config + + async def request_sanitizer( + request: Json, + context: LlmSanitizeRequestContext, + ) -> Json: + codec = context.resolve_codec() + assert codec is not None + annotated = await codec.decode(request) + calls.append(("request_decode", annotated)) + return await codec.encode(annotated, request) + + async def response_sanitizer( + response: Json, + context: LlmSanitizeResponseContext, + ) -> Json: + codec = context.resolve_codec() + assert codec is not None + calls.append(("response_decode", await codec.decode(response))) + return response + + ctx.register_llm_sanitize_request_guardrail("request", request_sanitizer) + ctx.register_llm_sanitize_response_guardrail("response", response_sanitizer) + + host = RecordingHostStub() + service = _service(CodecSanitizerPlugin(), host) + await _register(service) + + request_payload = _llm_payload( + request={"headers": {}, "content": {"model": "runtime-model"}}, + ) + request_payload.request_sanitize_context.CopyFrom( + pb.LlmSanitizeRequestContext( + codec=pb.LlmCodecIdentity( + kind=pb.LLM_CODEC_KIND_RUNTIME, + id="com.example.runtime", + ), + codec_capability_id="request-capability", + ) + ) + request_result = await service.Invoke( + _invoke_request( + "request", + pb.LLM_SANITIZE_REQUEST_GUARDRAIL, + llm=request_payload, + ), + AbortContext(), + ) + assert request_result.WhichOneof("result") == "json", request_result + + response_payload = _llm_payload(response={"model": "opaque-model", "message": "secret"}) + response_payload.response_sanitize_context.CopyFrom( + pb.LlmSanitizeResponseContext( + codec=pb.LlmCodecIdentity(kind=pb.LLM_CODEC_KIND_OPAQUE), + codec_capability_id="response-capability", + ) + ) + response_result = await service.Invoke( + _invoke_request( + "response", + pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, + llm=response_payload, + ), + AbortContext(), + ) + assert response_result.WhichOneof("result") == "json", response_result + assert calls == [ + ("request_decode", {"messages": [], "model": "runtime-model"}), + ( + "response_decode", + {"model": "opaque-model", "message": "secret"}, + ), + ] + assert [request.codec_capability_id for request in host.requests if hasattr(request, "codec_capability_id")] == [ + "request-capability", + "request-capability", + "response-capability", + ] + + +@pytest.mark.parametrize( + ("failure", "surface", "registration_name"), + [ + ("DecodeLlmCodecRequest", pb.LLM_SANITIZE_REQUEST_GUARDRAIL, "request"), + ("EncodeLlmCodecRequest", pb.LLM_SANITIZE_REQUEST_GUARDRAIL, "request"), + ("DecodeLlmCodecResponse", pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, "response"), + ], +) +async def test_llm_sanitizer_codec_rpc_failures_return_invocation_errors( + failure: str, + surface: int, + registration_name: str, +): + class CodecFailurePlugin(WorkerPlugin): + plugin_id = "tests.llm_codec_failure" + + def register(self, ctx: PluginContext, config: Json) -> None: + del config + + async def request_sanitizer( + request: Json, + context: LlmSanitizeRequestContext, + ) -> Json: + codec = context.resolve_codec() + assert codec is not None + annotated = await codec.decode(request) + return await codec.encode(annotated, request) + + async def response_sanitizer( + response: Json, + context: LlmSanitizeResponseContext, + ) -> Json: + codec = context.resolve_codec() + assert codec is not None + await codec.decode(response) + return response + + ctx.register_llm_sanitize_request_guardrail("request", request_sanitizer) + ctx.register_llm_sanitize_response_guardrail("response", response_sanitizer) + + host = RecordingHostStub() + host.failures[failure] = "error" + service = _service(CodecFailurePlugin(), host) + await _register(service) + + if surface == pb.LLM_SANITIZE_REQUEST_GUARDRAIL: + payload = _llm_payload(request={"headers": {}, "content": {"model": "test"}}) + payload.request_sanitize_context.CopyFrom( + pb.LlmSanitizeRequestContext( + codec=pb.LlmCodecIdentity(kind=pb.LLM_CODEC_KIND_OPAQUE), + codec_capability_id="request-capability", + ) + ) + else: + payload = _llm_payload(response={"model": "test", "message": "secret"}) + payload.response_sanitize_context.CopyFrom( + pb.LlmSanitizeResponseContext( + codec=pb.LlmCodecIdentity(kind=pb.LLM_CODEC_KIND_OPAQUE), + codec_capability_id="response-capability", + ) + ) + + result = await service.Invoke( + _invoke_request(registration_name, surface, llm=payload), + AbortContext(), + ) + assert result.WhichOneof("result") == "error" + assert f"{failure} failed" in result.error.message + + @pytest.mark.parametrize( "surface", [ diff --git a/python/tests/test_builtin_codecs.py b/python/tests/test_builtin_codecs.py index 99beb9054..156122583 100644 --- a/python/tests/test_builtin_codecs.py +++ b/python/tests/test_builtin_codecs.py @@ -394,7 +394,8 @@ def test_manual_call_end_response_codec_uses_sanitized_payload(self): def capture(event): captured_events.append(event) - def sanitize_response(response): + def sanitize_response(response, context): + del context return { "id": "chatcmpl-sanitized", "model": "gpt-4", diff --git a/python/tests/test_llm.py b/python/tests/test_llm.py index f2575c9aa..9a3662be4 100644 --- a/python/tests/test_llm.py +++ b/python/tests/test_llm.py @@ -22,6 +22,7 @@ scope, subscribers, ) +from nemo_relay.codecs import OpenAIChatCodec def make_request(): @@ -115,8 +116,22 @@ async def func(request): class TestLLMGuardrails: + @pytest.mark.parametrize( + ("register", "callback"), + [ + (guardrails.register_llm_sanitize_request, lambda request: request), + (guardrails.register_llm_sanitize_response, lambda response: response), + (guardrails.register_llm_sanitize_request, object()), + (guardrails.register_llm_sanitize_response, object()), + ], + ) + def test_sanitizer_registration_rejects_legacy_or_uninspectable_callbacks(self, register, callback): + with pytest.raises(TypeError, match="payload, context"): + register("py_llm_invalid_signature", 1, callback) + def test_sanitize_request_guardrail(self): - def sanitizer(request): + def sanitizer(request, context): + del context # request is an LLMRequest object; must return a new LLMRequest headers = request.headers headers["X-Sanitized"] = "true" @@ -126,7 +141,8 @@ def sanitizer(request): guardrails.deregister_llm_sanitize_request("py_llm_san_req") def test_sanitize_response_guardrail(self): - def sanitizer(response): + def sanitizer(response, context): + del context # response is a plain dict response["cleaned"] = True return response @@ -134,6 +150,112 @@ def sanitizer(response): guardrails.register_llm_sanitize_response("py_llm_san_resp", 1, sanitizer) guardrails.deregister_llm_sanitize_response("py_llm_san_resp") + def test_sanitizers_receive_a_structured_codec_context(self): + request_contexts = [] + response_contexts = [] + + def sanitize_request(request, context): + request_contexts.append(context) + return request + + def sanitize_response(response, context): + response_contexts.append(context) + return response + + guardrails.register_llm_sanitize_request("py_llm_structured_context_request", 1, sanitize_request) + guardrails.register_llm_sanitize_response("py_llm_structured_context_response", 1, sanitize_response) + try: + handle = llm.call("py_llm_structured_context", make_request()) + llm.call_end(handle, {"response": "ok"}) + finally: + guardrails.deregister_llm_sanitize_request("py_llm_structured_context_request") + guardrails.deregister_llm_sanitize_response("py_llm_structured_context_response") + + assert len(request_contexts) == 1 + assert len(response_contexts) == 1 + for context in [*request_contexts, *response_contexts]: + assert context.codec.kind == "none" + assert context.codec.id is None + + async def test_sanitizers_resolve_active_builtin_codecs(self): + request_codec_used = False + response_codec_used = False + + def sanitize_request(request, context): + nonlocal request_codec_used + assert context.codec.kind == "builtin" + assert context.codec.id == "openai_chat" + codec = context.resolve_codec() + assert codec is not None + annotated = codec.decode(request) + request_codec_used = True + return codec.encode(annotated, request) + + def sanitize_response(response, context): + nonlocal response_codec_used + assert context.codec.kind == "builtin" + assert context.codec.id == "openai_chat" + codec = context.resolve_codec() + assert codec is not None + assert codec.decode_response(response).model == "test-model" + response_codec_used = True + return response + + codec = OpenAIChatCodec() + response = { + "id": "chatcmpl-python", + "model": "test-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}}], + } + guardrails.register_llm_sanitize_request("py_llm_builtin_context_request", 1, sanitize_request) + guardrails.register_llm_sanitize_response("py_llm_builtin_context_response", 1, sanitize_response) + try: + result = await llm.execute( + "py_llm_builtin_context", + make_request(), + lambda request: response, + codec=codec, + response_codec=codec, + ) + finally: + guardrails.deregister_llm_sanitize_request("py_llm_builtin_context_request") + guardrails.deregister_llm_sanitize_response("py_llm_builtin_context_response") + + assert result == response + assert request_codec_used is True + assert response_codec_used is True + + def test_none_omits_payload_and_short_circuits_later_sanitizers(self): + events = [] + later_called = False + + def omit(request, context): + return None + + def later(request, context): + nonlocal later_called + later_called = True + return request + + subscribers.register("py_llm_none_sanitize_sub", events.append) + guardrails.register_llm_sanitize_request("py_llm_none_sanitize_first", 1, omit) + guardrails.register_llm_sanitize_request("py_llm_none_sanitize_later", 2, later) + try: + handle = llm.call("py_llm_none_sanitize", make_request()) + llm.call_end(handle, {"ok": True}) + finally: + guardrails.deregister_llm_sanitize_request("py_llm_none_sanitize_first") + guardrails.deregister_llm_sanitize_request("py_llm_none_sanitize_later") + try: + subscribers.flush() + finally: + subscribers.deregister("py_llm_none_sanitize_sub") + + start = _llm_event(events, "py_llm_none_sanitize", "start") + assert start.data is None + assert start.annotated_request is None + assert later_called is False + def test_conditional_execution_guardrail(self): def checker(request): return None @@ -148,18 +270,18 @@ def test_conditional_execution_direct(self): guardrails.deregister_llm_conditional_execution("py_llm_cond_direct") def test_duplicate_raises(self): - guardrails.register_llm_sanitize_request("py_llm_dup", 1, lambda r: r) + guardrails.register_llm_sanitize_request("py_llm_dup", 1, lambda r, context: r) with pytest.raises(RuntimeError): - guardrails.register_llm_sanitize_request("py_llm_dup", 1, lambda r: r) + guardrails.register_llm_sanitize_request("py_llm_dup", 1, lambda r, context: r) guardrails.deregister_llm_sanitize_request("py_llm_dup") - def test_sanitize_request_callable_error_falls_back_to_original_input(self): + def test_sanitize_request_callable_error_omits_observability_input(self): events = [] subscribers.register("py_llm_sanitize_req_sub", lambda event: events.append(event)) guardrails.register_llm_sanitize_request( "py_llm_sanitize_req_fail", 1, - lambda request: raise_runtime_error("boom"), + lambda request, context: raise_runtime_error("boom"), ) try: request = make_request() @@ -173,16 +295,16 @@ def test_sanitize_request_callable_error_falls_back_to_original_input(self): subscribers.deregister("py_llm_sanitize_req_sub") start = _llm_event(events, "llm_sanitize_req_fail", "start") - request = make_request() - assert start.data == {"headers": request.headers, "content": request.content} + assert start.data is None + assert start.annotated_request is None - def test_sanitize_request_invalid_return_falls_back_to_original_input(self): + def test_sanitize_request_invalid_return_omits_observability_input(self): events = [] subscribers.register("py_llm_sanitize_req_bad_sub", lambda event: events.append(event)) guardrails.register_llm_sanitize_request( "py_llm_sanitize_req_bad", 1, - cast(guardrails.LlmSanitizeRequestGuardrail, lambda request: object()), + cast(guardrails.LlmSanitizeRequestGuardrail, lambda request, context: object()), ) try: request = make_request() @@ -196,16 +318,16 @@ def test_sanitize_request_invalid_return_falls_back_to_original_input(self): subscribers.deregister("py_llm_sanitize_req_bad_sub") start = _llm_event(events, "llm_sanitize_req_bad", "start") - request = make_request() - assert start.data == {"headers": request.headers, "content": request.content} + assert start.data is None + assert start.annotated_request is None - def test_sanitize_response_callable_error_falls_back_to_original_output(self): + def test_sanitize_response_callable_error_omits_observability_output(self): events = [] subscribers.register("py_llm_sanitize_resp_sub", lambda event: events.append(event)) guardrails.register_llm_sanitize_response( "py_llm_sanitize_resp_fail", 1, - lambda response: raise_runtime_error("boom"), + lambda response, context: raise_runtime_error("boom"), ) try: handle = llm.call("llm_sanitize_resp_fail", make_request()) @@ -218,15 +340,16 @@ def test_sanitize_response_callable_error_falls_back_to_original_output(self): subscribers.deregister("py_llm_sanitize_resp_sub") end = _llm_event(events, "llm_sanitize_resp_fail", "end") - assert end.data == {"ok": True} + assert end.data is None + assert end.annotated_response is None - def test_sanitize_response_invalid_return_falls_back_to_original_output(self): + def test_sanitize_response_invalid_return_omits_observability_output(self): events = [] subscribers.register("py_llm_sanitize_resp_bad_sub", lambda event: events.append(event)) guardrails.register_llm_sanitize_response( "py_llm_sanitize_resp_bad", 1, - cast(guardrails.LlmSanitizeResponseGuardrail, lambda response: object()), + cast(guardrails.LlmSanitizeResponseGuardrail, lambda response, context: object()), ) try: handle = llm.call("llm_sanitize_resp_bad", make_request()) @@ -239,7 +362,29 @@ def test_sanitize_response_invalid_return_falls_back_to_original_output(self): subscribers.deregister("py_llm_sanitize_resp_bad_sub") end = _llm_event(events, "llm_sanitize_resp_bad", "end") - assert end.data == {"ok": True} + assert end.data is None + assert end.annotated_response is None + + def test_sanitize_response_guardrail_accepts_scalar_json_payloads(self): + events = [] + subscribers.register("py_llm_sanitize_scalar_sub", lambda event: events.append(event)) + guardrails.register_llm_sanitize_response( + "py_llm_sanitize_scalar", + 1, + lambda response, context: f"sanitized:{response}", + ) + try: + handle = llm.call("llm_sanitize_scalar", make_request()) + llm.call_end(handle, "raw-response") + finally: + guardrails.deregister_llm_sanitize_response("py_llm_sanitize_scalar") + try: + subscribers.flush() + finally: + subscribers.deregister("py_llm_sanitize_scalar_sub") + + end = _llm_event(events, "llm_sanitize_scalar", "end") + assert end.data == "sanitized:raw-response" def test_deregister_nonexistent(self): assert not guardrails.deregister_llm_sanitize_request("nope") diff --git a/python/tests/test_scope_local.py b/python/tests/test_scope_local.py index 654400805..d20c7382d 100644 --- a/python/tests/test_scope_local.py +++ b/python/tests/test_scope_local.py @@ -576,6 +576,20 @@ def test_deregister_nonexistent_returns_false(self): class TestScopeLocalLlmWrappers: + @pytest.mark.parametrize( + ("register", "callback"), + [ + (scope_local.register_llm_sanitize_request, lambda request: request), + (scope_local.register_llm_sanitize_response, lambda response: response), + (scope_local.register_llm_sanitize_request, object()), + (scope_local.register_llm_sanitize_response, object()), + ], + ) + def test_sanitizer_registration_rejects_legacy_or_uninspectable_callbacks(self, register, callback): + with scope.scope("invalid_scope_local_sanitizer", ScopeType.Agent) as handle: + with pytest.raises(TypeError, match="payload, context"): + register(handle, "invalid_scope_local_sanitizer", 1, callback) + def test_register_and_deregister_scope_local_wrappers(self): """Scope-local wrapper functions round-trip through the native API for both tool and LLM middleware.""" request = LLMRequest({}, {"messages": [], "model": "scope-local"}) @@ -602,10 +616,20 @@ async def stream_intercept(request_inner, next_fn): ) assert scope_local.deregister_tool_execution(handle, "sl_tool_exec_cov") is True - scope_local.register_llm_sanitize_request(handle, "sl_llm_req_cov", 1, lambda req: req) + scope_local.register_llm_sanitize_request( + handle, + "sl_llm_req_cov", + 1, + lambda req, _context: req, + ) assert scope_local.deregister_llm_sanitize_request(handle, "sl_llm_req_cov") is True - scope_local.register_llm_sanitize_response(handle, "sl_llm_resp_cov", 1, lambda response: response) + scope_local.register_llm_sanitize_response( + handle, + "sl_llm_resp_cov", + 1, + lambda response, _context: response, + ) assert scope_local.deregister_llm_sanitize_response(handle, "sl_llm_resp_cov") is True scope_local.register_llm_conditional_execution(handle, "sl_llm_cond_cov", 1, lambda req: None) @@ -635,7 +659,8 @@ async def test_scope_local_llm_sanitize_request_rewrites_event_input(self): events = [] request = LLMRequest({}, {"messages": [], "model": "scope-local"}) - def sanitize_request(req): + def sanitize_request(req, context): + del context return LLMRequest({"X-Scope-Local": "yes"}, req.content) with scope.scope("sl_llm_sanitize_scope", ScopeType.Agent) as handle: