nemo_flow/api/runtime/

state.rs

// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Process-global runtime state and middleware-chain builders.
//!
//! [`NemoFlowContextState`] owns the registries and helper methods that power
//! the public scope, tool, and LLM APIs. Advanced integrations can use this
//! type directly to register middleware, attach runtime extensions, and build
//! the resolved callback chains that the higher-level API layer executes.

use std::any::Any;
use std::collections::HashMap;
use std::sync::Arc;

use crate::api::event::{
    BaseEvent, CategoryProfile, Event, EventCategory, MarkEvent, ScopeCategory, ScopeEvent,
    llm_attributes_to_strings, scope_attributes_to_strings, tool_attributes_to_strings,
};
use crate::api::llm::{CreateLlmHandleParams, EndLlmHandleParams};
use crate::api::llm::{LlmHandle, LlmRequest};
use crate::api::registry::{ExecutionIntercept, GuardrailEntry, Intercept};
use crate::api::runtime::callbacks::{
    EventSubscriberFn, LlmConditionalFn, LlmExecutionFn, LlmExecutionNextFn, LlmRequestInterceptFn,
    LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionFn, LlmStreamExecutionNextFn,
    LlmStreamExecutionRegistryRefs, ToolConditionalFn, ToolExecutionFn, ToolExecutionNextFn,
    ToolInterceptFn, ToolSanitizeFn,
};
use crate::api::scope::{CreateScopeHandleParams, EndScopeHandleParams, ScopeHandle};
use crate::api::tool::ToolHandle;
use crate::api::tool::{CreateToolHandleParams, EndToolHandleParams};
use crate::codec::request::AnnotatedLlmRequest;
use crate::codec::response::AnnotatedLlmResponse;
use crate::context::registries::{
    merge_execution_intercept_callables, merge_guardrail_entries, merge_intercept_entries,
};
use crate::json::{Json, merge_json};
use crate::registry::SortedRegistry;
use chrono::{Duration, Utc};

/// Process-global runtime state backing middleware and event emission.
///
/// The public API layer stores one shared instance of this type for the
/// process. It contains global middleware registries, lifecycle subscribers,
/// and arbitrary extension slots used by bindings or integrations.
pub struct NemoFlowContextState {
    /// Global tool request sanitizers applied to emitted tool-start payloads.
    pub tool_sanitize_request_guardrails: SortedRegistry<GuardrailEntry<ToolSanitizeFn>>,
    /// Global tool response sanitizers applied to emitted tool-end payloads.
    pub tool_sanitize_response_guardrails: SortedRegistry<GuardrailEntry<ToolSanitizeFn>>,
    /// Global tool guardrails that can reject execution before the callback runs.
    pub tool_conditional_execution_guardrails: SortedRegistry<GuardrailEntry<ToolConditionalFn>>,
    /// Global tool request intercepts that can rewrite arguments before execution.
    pub tool_request_intercepts: SortedRegistry<Intercept<ToolInterceptFn>>,
    /// Global tool execution intercepts that wrap or replace callback execution.
    pub tool_execution_intercepts: SortedRegistry<ExecutionIntercept<ToolExecutionFn>>,
    /// Global LLM request sanitizers applied to emitted LLM-start payloads.
    pub llm_sanitize_request_guardrails: SortedRegistry<GuardrailEntry<LlmSanitizeRequestFn>>,
    /// Global LLM response sanitizers applied to emitted LLM-end payloads.
    pub llm_sanitize_response_guardrails: SortedRegistry<GuardrailEntry<LlmSanitizeResponseFn>>,
    /// Global LLM guardrails that can reject execution before the provider callback runs.
    pub llm_conditional_execution_guardrails: SortedRegistry<GuardrailEntry<LlmConditionalFn>>,
    /// Global LLM request intercepts that can rewrite or annotate requests.
    pub llm_request_intercepts: SortedRegistry<Intercept<LlmRequestInterceptFn>>,
    /// Global non-streaming LLM execution intercepts that wrap callback execution.
    pub llm_execution_intercepts: SortedRegistry<ExecutionIntercept<LlmExecutionFn>>,
    /// Global streaming LLM execution intercepts that wrap stream-producing callbacks.
    pub llm_stream_execution_intercepts: SortedRegistry<ExecutionIntercept<LlmStreamExecutionFn>>,
    /// Global lifecycle subscribers notified after runtime events are emitted.
    pub event_subscribers: HashMap<String, EventSubscriberFn>,
    /// Arbitrary binding- or integration-specific runtime extensions.
    pub extensions: HashMap<String, Box<dyn Any + Send + Sync>>,
}

