Skip to main content

rig_core/serve/
recorder.rs

1//! Runtime-independent recording of handler dispatches.
2//!
3//! ```
4//! use rig_core::serve::DispatchContext;
5//!
6//! let origin = DispatchContext::default();
7//! assert!(origin.parent.is_none());
8//! ```
9
10use std::sync::Arc;
11
12use crate::{
13    effect::{Answered, EffectId, EffectKind, HandlerDescriptor, HandlerKey},
14    error::ErrorReport,
15    streaming::{Item, StreamEvent},
16    wasm_compat::{WasmCompatSend, WasmCompatSync},
17};
18
19/// Recorded parent dispatch, stable program scope identifier and purpose,
20/// without live runtime handles.
21#[derive(Clone, Debug, Default, PartialEq, Eq)]
22pub struct DispatchContext {
23    /// The dispatch this one was made from, if a handler made it.
24    pub parent: Option<EffectId>,
25    /// The scope of the program that made it, if its dispatcher was scoped.
26    pub scope: Option<std::sync::Arc<str>>,
27    /// What it was made for, such as `"turn"` or `"summary"`, if its
28    /// dispatcher said.
29    pub purpose: Option<std::borrow::Cow<'static, str>>,
30}
31
32/// What a driver tells about the dispatches it serves. A driver calls
33/// [`handlers`](Self::handlers) once when recording starts,
34/// [`begin`](Self::begin) as each dispatch is handed to its handler,
35/// [`event`](Self::event) for every streamed event when
36/// [`keep_events`](Self::keep_events) says so and for the first otherwise,
37/// and
38/// [`resolve`](Self::resolve) when the outcome is known. A recorder is
39/// shared between the driver and its owner, so every method takes `&self`;
40/// it rides in the dispatch observer and uses the platform compatibility bounds.
41pub trait Recorder: WasmCompatSend + WasmCompatSync + 'static {
42    /// Handlers the driver serves: those registered when recording started,
43    /// then each one installed later, as it is installed. A key described
44    /// again is the same handler re-registered; the latest description
45    /// stands.
46    fn handlers(&self, handlers: Vec<HandlerDescriptor>);
47    /// A dispatch begins: its id, the key it was routed to, the effect, and
48    /// where it came from (its parent and scope).
49    fn begin(&self, id: EffectId, key: HandlerKey, kind: EffectKind, origin: DispatchContext);
50    /// Removes a begun dispatch decided by a layer before handler execution.
51    /// Replay reruns layer decisions rather than recording them as handler outcomes.
52    fn discard(&self, id: EffectId);
53    /// Replaces the recorded request with a same-family layer patch, so it
54    /// reflects the request served by the innermost handler.
55    fn patch(&self, id: EffectId, kind: EffectKind);
56    /// Whether streamed items are wanted verbatim ([`Self::event`]).
57    fn keep_events(&self) -> bool;
58    /// One streamed item of `id`: each one when
59    /// [`keep_events`](Self::keep_events) says so, else only the first, which
60    /// a timed recorder stamps.
61    fn event(&self, id: EffectId, item: &Item<StreamEvent>);
62    /// An error item at its original position in a kept stream. Unlike the
63    /// folded outcome, this includes errors after an earlier terminal item.
64    fn stream_error(&self, _id: EffectId, _error: &ErrorReport) {}
65    /// Who the streamed reply of `id` is from, before its first item, when
66    /// [`keep_events`](Self::keep_events) says so.
67    fn origin(&self, id: EffectId, origin: &crate::message::Origin);
68    /// Explicitly published output, a tool's or a layer's facts, delivered
69    /// before `resolve`. A driver snapshots it without consuming the
70    /// caller's published context.
71    fn tool_output(&self, id: EffectId, output: crate::tool::ToolResultContext);
72    /// The outcome of `id`.
73    fn resolve(&self, id: EffectId, outcome: Result<Answered, ErrorReport>);
74}
75
76/// The [`Observe`](super::Observe) a recording driver installs on each
77/// dispatch ([`Dispatch::recorded_by`](super::Dispatch::recorded_by)): it
78/// tells `recorder` everything about the dispatch `id`, and the tool output
79/// its handler published, if any, before the outcome.
80pub(crate) struct RecordingObserver {
81    pub(crate) recorder: Arc<dyn Recorder + Send + Sync>,
82    pub(crate) id: EffectId,
83    pub(crate) published: Option<Arc<crate::tool::PublishedContext>>,
84    /// Whether an item streamed yet.
85    pub(crate) streamed: bool,
86}
87
88impl super::Observe for RecordingObserver {
89    fn outcome(&mut self, outcome: Result<Answered, ErrorReport>) {
90        if let Some(output) = self
91            .published
92            .as_ref()
93            .and_then(|published| published.result_context())
94        {
95            self.recorder.tool_output(self.id, output);
96        }
97        self.recorder.resolve(self.id, outcome);
98    }
99
100    fn keep_events(&self) -> bool {
101        self.recorder.keep_events()
102    }
103
104    fn event(&mut self, item: &Item<StreamEvent>) {
105        // Past the first, an item matters only to a recorder that keeps them.
106        if !std::mem::replace(&mut self.streamed, true) || self.recorder.keep_events() {
107            self.recorder.event(self.id, item);
108        }
109    }
110
111    fn stream_error(&mut self, error: &ErrorReport) {
112        self.recorder.stream_error(self.id, error);
113    }
114
115    fn origin(&mut self, origin: &crate::message::Origin) {
116        self.recorder.origin(self.id, origin);
117    }
118
119    /// A layer decided the dispatch before any handler saw it: the record
120    /// is dropped, since replay reruns the layers.
121    fn discard(&mut self, _layer: &str) {
122        self.recorder.discard(self.id);
123    }
124
125    fn patch(&mut self, kind: &EffectKind) {
126        self.recorder.patch(self.id, kind.clone());
127    }
128}
129
130/// `work`, the task that drives the dispatch `id`, with a panic in it (in
131/// the handler or in its reply's stream) turned into an internal error that
132/// `recorder`, if any, also records as the dispatch's outcome. Unwinding
133/// drops the dispatch's observer, which records a cancellation; this
134/// replaces it.
135pub fn catch_panics<T>(
136    recorder: Option<Arc<dyn Recorder + Send + Sync>>,
137    id: EffectId,
138    work: impl Future<Output = Result<T, ErrorReport>>,
139) -> impl Future<Output = Result<T, ErrorReport>> {
140    use futures::FutureExt;
141    async move {
142        std::panic::AssertUnwindSafe(work)
143            .catch_unwind()
144            .await
145            .unwrap_or_else(|panic| {
146                let message = panic
147                    .downcast_ref::<&str>()
148                    .copied()
149                    .or_else(|| panic.downcast_ref::<String>().map(String::as_str))
150                    .unwrap_or("no message");
151                let report = ErrorReport::new(
152                    crate::error::ErrorKind::Internal,
153                    format!("panicked: {message}"),
154                );
155                if let Some(recorder) = recorder {
156                    recorder.resolve(id, Err(report.clone()));
157                }
158                Err(report)
159            })
160    }
161}