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    first_item_ms: Option<u64>,
64    finished_ms: Option<u64>,
65}
66
67impl Slot {
68    /// The record, once the dispatch resolved.
69    fn into_record(self) -> Option<EffectRecord> {
70        Some(EffectRecord {
71            published: self.tool_output,
72            id: self.id,
73            key: self.key,
74            kind: self.kind,
75            outcome: self.outcome?,
76            events: self
77                .events
78                .and_then(|items| StreamRecord::from_items(items).ok()),
79            stream_origin: self.stream_origin,
80            parent: self.origin.parent,
81            scope: self.origin.scope,
82            purpose: self.origin.purpose,
83            turn: None,
84            started_ms: self.started_ms,
85            first_item_ms: self.first_item_ms,
86            finished_ms: self.finished_ms,
87            reviewed: None,
88            disposition: Disposition::Ran,
89            decided_by: None,
90        })
91    }
92}
93
94impl RecordBuffer {
95    /// An empty buffer that records a stream as its folded outcome and keeps
96    /// no times.
97    pub fn new() -> Self {
98        Self::default()
99    }
100
101    /// This buffer keeping a streamed dispatch's items verbatim in its
102    /// record ([`EffectRecord::events`]), so a replay re-emits the original
103    /// delta boundaries.
104    #[must_use]
105    pub fn keeping_stream_events(mut self) -> Self {
106        self.keep_events = true;
107        self
108    }
109
110    /// This buffer stamping each record with the wall-clock time its
111    /// dispatch began ([`EffectRecord::started_ms`]), streamed its first
112    /// item ([`EffectRecord::first_item_ms`]) and resolved
113    /// ([`EffectRecord::finished_ms`]). Without it records carry no times,
114    /// so a recording is the same on every run.
115    #[must_use]
116    pub fn timed(mut self) -> Self {
117        self.timed = true;
118        self
119    }
120
121    /// Takes the record of `id` once it resolved; `None` while it is in
122    /// flight, or when the buffer has none.
123    pub fn take(&self, id: EffectId) -> Option<EffectRecord> {
124        let mut slots = self.slots();
125        let position = slots
126            .iter()
127            .rposition(|slot| slot.id == id && slot.outcome.is_some())?;
128        slots.remove(position).into_record()
129    }
130
131    /// Takes every resolved record, in the order the dispatches began,
132    /// leaving the ones in flight.
133    pub fn take_resolved(&self) -> Vec<EffectRecord> {
134        self.take_resolved_where(|_| true)
135    }
136
137    /// Takes every resolved record whose id `take` accepts, in the order
138    /// the dispatches began, leaving the rest.
139    pub fn take_resolved_where(&self, mut take: impl FnMut(EffectId) -> bool) -> Vec<EffectRecord> {
140        let mut slots = self.slots();
141        let (taken, left) = std::mem::take(&mut *slots)
142            .into_iter()
143            .partition(|slot| slot.outcome.is_some() && take(slot.id));
144        *slots = left;
145        drop(slots);
146        taken.into_iter().filter_map(Slot::into_record).collect()
147    }
148
149    /// A copy of every resolved record, in the order the dispatches began.
150    pub fn records(&self) -> Vec<EffectRecord> {
151        let slots = self.slots();
152        let resolved: Vec<Slot> = slots
153            .iter()
154            .filter(|slot| slot.outcome.is_some())
155            .cloned()
156            .collect();
157        drop(slots);
158        resolved.into_iter().filter_map(Slot::into_record).collect()
159    }
160
161    /// Whether the buffer holds a record of `id`, resolved or not.
162    pub fn contains(&self, id: EffectId) -> bool {
163        self.slots().iter().any(|slot| slot.id == id)
164    }
165
166    /// Drops the record of `id`, resolved or not; the key its dispatch was
167    /// routed to, or `None` when the buffer held none.
168    pub fn remove(&self, id: EffectId) -> Option<HandlerKey> {
169        let mut slots = self.slots();
170        let position = slots.iter().rposition(|slot| slot.id == id)?;
171        Some(slots.remove(position).key)
172    }
173
174    /// How many streamed items the record of `id` keeps so far, or `None`
175    /// when it keeps none.
176    pub fn kept_events(&self, id: EffectId) -> Option<usize> {
177        let slots = self.slots();
178        let slot = slots.iter().rev().find(|slot| slot.id == id)?;
179        slot.events.as_ref().map(Vec::len)
180    }
181
182    /// Records resolved and not yet taken.
183    pub fn resolved(&self) -> usize {
184        let slots = self.slots();
185        slots.iter().filter(|slot| slot.outcome.is_some()).count()
186    }
187
188    /// Records begun and not yet resolved.
189    pub fn in_flight(&self) -> usize {
190        let slots = self.slots();
191        slots.iter().filter(|slot| slot.outcome.is_none()).count()
192    }
193
194    fn slots(&self) -> MutexGuard<'_, Vec<Slot>> {
195        self.slots.lock().unwrap_or_else(PoisonError::into_inner)
196    }
197
198    /// Edits the newest record of `id` that `keep` accepts. Records in
199    /// flight are usually the newest, so the search starts there.
200    fn with_slot(&self, id: EffectId, keep: impl Fn(&Slot) -> bool, edit: impl FnOnce(&mut Slot)) {
201        let mut slots = self.slots();
202        if let Some(slot) = slots
203            .iter_mut()
204            .rev()
205            .find(|slot| slot.id == id && keep(slot))
206        {
207            edit(slot);
208        }
209    }
210
211    /// Milliseconds since the Unix epoch, when the buffer keeps times.
212    fn now(&self) -> Option<u64> {
213        self.timed.then(crate::now_ms)
214    }
215}
216
217impl std::fmt::Debug for RecordBuffer {
218    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
219        f.debug_struct("RecordBuffer")
220            .field("resolved", &self.resolved())
221            .field("in_flight", &self.in_flight())
222            .finish()
223    }
224}
225
226impl Recorder for RecordBuffer {
227    fn handlers(&self, _handlers: Vec<HandlerDescriptor>) {}
228
229    fn begin(&self, id: EffectId, key: HandlerKey, kind: EffectKind, origin: DispatchContext) {
230        let events = (self.keep_events && kind.streams()).then(Vec::new);
231        let started_ms = self.now();
232        self.slots().push(Slot {
233            id,
234            key,
235            kind,
236            origin,
237            outcome: None,
238            events,
239            stream_origin: None,
240            tool_output: None,
241            started_ms,
242            first_item_ms: None,
243            finished_ms: None,
244        });
245    }
246
247    fn discard(&self, id: EffectId) {
248        self.remove(id);
249    }
250
251    fn patch(&self, id: EffectId, kind: EffectKind) {
252        self.with_slot(id, |_| true, |slot| slot.kind = kind);
253    }
254
255    fn keep_events(&self) -> bool {
256        self.keep_events
257    }
258
259    fn event(&self, id: EffectId, item: &Item<StreamEvent>) {
260        if !self.keep_events && !self.timed {
261            return;
262        }
263        self.with_slot(
264            id,
265            |_| true,
266            |slot| {
267                slot.first_item_ms = slot.first_item_ms.or_else(|| self.now());
268                if let Some(events) = &mut slot.events {
269                    events.push(item.clone());
270                }
271            },
272        );
273    }
274
275    fn origin(&self, id: EffectId, origin: &crate::message::Origin) {
276        let keeps = |slot: &Slot| slot.events.is_some();
277        self.with_slot(id, keeps, |slot| {
278            slot.stream_origin = Some(origin.clone());
279        });
280    }
281
282    fn tool_output(&self, id: EffectId, output: ToolResultContext) {
283        let open = |slot: &Slot| slot.outcome.is_none();
284        self.with_slot(id, open, |slot| slot.tool_output = Some(output));
285    }
286
287    fn resolve(&self, id: EffectId, outcome: Result<Answered, ErrorReport>) {
288        let finished_ms = self.now();
289        self.with_slot(
290            id,
291            |_| true,
292            |slot| {
293                slot.outcome = Some(outcome);
294                slot.finished_ms = finished_ms;
295            },
296        );
297    }
298}
299
300#[cfg(test)]
301mod tests;
rig_core/serve/record_buffer.rs

rig_core/serve/
record_buffer.rs