impl NemoFlowContextState {
    /// Create an empty runtime state with no registered middleware.
    ///
    /// # Returns
    /// A [`NemoFlowContextState`] with empty registries, no subscribers, and no
    /// extensions.
    pub fn new() -> Self {
        Self {
            tool_sanitize_request_guardrails: SortedRegistry::new(|entry| entry.priority),
            tool_sanitize_response_guardrails: SortedRegistry::new(|entry| entry.priority),
            tool_conditional_execution_guardrails: SortedRegistry::new(|entry| entry.priority),
            tool_request_intercepts: SortedRegistry::new(|entry| entry.priority),
            tool_execution_intercepts: SortedRegistry::new(|entry| entry.priority),
            llm_sanitize_request_guardrails: SortedRegistry::new(|entry| entry.priority),
            llm_sanitize_response_guardrails: SortedRegistry::new(|entry| entry.priority),
            llm_conditional_execution_guardrails: SortedRegistry::new(|entry| entry.priority),
            llm_request_intercepts: SortedRegistry::new(|entry| entry.priority),
            llm_execution_intercepts: SortedRegistry::new(|entry| entry.priority),
            llm_stream_execution_intercepts: SortedRegistry::new(|entry| entry.priority),
            event_subscribers: HashMap::new(),
            extensions: HashMap::new(),
        }
    }

    /// Store an arbitrary runtime extension under `key`.
    ///
    /// Extensions let bindings or integrations attach shared state to the
    /// process-global runtime without adding new first-class fields.
    ///
    /// # Parameters
    /// - `key`: Stable identifier for the extension slot.
    /// - `value`: Typed extension value to store.
    pub fn set_extension<T: Any + Send + Sync>(&mut self, key: impl Into<String>, value: T) {
        self.extensions.insert(key.into(), Box::new(value));
    }

    /// Borrow a typed runtime extension by key.
    ///
    /// # Parameters
    /// - `key`: Extension slot name.
    ///
    /// # Returns
    /// `Some(&T)` when an extension exists under `key` with the requested type
    /// and `None` otherwise.
    pub fn get_extension<T: Any + Send + Sync>(&self, key: &str) -> Option<&T> {
        self.extensions
            .get(key)
            .and_then(|value| value.downcast_ref::<T>())
    }

    /// Mutably borrow a typed runtime extension by key.
    ///
    /// # Parameters
    /// - `key`: Extension slot name.
    ///
    /// # Returns
    /// `Some(&mut T)` when an extension exists under `key` with the requested
    /// type and `None` otherwise.
    pub fn get_extension_mut<T: Any + Send + Sync>(&mut self, key: &str) -> Option<&mut T> {
        self.extensions
            .get_mut(key)
            .and_then(|value| value.downcast_mut::<T>())
    }

    /// Remove a runtime extension by key.
    ///
    /// # Parameters
    /// - `key`: Extension slot name.
    ///
    /// # Returns
    /// `true` when an extension was removed and `false` when no extension was
    /// stored under `key`.
    pub fn remove_extension(&mut self, key: &str) -> bool {
        self.extensions.remove(key).is_some()
    }

    /// Combine global and scope-local subscribers into one delivery list.
    ///
    /// # Parameters
    /// - `scope_local_subscribers`: Subscribers collected from the active scope
    ///   stack.
    ///
    /// # Returns
    /// A vector containing all global subscribers followed by the provided
    /// scope-local subscribers.
    pub fn collect_event_subscribers(
        &self,
        scope_local_subscribers: &[EventSubscriberFn],
    ) -> Vec<EventSubscriberFn> {
        let mut subscribers =
            Vec::with_capacity(self.event_subscribers.len() + scope_local_subscribers.len());
        subscribers.extend(self.event_subscribers.values().cloned());
        subscribers.extend(scope_local_subscribers.iter().cloned());
        subscribers
    }

