rig_compose/
trace.rs

1//! Dispatch trace records.
2//!
3//! [`DispatchTrace`] is an additive, deterministic record of what happened
4//! during a [`dispatch_tool_invocations_with_trace`](crate::dispatch_tool_invocations_with_trace)
5//! call: each hook's decision, the eventual invocation outcome, errors, and
6//! reservation cleanups. Hosts use it to explain a tool-loop replay or to
7//! export structured policy traces without depending on a concrete tracing
8//! backend.
9//!
10//! # Example
11//!
12//! ```no_run
13//! use rig_compose::{
14//!     DispatchTrace, ToolRegistry, dispatch_tool_invocations_with_trace,
15//! };
16//!
17//! # async fn run(registry: ToolRegistry, invocations: Vec<rig_compose::ToolInvocation>) -> Result<(), rig_compose::KernelError> {
18//! let trace = DispatchTrace::new();
19//! let _results =
20//!     dispatch_tool_invocations_with_trace(&registry, &invocations, &[], &trace).await?;
21//! for event in trace.events() {
22//!     tracing::debug!(?event, "dispatch trace event");
23//! }
24//! # Ok(()) }
25//! ```
26
27use std::sync::{Arc, Mutex};
28
29use crate::normalizer::ToolDispatchAction;
30
31/// One observable event within a dispatch trace. Events are pushed in order;
32/// `invocation_index` and `hook_index` refer to positions in the slices passed
33/// to
34/// [`dispatch_tool_invocations_with_trace`](crate::dispatch_tool_invocations_with_trace).
35#[derive(Debug, Clone, PartialEq)]
36pub enum DispatchTraceEvent {
37    /// A hook's `before_invocation` returned a decision.
38    HookBefore {
39        invocation_index: usize,
40        hook_index: usize,
41        decision: TracedAction,
42    },
43    /// A hook's `before_invocation` returned an error and aborted dispatch.
44    HookBeforeError {
45        invocation_index: usize,
46        hook_index: usize,
47        message: String,
48    },
49    /// A hook's `on_invocation_error` was notified so it could release
50    /// resources reserved in `before_invocation`.
51    HookCleanup {
52        invocation_index: usize,
53        hook_index: usize,
54    },
55    /// A hook's `after_invocation_with_outcome` ran successfully.
56    HookAfter {
57        invocation_index: usize,
58        hook_index: usize,
59    },
60    /// The invocation finished with a final outcome.
61    InvocationOutcome {
62        invocation_index: usize,
63        outcome: TracedOutcome,
64    },
65}
66
67/// Lightweight projection of [`ToolDispatchAction`] suitable for trace
68/// recording. Decouples the trace shape from the runtime action enum so adding
69/// fields to [`ToolDispatchAction`] does not silently widen trace records.
70#[derive(Debug, Clone, PartialEq)]
71pub enum TracedAction {
72    Continue,
73    Skip { reason: Option<String> },
74    Terminate { reason: String },
75}
76
77impl From<&ToolDispatchAction> for TracedAction {
78    fn from(value: &ToolDispatchAction) -> Self {
79        match value {
80            ToolDispatchAction::Continue => TracedAction::Continue,
81            ToolDispatchAction::Skip { reason, .. } => TracedAction::Skip {
82                reason: reason.clone(),
83            },
84            ToolDispatchAction::Terminate { reason } => TracedAction::Terminate {
85                reason: reason.clone(),
86            },
87        }
88    }
89}
90
91/// Final outcome recorded for one invocation.
92#[derive(Debug, Clone, PartialEq)]
93pub enum TracedOutcome {
94    /// The tool body ran and produced a result.
95    Completed,
96    /// A hook synthesised a skip result.
97    Skipped { reason: Option<String> },
98    /// A hook stopped the dispatch loop.
99    Terminated { reason: String },
100    /// The tool body or a hook returned an error.
101    Failed { message: String },
102}
103
104/// Append-only collector for [`DispatchTraceEvent`] values.
105///
106/// Clones share the same underlying buffer (Arc + Mutex), so a trace handed to
107/// the dispatcher and a trace inspected by the host see the same events.
108#[derive(Debug, Default, Clone)]
109pub struct DispatchTrace {
110    inner: Arc<Mutex<Vec<DispatchTraceEvent>>>,
111}
112
113impl DispatchTrace {
114    /// Create an empty trace.
115    pub fn new() -> Self {
116        Self::default()
117    }
118
119    /// Append an event. Used by the dispatcher; hosts normally read via
120    /// [`Self::events`].
121    pub(crate) fn push(&self, event: DispatchTraceEvent) {
122        if let Ok(mut guard) = self.inner.lock() {
123            guard.push(event);
124        }
125    }
126
127    /// Snapshot the events recorded so far.
128    pub fn events(&self) -> Vec<DispatchTraceEvent> {
129        self.inner
130            .lock()
131            .map(|guard| guard.clone())
132            .unwrap_or_default()
133    }
134
135    /// Number of recorded events.
136    pub fn len(&self) -> usize {
137        self.inner.lock().map(|g| g.len()).unwrap_or(0)
138    }
139
140    /// Whether no events have been recorded.
141    pub fn is_empty(&self) -> bool {
142        self.len() == 0
143    }
144}
rig_compose/trace.rs

rig_compose/
trace.rs
