rig_ecs/bus/
record.rs

1//! Effect recording and cancellation observers for the world bus.
2//!
3//! ```
4//! use rig_ecs::bus::record::DeliveryBatch;
5//! let batch = DeliveryBatch::default();
6//! assert_eq!(batch.0, 0);
7//! ```
8
9use std::sync::Arc;
10
11use bevy_ecs::{
12    lifecycle::{Insert, Remove},
13    prelude::*,
14};
15use rig_core::{
16    effect::{Delivery, DeliveryKind, EffectId, EffectKind, HandlerDescriptor, Outcome},
17    error::ErrorReport,
18    serve::{Recorder, cancelled},
19    streaming::{Item, Relayed, StreamEvent},
20};
21
22use super::{
23    effect::{EffectOutcome, InFlight, Issued},
24    handlers::Bound,
25};
26use bevy_ecs::component::Component;
27
28/// The world's recorder. `Dispatch` opens a record as it takes an effect,
29/// `Collect` closes it as the outcome lands, a despawn in flight closes it
30/// as cancelled, and every `Bound` insert describes its handler. Install
31/// with [`Recording::install`] (which also describes the handlers bound so
32/// far) or as a plain resource before any handler is bound. It derefs to
33/// the [`Recorder`] it records into.
34#[derive(Resource, Clone)]
35pub struct Recording(Arc<dyn Recorder + Send + Sync>);
36
37impl Recording {
38    /// A recording into `recorder`.
39    pub fn new(recorder: impl Recorder + Send + Sync) -> Self {
40        Self(Arc::new(recorder))
41    }
42
43    /// Begin delivery tracking, describe currently bound handlers in key order,
44    /// and install `recorder` as the world's recording resource.
45    pub fn install(world: &mut World, recorder: impl Recorder + Send + Sync) {
46        let recording = Self::new(recorder);
47        recording.begin_delivery_tracking();
48        let mut bound: Vec<HandlerDescriptor> = world
49            .query::<&Bound>()
50            .iter(world)
51            .map(|bound| bound.descriptor.clone())
52            .collect();
53        bound.sort_by(|a, b| a.key.cmp(&b.key));
54        recording.handlers(bound);
55        world.insert_resource(recording);
56    }
57}
58
59impl std::ops::Deref for Recording {
60    type Target = dyn Recorder + Send + Sync;
61
62    fn deref(&self) -> &Self::Target {
63        &*self.0
64    }
65}
66
67/// The current delivery batch. Advances once per schedule pass.
68#[derive(Resource, Default)]
69pub struct DeliveryBatch(pub u64);
70
71/// Begin the next pass's observation group.
72pub fn begin_delivery_pass(mut batch: ResMut<DeliveryBatch>) {
73    batch.0 += 1;
74}
75
76/// Record visibility when the outcome is inserted, not later when a query
77/// happens to visit it. This also captures answers from world-served handlers.
78pub fn record_outcome(
79    added: On<Add, EffectOutcome>,
80    issued: Query<(&Issued, Has<super::collect::CollectedOutcome>), With<InFlight>>,
81    recording: Option<Res<Recording>>,
82    batch: Res<DeliveryBatch>,
83) {
84    if let Some(recording) = recording
85        && let Ok((Issued(id), collected)) = issued.get(added.event().entity)
86    {
87        if !collected {
88            recording.unsupported_delivery("an in-flight EffectOutcome bypassed Collect; submit world answers with WorldOutcome or typed Answer instead");
89        }
90        recording.delivery(Delivery {
91            batch: batch.0,
92            id: *id,
93            kind: DeliveryKind::Outcome,
94        });
95    }
96}
97
98/// Shared in-flight observation state containing the original handler answer
99/// and whether a layer discarded dispatch. Not serialized; layer verdicts do not
100/// replace the recorded answer.
101#[derive(Component, Clone, Default)]
102pub struct Observed(pub Arc<ObservedState>);
103
104/// The dispatch's replacement slot (`Dispatch::replaced_by`): the layer
105/// whose verdict the consumer's answer is from, once one said so. Runtime-
106/// only; removed with the in-flight markers at collection.
107#[derive(Component)]
108pub struct ReplacedBy(pub Arc<std::sync::Mutex<Option<String>>>);
109
110impl ReplacedBy {
111    /// The layer named, if any.
112    pub fn layer(&self) -> Option<String> {
113        self.0
114            .lock()
115            .unwrap_or_else(std::sync::PoisonError::into_inner)
116            .clone()
117    }
118}
119
120/// The observer's slots.
121#[derive(Default)]
122pub struct ObservedState {
123    state: std::sync::Mutex<Observation>,
124}
125
126#[derive(Default)]
127struct Observation {
128    outcome: Option<Result<Outcome, ErrorReport>>,
129    discarded: bool,
130    closed: bool,
131}
132
133impl ObservedState {
134    fn lock(&self) -> std::sync::MutexGuard<'_, Observation> {
135        self.state
136            .lock()
137            .unwrap_or_else(std::sync::PoisonError::into_inner)
138    }
139
140    /// Close observation and take the original answer. No later worker callback
141    /// can modify recording after this boundary, even if its poll was in progress.
142    pub fn take_outcome(&self) -> Option<Result<Outcome, ErrorReport>> {
143        let mut state = self.lock();
144        state.closed = true;
145        state.outcome.take()
146    }
147
148    /// Whether the original handler answer has been observed.
149    pub fn has_outcome(&self) -> bool {
150        self.lock().outcome.is_some()
151    }
152
153    /// Whether a layer discarded the dispatch.
154    pub fn is_discarded(&self) -> bool {
155        self.lock().discarded
156    }
157}
158
159/// Records original answers, stream items, and layer patches or discards for a
160/// task-served dispatch. Once observation closes, worker callbacks cannot mutate
161/// the recording.
162pub struct WorldObserver {
163    /// Explicit host operation, already bound to this dispatch's subject.
164    pub adapter: Option<rig_core::observe::AdapterContext>,
165    /// Tool output shared with the caller, read without consuming it.
166    pub published: Option<Arc<rig_core::tool::PublishedContext>>,
167    /// The dispatch.
168    pub id: EffectId,
169    /// The world's recorder, if any.
170    pub recording: Option<Recording>,
171    /// The shared slots.
172    pub observed: Arc<ObservedState>,
173    /// The world's witness with the current dispatch subject.
174    pub witness: Option<(super::witness::Witnessing, rig_core::observe::Subject)>,
175}
176
177impl WorldObserver {
178    fn record_answer(&self, state: &mut Observation, outcome: &Result<Outcome, ErrorReport>) {
179        if let (Some(recording), Some(output)) = (
180            &self.recording,
181            self.published
182                .as_ref()
183                .and_then(|published| published.result_context()),
184        ) {
185            recording.tool_output(self.id, output);
186        }
187        state.outcome = Some(outcome.clone());
188    }
189
190    fn record_item(&self, item: &Result<Relayed, ErrorReport>) {
191        if let Some(recording) = &self.recording
192            && recording.keep_events()
193        {
194            match item {
195                Ok(Relayed::Origin(origin)) => recording.origin(self.id, origin),
196                Ok(Relayed::Item(item)) => recording.event(self.id, item),
197                Ok(Relayed::Done(_)) => {}
198                Err(error) => recording.stream_error(self.id, error),
199            }
200        }
201    }
202}
203
204impl rig_core::serve::Observe for WorldObserver {
205    fn adapter_context(&self) -> Option<rig_core::observe::AdapterContext> {
206        if let Some(context) = &self.adapter {
207            return Some(context.clone());
208        }
209        let (witness, subject) = self.witness.as_ref()?;
210        let mut subject = subject.clone();
211        subject.effect = Some(self.id);
212        Some(rig_core::observe::AdapterContext::new(
213            witness.sink().clone(),
214            subject,
215            format!("effect/{}", self.id),
216        ))
217    }
218
219    fn outcome(&mut self, outcome: &Result<Outcome, ErrorReport>) {
220        let mut state = self.observed.lock();
221        if !state.closed {
222            self.record_answer(&mut state, outcome);
223        }
224    }
225
226    fn keep_events(&self) -> bool {
227        self.recording
228            .as_ref()
229            .is_some_and(|recording| recording.keep_events())
230    }
231
232    fn stream_item(
233        &mut self,
234        item: &Result<Relayed, ErrorReport>,
235        outcome: Option<&Result<Outcome, ErrorReport>>,
236    ) {
237        // Source polling and folding happen before this lock. Only observation
238        // is atomic with cancellation, so arbitrary handler work cannot hold it.
239        let mut state = self.observed.lock();
240        if state.closed {
241            return;
242        }
243        self.record_item(item);
244        if let Some(outcome) = outcome {
245            self.record_answer(&mut state, outcome);
246        }
247    }
248
249    fn origin(&mut self, origin: &rig_core::message::Origin) {
250        let state = self.observed.lock();
251        if !state.closed
252            && let Some(recording) = &self.recording
253        {
254            recording.origin(self.id, origin);
255        }
256    }
257
258    fn event(&mut self, item: &Item<StreamEvent>) {
259        let state = self.observed.lock();
260        if !state.closed
261            && let Some(recording) = &self.recording
262        {
263            recording.event(self.id, item);
264        }
265    }
266
267    fn stream_error(&mut self, error: &ErrorReport) {
268        let state = self.observed.lock();
269        if !state.closed
270            && let Some(recording) = &self.recording
271        {
272            recording.stream_error(self.id, error);
273        }
274    }
275
276    fn discard(&mut self, layer: &str) {
277        let mut state = self.observed.lock();
278        if !state.closed {
279            state.discarded = true;
280            if let Some(recording) = &self.recording {
281                recording.discard(self.id);
282            }
283            if let Some((witness, subject)) = &self.witness {
284                let mut subject = subject.clone();
285                subject.effect = Some(self.id);
286                witness.emit(
287                    subject,
288                    rig_core::observe::Stage::Handler,
289                    rig_core::observe::Emitter::named(layer),
290                    rig_core::observe::Action::Denied {
291                        reason: rig_core::observe::Reason::code("layer_discarded"),
292                    },
293                );
294            }
295        }
296    }
297
298    fn patch(&mut self, kind: &EffectKind) {
299        let state = self.observed.lock();
300        if !state.closed
301            && let Some(recording) = &self.recording
302        {
303            recording.patch(self.id, kind.clone());
304        }
305    }
306}
307
308/// Close recording when an effect leaves flight before settlement.
309/// Preserve any original handler answer and tool output; otherwise record
310/// cancellation. An observed answer cancelled before delivery gets a cancellation
311/// boundary, not an invented outcome delivery. Discarded dispatches stay unrecorded.
312pub fn record_cancelled(
313    removed: On<Remove, InFlight>,
314    effects: Query<(
315        &Issued,
316        Option<&EffectOutcome>,
317        Has<super::collect::CollectedOutcome>,
318        Option<&super::effect::Publishing>,
319        Option<&super::effect::ToolOutputs>,
320        Option<&Observed>,
321    )>,
322    recording: Option<Res<Recording>>,
323    batch: Res<DeliveryBatch>,
324) {
325    let Some(recording) = recording else {
326        return;
327    };
328    let Ok((Issued(id), outcome, collected, publishing, outputs, observed)) =
329        effects.get(removed.event().entity)
330    else {
331        return;
332    };
333    if observed.is_some_and(|observed| observed.0.is_discarded()) {
334        return;
335    }
336    match outcome {
337        // Landed but not settled: `settle` retires `CollectedOutcome` before
338        // it retires `InFlight`, so its presence here means settlement never
339        // ran (the entity is being despawned with its answer on it).
340        Some(EffectOutcome(landed)) if collected => {
341            let original = observed.and_then(|observed| observed.0.take_outcome());
342            if observed.is_none()
343                && let Some(outputs) = outputs
344            {
345                recording.tool_output(*id, outputs.0.result_context());
346            }
347            recording.resolve(*id, original.unwrap_or_else(|| landed.clone()));
348        }
349        Some(_) => {}
350        None => {
351            let original = observed.and_then(|observed| observed.0.take_outcome());
352            let output = publishing
353                .and_then(|published| published.0.result_context())
354                .or_else(|| outputs.map(|outputs| outputs.0.result_context()));
355            if let Some(output) = output {
356                recording.tool_output(*id, output);
357            }
358            if original.as_ref().is_some_and(|answer| {
359                !answer
360                    .as_ref()
361                    .is_err_and(|error| error.kind == rig_core::error::ErrorKind::Cancelled)
362            }) {
363                recording.delivery(Delivery {
364                    batch: batch.0,
365                    id: *id,
366                    kind: DeliveryKind::Cancelled,
367                });
368            }
369            recording.resolve(*id, original.unwrap_or_else(|| Err(cancelled())));
370        }
371    }
372}
373
374/// An in-flight effect losing `InFlight` without an outcome, seen by the
375/// witness: a cancelled dispatch, whether or not a record is kept.
376pub fn witness_cancelled(
377    removed: On<Remove, InFlight>,
378    effects: Query<(Has<EffectOutcome>, Option<&Observed>), With<Issued>>,
379    subjects: super::witness::Subjects,
380    witness: Option<Res<super::witness::Witnessing>>,
381) {
382    let entity = removed.event().entity;
383    let Some(witness) = witness else {
384        return;
385    };
386    let Ok((answered, observed)) = effects.get(entity) else {
387        return;
388    };
389    if answered || observed.is_some_and(|observed| observed.0.is_discarded()) {
390        return;
391    }
392    let observation = rig_core::observe::Observation::new(
393        subjects.of(entity),
394        rig_core::observe::Stage::Collect,
395        super::witness::bus_emitter(),
396        rig_core::observe::Action::Cancelled {
397            reason: rig_core::observe::Reason::from_report(&cancelled()),
398        },
399    );
400    witness.observe(observation);
401}
402
403/// Describe an inserted or replaced [`Bound`] component to the installed recorder.
404pub fn record_bound(
405    inserted: On<Insert, Bound>,
406    bound: Query<&Bound>,
407    recording: Option<Res<Recording>>,
408) {
409    let Some(recording) = recording else {
410        return;
411    };
412    if let Ok(bound) = bound.get(inserted.event().entity) {
413        recording.handlers(vec![bound.descriptor.clone()]);
414    }
415}
416
417#[cfg(test)]
418mod tests;
rig_ecs/bus/record.rs

rig_ecs/bus/
record.rs