    /// Deliver an event to every subscriber in order.
    ///
    /// # Parameters
    /// - `event`: Fully constructed lifecycle event to deliver.
    /// - `subscribers`: Subscribers that should observe the event.
    pub fn emit_event(event: &Event, subscribers: &[EventSubscriberFn]) {
        for subscriber in subscribers {
            subscriber(event);
        }
    }

    /// Build a standalone mark event.
    ///
    /// # Parameters
    /// - `params`: A pre-built [`MarkEvent`] to wrap in an [`Event`].
    ///
    /// # Returns
    /// A mark [`Event`] containing the provided [`MarkEvent`].
    pub fn create_event(&self, params: MarkEvent) -> Event {
        Event::Mark(params)
    }

    /// Create a new scope handle.
    ///
    /// # Parameters
    /// - `name`: Human-readable scope name.
    /// - `parent_uuid`: Optional parent scope UUID.
    /// - `scope_type`: Semantic category of the scope.
    /// - `attributes`: Scope attribute bitflags.
    /// - `data`: Optional application payload stored on the handle.
    /// - `metadata`: Optional metadata stored on the handle.
    /// - `timestamp`: Optional handle start time. When omitted, the current
    ///   UTC time is used.
    ///
    /// # Returns
    /// A new [`ScopeHandle`] with a fresh UUID.
    pub fn create_scope_handle(&self, params: CreateScopeHandleParams<'_>) -> ScopeHandle {
        ScopeHandle::builder()
            .name(params.name)
            .scope_type(params.scope_type)
            .started_at(params.timestamp.unwrap_or_else(Utc::now))
            .attributes(params.attributes)
            .parent_uuid_opt(params.parent_uuid)
            .data_opt(params.data)
            .metadata_opt(params.metadata)
            .build()
    }

    /// Build a scope-start event from a handle.
    ///
    /// # Parameters
    /// - `handle`: Scope handle to serialize into an event.
    /// - `data`: Optional semantic input payload exported on the start event.
    ///
    /// # Returns
    /// A scope-start [`Event`] derived from the provided handle.
    pub fn build_scope_start_event(&self, handle: &ScopeHandle, data: Option<Json>) -> Event {
        Event::Scope(ScopeEvent::new(
            BaseEvent::builder()
                .parent_uuid_opt(handle.parent_uuid)
                .uuid(handle.uuid)
                .timestamp(handle.started_at)
                .name(handle.name.as_str())
                .data_opt(data)
                .metadata_opt(handle.metadata.clone())
                .build(),
            ScopeCategory::Start,
            scope_attributes_to_strings(handle.attributes),
            EventCategory::from(handle.scope_type),
            None,
        ))
    }

    /// Build a scope-end event from a handle.
    ///
    /// # Parameters
    /// - `handle`: Scope handle to serialize into an event.
    /// - `data`: Optional data payload returned from the scope.
    ///
    /// # Returns
    /// A scope-end [`Event`] derived from the provided handle.
    pub fn end_scope_handle(&self, handle: &ScopeHandle, data: Option<Json>) -> Event {
        self.build_scope_end_event(
            EndScopeHandleParams::builder()
                .handle(handle)
                .data_opt(data)
                .build(),
        )
    }

    /// Build a scope-end event from builder parameters.
    ///
    /// # Parameters
    /// - `params`: Scope end-event builder parameters.
    ///
    /// # Returns
    /// A scope-end [`Event`] derived from the provided parameters.
    pub fn build_scope_end_event(&self, params: EndScopeHandleParams<'_>) -> Event {
        let handle = params.handle;
        Event::Scope(ScopeEvent::new(
            BaseEvent::builder()
                .parent_uuid_opt(handle.parent_uuid)
                .uuid(handle.uuid)
                .timestamp(
                    params
                        .timestamp
                        .unwrap_or_else(|| end_timestamp_after(handle.started_at)),
                )
                .name(handle.name.as_str())
                .data_opt(params.data)
                .metadata_opt(handle.metadata.clone())
                .build(),
            ScopeCategory::End,
            scope_attributes_to_strings(handle.attributes),
            EventCategory::from(handle.scope_type),
            None,
        ))
    }

