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}