rig_core/serve/
record_buffer.rs

1//! A [`Recorder`] that collects whole [`EffectRecord`]s by effect id, for
2//! its owner to take.
3//!
4//! ```
5//! use rig_core::effect::{EffectId, EffectKind, Answered, tool_key};
6//! use rig_core::serve::{DispatchContext, RecordBuffer, Recorder};
7//! use rig_core::tool::ToolResult;
8//!
9//! let buffer = RecordBuffer::new();
10//! let id = EffectId::from_raw(1);
11//! let kind = EffectKind::ToolCall { name: "add".into(), args: "{}".into() };
12//! buffer.begin(id, tool_key("add"), kind, DispatchContext::default());
13//! assert_eq!(buffer.in_flight(), 1);
14//! buffer.resolve(id, Ok(Answered::ToolResult { result: ToolResult::skipped("no") }));
15//! let record = buffer.take(id).ok_or("resolved")?;
16//! assert!(record.outcome.is_ok());
17//! # Ok::<(), &str>(())
18//! ```
19
20use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
21
22use super::{DispatchContext, Recorder};
23use crate::{
24    effect::{
25        Answered, Disposition, EffectId, EffectKind, EffectRecord, HandlerDescriptor, HandlerKey,
26    },
27    error::ErrorReport,
28    streaming::{Item, StreamEvent, StreamRecord},
29    tool::ToolResultContext,
30};
31
32/// Collects the record of every dispatch it is told about, in the order the
33/// dispatches began, whatever order they resolve in. A dispatch recorded by
34/// it ([`Dispatch::recorded_by`](super::Dispatch::recorded_by)), and every
35/// dispatch nested in it, becomes one [`EffectRecord`] once resolved: a
36/// layer's [`discard`](Recorder::discard) drops it and a
37/// [`patch`](Recorder::patch) replaces its request, as an effect log keeps
38/// them. Clones share the records.
39///
40/// It keeps no handler descriptions and no delivery boundaries: an owner
41/// that writes a log keeps those itself. Its records are
42/// [`Disposition::Ran`] with no reviewed outcome; a host that answers,
43/// reviews or cancels effects itself sets those on the records it takes.
44#[derive(Clone, Default)]
45pub struct RecordBuffer {
46    slots: Arc<Mutex<Vec<Slot>>>,
47    keep_events: bool,
48    timed: bool,
49}
50
51/// One dispatch: opened when it begins, filled when it resolves.
52#[derive(Clone)]
53struct Slot {
54    id: EffectId,
55    key: HandlerKey,
56    kind: EffectKind,
57    origin: DispatchContext,
58    outcome: Option<Result<Answered, ErrorReport>>,
59    events: Option<Vec<Item<StreamEvent>>>,
60    stream_origin: Option<crate::message::Origin>,
61    tool_output: Option<ToolResultContext>,
62    started_ms: Option<u64>,
63    finished_ms: Option<u64>,
64}
65
66impl Slot {
67    /// The record, once the dispatch resolved.
68    fn into_record(self) -> Option<EffectRecord> {
69        Some(EffectRecord {
70            published: self.tool_output,
71            id: self.id,
72            key: self.key,
73            kind: self.kind,
74            outcome: self.outcome?,
75            events: self
76                .events
77                .and_then(|items| StreamRecord::from_items(items).ok()),
78            stream_origin: self.stream_origin,
79            parent: self.origin.parent,
80            scope: self.origin.scope,
81            purpose: self.origin.purpose,
82            turn: None,
83            started_ms: self.started_ms,
84            finished_ms: self.finished_ms,
85            reviewed: None,
86            disposition: Disposition::Ran,
87            decided_by: None,
88        })
89    }
90}
91
92impl RecordBuffer {
93    /// An empty buffer that records a stream as its folded outcome and keeps
94    /// no times.
95    pub fn new() -> Self {
96        Self::default()
97    }
98
99    /// This buffer keeping a streamed dispatch's items verbatim in its
100    /// record ([`EffectRecord::events`]), so a replay re-emits the original
101    /// delta boundaries.
102    #[must_use]
103    pub fn keeping_stream_events(mut self) -> Self {
104        self.keep_events = true;
105        self
106    }
107
108    /// This buffer stamping each record with the wall-clock time its
109    /// dispatch began ([`EffectRecord::started_ms`]) and resolved
110    /// ([`EffectRecord::finished_ms`]). Without it records carry no times,
111    /// so a recording is the same on every run.
112    #[must_use]
113    pub fn timed(mut self) -> Self {
114        self.timed = true;
115        self
116    }
117
118    /// Takes the record of `id` once it resolved; `None` while it is in
119    /// flight, or when the buffer has none.
120    pub fn take(&self, id: EffectId) -> Option<EffectRecord> {
121        let mut slots = self.slots();
122        let position = slots
123            .iter()
124            .rposition(|slot| slot.id == id && slot.outcome.is_some())?;
125        slots.remove(position).into_record()
126    }
127
128    /// Takes every resolved record, in the order the dispatches began,
129    /// leaving the ones in flight.
130    pub fn take_resolved(&self) -> Vec<EffectRecord> {
131        self.take_resolved_where(|_| true)
132    }
133
134    /// Takes every resolved record whose id `take` accepts, in the order
135    /// the dispatches began, leaving the rest.
136    pub fn take_resolved_where(&self, mut take: impl FnMut(EffectId) -> bool) -> Vec<EffectRecord> {
137        let mut slots = self.slots();
138        let (taken, left) = std::mem::take(&mut *slots)
139            .into_iter()
140            .partition(|slot| slot.outcome.is_some() && take(slot.id));
141        *slots = left;
142        drop(slots);
143        taken.into_iter().filter_map(Slot::into_record).collect()
144    }
145
146    /// A copy of every resolved record, in the order the dispatches began.
147    pub fn records(&self) -> Vec<EffectRecord> {
148        let slots = self.slots();
149        let resolved: Vec<Slot> = slots
150            .iter()
151            .filter(|slot| slot.outcome.is_some())
152            .cloned()
153            .collect();
154        drop(slots);
155        resolved.into_iter().filter_map(Slot::into_record).collect()
156    }
157
158    /// Whether the buffer holds a record of `id`, resolved or not.
159    pub fn contains(&self, id: EffectId) -> bool {
160        self.slots().iter().any(|slot| slot.id == id)
161    }
162
163    /// Drops the record of `id`, resolved or not; the key its dispatch was
164    /// routed to, or `None` when the buffer held none.
165    pub fn remove(&self, id: EffectId) -> Option<HandlerKey> {
166        let mut slots = self.slots();
167        let position = slots.iter().rposition(|slot| slot.id == id)?;
168        Some(slots.remove(position).key)
169    }
170
171    /// How many streamed items the record of `id` keeps so far, or `None`
172    /// when it keeps none.
173    pub fn kept_events(&self, id: EffectId) -> Option<usize> {
174        let slots = self.slots();
175        let slot = slots.iter().rev().find(|slot| slot.id == id)?;
176        slot.events.as_ref().map(Vec::len)
177    }
178
179    /// Records resolved and not yet taken.
180    pub fn resolved(&self) -> usize {
181        let slots = self.slots();
182        slots.iter().filter(|slot| slot.outcome.is_some()).count()
183    }
184
185    /// Records begun and not yet resolved.
186    pub fn in_flight(&self) -> usize {
187        let slots = self.slots();
188        slots.iter().filter(|slot| slot.outcome.is_none()).count()
189    }
190
191    fn slots(&self) -> MutexGuard<'_, Vec<Slot>> {
192        self.slots.lock().unwrap_or_else(PoisonError::into_inner)
193    }
194
195    /// Edits the newest record of `id` that `keep` accepts. Records in
196    /// flight are usually the newest, so the search starts there.
197    fn with_slot(&self, id: EffectId, keep: impl Fn(&Slot) -> bool, edit: impl FnOnce(&mut Slot)) {
198        let mut slots = self.slots();
199        if let Some(slot) = slots
200            .iter_mut()
201            .rev()
202            .find(|slot| slot.id == id && keep(slot))
203        {
204            edit(slot);
205        }
206    }
207
208    /// Milliseconds since the Unix epoch, when the buffer keeps times.
209    fn now(&self) -> Option<u64> {
210        use web_time::{SystemTime, UNIX_EPOCH};
211        if !self.timed {
212            return None;
213        }
214        let since = SystemTime::now().duration_since(UNIX_EPOCH).ok()?;
215        Some(u64::try_from(since.as_millis()).unwrap_or(u64::MAX))
216    }
217}
218
219impl std::fmt::Debug for RecordBuffer {
220    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
221        f.debug_struct("RecordBuffer")
222            .field("resolved", &self.resolved())
223            .field("in_flight", &self.in_flight())
224            .finish()
225    }
226}
227
228impl Recorder for RecordBuffer {
229    fn handlers(&self, _handlers: Vec<HandlerDescriptor>) {}
230
231    fn begin(&self, id: EffectId, key: HandlerKey, kind: EffectKind, origin: DispatchContext) {
232        let events = (self.keep_events && kind.streams()).then(Vec::new);
233        let started_ms = self.now();
234        self.slots().push(Slot {
235            id,
236            key,
237            kind,
238            origin,
239            outcome: None,
240            events,
241            stream_origin: None,
242            tool_output: None,
243            started_ms,
244            finished_ms: None,
245        });
246    }
247
248    fn discard(&self, id: EffectId) {
249        self.remove(id);
250    }
251
252    fn patch(&self, id: EffectId, kind: EffectKind) {
253        self.with_slot(id, |_| true, |slot| slot.kind = kind);
254    }
255
256    fn keep_events(&self) -> bool {
257        self.keep_events
258    }
259
260    fn event(&self, id: EffectId, item: &Item<StreamEvent>) {
261        self.with_slot(
262            id,
263            |_| true,
264            |slot| {
265                if let Some(events) = &mut slot.events {
266                    events.push(item.clone());
267                }
268            },
269        );
270    }
271
272    fn origin(&self, id: EffectId, origin: &crate::message::Origin) {
273        let keeps = |slot: &Slot| slot.events.is_some();
274        self.with_slot(id, keeps, |slot| {
275            slot.stream_origin = Some(origin.clone());
276        });
277    }
278
279    fn tool_output(&self, id: EffectId, output: ToolResultContext) {
280        let open = |slot: &Slot| slot.outcome.is_none();
281        self.with_slot(id, open, |slot| slot.tool_output = Some(output));
282    }
283
284    fn resolve(&self, id: EffectId, outcome: Result<Answered, ErrorReport>) {
285        let finished_ms = self.now();
286        self.with_slot(
287            id,
288            |_| true,
289            |slot| {
290                slot.outcome = Some(outcome);
291                slot.finished_ms = finished_ms;
292            },
293        );
294    }
295}
296
297#[cfg(test)]
298mod tests;
rig_core/serve/record_buffer.rs

rig_core/serve/
record_buffer.rs