    /// Create a new tool handle.
    ///
    /// # Parameters
    /// - `name`: Tool name recorded on emitted events.
    /// - `parent_uuid`: Optional parent scope UUID.
    /// - `attributes`: Tool attribute bitflags.
    /// - `data`: Optional application payload stored on the handle.
    /// - `metadata`: Optional metadata stored on the handle.
    /// - `tool_call_id`: Optional provider-specific correlation identifier.
    /// - `timestamp`: Optional handle start time. When omitted, the current
    ///   UTC time is used.
    ///
    /// # Returns
    /// A new [`ToolHandle`] with a fresh UUID.
    pub fn create_tool_handle(&self, params: CreateToolHandleParams<'_>) -> ToolHandle {
        ToolHandle::builder()
            .name(params.name)
            .started_at(params.timestamp.unwrap_or_else(Utc::now))
            .attributes(params.attributes)
            .parent_uuid_opt(params.parent_uuid)
            .data_opt(params.data)
            .metadata_opt(params.metadata)
            .tool_call_id_opt(params.tool_call_id)
            .build()
    }

    /// Build a tool-start event from a handle.
    ///
    /// # Parameters
    /// - `handle`: Tool handle to serialize into an event.
    /// - `data`: Optional tool input payload.
    ///
    /// # Returns
    /// A tool-start [`Event`] derived from the provided handle.
    pub fn build_tool_start_event(&self, handle: &ToolHandle, data: Option<Json>) -> Event {
        Event::Scope(ScopeEvent::new(
            BaseEvent::builder()
                .parent_uuid_opt(handle.parent_uuid)
                .uuid(handle.uuid)
                .timestamp(handle.started_at)
                .name(handle.name.as_str())
                .data_opt(data)
                .metadata_opt(handle.metadata.clone())
                .build(),
            ScopeCategory::Start,
            tool_attributes_to_strings(handle.attributes),
            EventCategory::tool(),
            Some(
                CategoryProfile::builder()
                    .tool_call_id_opt(handle.tool_call_id.clone())
                    .build(),
            ),
        ))
    }

    /// Build a tool-end event from a handle and optional overrides.
    ///
    /// # Parameters
    /// - `handle`: Tool handle to serialize into an event.
    /// - `data`: Optional end-event data payload.
    /// - `metadata`: Optional metadata payload merged over `handle.metadata`.
    ///
    /// # Returns
    /// A tool-end [`Event`] derived from the provided handle.
    pub fn end_tool_handle(
        &self,
        handle: &ToolHandle,
        data: Option<Json>,
        metadata: Option<Json>,
    ) -> Event {
        self.build_tool_end_event(
            EndToolHandleParams::builder()
                .handle(handle)
                .data_opt(data)
                .metadata_opt(metadata)
                .build(),
        )
    }

    /// Build a tool-end event from builder parameters.
    ///
    /// The `metadata` payload is merged over the metadata already stored on
    /// the handle.
    ///
    /// # Parameters
    /// - `params`: Tool end-event builder parameters.
    ///
    /// # Returns
    /// A tool-end [`Event`] derived from the provided parameters.
    pub fn build_tool_end_event(&self, params: EndToolHandleParams<'_>) -> Event {
        let handle = params.handle;
        Event::Scope(ScopeEvent::new(
            BaseEvent::builder()
                .parent_uuid_opt(handle.parent_uuid)
                .uuid(handle.uuid)
                .timestamp(
                    params
                        .timestamp
                        .unwrap_or_else(|| end_timestamp_after(handle.started_at)),
                )
                .name(handle.name.as_str())
                .data_opt(params.data)
                .metadata_opt(merge_json(handle.metadata.clone(), params.metadata))
                .build(),
            ScopeCategory::End,
            tool_attributes_to_strings(handle.attributes),
            EventCategory::tool(),
            Some(
                CategoryProfile::builder()
                    .tool_call_id_opt(handle.tool_call_id.clone())
                    .build(),
            ),
        ))
    }

