nemo_relay/api/runtime/
state.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Process-global runtime state and middleware-chain builders.
5//!
6//! [`NemoRelayContextState`] owns the registries and helper methods that power
7//! the public scope, tool, and LLM APIs. Advanced integrations can use this
8//! type directly to register middleware, attach runtime extensions, and build
9//! the resolved callback chains that the higher-level API layer executes.
10
11use std::any::Any;
12use std::collections::HashMap;
13use std::future::Future;
14use std::panic::AssertUnwindSafe;
15use std::pin::Pin;
16use std::sync::atomic::{AtomicUsize, Ordering};
17use std::sync::{Arc, Mutex};
18use std::task::{Context, Poll};
19
20use futures_util::{FutureExt, Stream};
21
22use crate::api::event::{
23    BaseEvent, CategoryProfile, Event, EventCategory, EventSanitizeFields, MarkEvent,
24    ScopeCategory, ScopeEvent, llm_attributes_to_strings, scope_attributes_to_strings,
25    tool_attributes_to_strings,
26};
27use crate::api::llm::{CreateLlmHandleParams, EndLlmHandleParams};
28use crate::api::llm::{LlmHandle, LlmRequest};
29use crate::api::registry::{ExecutionIntercept, Guardrail, Intercept};
30use crate::api::runtime::ScopeStackHandle;
31use crate::api::runtime::callbacks::{
32    EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionFn, LlmExecutionNextFn,
33    LlmJsonStream, LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn,
34    LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionFn,
35    LlmStreamExecutionNextFn, LlmStreamExecutionRegistryRefs, LlmStreamInner, ToolConditionalFn,
36    ToolExecutionFn, ToolExecutionNextFn, ToolExecutionOutcomeNextFn, ToolInterceptFn,
37    ToolSanitizeFn,
38};
39use crate::api::runtime::continuation_context::{
40    MiddlewareContinuationContext, MiddlewareContinuationGuard, MiddlewareContinuationLease,
41};
42use crate::api::runtime::subscriber_dispatcher;
43use crate::api::scope::{CreateScopeHandleParams, EndScopeHandleParams, ScopeHandle, ScopeType};
44use crate::api::shared::snapshot_event_sanitizers;
45use crate::api::tool::ToolHandle;
46use crate::api::tool::{
47    CreateToolHandleParams, EndToolHandleParams, ToolExecutionInterceptOutcome,
48};
49use crate::codec::request::AnnotatedLlmRequest;
50use crate::codec::response::AnnotatedLlmResponse;
51use crate::context::registries::{
52    merge_execution_intercept_callables, merge_guardrail_entries, merge_intercept_entries,
53};
54use crate::error::FlowError;
55use crate::json::{Json, merge_json};
56use crate::registry::SortedRegistry;
57use chrono::{Duration, Utc};
58use serde_json::json;
59use uuid::Uuid;
60
61struct ContinuationGuardedLlmStream {
62    inner: LlmJsonStream,
63    guard: Option<MiddlewareContinuationGuard>,
64}
65
66struct ContextualizedLlmStream {
67    inner: LlmJsonStream,
68    context: MiddlewareContinuationContext,
69}
70
71impl Stream for ContextualizedLlmStream {
72    type Item = crate::error::Result<Json>;
73
74    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
75        let this = self.get_mut();
76        let context = this.context.clone();
77        let inner = &mut this.inner;
78        let future = context.run(futures_util::future::poll_fn(|inner_cx| {
79            Pin::new(&mut *inner).poll_next(inner_cx)
80        }));
81        tokio::pin!(future);
82        future.poll(cx)
83    }
84}
85
86impl LlmStreamInner for ContextualizedLlmStream {
87    fn terminalize(self: Pin<&mut Self>) {
88        self.get_mut().inner.terminalize();
89    }
90
91    fn close(
92        self: Pin<&mut Self>,
93    ) -> Pin<Box<dyn Future<Output = crate::error::Result<()>> + Send + '_>> {
94        Box::pin(async move {
95            let this = self.get_mut();
96            let context = this.context.clone();
97            context.run(this.inner.close()).await
98        })
99    }
100}
101
102pub(crate) fn contextualize_stream(
103    stream: LlmJsonStream,
104    context: MiddlewareContinuationContext,
105) -> LlmJsonStream {
106    LlmJsonStream::from_closeable(ContextualizedLlmStream {
107        inner: stream,
108        context,
109    })
110}
111
112impl Stream for ContinuationGuardedLlmStream {
113    type Item = crate::error::Result<Json>;
114
115    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
116        let this = self.get_mut();
117        let result = Pin::new(&mut this.inner).poll_next(cx);
118        if matches!(&result, Poll::Ready(None)) {
119            this.guard.take();
120        }
121        result
122    }
123}
124
125impl LlmStreamInner for ContinuationGuardedLlmStream {
126    fn terminalize(self: Pin<&mut Self>) {
127        let this = self.get_mut();
128        this.guard.take();
129        this.inner.terminalize();
130    }
131
132    fn close(
133        self: Pin<&mut Self>,
134    ) -> Pin<Box<dyn Future<Output = crate::error::Result<()>> + Send + '_>> {
135        Box::pin(async move {
136            let this = self.get_mut();
137            let guard = this.guard.take();
138            let result = this.inner.close().await;
139            drop(guard);
140            result
141        })
142    }
143}
144
145fn guard_stream_continuation(
146    stream: LlmJsonStream,
147    guard: MiddlewareContinuationGuard,
148) -> LlmJsonStream {
149    LlmJsonStream::from_closeable(ContinuationGuardedLlmStream {
150        inner: stream,
151        guard: Some(guard),
152    })
153}
154
155struct GuardrailScopeCompletion<'a> {
156    handle: Option<ScopeHandle>,
157    subscribers: &'a [EventSubscriberFn],
158    scope_stack: ScopeStackHandle,
159    pending_publication: Option<subscriber_dispatcher::PendingPublication>,
160}
161
162impl GuardrailScopeCompletion<'_> {
163    fn new(
164        handle: ScopeHandle,
165        subscribers: &[EventSubscriberFn],
166        scope_stack: ScopeStackHandle,
167    ) -> GuardrailScopeCompletion<'_> {
168        GuardrailScopeCompletion {
169            handle: Some(handle),
170            subscribers,
171            scope_stack,
172            pending_publication: (!subscribers.is_empty())
173                .then(subscriber_dispatcher::register_pending_publication)
174                .flatten(),
175        }
176    }
177
178    fn finish(mut self, output: Json) {
179        let handle = self.handle.take().expect("guardrail scope handle");
180        NemoRelayContextState::emit_guardrail_scope_end(
181            &handle,
182            output,
183            self.subscribers,
184            self.scope_stack.clone(),
185        );
186        drop(self.pending_publication.take());
187    }
188}
189
190impl Drop for GuardrailScopeCompletion<'_> {
191    fn drop(&mut self) {
192        let Some(handle) = self.handle.take() else {
193            return;
194        };
195        NemoRelayContextState::emit_guardrail_scope_end(
196            &handle,
197            json!({
198                "allowed": false,
199                "cancelled": true,
200                "error": "guardrail evaluation cancelled",
201            }),
202            self.subscribers,
203            self.scope_stack.clone(),
204        );
205        drop(self.pending_publication.take());
206    }
207}
208
209/// Process-global runtime state backing middleware and event emission.
210///
211/// The public API layer stores one shared instance of this type for the
212/// process. It contains global middleware registries, lifecycle subscribers,
213/// and arbitrary extension slots used by bindings or integrations.
214pub struct NemoRelayContextState {
215    /// Global mark event field sanitizers.
216    pub(crate) mark_sanitize_guardrails: SortedRegistry<Guardrail<EventSanitizeFn>>,
217    /// Global scope-start event field sanitizers.
218    pub(crate) scope_sanitize_start_guardrails: SortedRegistry<Guardrail<EventSanitizeFn>>,
219    /// Global scope-end event field sanitizers.
220    pub(crate) scope_sanitize_end_guardrails: SortedRegistry<Guardrail<EventSanitizeFn>>,
221    /// Global tool request sanitizers applied to emitted tool-start payloads.
222    pub(crate) tool_sanitize_request_guardrails: SortedRegistry<Guardrail<ToolSanitizeFn>>,
223    /// Global tool response sanitizers applied to emitted tool-end payloads.
224    pub(crate) tool_sanitize_response_guardrails: SortedRegistry<Guardrail<ToolSanitizeFn>>,
225    /// Global tool guardrails that can reject execution before the callback runs.
226    pub(crate) tool_conditional_execution_guardrails: SortedRegistry<Guardrail<ToolConditionalFn>>,
227    /// Global tool request intercepts that can rewrite arguments before execution.
228    pub(crate) tool_request_intercepts: SortedRegistry<Intercept<ToolInterceptFn>>,
229    /// Global tool execution intercepts that wrap or replace callback execution.
230    pub(crate) tool_execution_intercepts: SortedRegistry<ExecutionIntercept<ToolExecutionFn>>,
231    /// Global LLM request sanitizers applied to emitted LLM-start payloads.
232    pub(crate) llm_sanitize_request_guardrails: SortedRegistry<Guardrail<LlmSanitizeRequestFn>>,
233    /// Global LLM response sanitizers applied to emitted LLM-end payloads.
234    pub(crate) llm_sanitize_response_guardrails: SortedRegistry<Guardrail<LlmSanitizeResponseFn>>,
235    /// Global LLM guardrails that can reject execution before the provider callback runs.
236    pub(crate) llm_conditional_execution_guardrails: SortedRegistry<Guardrail<LlmConditionalFn>>,
237    /// Global LLM request intercepts that can rewrite or annotate requests.
238    pub(crate) llm_request_intercepts: SortedRegistry<Intercept<LlmRequestInterceptFn>>,
239    /// Global non-streaming LLM execution intercepts that wrap callback execution.
240    pub(crate) llm_execution_intercepts: SortedRegistry<ExecutionIntercept<LlmExecutionFn>>,
241    /// Global streaming LLM execution intercepts that wrap stream-producing callbacks.
242    pub(crate) llm_stream_execution_intercepts:
243        SortedRegistry<ExecutionIntercept<LlmStreamExecutionFn>>,
244    /// Global lifecycle subscribers notified after runtime events are emitted.
245    pub(crate) event_subscribers: HashMap<String, EventSubscriberFn>,
246    /// Whether LLM start events retain complete sanitized request payloads.
247    pub(crate) observability_full_payloads_enabled: bool,
248    /// Arbitrary binding- or integration-specific runtime extensions.
249    pub(crate) extensions: HashMap<String, Box<dyn Any + Send + Sync>>,
250}
251
252impl NemoRelayContextState {
253    /// Create an empty runtime state with no registered middleware.
254    ///
255    /// # Returns
256    /// A [`NemoRelayContextState`] with empty registries, no subscribers, and no
257    /// extensions.
258    pub fn new() -> Self {
259        Self {
260            mark_sanitize_guardrails: SortedRegistry::new(),
261            scope_sanitize_start_guardrails: SortedRegistry::new(),
262            scope_sanitize_end_guardrails: SortedRegistry::new(),
263            tool_sanitize_request_guardrails: SortedRegistry::new(),
264            tool_sanitize_response_guardrails: SortedRegistry::new(),
265            tool_conditional_execution_guardrails: SortedRegistry::new(),
266            tool_request_intercepts: SortedRegistry::new(),
267            tool_execution_intercepts: SortedRegistry::new(),
268            llm_sanitize_request_guardrails: SortedRegistry::new(),
269            llm_sanitize_response_guardrails: SortedRegistry::new(),
270            llm_conditional_execution_guardrails: SortedRegistry::new(),
271            llm_request_intercepts: SortedRegistry::new(),
272            llm_execution_intercepts: SortedRegistry::new(),
273            llm_stream_execution_intercepts: SortedRegistry::new(),
274            event_subscribers: HashMap::new(),
275            observability_full_payloads_enabled: false,
276            extensions: HashMap::new(),
277        }
278    }
279
280    /// Store an arbitrary runtime extension under `key`.
281    ///
282    /// Extensions let bindings or integrations attach shared state to the
283    /// process-global runtime without adding new first-class fields.
284    ///
285    /// # Parameters
286    /// - `key`: Stable identifier for the extension slot.
287    /// - `value`: Typed extension value to store.
288    pub fn set_extension<T: Any + Send + Sync>(&mut self, key: impl Into<String>, value: T) {
289        self.extensions.insert(key.into(), Box::new(value));
290    }
291
292    /// Borrow a typed runtime extension by key.
293    ///
294    /// # Parameters
295    /// - `key`: Extension slot name.
296    ///
297    /// # Returns
298    /// `Some(&T)` when an extension exists under `key` with the requested type
299    /// and `None` otherwise.
300    pub fn get_extension<T: Any + Send + Sync>(&self, key: &str) -> Option<&T> {
301        self.extensions
302            .get(key)
303            .and_then(|value| value.downcast_ref::<T>())
304    }
305
306    /// Mutably borrow a typed runtime extension by key.
307    ///
308    /// # Parameters
309    /// - `key`: Extension slot name.
310    ///
311    /// # Returns
312    /// `Some(&mut T)` when an extension exists under `key` with the requested
313    /// type and `None` otherwise.
314    pub fn get_extension_mut<T: Any + Send + Sync>(&mut self, key: &str) -> Option<&mut T> {
315        self.extensions
316            .get_mut(key)
317            .and_then(|value| value.downcast_mut::<T>())
318    }
319
320    /// Remove a runtime extension by key.
321    ///
322    /// # Parameters
323    /// - `key`: Extension slot name.
324    ///
325    /// # Returns
326    /// `true` when an extension was removed and `false` when no extension was
327    /// stored under `key`.
328    pub fn remove_extension(&mut self, key: &str) -> bool {
329        self.extensions.remove(key).is_some()
330    }
331
332    /// Combine global and scope-local subscribers into one delivery list.
333    ///
334    /// # Parameters
335    /// - `scope_local_subscribers`: Subscribers collected from the active scope
336    ///   stack.
337    ///
338    /// # Returns
339    /// A vector containing all global subscribers followed by the provided
340    /// scope-local subscribers.
341    pub(crate) fn collect_event_subscribers(
342        &self,
343        scope_local_subscribers: &[EventSubscriberFn],
344    ) -> Vec<EventSubscriberFn> {
345        let mut subscribers =
346            Vec::with_capacity(self.event_subscribers.len() + scope_local_subscribers.len());
347        subscribers.extend(self.event_subscribers.values().cloned());
348        subscribers.extend(scope_local_subscribers.iter().cloned());
349        subscribers
350    }
351
352    /// Deliver an event to every subscriber in order.
353    ///
354    /// # Parameters
355    /// - `event`: Fully constructed lifecycle event to deliver.
356    /// - `subscribers`: Subscribers that should observe the event.
357    #[cfg(test)]
358    pub(crate) fn emit_event(event: &Event, subscribers: &[EventSubscriberFn]) {
359        let _ = subscriber_dispatcher::dispatch_event(event, subscribers);
360    }
361
362    /// Build a standalone mark event.
363    ///
364    /// # Parameters
365    /// - `params`: A pre-built [`MarkEvent`] to wrap in an [`Event`].
366    ///
367    /// # Returns
368    /// A mark [`Event`] containing the provided [`MarkEvent`].
369    pub fn create_event(&self, params: MarkEvent) -> Event {
370        Event::Mark(params)
371    }
372
373    /// Create a new scope handle.
374    ///
375    /// # Parameters
376    /// - `name`: Human-readable scope name.
377    /// - `parent_uuid`: Optional parent scope UUID.
378    /// - `scope_type`: Semantic category of the scope.
379    /// - `attributes`: Scope attribute bitflags.
380    /// - `data`: Optional application payload stored on the handle.
381    /// - `metadata`: Optional metadata stored on the handle.
382    /// - `timestamp`: Optional handle start time. When omitted, the current
383    ///   UTC time is used.
384    ///
385    /// # Returns
386    /// A new [`ScopeHandle`] with a fresh UUID.
387    pub fn create_scope_handle(&self, params: CreateScopeHandleParams<'_>) -> ScopeHandle {
388        ScopeHandle::builder()
389            .name(params.name)
390            .scope_type(params.scope_type)
391            .started_at(params.timestamp.unwrap_or_else(Utc::now))
392            .attributes(params.attributes)
393            .parent_uuid_opt(params.parent_uuid)
394            .data_opt(params.data)
395            .metadata_opt(params.metadata)
396            .build()
397    }
398
399    /// Build a scope-start event from a handle.
400    ///
401    /// # Parameters
402    /// - `handle`: Scope handle to serialize into an event.
403    /// - `data`: Optional semantic input payload exported on the start event.
404    ///
405    /// # Returns
406    /// A scope-start [`Event`] derived from the provided handle.
407    pub fn build_scope_start_event(&self, handle: &ScopeHandle, data: Option<Json>) -> Event {
408        Event::Scope(ScopeEvent::new(
409            BaseEvent::builder()
410                .parent_uuid_opt(handle.parent_uuid)
411                .uuid(handle.uuid)
412                .timestamp(handle.started_at)
413                .name(handle.name.as_str())
414                .data_opt(data)
415                .metadata_opt(handle.metadata.clone())
416                .build(),
417            ScopeCategory::Start,
418            scope_attributes_to_strings(handle.attributes),
419            EventCategory::from(handle.scope_type),
420            None,
421        ))
422    }
423
424    /// Build a scope-end event from a handle.
425    ///
426    /// # Parameters
427    /// - `handle`: Scope handle to serialize into an event.
428    /// - `data`: Optional data payload returned from the scope.
429    /// - `metadata`: Optional metadata payload merged over `handle.metadata`.
430    ///
431    /// # Returns
432    /// A scope-end [`Event`] derived from the provided handle.
433    pub fn end_scope_handle(
434        &self,
435        handle: &ScopeHandle,
436        data: Option<Json>,
437        metadata: Option<Json>,
438    ) -> Event {
439        self.build_scope_end_event(
440            EndScopeHandleParams::builder()
441                .handle(handle)
442                .data_opt(data)
443                .metadata_opt(metadata)
444                .build(),
445        )
446    }
447
448    /// Build a scope-end event from builder parameters.
449    ///
450    /// The `metadata` payload is merged over the metadata already stored on
451    /// the handle.
452    ///
453    /// # Parameters
454    /// - `params`: Scope end-event builder parameters.
455    ///
456    /// # Returns
457    /// A scope-end [`Event`] derived from the provided parameters.
458    pub fn build_scope_end_event(&self, params: EndScopeHandleParams<'_>) -> Event {
459        let handle = params.handle;
460        Event::Scope(ScopeEvent::new(
461            BaseEvent::builder()
462                .parent_uuid_opt(handle.parent_uuid)
463                .uuid(handle.uuid)
464                .timestamp(
465                    params
466                        .timestamp
467                        .unwrap_or_else(|| end_timestamp_after(handle.started_at)),
468                )
469                .name(handle.name.as_str())
470                .data_opt(params.data)
471                .metadata_opt(merge_json(handle.metadata.clone(), params.metadata))
472                .build(),
473            ScopeCategory::End,
474            scope_attributes_to_strings(handle.attributes),
475            EventCategory::from(handle.scope_type),
476            None,
477        ))
478    }
479
480    /// Create a new tool handle.
481    ///
482    /// # Parameters
483    /// - `name`: Tool name recorded on emitted events.
484    /// - `parent_uuid`: Optional parent scope UUID.
485    /// - `attributes`: Tool attribute bitflags.
486    /// - `data`: Optional application payload stored on the handle.
487    /// - `metadata`: Optional metadata stored on the handle.
488    /// - `tool_call_id`: Optional provider-specific correlation identifier.
489    /// - `timestamp`: Optional handle start time. When omitted, the current
490    ///   UTC time is used.
491    ///
492    /// # Returns
493    /// A new [`ToolHandle`] with a fresh UUID.
494    pub fn create_tool_handle(&self, params: CreateToolHandleParams<'_>) -> ToolHandle {
495        ToolHandle::builder()
496            .name(params.name)
497            .started_at(params.timestamp.unwrap_or_else(Utc::now))
498            .attributes(params.attributes)
499            .parent_uuid_opt(params.parent_uuid)
500            .data_opt(params.data)
501            .metadata_opt(params.metadata)
502            .tool_call_id_opt(params.tool_call_id)
503            .build()
504    }
505
506    /// Build a tool-start event from a handle.
507    ///
508    /// # Parameters
509    /// - `handle`: Tool handle to serialize into an event.
510    /// - `data`: Optional tool input payload.
511    ///
512    /// # Returns
513    /// A tool-start [`Event`] derived from the provided handle.
514    pub fn build_tool_start_event(&self, handle: &ToolHandle, data: Option<Json>) -> Event {
515        Event::Scope(ScopeEvent::new(
516            BaseEvent::builder()
517                .parent_uuid_opt(handle.parent_uuid)
518                .uuid(handle.uuid)
519                .timestamp(handle.started_at)
520                .name(handle.name.as_str())
521                .data_opt(data)
522                .metadata_opt(handle.metadata.clone())
523                .build(),
524            ScopeCategory::Start,
525            tool_attributes_to_strings(handle.attributes),
526            EventCategory::tool(),
527            Some(
528                CategoryProfile::builder()
529                    .tool_call_id_opt(handle.tool_call_id.clone())
530                    .build(),
531            ),
532        ))
533    }
534
535    /// Build a tool-end event from a handle and optional overrides.
536    ///
537    /// # Parameters
538    /// - `handle`: Tool handle to serialize into an event.
539    /// - `data`: Optional end-event data payload.
540    /// - `metadata`: Optional metadata payload merged over `handle.metadata`.
541    ///
542    /// # Returns
543    /// A tool-end [`Event`] derived from the provided handle.
544    pub fn end_tool_handle(
545        &self,
546        handle: &ToolHandle,
547        data: Option<Json>,
548        metadata: Option<Json>,
549    ) -> Event {
550        self.build_tool_end_event(
551            EndToolHandleParams::builder()
552                .handle(handle)
553                .data_opt(data)
554                .metadata_opt(metadata)
555                .build(),
556        )
557    }
558
559    /// Build a tool-end event from builder parameters.
560    ///
561    /// The `metadata` payload is merged over the metadata already stored on
562    /// the handle.
563    ///
564    /// # Parameters
565    /// - `params`: Tool end-event builder parameters.
566    ///
567    /// # Returns
568    /// A tool-end [`Event`] derived from the provided parameters.
569    pub fn build_tool_end_event(&self, params: EndToolHandleParams<'_>) -> Event {
570        let handle = params.handle;
571        Event::Scope(ScopeEvent::new(
572            BaseEvent::builder()
573                .parent_uuid_opt(handle.parent_uuid)
574                .uuid(handle.uuid)
575                .timestamp(
576                    params
577                        .timestamp
578                        .unwrap_or_else(|| end_timestamp_after(handle.started_at)),
579                )
580                .name(handle.name.as_str())
581                .data_opt(params.data)
582                .metadata_opt(merge_json(handle.metadata.clone(), params.metadata))
583                .build(),
584            ScopeCategory::End,
585            tool_attributes_to_strings(handle.attributes),
586            EventCategory::tool(),
587            Some(
588                CategoryProfile::builder()
589                    .tool_call_id_opt(handle.tool_call_id.clone())
590                    .build(),
591            ),
592        ))
593    }
594
595    /// Create a new LLM handle.
596    ///
597    /// # Parameters
598    /// - `name`: Logical provider or model family name. Gateway-managed LLM
599    ///   calls use provider route names such as `anthropic.messages`, which
600    ///   become the emitted event name.
601    /// - `parent_uuid`: Optional parent scope UUID.
602    /// - `attributes`: LLM attribute bitflags.
603    /// - `data`: Optional application payload stored on the handle.
604    /// - `metadata`: Optional metadata stored on the handle.
605    /// - `model_name`: Optional normalized model name stored on the handle.
606    /// - `timestamp`: Optional handle start time. When omitted, the current
607    ///   UTC time is used.
608    ///
609    /// # Returns
610    /// A new [`LlmHandle`] with a fresh UUID.
611    pub fn create_llm_handle(&self, params: CreateLlmHandleParams<'_>) -> LlmHandle {
612        LlmHandle::builder()
613            .name(params.name)
614            .started_at(params.timestamp.unwrap_or_else(Utc::now))
615            .attributes(params.attributes)
616            .parent_uuid_opt(params.parent_uuid)
617            .data_opt(params.data)
618            .metadata_opt(params.metadata)
619            .model_name_opt(params.model_name)
620            .build()
621    }
622
623    /// Build an LLM-start event from a handle.
624    ///
625    /// # Parameters
626    /// - `handle`: LLM handle to serialize into an event.
627    /// - `data`: Sanitized LLM request payload.
628    /// - `annotated_request`: Optional normalized request annotation.
629    ///
630    /// # Returns
631    /// An LLM-start [`Event`] derived from the provided handle.
632    pub fn build_llm_start_event(
633        &self,
634        handle: &LlmHandle,
635        data: Option<Json>,
636        annotated_request: Option<Arc<AnnotatedLlmRequest>>,
637    ) -> Event {
638        Event::Scope(ScopeEvent::new(
639            BaseEvent::builder()
640                .parent_uuid_opt(handle.parent_uuid)
641                .uuid(handle.uuid)
642                .timestamp(handle.started_at)
643                .name(handle.name.as_str())
644                .data_opt(data)
645                .metadata_opt(handle.metadata.clone())
646                .build(),
647            ScopeCategory::Start,
648            llm_attributes_to_strings(handle.attributes),
649            EventCategory::llm(),
650            Some(
651                CategoryProfile::builder()
652                    .model_name_opt(handle.model_name.clone())
653                    .annotated_request_opt(annotated_request)
654                    .build(),
655            ),
656        ))
657    }
658
659    /// Build an LLM-end event from a handle and optional overrides.
660    ///
661    /// # Parameters
662    /// - `handle`: LLM handle to serialize into an event.
663    /// - `data`: Sanitized LLM response payload.
664    /// - `metadata`: Optional metadata payload merged over `handle.metadata`.
665    /// - `annotated_response`: Optional normalized response annotation.
666    ///
667    /// # Returns
668    /// An LLM-end [`Event`] derived from the provided handle.
669    pub fn end_llm_handle(
670        &self,
671        handle: &LlmHandle,
672        data: Option<Json>,
673        metadata: Option<Json>,
674        annotated_response: Option<Arc<AnnotatedLlmResponse>>,
675    ) -> Event {
676        self.build_llm_end_event(
677            EndLlmHandleParams::builder()
678                .handle(handle)
679                .data_opt(data)
680                .metadata_opt(metadata)
681                .annotated_response_opt(annotated_response)
682                .build(),
683        )
684    }
685
686    /// Build an LLM-end event from builder parameters.
687    ///
688    /// The `metadata` payload is merged over the metadata already stored on
689    /// the handle.
690    ///
691    /// # Parameters
692    /// - `params`: LLM end-event builder parameters.
693    ///
694    /// # Returns
695    /// An LLM-end [`Event`] derived from the provided parameters.
696    pub fn build_llm_end_event(&self, params: EndLlmHandleParams<'_>) -> Event {
697        let handle = params.handle;
698        Event::Scope(ScopeEvent::new(
699            BaseEvent::builder()
700                .parent_uuid_opt(handle.parent_uuid)
701                .uuid(handle.uuid)
702                .timestamp(
703                    params
704                        .timestamp
705                        .unwrap_or_else(|| end_timestamp_after(handle.started_at)),
706                )
707                .name(handle.name.as_str())
708                .data_opt(params.data)
709                .metadata_opt(merge_json(handle.metadata.clone(), params.metadata))
710                .build(),
711            ScopeCategory::End,
712            llm_attributes_to_strings(handle.attributes),
713            EventCategory::llm(),
714            Some(
715                CategoryProfile::builder()
716                    .model_name_opt(handle.model_name.clone())
717                    .annotated_response_opt(params.annotated_response)
718                    .build(),
719            ),
720        ))
721    }
722
723    fn emit_guardrail_scope_start(
724        name: &str,
725        parent_uuid: Option<Uuid>,
726        metadata: Option<Json>,
727        input: Json,
728        subscribers: &[EventSubscriberFn],
729        scope_stack: ScopeStackHandle,
730    ) -> ScopeHandle {
731        let handle = ScopeHandle::builder()
732            .name(name)
733            .scope_type(ScopeType::Guardrail)
734            .parent_uuid_opt(parent_uuid)
735            .metadata_opt(metadata)
736            .build();
737        let event = Event::Scope(ScopeEvent::new(
738            BaseEvent::builder()
739                .parent_uuid_opt(handle.parent_uuid)
740                .uuid(handle.uuid)
741                .timestamp(handle.started_at)
742                .name(handle.name.as_str())
743                .data(input)
744                .metadata_opt(handle.metadata.clone())
745                .build(),
746            ScopeCategory::Start,
747            scope_attributes_to_strings(handle.attributes),
748            EventCategory::from(handle.scope_type),
749            None,
750        ));
751        let sanitizers = snapshot_event_sanitizers(&event, &scope_stack).unwrap_or_default();
752        subscriber_dispatcher::dispatch_sanitized_event(
753            event,
754            sanitizers,
755            subscribers,
756            scope_stack,
757        );
758        handle
759    }
760
761    fn emit_guardrail_scope_end(
762        handle: &ScopeHandle,
763        output: Json,
764        subscribers: &[EventSubscriberFn],
765        scope_stack: ScopeStackHandle,
766    ) {
767        let event = Event::Scope(ScopeEvent::new(
768            BaseEvent::builder()
769                .parent_uuid_opt(handle.parent_uuid)
770                .uuid(handle.uuid)
771                .timestamp(end_timestamp_after(handle.started_at))
772                .name(handle.name.as_str())
773                .data(output)
774                .metadata_opt(handle.metadata.clone())
775                .build(),
776            ScopeCategory::End,
777            scope_attributes_to_strings(handle.attributes),
778            EventCategory::from(handle.scope_type),
779            None,
780        ));
781        let sanitizers = snapshot_event_sanitizers(&event, &scope_stack).unwrap_or_default();
782        subscriber_dispatcher::dispatch_sanitized_event(
783            event,
784            sanitizers,
785            subscribers,
786            scope_stack,
787        );
788    }
789
790    /// Snapshot event sanitizer entries in priority order.
791    pub(crate) fn event_sanitize_entries(
792        global: &SortedRegistry<Guardrail<EventSanitizeFn>>,
793        scope_locals: &[&SortedRegistry<Guardrail<EventSanitizeFn>>],
794    ) -> Vec<Guardrail<EventSanitizeFn>> {
795        merge_guardrail_entries(global, scope_locals)
796            .into_iter()
797            .cloned()
798            .collect()
799    }
800
801    /// Apply an event sanitizer snapshot to the mutable observability fields.
802    pub(crate) async fn event_sanitize_snapshot_chain(
803        mut event: Event,
804        entries: &[Guardrail<EventSanitizeFn>],
805    ) -> Event {
806        for entry in entries {
807            let fields = event.sanitize_fields();
808            let callback = Arc::clone(&entry.payload);
809            let context = Arc::new(event);
810            let callback_context = Arc::clone(&context);
811            let outcome = AssertUnwindSafe(async move { callback(callback_context, fields).await })
812                .catch_unwind()
813                .await;
814            event = Arc::try_unwrap(context).unwrap_or_else(|context| (*context).clone());
815            match outcome {
816                Ok(Ok(fields)) => event.apply_sanitize_fields(fields),
817                Ok(Err(_error)) => {
818                    log::error!(
819                        target: "nemo_relay.runtime",
820                        event = "event_sanitizer_failed",
821                        sanitizer = entry.name.as_str(),
822                        event_name = event.name();
823                        "Event sanitizer failed; clearing observability fields"
824                    );
825                    event.apply_sanitize_fields(EventSanitizeFields::default());
826                    break;
827                }
828                Err(_) => {
829                    log::error!(
830                        target: "nemo_relay.runtime",
831                        event = "event_sanitizer_panicked",
832                        sanitizer = entry.name.as_str(),
833                        event_name = event.name();
834                        "Event sanitizer panicked; clearing observability fields"
835                    );
836                    event.apply_sanitize_fields(EventSanitizeFields::default());
837                    break;
838                }
839            }
840        }
841        event
842    }
843
844    /// Snapshot tool request sanitizers in priority order.
845    ///
846    /// # Parameters
847    /// - `scope_locals`: Scope-local sanitizer registries collected from the
848    ///   active scope stack.
849    ///
850    /// # Returns
851    /// Named sanitizer snapshots that can be evaluated after registry locks
852    /// are released.
853    pub(crate) fn tool_sanitize_request_entries(
854        &self,
855        scope_locals: &[&SortedRegistry<Guardrail<ToolSanitizeFn>>],
856    ) -> Vec<Guardrail<ToolSanitizeFn>> {
857        merge_guardrail_entries(&self.tool_sanitize_request_guardrails, scope_locals)
858            .into_iter()
859            .cloned()
860            .collect()
861    }
862
863    /// Run a snapshot of tool request sanitizers in priority order.
864    ///
865    /// # Parameters
866    /// - `name`: Tool name associated with the request.
867    /// - `args`: Raw tool arguments to sanitize for observability.
868    /// - `entries`: Sanitizer snapshots to evaluate.
869    ///
870    /// # Returns
871    /// The sanitized JSON payload after every provided guardrail has run, or
872    /// `None` when a sanitizer failure omits the payload.
873    pub(crate) async fn tool_sanitize_request_snapshot_chain(
874        name: &str,
875        args: Json,
876        entries: &[Guardrail<ToolSanitizeFn>],
877    ) -> Option<Json> {
878        let mut value = Some(args);
879        for entry in entries {
880            if let Some(current) = value.take() {
881                let callback = Arc::clone(&entry.payload);
882                let callback_name = name.to_string();
883                match AssertUnwindSafe(async move { callback(callback_name, current).await })
884                    .catch_unwind()
885                    .await
886                {
887                    Ok(Ok(next)) => value = Some(next),
888                    Ok(Err(_error)) => log::error!(
889                        target: "nemo_relay.runtime",
890                        event = "tool_request_sanitizer_failed",
891                        sanitizer = entry.name.as_str(),
892                        tool_name = name;
893                        "Tool request sanitizer failed; omitting the observability payload"
894                    ),
895                    Err(_) => log::error!(
896                        target: "nemo_relay.runtime",
897                        event = "tool_request_sanitizer_panicked",
898                        sanitizer = entry.name.as_str(),
899                        tool_name = name;
900                        "Tool request sanitizer panicked; omitting the observability payload"
901                    ),
902                }
903            }
904        }
905        value
906    }
907
908    /// Snapshot tool response sanitizers in priority order.
909    ///
910    /// # Parameters
911    /// - `scope_locals`: Scope-local sanitizer registries collected from the
912    ///   active scope stack.
913    ///
914    /// # Returns
915    /// Named sanitizer snapshots that can be evaluated after registry locks
916    /// are released.
917    pub(crate) fn tool_sanitize_response_entries(
918        &self,
919        scope_locals: &[&SortedRegistry<Guardrail<ToolSanitizeFn>>],
920    ) -> Vec<Guardrail<ToolSanitizeFn>> {
921        merge_guardrail_entries(&self.tool_sanitize_response_guardrails, scope_locals)
922            .into_iter()
923            .cloned()
924            .collect()
925    }
926
927    /// Run a snapshot of tool response sanitizers in priority order.
928    ///
929    /// # Parameters
930    /// - `name`: Tool name associated with the response.
931    /// - `result`: Raw tool result to sanitize for observability.
932    /// - `entries`: Sanitizer snapshots to evaluate.
933    ///
934    /// # Returns
935    /// The sanitized JSON payload after every provided guardrail has run, or
936    /// `None` when a sanitizer failure omits the payload.
937    pub(crate) async fn tool_sanitize_response_snapshot_chain(
938        name: &str,
939        result: Json,
940        entries: &[Guardrail<ToolSanitizeFn>],
941    ) -> Option<Json> {
942        let mut value = Some(result);
943        for entry in entries {
944            if let Some(current) = value.take() {
945                let callback = Arc::clone(&entry.payload);
946                let callback_name = name.to_string();
947                match AssertUnwindSafe(async move { callback(callback_name, current).await })
948                    .catch_unwind()
949                    .await
950                {
951                    Ok(Ok(next)) => value = Some(next),
952                    Ok(Err(_error)) => log::error!(
953                        target: "nemo_relay.runtime",
954                        event = "tool_response_sanitizer_failed",
955                        sanitizer = entry.name.as_str(),
956                        tool_name = name;
957                        "Tool response sanitizer failed; omitting the observability payload"
958                    ),
959                    Err(_) => log::error!(
960                        target: "nemo_relay.runtime",
961                        event = "tool_response_sanitizer_panicked",
962                        sanitizer = entry.name.as_str(),
963                        tool_name = name;
964                        "Tool response sanitizer panicked; omitting the observability payload"
965                    ),
966                }
967            }
968        }
969        value
970    }
971
972    /// Snapshot tool conditional-execution guardrails in priority order.
973    ///
974    /// # Parameters
975    /// - `scope_locals`: Scope-local conditional guardrail registries collected
976    ///   from the active scope stack.
977    ///
978    /// # Returns
979    /// Named guardrail snapshots that can be evaluated after registry locks
980    /// are released.
981    pub(crate) fn tool_conditional_execution_entries(
982        &self,
983        scope_locals: &[&SortedRegistry<Guardrail<ToolConditionalFn>>],
984    ) -> Vec<Guardrail<ToolConditionalFn>> {
985        merge_guardrail_entries(&self.tool_conditional_execution_guardrails, scope_locals)
986            .into_iter()
987            .cloned()
988            .collect()
989    }
990
991    /// Evaluate a snapshot of tool conditional-execution guardrails in priority order.
992    ///
993    /// This function emits guardrail scope start/end events while evaluating
994    /// the provided entries. Callers should pass entries snapped from the
995    /// global and scope-local registries so subscriber callbacks run without
996    /// registry locks held. If `entries` is empty, no guardrail scopes are
997    /// emitted. Guardrail start events identify the guardrail and target but
998    /// intentionally omit raw tool arguments from their event data.
999    ///
1000    /// # Parameters
1001    /// - `name`: Tool name associated with the request.
1002    /// - `args`: Tool arguments to validate.
1003    /// - `entries`: Borrowed conditional guardrail snapshots to evaluate.
1004    /// - `subscribers`: Event subscribers that should observe guardrail scope
1005    ///   start/end events.
1006    /// - `parent_uuid`: Optional parent scope UUID for emitted guardrail
1007    ///   scopes.
1008    /// - `metadata`: Optional metadata attached to emitted guardrail scopes.
1009    ///
1010    /// # Returns
1011    /// A [`Result`](crate::error::Result) containing `Ok(None)` when execution
1012    /// is allowed or `Ok(Some(reason))` when a guardrail rejects the call.
1013    ///
1014    /// # Errors
1015    /// Propagates any error returned by a guardrail callback after emitting the
1016    /// corresponding guardrail scope end event.
1017    pub(crate) async fn tool_conditional_execution_snapshot_chain(
1018        name: &str,
1019        args: &Json,
1020        entries: &[Guardrail<ToolConditionalFn>],
1021        subscribers: &[EventSubscriberFn],
1022        parent_uuid: Option<Uuid>,
1023        metadata: Option<Json>,
1024    ) -> crate::error::Result<Option<String>> {
1025        for entry in entries {
1026            let scope_stack = super::current_scope_stack();
1027            let handle = Self::emit_guardrail_scope_start(
1028                &entry.name,
1029                parent_uuid,
1030                metadata.clone(),
1031                json!({
1032                    "kind": "tool_conditional_execution",
1033                    "target_name": name,
1034                }),
1035                subscribers,
1036                scope_stack.clone(),
1037            );
1038            let completion = GuardrailScopeCompletion::new(handle, subscribers, scope_stack);
1039            let callback = Arc::clone(&entry.payload);
1040            let callback_name = name.to_string();
1041            let callback_args = args.clone();
1042            let result =
1043                match AssertUnwindSafe(async move { callback(callback_name, callback_args).await })
1044                    .catch_unwind()
1045                    .await
1046                {
1047                    Ok(result) => result,
1048                    Err(_) => Err(FlowError::Internal(format!(
1049                        "tool conditional guardrail '{}' panicked",
1050                        entry.name
1051                    ))),
1052                };
1053            let output = match &result {
1054                Ok(Some(reason)) => json!({
1055                    "allowed": false,
1056                    "rejected": true,
1057                    "rejection_reason": reason,
1058                }),
1059                Ok(None) => json!({
1060                    "allowed": true,
1061                    "rejected": false,
1062                }),
1063                Err(error) => json!({
1064                    "allowed": false,
1065                    "error": error.to_string(),
1066                }),
1067            };
1068            completion.finish(output);
1069            if let Some(error) = result? {
1070                return Ok(Some(error));
1071            }
1072        }
1073        Ok(None)
1074    }
1075
1076    /// Snapshot tool request intercepts in priority order.
1077    ///
1078    /// # Parameters
1079    /// - `scope_locals`: Scope-local request intercept registries collected
1080    ///   from the active scope stack.
1081    ///
1082    /// # Returns
1083    /// Named intercept snapshots that can be evaluated after registry locks
1084    /// are released.
1085    pub(crate) fn tool_request_intercept_entries(
1086        &self,
1087        scope_locals: &[&SortedRegistry<Intercept<ToolInterceptFn>>],
1088    ) -> Vec<Intercept<ToolInterceptFn>> {
1089        merge_intercept_entries(&self.tool_request_intercepts, scope_locals)
1090            .into_iter()
1091            .cloned()
1092            .collect()
1093    }
1094
1095    /// Run a snapshot of tool request intercepts in priority order.
1096    ///
1097    /// # Parameters
1098    /// - `name`: Tool name associated with the request.
1099    /// - `args`: Tool arguments to pass through the intercept chain.
1100    /// - `entries`: Intercept snapshots to evaluate.
1101    ///
1102    /// # Returns
1103    /// A [`Result`] containing the final JSON argument payload.
1104    ///
1105    /// # Errors
1106    /// Propagates any error returned by an intercept callback.
1107    ///
1108    /// # Notes
1109    /// If an intercept entry has `break_chain` enabled, later intercepts are
1110    /// skipped after that entry runs.
1111    pub(crate) async fn tool_request_intercepts_snapshot_chain(
1112        name: &str,
1113        args: Json,
1114        entries: &[Intercept<ToolInterceptFn>],
1115    ) -> crate::error::Result<Json> {
1116        let mut value = args;
1117        for entry in entries {
1118            let callback = Arc::clone(&entry.payload.callable);
1119            let callback_name = name.to_string();
1120            value = match AssertUnwindSafe(async move { callback(callback_name, value).await })
1121                .catch_unwind()
1122                .await
1123            {
1124                Ok(result) => result?,
1125                Err(_) => {
1126                    return Err(FlowError::Internal(format!(
1127                        "tool request intercept '{}' panicked",
1128                        entry.name
1129                    )));
1130                }
1131            };
1132            if entry.payload.break_chain {
1133                break;
1134            }
1135        }
1136        Ok(value)
1137    }
1138
1139    /// Build the composed tool execution continuation chain.
1140    ///
1141    /// # Parameters
1142    /// - `name`: Tool name passed into each execution intercept.
1143    /// - `default_fn`: Base tool callback that should run after all intercepts.
1144    /// - `scope_locals`: Scope-local execution intercept registries collected
1145    ///   from the active scope stack.
1146    ///
1147    /// # Returns
1148    /// A composed [`ToolExecutionOutcomeNextFn`] that wraps `default_fn` in
1149    /// every matching execution intercept.
1150    pub(crate) fn tool_build_execution_chain(
1151        &self,
1152        name: &str,
1153        default_fn: ToolExecutionNextFn,
1154        scope_locals: &[&SortedRegistry<ExecutionIntercept<ToolExecutionFn>>],
1155    ) -> ToolExecutionOutcomeNextFn {
1156        let matching =
1157            merge_execution_intercept_callables(&self.tool_execution_intercepts, scope_locals);
1158        let mut next: ToolExecutionOutcomeNextFn = Arc::new(move |args| {
1159            let default_fn = default_fn.clone();
1160            Box::pin(async move {
1161                default_fn(args)
1162                    .await
1163                    .map(ToolExecutionInterceptOutcome::new)
1164            })
1165        });
1166        let name = name.to_string();
1167        for (callable, _) in matching.into_iter().rev() {
1168            let current_next = next.clone();
1169            let current_name = name.clone();
1170            next = Arc::new(move |args| {
1171                let callable = callable.clone();
1172                let current_name = current_name.clone();
1173                let (continuation, continuation_guard) = MiddlewareContinuationLease::capture();
1174                let next_sequence = Arc::new(AtomicUsize::new(0));
1175                let downstream_marks = Arc::new(Mutex::new(Vec::new()));
1176                let raw_next: ToolExecutionNextFn = {
1177                    let current_next = current_next.clone();
1178                    let continuation = continuation.clone();
1179                    let next_sequence = next_sequence.clone();
1180                    let downstream_marks = downstream_marks.clone();
1181                    Arc::new(move |args| {
1182                        let sequence = next_sequence.fetch_add(1, Ordering::Relaxed);
1183                        let current_next = current_next.clone();
1184                        let invocation = continuation.begin();
1185                        let downstream_marks = downstream_marks.clone();
1186                        Box::pin(async move {
1187                            let outcome = invocation?.invoke(move || current_next(args)).await?;
1188                            downstream_marks
1189                                .lock()
1190                                .expect("tool pending mark accumulator lock poisoned")
1191                                .push((sequence, outcome.pending_marks));
1192                            Ok(outcome.result)
1193                        })
1194                    })
1195                };
1196                Box::pin(async move {
1197                    let outcome = callable(&current_name, args, raw_next).await;
1198                    drop(continuation_guard);
1199                    let mut outcome = outcome?;
1200                    let mut downstream_batches = std::mem::take(
1201                        &mut *downstream_marks
1202                            .lock()
1203                            .expect("tool pending mark accumulator lock poisoned"),
1204                    );
1205                    downstream_batches.sort_by_key(|(sequence, _)| *sequence);
1206                    let mut marks = downstream_batches
1207                        .into_iter()
1208                        .flat_map(|(_, marks)| marks)
1209                        .collect::<Vec<_>>();
1210                    marks.append(&mut outcome.pending_marks);
1211                    outcome.pending_marks = marks;
1212                    Ok(outcome)
1213                })
1214            });
1215        }
1216        next
1217    }
1218
1219    /// Snapshot LLM request sanitizers in priority order.
1220    ///
1221    /// # Parameters
1222    /// - `scope_locals`: Scope-local sanitizer registries collected from the
1223    ///   active scope stack.
1224    ///
1225    /// # Returns
1226    /// Named sanitizer snapshots that can be evaluated after registry locks
1227    /// are released.
1228    pub(crate) fn llm_sanitize_request_entries(
1229        &self,
1230        scope_locals: &[&SortedRegistry<Guardrail<LlmSanitizeRequestFn>>],
1231    ) -> Vec<Guardrail<LlmSanitizeRequestFn>> {
1232        merge_guardrail_entries(&self.llm_sanitize_request_guardrails, scope_locals)
1233            .into_iter()
1234            .cloned()
1235            .collect()
1236    }
1237
1238    /// Run a snapshot of LLM request sanitizers in priority order.
1239    ///
1240    /// # Parameters
1241    /// - `request`: Raw LLM request to sanitize for observability.
1242    /// - `entries`: Sanitizer snapshots to evaluate.
1243    ///
1244    /// # Returns
1245    /// The sanitized [`LlmRequest`] after every provided guardrail has run, or
1246    /// `None` when a sanitizer errors or panics.
1247    pub(crate) async fn llm_sanitize_request_snapshot_chain(
1248        request: LlmRequest,
1249        context: LlmSanitizeRequestContext,
1250        entries: &[Guardrail<LlmSanitizeRequestFn>],
1251    ) -> Option<LlmRequest> {
1252        let mut value = Some(request);
1253        for entry in entries {
1254            if let Some(current) = value.take() {
1255                let callback = Arc::clone(&entry.payload);
1256                let callback_value = current.clone();
1257                let callback_context = context.clone();
1258                match AssertUnwindSafe(
1259                    async move { callback(callback_value, callback_context).await },
1260                )
1261                .catch_unwind()
1262                .await
1263                {
1264                    Ok(Ok(next)) => value = next,
1265                    Ok(Err(_error)) => {
1266                        log::error!(
1267                            target: "nemo_relay.runtime",
1268                            event = "llm_request_sanitizer_failed",
1269                            sanitizer = entry.name.as_str();
1270                            "LLM request sanitizer failed; omitting the observability payload"
1271                        );
1272                    }
1273                    Err(_) => {
1274                        log::error!(
1275                            target: "nemo_relay.runtime",
1276                            event = "llm_request_sanitizer_panicked",
1277                            sanitizer = entry.name.as_str();
1278                            "LLM request sanitizer panicked; omitting the observability payload"
1279                        );
1280                    }
1281                }
1282            }
1283        }
1284        value
1285    }
1286
1287    /// Snapshot LLM response sanitizers in priority order.
1288    ///
1289    /// # Parameters
1290    /// - `scope_locals`: Scope-local sanitizer registries collected from the
1291    ///   active scope stack.
1292    ///
1293    /// # Returns
1294    /// Named sanitizer snapshots that can be evaluated after registry locks
1295    /// are released.
1296    pub(crate) fn llm_sanitize_response_entries(
1297        &self,
1298        scope_locals: &[&SortedRegistry<Guardrail<LlmSanitizeResponseFn>>],
1299    ) -> Vec<Guardrail<LlmSanitizeResponseFn>> {
1300        merge_guardrail_entries(&self.llm_sanitize_response_guardrails, scope_locals)
1301            .into_iter()
1302            .cloned()
1303            .collect()
1304    }
1305
1306    /// Run a snapshot of LLM response sanitizers in priority order.
1307    ///
1308    /// # Parameters
1309    /// - `response`: Raw response payload to sanitize for observability.
1310    /// - `entries`: Sanitizer snapshots to evaluate.
1311    ///
1312    /// # Returns
1313    /// The sanitized response payload after every provided guardrail has run,
1314    /// or `None` when a sanitizer errors or panics.
1315    pub(crate) async fn llm_sanitize_response_snapshot_chain(
1316        response: Json,
1317        context: LlmSanitizeResponseContext,
1318        entries: &[Guardrail<LlmSanitizeResponseFn>],
1319    ) -> Option<Json> {
1320        let mut value = Some(response);
1321        for entry in entries {
1322            if let Some(current) = value.take() {
1323                let callback = Arc::clone(&entry.payload);
1324                let callback_value = current.clone();
1325                let callback_context = context.clone();
1326                match AssertUnwindSafe(
1327                    async move { callback(callback_value, callback_context).await },
1328                )
1329                .catch_unwind()
1330                .await
1331                {
1332                    Ok(Ok(next)) => value = next,
1333                    Ok(Err(_error)) => {
1334                        log::error!(
1335                            target: "nemo_relay.runtime",
1336                            event = "llm_response_sanitizer_failed",
1337                            sanitizer = entry.name.as_str();
1338                            "LLM response sanitizer failed; omitting the observability payload"
1339                        );
1340                    }
1341                    Err(_) => {
1342                        log::error!(
1343                            target: "nemo_relay.runtime",
1344                            event = "llm_response_sanitizer_panicked",
1345                            sanitizer = entry.name.as_str();
1346                            "LLM response sanitizer panicked; omitting the observability payload"
1347                        );
1348                    }
1349                }
1350            }
1351        }
1352        value
1353    }
1354
1355    /// Snapshot LLM conditional-execution guardrails in priority order.
1356    ///
1357    /// # Parameters
1358    /// - `scope_locals`: Scope-local conditional guardrail registries collected
1359    ///   from the active scope stack.
1360    ///
1361    /// # Returns
1362    /// Named guardrail snapshots that can be evaluated after registry locks
1363    /// are released.
1364    pub(crate) fn llm_conditional_execution_entries(
1365        &self,
1366        scope_locals: &[&SortedRegistry<Guardrail<LlmConditionalFn>>],
1367    ) -> Vec<Guardrail<LlmConditionalFn>> {
1368        merge_guardrail_entries(&self.llm_conditional_execution_guardrails, scope_locals)
1369            .into_iter()
1370            .cloned()
1371            .collect()
1372    }
1373
1374    /// Evaluate a snapshot of LLM conditional-execution guardrails in priority order.
1375    ///
1376    /// This function emits guardrail scope start/end events while evaluating
1377    /// the provided entries. Callers should pass entries snapped from the
1378    /// global and scope-local registries so subscriber callbacks run without
1379    /// registry locks held. If `entries` is empty, no guardrail scopes are
1380    /// emitted. Guardrail start events identify the guardrail but intentionally
1381    /// omit raw LLM requests from their event data.
1382    ///
1383    /// # Parameters
1384    /// - `request`: LLM request to validate.
1385    /// - `entries`: Borrowed conditional guardrail snapshots to evaluate.
1386    /// - `subscribers`: Event subscribers that should observe guardrail scope
1387    ///   start/end events.
1388    /// - `parent_uuid`: Optional parent scope UUID for emitted guardrail
1389    ///   scopes.
1390    /// - `metadata`: Optional metadata attached to emitted guardrail scopes.
1391    ///
1392    /// # Returns
1393    /// A [`Result`](crate::error::Result) containing `Ok(None)` when execution
1394    /// is allowed or `Ok(Some(reason))` when a guardrail rejects the call.
1395    ///
1396    /// # Errors
1397    /// Propagates any error returned by a guardrail callback after emitting the
1398    /// corresponding guardrail scope end event.
1399    pub(crate) async fn llm_conditional_execution_snapshot_chain(
1400        request: &LlmRequest,
1401        entries: &[Guardrail<LlmConditionalFn>],
1402        subscribers: &[EventSubscriberFn],
1403        parent_uuid: Option<Uuid>,
1404        metadata: Option<Json>,
1405    ) -> crate::error::Result<Option<String>> {
1406        for entry in entries {
1407            let scope_stack = super::current_scope_stack();
1408            let handle = Self::emit_guardrail_scope_start(
1409                &entry.name,
1410                parent_uuid,
1411                metadata.clone(),
1412                json!({
1413                    "kind": "llm_conditional_execution",
1414                }),
1415                subscribers,
1416                scope_stack.clone(),
1417            );
1418            let completion = GuardrailScopeCompletion::new(handle, subscribers, scope_stack);
1419            let callback = Arc::clone(&entry.payload);
1420            let callback_request = request.clone();
1421            let result = match AssertUnwindSafe(async move { callback(callback_request).await })
1422                .catch_unwind()
1423                .await
1424            {
1425                Ok(result) => result,
1426                Err(_) => Err(FlowError::Internal(format!(
1427                    "LLM conditional guardrail '{}' panicked",
1428                    entry.name
1429                ))),
1430            };
1431            let output = match &result {
1432                Ok(Some(reason)) => json!({
1433                    "allowed": false,
1434                    "rejected": true,
1435                    "rejection_reason": reason,
1436                }),
1437                Ok(None) => json!({
1438                    "allowed": true,
1439                    "rejected": false,
1440                }),
1441                Err(error) => json!({
1442                    "allowed": false,
1443                    "error": error.to_string(),
1444                }),
1445            };
1446            completion.finish(output);
1447            if let Some(error) = result? {
1448                return Ok(Some(error));
1449            }
1450        }
1451        Ok(None)
1452    }
1453
1454    /// Snapshot LLM request intercepts in priority order.
1455    ///
1456    /// # Parameters
1457    /// - `scope_locals`: Scope-local request intercept registries collected
1458    ///   from the active scope stack.
1459    ///
1460    /// # Returns
1461    /// Named intercept snapshots that can be evaluated after registry locks
1462    /// are released.
1463    pub(crate) fn llm_request_intercept_entries(
1464        &self,
1465        scope_locals: &[&SortedRegistry<Intercept<LlmRequestInterceptFn>>],
1466    ) -> Vec<Intercept<LlmRequestInterceptFn>> {
1467        merge_intercept_entries(&self.llm_request_intercepts, scope_locals)
1468            .into_iter()
1469            .cloned()
1470            .collect()
1471    }
1472
1473    /// Run a snapshot of LLM request intercepts in priority order.
1474    ///
1475    /// # Parameters
1476    /// - `name`: Logical provider or model family name.
1477    /// - `request`: LLM request to pass through the intercept chain.
1478    /// - `annotated`: Optional normalized request annotation to carry through
1479    ///   the chain.
1480    /// - `entries`: Intercept snapshots to evaluate.
1481    /// - `codec_active`: Whether request content is owned by the normalized
1482    ///   annotation and must remain unchanged by callbacks.
1483    ///
1484    /// # Returns
1485    /// A [`Result`] containing the final request and annotation pair.
1486    ///
1487    /// # Errors
1488    /// Propagates any error returned by an intercept callback.
1489    ///
1490    /// # Notes
1491    /// If an intercept entry has `break_chain` enabled, later intercepts are
1492    /// skipped after that entry runs.
1493    pub(crate) async fn llm_request_intercepts_snapshot_chain(
1494        name: &str,
1495        request: LlmRequest,
1496        annotated: Option<AnnotatedLlmRequest>,
1497        entries: &[Intercept<LlmRequestInterceptFn>],
1498        codec_active: bool,
1499    ) -> crate::error::Result<crate::api::llm::LlmRequestInterceptOutcome> {
1500        Self::llm_request_intercepts_snapshot_chain_with_recorder(
1501            name,
1502            request,
1503            annotated,
1504            entries,
1505            codec_active,
1506            None,
1507        )
1508        .await
1509    }
1510
1511    /// Run a request-intercept snapshot while ingesting optimization evidence
1512    /// directly into the managed call's bounded accumulator.
1513    pub(crate) async fn llm_request_intercepts_snapshot_chain_with_recorder(
1514        name: &str,
1515        request: LlmRequest,
1516        annotated: Option<AnnotatedLlmRequest>,
1517        entries: &[Intercept<LlmRequestInterceptFn>],
1518        codec_active: bool,
1519        optimization_recorder: Option<&crate::api::optimization::LlmOptimizationRecorder>,
1520    ) -> crate::error::Result<crate::api::llm::LlmRequestInterceptOutcome> {
1521        let mut request_value = request;
1522        let mut annotated_value = annotated;
1523        let mut pending_marks = Vec::new();
1524        let mut optimization_contributions = Vec::new();
1525        for entry in entries {
1526            let input_content = request_value.content.clone();
1527            let callback = Arc::clone(&entry.payload.callable);
1528            let callback_name = name.to_string();
1529            let outcome = match AssertUnwindSafe(async move {
1530                callback(callback_name, request_value, annotated_value).await
1531            })
1532            .catch_unwind()
1533            .await
1534            {
1535                Ok(result) => result?,
1536                Err(_) => {
1537                    return Err(FlowError::Internal(format!(
1538                        "LLM request intercept '{}' panicked",
1539                        entry.name
1540                    )));
1541                }
1542            };
1543            if codec_active && outcome.request.content != input_content {
1544                return Err(crate::error::FlowError::InvalidArgument(format!(
1545                    "LLM request intercept '{}' changed request.content while a request codec is active; modify annotated_request instead",
1546                    entry.name
1547                )));
1548            }
1549            if codec_active && outcome.annotated_request.is_none() {
1550                return Err(crate::error::FlowError::InvalidArgument(format!(
1551                    "LLM request intercept '{}' omitted annotated_request while a request codec is active",
1552                    entry.name
1553                )));
1554            }
1555            request_value = outcome.request;
1556            annotated_value = outcome.annotated_request;
1557            pending_marks.extend(outcome.pending_marks);
1558            if let Some(recorder) = optimization_recorder {
1559                recorder.record_all(outcome.optimization_contributions);
1560            } else {
1561                optimization_contributions.extend(outcome.optimization_contributions);
1562            }
1563            if entry.payload.break_chain {
1564                break;
1565            }
1566        }
1567        Ok(crate::api::llm::LlmRequestInterceptOutcome {
1568            request: request_value,
1569            annotated_request: annotated_value,
1570            pending_marks,
1571            optimization_contributions,
1572        })
1573    }
1574
1575    /// Build the composed non-streaming LLM execution continuation chain.
1576    ///
1577    /// # Parameters
1578    /// - `name`: Logical provider or model family name passed into each
1579    ///   execution intercept.
1580    /// - `default_fn`: Base provider callback that should run after all
1581    ///   intercepts.
1582    /// - `scope_locals`: Scope-local execution intercept registries collected
1583    ///   from the active scope stack.
1584    ///
1585    /// # Returns
1586    /// A composed [`LlmExecutionNextFn`] that wraps `default_fn` in every
1587    /// matching execution intercept.
1588    pub(crate) fn llm_build_execution_chain(
1589        &self,
1590        name: &str,
1591        default_fn: LlmExecutionNextFn,
1592        scope_locals: &[&SortedRegistry<ExecutionIntercept<LlmExecutionFn>>],
1593    ) -> LlmExecutionNextFn {
1594        let matching =
1595            merge_execution_intercept_callables(&self.llm_execution_intercepts, scope_locals);
1596        let mut next = default_fn;
1597        let name = name.to_string();
1598        for (callable, _) in matching.into_iter().rev() {
1599            let current_next = next.clone();
1600            let current_name = name.clone();
1601            next = Arc::new(move |request| {
1602                let callable = callable.clone();
1603                let current_next = current_next.clone();
1604                let current_name = current_name.clone();
1605                Box::pin(async move {
1606                    let (continuation, continuation_guard) = MiddlewareContinuationLease::capture();
1607                    let raw_next: LlmExecutionNextFn = Arc::new(move |request| {
1608                        let invocation = continuation.begin();
1609                        let current_next = current_next.clone();
1610                        Box::pin(
1611                            async move { invocation?.invoke(move || current_next(request)).await },
1612                        )
1613                    });
1614                    let result = callable(&current_name, request, raw_next).await;
1615                    drop(continuation_guard);
1616                    result
1617                })
1618            });
1619        }
1620        next
1621    }
1622
1623    /// Build the composed streaming LLM execution continuation chain.
1624    ///
1625    /// # Parameters
1626    /// - `name`: Logical provider or model family name passed into each
1627    ///   execution intercept.
1628    /// - `default_fn`: Base stream-producing callback that should run after all
1629    ///   intercepts.
1630    /// - `scope_locals`: Scope-local execution intercept registries collected
1631    ///   from the active scope stack.
1632    ///
1633    /// # Returns
1634    /// A composed [`LlmStreamExecutionNextFn`] that wraps `default_fn` in every
1635    /// matching execution intercept.
1636    pub(crate) fn llm_stream_build_execution_chain(
1637        &self,
1638        name: &str,
1639        default_fn: LlmStreamExecutionNextFn,
1640        scope_locals: LlmStreamExecutionRegistryRefs<'_>,
1641    ) -> LlmStreamExecutionNextFn {
1642        let matching = merge_execution_intercept_callables(
1643            &self.llm_stream_execution_intercepts,
1644            scope_locals,
1645        );
1646        let mut next = default_fn;
1647        let name = name.to_string();
1648        for (callable, _) in matching.into_iter().rev() {
1649            let current_next = next.clone();
1650            let current_name = name.clone();
1651            next = Arc::new(move |request| {
1652                let callable = callable.clone();
1653                let current_next = current_next.clone();
1654                let current_name = current_name.clone();
1655                Box::pin(async move {
1656                    let (continuation, continuation_guard) = MiddlewareContinuationLease::capture();
1657                    let raw_next: LlmStreamExecutionNextFn = Arc::new(move |request| {
1658                        let invocation = continuation.begin();
1659                        let current_next = current_next.clone();
1660                        Box::pin(async move {
1661                            let invocation = invocation?;
1662                            let context = invocation.context().clone();
1663                            let stream = invocation.invoke(move || current_next(request)).await?;
1664                            Ok(contextualize_stream(stream, context))
1665                        })
1666                    });
1667                    let result = callable(&current_name, request, raw_next).await;
1668                    result.map(|stream| guard_stream_continuation(stream, continuation_guard))
1669                })
1670            });
1671        }
1672        next
1673    }
1674}
1675
1676fn end_timestamp_after(started_at: chrono::DateTime<Utc>) -> chrono::DateTime<Utc> {
1677    let now = Utc::now();
1678    std::cmp::max(now, started_at + Duration::microseconds(1))
1679}
1680
1681impl Default for NemoRelayContextState {
1682    fn default() -> Self {
1683        Self::new()
1684    }
1685}
1686
1687#[cfg(test)]
1688#[path = "../../../tests/unit/runtime_state_tests.rs"]
1689mod tests;
nemo_relay/api/runtime/state.rs

nemo_relay/api/runtime/
state.rs
