nemo_relay/api/
tool.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use serde_json::json;
5
6use crate::api::runtime::NemoRelayContextState;
7use crate::api::runtime::ToolExecutionNextFn;
8use crate::api::runtime::current_scope_stack;
9use crate::api::runtime::global_context;
10use crate::api::scope::event;
11use crate::api::scope::{EmitMarkEventParams, ScopeHandle};
12use crate::api::shared::{
13    ensure_runtime_owner, metadata_with_otel_status, resolve_parent_uuid,
14    snapshot_event_subscribers,
15};
16use crate::error::{FlowError, Result};
17use crate::json::Json;
18use bitflags::bitflags;
19use chrono::{DateTime, Utc};
20use serde::{Deserialize, Serialize};
21use typed_builder::TypedBuilder;
22use uuid::Uuid;
23
24bitflags! {
25    /// Bitflags that modify tool-call behavior and observability.
26    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
27    pub struct ToolAttributes: u32 {
28        /// Marks the tool as executing out-of-process.
29        const REMOTE = 0b01;
30    }
31}
32
33/// Runtime-owned handle identifying an active or completed tool call.
34#[derive(Debug, Clone, Serialize, Deserialize, TypedBuilder)]
35#[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))]
36pub struct ToolHandle {
37    /// Unique tool-call identifier.
38    #[builder(default = Uuid::now_v7())]
39    pub uuid: Uuid,
40    /// Timestamp captured when the tool handle was created.
41    #[builder(default = Utc::now())]
42    pub started_at: DateTime<Utc>,
43    /// Tool name recorded on lifecycle events.
44    #[builder(setter(into))]
45    pub name: String,
46    /// Optional application payload stored on the handle.
47    #[builder(default)]
48    pub data: Option<Json>,
49    /// Optional metadata attached to the tool span.
50    #[builder(default)]
51    pub metadata: Option<Json>,
52    /// Tool behavior flags.
53    #[builder(default = ToolAttributes::empty())]
54    pub attributes: ToolAttributes,
55    /// UUID of the parent scope, if any.
56    #[builder(default)]
57    pub parent_uuid: Option<Uuid>,
58    /// Optional provider-specific tool-call correlation identifier.
59    #[builder(default, setter(into))]
60    pub tool_call_id: Option<String>,
61}
62
63/// Builder parameters for [`NemoRelayContextState::create_tool_handle`].
64#[derive(Debug, Clone, TypedBuilder)]
65#[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))]
66pub struct CreateToolHandleParams<'a> {
67    /// Tool name recorded on emitted events.
68    pub name: &'a str,
69    /// Optional parent scope UUID.
70    #[builder(default)]
71    pub parent_uuid: Option<uuid::Uuid>,
72    /// Tool attribute bitflags.
73    #[builder(default = ToolAttributes::empty())]
74    pub attributes: ToolAttributes,
75    /// Optional application payload stored on the handle.
76    #[builder(default)]
77    pub data: Option<Json>,
78    /// Optional metadata stored on the handle.
79    #[builder(default)]
80    pub metadata: Option<Json>,
81    /// Optional provider-specific correlation identifier.
82    #[builder(default, setter(into))]
83    pub tool_call_id: Option<String>,
84    /// Optional timestamp captured as the handle start time and reused by the
85    /// emitted start event. When omitted, the current UTC time is used.
86    #[builder(default)]
87    pub timestamp: Option<DateTime<Utc>>,
88}
89
90/// Builder parameters for [`NemoRelayContextState::build_tool_end_event`].
91#[derive(Debug, Clone, TypedBuilder)]
92#[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))]
93pub struct EndToolHandleParams<'a> {
94    /// Tool handle to serialize into the emitted end event.
95    pub handle: &'a ToolHandle,
96    /// Optional data payload merged over the handle data.
97    #[builder(default)]
98    pub data: Option<Json>,
99    /// Optional metadata payload merged over the handle metadata.
100    #[builder(default)]
101    pub metadata: Option<Json>,
102    /// Optional timestamp recorded on the emitted end event. When omitted, the
103    /// runtime records the current UTC time, or one microsecond after the
104    /// handle start time if the current time is not later.
105    #[builder(default)]
106    pub timestamp: Option<DateTime<Utc>>,
107}
108
109/// Builder parameters for [`tool_call`].
110#[derive(TypedBuilder)]
111#[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))]
112pub struct ToolCallParams<'a> {
113    /// Tool name recorded on the emitted lifecycle event.
114    pub name: &'a str,
115    /// Raw tool arguments associated with the span.
116    pub args: Json,
117    /// Optional explicit parent scope.
118    #[builder(default)]
119    pub parent: Option<&'a ScopeHandle>,
120    /// Tool attribute bitflags applied to the span.
121    #[builder(default = ToolAttributes::empty())]
122    pub attributes: ToolAttributes,
123    /// Optional application payload stored on the handle but not emitted as
124    /// Agent Trajectory Observability Format (ATOF) data.
125    #[builder(default)]
126    pub data: Option<Json>,
127    /// Optional JSON metadata recorded on the start event.
128    #[builder(default)]
129    pub metadata: Option<Json>,
130    /// Optional provider-specific correlation identifier.
131    #[builder(default, setter(into))]
132    pub tool_call_id: Option<String>,
133    /// Optional timestamp captured as the handle start time and reused by the
134    /// emitted start event. When omitted, the current UTC time is used.
135    #[builder(default)]
136    pub timestamp: Option<DateTime<Utc>>,
137}
138
139/// Builder parameters for [`tool_call_execute`].
140#[derive(TypedBuilder)]
141#[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))]
142pub struct ToolCallExecuteParams {
143    /// Tool name recorded on emitted lifecycle events.
144    #[builder(setter(into))]
145    pub name: String,
146    /// Raw tool arguments passed into the managed pipeline.
147    pub args: Json,
148    /// Tool callback or execution continuation.
149    pub func: ToolExecutionNextFn,
150    /// Optional explicit parent scope for the emitted tool span.
151    #[builder(default)]
152    pub parent: Option<ScopeHandle>,
153    /// Tool attribute bitflags applied to the managed span.
154    #[builder(default = ToolAttributes::empty())]
155    pub attributes: ToolAttributes,
156    /// Optional application payload stored on the handle but not emitted as
157    /// Agent Trajectory Observability Format (ATOF) data.
158    #[builder(default)]
159    pub data: Option<Json>,
160    /// Optional JSON metadata recorded on emitted events.
161    #[builder(default)]
162    pub metadata: Option<Json>,
163}
164
165/// Builder parameters for [`tool_call_end`].
166#[derive(TypedBuilder)]
167#[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))]
168pub struct ToolCallEndParams<'a> {
169    /// Tool handle to close.
170    pub handle: &'a ToolHandle,
171    /// Raw tool result associated with the end event.
172    pub result: Json,
173    /// Optional application payload retained for compatibility; Agent
174    /// Trajectory Observability Format (ATOF) data is the result.
175    #[builder(default)]
176    pub data: Option<Json>,
177    /// Optional JSON metadata recorded on the end event.
178    #[builder(default)]
179    pub metadata: Option<Json>,
180    /// Optional timestamp recorded on the emitted end event. When omitted, the
181    /// runtime records the current UTC time, or one microsecond after the
182    /// handle start time if the current time is not later.
183    #[builder(default)]
184    pub timestamp: Option<DateTime<Utc>>,
185}
186
187/// Start a manual tool lifecycle span.
188///
189/// This emits a tool-start event after applying sanitize-request guardrails to
190/// the payload recorded for observability.
191///
192/// # Parameters
193/// - `name`: Tool name recorded on the emitted lifecycle event.
194/// - `args`: Raw tool arguments associated with the span.
195/// - `parent`: Optional explicit parent scope.
196/// - `attributes`: Tool attribute bitflags applied to the span.
197/// - `data`: Optional application payload stored on the returned handle. The
198///   emitted start event data is the sanitized `args` payload.
199/// - `metadata`: Optional JSON metadata recorded on the start event.
200/// - `tool_call_id`: Optional provider-specific correlation identifier.
201/// - `timestamp`: Optional timestamp recorded as the handle start time and on
202///   the emitted start event. When `None`, the current UTC time is used.
203///
204/// # Returns
205/// A [`Result`] containing the created [`ToolHandle`].
206///
207/// # Errors
208/// Returns an error when the runtime owner check fails or when internal state
209/// cannot be read safely.
210///
211/// # Notes
212/// Sanitize-request guardrails affect only the emitted start-event payload, not
213/// the caller-owned `args` value.
214pub fn tool_call(params: ToolCallParams<'_>) -> Result<ToolHandle> {
215    ensure_runtime_owner()?;
216    let parent_uuid = resolve_parent_uuid(params.parent);
217    let (handle, event, subscribers) = {
218        let scope_stack = current_scope_stack();
219        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
220        let scope_locals = scope_guard.collect_scope_local_registries(|registries| {
221            &registries.tool_sanitize_request_guardrails
222        });
223        let scope_subscribers = scope_guard.collect_scope_local_subscribers();
224        let subscribers = snapshot_event_subscribers(scope_subscribers)?;
225        let context = global_context();
226        let state = context
227            .read()
228            .map_err(|error| FlowError::Internal(error.to_string()))?;
229
230        let sanitized_args =
231            state.tool_sanitize_request_chain(params.name, params.args, &scope_locals);
232        let handle_params = CreateToolHandleParams::builder()
233            .name(params.name)
234            .parent_uuid_opt(parent_uuid)
235            .attributes(params.attributes)
236            .data_opt(params.data)
237            .metadata_opt(params.metadata)
238            .tool_call_id_opt(params.tool_call_id)
239            .timestamp_opt(params.timestamp)
240            .build();
241        let handle = state.create_tool_handle(handle_params);
242        let event = state.build_tool_start_event(&handle, Some(sanitized_args));
243        (handle, event, subscribers)
244    };
245    NemoRelayContextState::emit_event(&event, &subscribers);
246    Ok(handle)
247}
248
249/// Finish a manual tool lifecycle span.
250///
251/// This emits a tool-end event for a handle previously returned by
252/// [`tool_call`].
253///
254/// # Parameters
255/// - `handle`: Tool handle to close.
256/// - `result`: Raw tool result associated with the end event.
257/// - `data`: Optional application payload retained for compatibility. The
258///   emitted end event data is the sanitized `result` unless it sanitizes to
259///   JSON null, in which case this payload is used.
260/// - `metadata`: Optional JSON metadata recorded on the end event.
261/// - `timestamp`: Optional timestamp recorded on the emitted end event. When
262///   `None`, the runtime uses the current UTC time, or one microsecond after
263///   the handle start time if the current time is not later.
264///
265/// # Returns
266/// A [`Result`] that is `Ok(())` when the end event has been emitted.
267///
268/// # Errors
269/// Returns an error when the runtime owner check fails or when internal state
270/// cannot be read safely.
271///
272/// # Notes
273/// Sanitize-response guardrails affect only the emitted end-event payload, not
274/// the caller-owned `result` value.
275pub fn tool_call_end(params: ToolCallEndParams<'_>) -> Result<()> {
276    ensure_runtime_owner()?;
277    let (event, subscribers) = {
278        let scope_stack = current_scope_stack();
279        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
280        let scope_locals = scope_guard.collect_scope_local_registries(|registries| {
281            &registries.tool_sanitize_response_guardrails
282        });
283        let scope_subscribers = scope_guard.collect_scope_local_subscribers();
284        let subscribers = snapshot_event_subscribers(scope_subscribers)?;
285        let context = global_context();
286        let state = context
287            .read()
288            .map_err(|error| FlowError::Internal(error.to_string()))?;
289
290        let sanitized_result =
291            state.tool_sanitize_response_chain(&params.handle.name, params.result, &scope_locals);
292        let data = if sanitized_result.is_null() {
293            params.data
294        } else {
295            Some(sanitized_result)
296        };
297        let event = state.build_tool_end_event(
298            EndToolHandleParams::builder()
299                .handle(params.handle)
300                .data_opt(data)
301                .metadata_opt(params.metadata)
302                .timestamp_opt(params.timestamp)
303                .build(),
304        );
305        (event, subscribers)
306    };
307    NemoRelayContextState::emit_event(&event, &subscribers);
308    Ok(())
309}
310
311fn emit_tool_end_without_output(handle: &ToolHandle, metadata: Option<Json>) -> Result<()> {
312    ensure_runtime_owner()?;
313    let (event, subscribers) = {
314        let scope_stack = current_scope_stack();
315        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
316        let scope_subscribers = scope_guard.collect_scope_local_subscribers();
317        let subscribers = snapshot_event_subscribers(scope_subscribers)?;
318        let context = global_context();
319        let state = context
320            .read()
321            .map_err(|error| FlowError::Internal(error.to_string()))?;
322        let event = state.end_tool_handle(handle, handle.data.clone(), metadata);
323        (event, subscribers)
324    };
325    NemoRelayContextState::emit_event(&event, &subscribers);
326    Ok(())
327}
328
329/// Execute a tool call through the managed middleware pipeline.
330///
331/// This runs conditional-execution guardrails, request intercepts,
332/// sanitize-request guardrails, execution intercepts, the tool callback, and
333/// sanitize-response guardrails in the runtime-defined order.
334///
335/// # Parameters
336/// - `name`: Tool name recorded on emitted lifecycle events.
337/// - `args`: Raw tool arguments passed into the managed pipeline.
338/// - `func`: Tool callback or execution continuation.
339/// - `parent`: Optional explicit parent scope for the emitted tool span.
340/// - `attributes`: Tool attribute bitflags applied to the managed span.
341/// - `data`: Optional application payload stored on the managed tool handle.
342///   It may be used on failure end events that have no output payload.
343/// - `metadata`: Optional JSON metadata recorded on emitted events.
344///
345/// # Returns
346/// A [`Result`] containing the raw tool result returned by the callback or an
347/// execution intercept.
348///
349/// # Errors
350/// Returns [`FlowError::GuardrailRejected`] when conditional-execution
351/// guardrails block the call, or any error raised by request intercepts,
352/// execution intercepts, or the callback itself.
353///
354/// # Notes
355/// When execution fails after the start event has been emitted, the runtime
356/// still emits a tool-end event without an output payload.
357pub async fn tool_call_execute(params: ToolCallExecuteParams) -> Result<Json> {
358    let ToolCallExecuteParams {
359        name,
360        args,
361        func,
362        parent,
363        attributes,
364        data,
365        metadata,
366    } = params;
367    ensure_runtime_owner()?;
368    {
369        let (entries, subscribers, parent_uuid, guardrail_metadata) = {
370            let scope_stack = current_scope_stack();
371            let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
372            let scope_locals = scope_guard.collect_scope_local_registries(|registries| {
373                &registries.tool_conditional_execution_guardrails
374            });
375            let scope_subscribers = scope_guard.collect_scope_local_subscribers();
376            let context = global_context();
377            let state = context
378                .read()
379                .map_err(|error| FlowError::Internal(error.to_string()))?;
380            let entries = state.tool_conditional_execution_entries(&scope_locals);
381            let subscribers = state.collect_event_subscribers(&scope_subscribers);
382            (
383                entries,
384                subscribers,
385                resolve_parent_uuid(parent.as_ref()),
386                metadata.clone(),
387            )
388        };
389        if let Some(error) = NemoRelayContextState::tool_conditional_execution_snapshot_chain(
390            &name,
391            &args,
392            &entries,
393            &subscribers,
394            parent_uuid,
395            guardrail_metadata,
396        )? {
397            let mut rejection_data = json!({});
398            if let Some(object) = rejection_data.as_object_mut() {
399                object.insert("rejected".into(), json!(true));
400                object.insert("rejection_reason".into(), json!(&error));
401            }
402            let _ = event(
403                EmitMarkEventParams::builder()
404                    .name(&name)
405                    .parent_opt(parent.as_ref())
406                    .data(rejection_data)
407                    .metadata_opt(metadata.clone())
408                    .build(),
409            );
410            return Err(FlowError::GuardrailRejected(error));
411        }
412    }
413
414    let intercepted_args = {
415        let scope_stack = current_scope_stack();
416        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
417        let scope_locals = scope_guard
418            .collect_scope_local_registries(|registries| &registries.tool_request_intercepts);
419        let context = global_context();
420        let state = context
421            .read()
422            .map_err(|error| FlowError::Internal(error.to_string()))?;
423        state.tool_request_intercepts_chain(&name, args, &scope_locals)?
424    };
425
426    let handle = tool_call(
427        ToolCallParams::builder()
428            .name(name.as_str())
429            .args(intercepted_args.clone())
430            .parent_opt(parent.as_ref())
431            .attributes(attributes)
432            .data_opt(data.clone())
433            .metadata_opt(metadata.clone())
434            .build(),
435    )?;
436
437    let execution = {
438        let scope_stack = current_scope_stack();
439        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
440        let scope_locals = scope_guard
441            .collect_scope_local_registries(|registries| &registries.tool_execution_intercepts);
442        let context = global_context();
443        let state = context
444            .read()
445            .map_err(|error| FlowError::Internal(error.to_string()))?;
446        state.tool_build_execution_chain(&name, func, &scope_locals)
447    };
448
449    match execution(intercepted_args).await {
450        Ok(result) => {
451            let end_metadata = metadata_with_otel_status(metadata, "OK", None);
452            tool_call_end(
453                ToolCallEndParams::builder()
454                    .handle(&handle)
455                    .result(result.clone())
456                    .data_opt(data)
457                    .metadata_opt(end_metadata)
458                    .build(),
459            )?;
460            Ok(result)
461        }
462        Err(error) => {
463            let end_metadata =
464                metadata_with_otel_status(metadata, "ERROR", Some(error.to_string()));
465            let _ = emit_tool_end_without_output(&handle, end_metadata);
466            Err(error)
467        }
468    }
469}
470
471/// Run only the tool request-intercept chain.
472///
473/// This applies the currently active global and scope-local request intercepts
474/// without emitting lifecycle events or invoking tool execution.
475///
476/// # Parameters
477/// - `name`: Tool name used when resolving the intercept chain.
478/// - `args`: Raw tool arguments to transform.
479///
480/// # Returns
481/// A [`Result`] containing the transformed JSON arguments.
482///
483/// # Errors
484/// Returns any error raised by the request-intercept chain.
485///
486/// # Notes
487/// Conditional guardrails and execution intercepts are not run by this helper.
488pub fn tool_request_intercepts(name: &str, args: Json) -> Result<Json> {
489    ensure_runtime_owner()?;
490    let scope_stack = current_scope_stack();
491    let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
492    let scope_locals = scope_guard
493        .collect_scope_local_registries(|registries| &registries.tool_request_intercepts);
494    let context = global_context();
495    let state = context
496        .read()
497        .map_err(|error| FlowError::Internal(error.to_string()))?;
498    state.tool_request_intercepts_chain(name, args, &scope_locals)
499}
500
501/// Run only the tool conditional-execution guardrail chain.
502///
503/// This evaluates whether a tool call should be allowed to proceed without
504/// invoking request intercepts or execution. Each evaluated guardrail emits an
505/// automatic guardrail scope start/end pair for observability.
506///
507/// # Parameters
508/// - `name`: Tool name used when resolving the guardrail chain.
509/// - `args`: Raw tool arguments to validate.
510///
511/// # Returns
512/// A [`Result`] that is `Ok(())` when all guardrails allow execution.
513///
514/// # Errors
515/// Returns [`FlowError::GuardrailRejected`] when a guardrail blocks execution,
516/// or any error raised by the guardrail chain itself.
517///
518/// # Notes
519/// This helper is useful for preflight checks when the caller needs the
520/// rejection result without starting a tool span. Guardrail scopes are still
521/// emitted for the conditional checks themselves.
522pub fn tool_conditional_execution(name: &str, args: &Json) -> Result<()> {
523    ensure_runtime_owner()?;
524    let (entries, subscribers, parent_uuid) = {
525        let scope_stack = current_scope_stack();
526        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
527        let scope_locals = scope_guard.collect_scope_local_registries(|registries| {
528            &registries.tool_conditional_execution_guardrails
529        });
530        let scope_subscribers = scope_guard.collect_scope_local_subscribers();
531        let context = global_context();
532        let state = context
533            .read()
534            .map_err(|error| FlowError::Internal(error.to_string()))?;
535        let entries = state.tool_conditional_execution_entries(&scope_locals);
536        let subscribers = state.collect_event_subscribers(&scope_subscribers);
537        (entries, subscribers, resolve_parent_uuid(None))
538    };
539    if let Some(error) = NemoRelayContextState::tool_conditional_execution_snapshot_chain(
540        name,
541        args,
542        &entries,
543        &subscribers,
544        parent_uuid,
545        None,
546    )? {
547        return Err(FlowError::GuardrailRejected(error));
548    }
549    Ok(())
550}
551
552#[cfg(test)]
553#[path = "../../tests/unit/tool_api_tests.rs"]
554mod tests;
nemo_relay/api/tool.rs

nemo_relay/api/
tool.rs