    /// Create a new LLM handle.
    ///
    /// # Parameters
    /// - `name`: Logical provider or model family name.
    /// - `parent_uuid`: Optional parent scope UUID.
    /// - `attributes`: LLM attribute bitflags.
    /// - `data`: Optional application payload stored on the handle.
    /// - `metadata`: Optional metadata stored on the handle.
    /// - `model_name`: Optional normalized model name stored on the handle.
    /// - `timestamp`: Optional handle start time. When omitted, the current
    ///   UTC time is used.
    ///
    /// # Returns
    /// A new [`LlmHandle`] with a fresh UUID.
    pub fn create_llm_handle(&self, params: CreateLlmHandleParams<'_>) -> LlmHandle {
        LlmHandle::builder()
            .name(params.name)
            .started_at(params.timestamp.unwrap_or_else(Utc::now))
            .attributes(params.attributes)
            .parent_uuid_opt(params.parent_uuid)
            .data_opt(params.data)
            .metadata_opt(params.metadata)
            .model_name_opt(params.model_name)
            .build()
    }

    /// Build an LLM-start event from a handle.
    ///
    /// # Parameters
    /// - `handle`: LLM handle to serialize into an event.
    /// - `data`: Sanitized LLM request payload.
    /// - `annotated_request`: Optional normalized request annotation.
    ///
    /// # Returns
    /// An LLM-start [`Event`] derived from the provided handle.
    pub fn build_llm_start_event(
        &self,
        handle: &LlmHandle,
        data: Option<Json>,
        annotated_request: Option<Arc<AnnotatedLlmRequest>>,
    ) -> Event {
        Event::Scope(ScopeEvent::new(
            BaseEvent::builder()
                .parent_uuid_opt(handle.parent_uuid)
                .uuid(handle.uuid)
                .timestamp(handle.started_at)
                .name(handle.name.as_str())
                .data_opt(data)
                .metadata_opt(handle.metadata.clone())
                .build(),
            ScopeCategory::Start,
            llm_attributes_to_strings(handle.attributes),
            EventCategory::llm(),
            Some(
                CategoryProfile::builder()
                    .model_name_opt(handle.model_name.clone())
                    .annotated_request_opt(annotated_request)
                    .build(),
            ),
        ))
    }

    /// Build an LLM-end event from a handle and optional overrides.
    ///
    /// # Parameters
    /// - `handle`: LLM handle to serialize into an event.
    /// - `data`: Sanitized LLM response payload.
    /// - `metadata`: Optional metadata payload merged over `handle.metadata`.
    /// - `annotated_response`: Optional normalized response annotation.
    ///
    /// # Returns
    /// An LLM-end [`Event`] derived from the provided handle.
    pub fn end_llm_handle(
        &self,
        handle: &LlmHandle,
        data: Option<Json>,
        metadata: Option<Json>,
        annotated_response: Option<Arc<AnnotatedLlmResponse>>,
    ) -> Event {
        self.build_llm_end_event(
            EndLlmHandleParams::builder()
                .handle(handle)
                .data_opt(data)
                .metadata_opt(metadata)
                .annotated_response_opt(annotated_response)
                .build(),
        )
    }

    /// Build an LLM-end event from builder parameters.
    ///
    /// The `metadata` payload is merged over the metadata already stored on
    /// the handle.
    ///
    /// # Parameters
    /// - `params`: LLM end-event builder parameters.
    ///
    /// # Returns
    /// An LLM-end [`Event`] derived from the provided parameters.
    pub fn build_llm_end_event(&self, params: EndLlmHandleParams<'_>) -> Event {
        let handle = params.handle;
        Event::Scope(ScopeEvent::new(
            BaseEvent::builder()
                .parent_uuid_opt(handle.parent_uuid)
                .uuid(handle.uuid)
                .timestamp(
                    params
                        .timestamp
                        .unwrap_or_else(|| end_timestamp_after(handle.started_at)),
                )
                .name(handle.name.as_str())
                .data_opt(params.data)
                .metadata_opt(merge_json(handle.metadata.clone(), params.metadata))
                .build(),
            ScopeCategory::End,
            llm_attributes_to_strings(handle.attributes),
            EventCategory::llm(),
            Some(
                CategoryProfile::builder()
                    .model_name_opt(handle.model_name.clone())
                    .annotated_response_opt(params.annotated_response)
                    .build(),
            ),
        ))
    }

