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

nemo_relay/api/runtime/
state.rs