    /// Run tool request sanitizers across global and scope-local registries.
    ///
    /// # Parameters
    /// - `name`: Tool name associated with the request.
    /// - `args`: Raw tool arguments to sanitize for observability.
    /// - `scope_locals`: Scope-local sanitizer registries collected from the
    ///   active scope stack.
    ///
    /// # Returns
    /// The sanitized JSON payload after every matching guardrail has run.
    pub fn tool_sanitize_request_chain(
        &self,
        name: &str,
        args: Json,
        scope_locals: &[&SortedRegistry<GuardrailEntry<ToolSanitizeFn>>],
    ) -> Json {
        let entries = merge_guardrail_entries(&self.tool_sanitize_request_guardrails, scope_locals);
        let mut value = args;
        for entry in entries {
            value = (entry.guardrail)(name, value);
        }
        value
    }

    /// Run tool response sanitizers across global and scope-local registries.
    ///
    /// # Parameters
    /// - `name`: Tool name associated with the response.
    /// - `result`: Raw tool result to sanitize for observability.
    /// - `scope_locals`: Scope-local sanitizer registries collected from the
    ///   active scope stack.
    ///
    /// # Returns
    /// The sanitized JSON payload after every matching guardrail has run.
    pub fn tool_sanitize_response_chain(
        &self,
        name: &str,
        result: Json,
        scope_locals: &[&SortedRegistry<GuardrailEntry<ToolSanitizeFn>>],
    ) -> Json {
        let entries =
            merge_guardrail_entries(&self.tool_sanitize_response_guardrails, scope_locals);
        let mut value = result;
        for entry in entries {
            value = (entry.guardrail)(name, value);
        }
        value
    }

    /// Evaluate tool conditional-execution guardrails in priority order.
    ///
    /// # Parameters
    /// - `name`: Tool name associated with the request.
    /// - `args`: Tool arguments to validate.
    /// - `scope_locals`: Scope-local conditional guardrail registries collected
    ///   from the active scope stack.
    ///
    /// # Returns
    /// A [`Result`] containing `Ok(None)` when execution is allowed or
    /// `Ok(Some(reason))` when a guardrail rejects the call.
    ///
    /// # Errors
    /// Propagates any error returned by a guardrail callback.
    pub fn tool_conditional_execution_chain(
        &self,
        name: &str,
        args: &Json,
        scope_locals: &[&SortedRegistry<GuardrailEntry<ToolConditionalFn>>],
    ) -> crate::error::Result<Option<String>> {
        let entries =
            merge_guardrail_entries(&self.tool_conditional_execution_guardrails, scope_locals);
        for entry in entries {
            if let Some(error) = (entry.guardrail)(name, args)? {
                return Ok(Some(error));
            }
        }
        Ok(None)
    }

    /// Run tool request intercepts in priority order.
    ///
    /// # Parameters
    /// - `name`: Tool name associated with the request.
    /// - `args`: Tool arguments to pass through the intercept chain.
    /// - `scope_locals`: Scope-local request intercept registries collected
    ///   from the active scope stack.
    ///
    /// # Returns
    /// A [`Result`] containing the final JSON argument payload.
    ///
    /// # Errors
    /// Propagates any error returned by an intercept callback.
    ///
    /// # Notes
    /// If an intercept entry has `break_chain` enabled, later intercepts are
    /// skipped after that entry runs.
    pub fn tool_request_intercepts_chain(
        &self,
        name: &str,
        args: Json,
        scope_locals: &[&SortedRegistry<Intercept<ToolInterceptFn>>],
    ) -> crate::error::Result<Json> {
        let entries = merge_intercept_entries(&self.tool_request_intercepts, scope_locals);
        let mut value = args;
        for entry in entries {
            value = (entry.callable)(name, value)?;
            if entry.break_chain {
                break;
            }
        }
        Ok(value)
    }

    /// Build the composed tool execution continuation chain.
    ///
    /// # Parameters
    /// - `name`: Tool name passed into each execution intercept.
    /// - `default_fn`: Base tool callback that should run after all intercepts.
    /// - `scope_locals`: Scope-local execution intercept registries collected
    ///   from the active scope stack.
    ///
    /// # Returns
    /// A composed [`ToolExecutionNextFn`] that wraps `default_fn` in every
    /// matching execution intercept.
    pub fn tool_build_execution_chain(
        &self,
        name: &str,
        default_fn: ToolExecutionNextFn,
        scope_locals: &[&SortedRegistry<ExecutionIntercept<ToolExecutionFn>>],
    ) -> ToolExecutionNextFn {
        let matching =
            merge_execution_intercept_callables(&self.tool_execution_intercepts, scope_locals);
        let mut next = default_fn;
        let name = name.to_string();
        for (callable, _) in matching.into_iter().rev() {
            let current_next = next.clone();
            let current_name = name.clone();
            next = Arc::new(move |args| callable(&current_name, args, current_next.clone()));
        }
        next
    }

    /// Run LLM request sanitizers across global and scope-local registries.
    ///
    /// # Parameters
    /// - `request`: Raw LLM request to sanitize for observability.
    /// - `scope_locals`: Scope-local sanitizer registries collected from the
    ///   active scope stack.
    ///
    /// # Returns
    /// The sanitized [`LlmRequest`] after every matching guardrail has run.
    pub fn llm_sanitize_request_chain(
        &self,
        request: LlmRequest,
        scope_locals: &[&SortedRegistry<GuardrailEntry<LlmSanitizeRequestFn>>],
    ) -> LlmRequest {
        let entries = merge_guardrail_entries(&self.llm_sanitize_request_guardrails, scope_locals);
        let mut value = request;
        for entry in entries {
            value = (entry.guardrail)(value);
        }
        value
    }

    /// Run LLM response sanitizers across global and scope-local registries.
    ///
    /// # Parameters
    /// - `response`: Raw response payload to sanitize for observability.
    /// - `scope_locals`: Scope-local sanitizer registries collected from the
    ///   active scope stack.
    ///
    /// # Returns
    /// The sanitized response payload after every matching guardrail has run.
    pub fn llm_sanitize_response_chain(
        &self,
        response: Json,
        scope_locals: &[&SortedRegistry<GuardrailEntry<LlmSanitizeResponseFn>>],
    ) -> Json {
        let entries = merge_guardrail_entries(&self.llm_sanitize_response_guardrails, scope_locals);
        let mut value = response;
        for entry in entries {
            value = (entry.guardrail)(value);
        }
        value
    }

    /// Evaluate LLM conditional-execution guardrails in priority order.
    ///
    /// # Parameters
    /// - `request`: LLM request to validate.
    /// - `scope_locals`: Scope-local conditional guardrail registries collected
    ///   from the active scope stack.
    ///
    /// # Returns
    /// A [`Result`] containing `Ok(None)` when execution is allowed or
    /// `Ok(Some(reason))` when a guardrail rejects the call.
    ///
    /// # Errors
    /// Propagates any error returned by a guardrail callback.
    pub fn llm_conditional_execution_chain(
        &self,
        request: &LlmRequest,
        scope_locals: &[&SortedRegistry<GuardrailEntry<LlmConditionalFn>>],
    ) -> crate::error::Result<Option<String>> {
        let entries =
            merge_guardrail_entries(&self.llm_conditional_execution_guardrails, scope_locals);
        for entry in entries {
            if let Some(error) = (entry.guardrail)(request)? {
                return Ok(Some(error));
            }
        }
        Ok(None)
    }

    /// Run LLM request intercepts in priority order.
    ///
    /// # Parameters
    /// - `name`: Logical provider or model family name.
    /// - `request`: LLM request to pass through the intercept chain.
    /// - `annotated`: Optional normalized request annotation to carry through
    ///   the chain.
    /// - `scope_locals`: Scope-local request intercept registries collected
    ///   from the active scope stack.
    ///
    /// # Returns
    /// A [`Result`] containing the final request and annotation pair.
    ///
    /// # Errors
    /// Propagates any error returned by an intercept callback.
    ///
    /// # Notes
    /// If an intercept entry has `break_chain` enabled, later intercepts are
    /// skipped after that entry runs.
    pub fn llm_request_intercepts_chain(
        &self,
        name: &str,
        request: LlmRequest,
        annotated: Option<AnnotatedLlmRequest>,
        scope_locals: &[&SortedRegistry<Intercept<LlmRequestInterceptFn>>],
    ) -> crate::error::Result<(LlmRequest, Option<AnnotatedLlmRequest>)> {
        let entries = merge_intercept_entries(&self.llm_request_intercepts, scope_locals);
        let mut request_value = request;
        let mut annotated_value = annotated;
        for entry in entries {
            let (new_request, new_annotated) =
                (entry.callable)(name, request_value, annotated_value)?;
            request_value = new_request;
            annotated_value = new_annotated;
            if entry.break_chain {
                break;
            }
        }
        Ok((request_value, annotated_value))
    }

    /// Build the composed non-streaming LLM execution continuation chain.
    ///
    /// # Parameters
    /// - `name`: Logical provider or model family name passed into each
    ///   execution intercept.
    /// - `default_fn`: Base provider callback that should run after all
    ///   intercepts.
    /// - `scope_locals`: Scope-local execution intercept registries collected
    ///   from the active scope stack.
    ///
    /// # Returns
    /// A composed [`LlmExecutionNextFn`] that wraps `default_fn` in every
    /// matching execution intercept.
    pub fn llm_build_execution_chain(
        &self,
        name: &str,
        default_fn: LlmExecutionNextFn,
        scope_locals: &[&SortedRegistry<ExecutionIntercept<LlmExecutionFn>>],
    ) -> LlmExecutionNextFn {
        let matching =
            merge_execution_intercept_callables(&self.llm_execution_intercepts, scope_locals);
        let mut next = default_fn;
        let name = name.to_string();
        for (callable, _) in matching.into_iter().rev() {
            let current_next = next.clone();
            let current_name = name.clone();
            next = Arc::new(move |request| callable(&current_name, request, current_next.clone()));
        }
        next
    }

    /// Build the composed streaming LLM execution continuation chain.
    ///
    /// # Parameters
    /// - `name`: Logical provider or model family name passed into each
    ///   execution intercept.
    /// - `default_fn`: Base stream-producing callback that should run after all
    ///   intercepts.
    /// - `scope_locals`: Scope-local execution intercept registries collected
    ///   from the active scope stack.
    ///
    /// # Returns
    /// A composed [`LlmStreamExecutionNextFn`] that wraps `default_fn` in every
    /// matching execution intercept.
    pub fn llm_stream_build_execution_chain(
        &self,
        name: &str,
        default_fn: LlmStreamExecutionNextFn,
        scope_locals: LlmStreamExecutionRegistryRefs<'_>,
    ) -> LlmStreamExecutionNextFn {
        let matching = merge_execution_intercept_callables(
            &self.llm_stream_execution_intercepts,
            scope_locals,
        );
        let mut next = default_fn;
        let name = name.to_string();
        for (callable, _) in matching.into_iter().rev() {
            let current_next = next.clone();
            let current_name = name.clone();
            next = Arc::new(move |request| callable(&current_name, request, current_next.clone()));
        }
        next
    }
}

fn end_timestamp_after(started_at: chrono::DateTime<Utc>) -> chrono::DateTime<Utc> {
    let now = Utc::now();
    if now > started_at {
        now
    } else {
        started_at + Duration::microseconds(1)
    }
}

impl Default for NemoFlowContextState {
    fn default() -> Self {
        Self::new()
    }
}