rig_ecs/
effects.rs

1//! The one dispatch path. Every effect entity ([`stage`](super::stage)),
2//! a model call, a tool call or a plugin's effect, runs through
3//! [`Effects`] under its effect id, recorded in the kernel's own
4//! [`RecordBuffer`], and its record is written as a
5//! [`EffectRecorded`](super::stage::EffectRecorded) when it settles. A plugin that
6//! keeps an effect log reads those, and continues the ids of the log it
7//! appends to with [`Effects::continue_after`].
8//!
9//! A streamed model call is recorded as the response it ended with, not
10//! its stream items: views follow the stream on the call ([`Partial`]),
11//! and a replay answers the call with its recorded outcome before it
12//! starts, in Gate, so nothing re-emits a stream.
13//!
14//! [`Partial`]: super::agent::Partial
15
16use std::sync::Arc;
17use std::sync::atomic::{AtomicU64, Ordering};
18
19use bevy_ecs::prelude::*;
20use bevy_tasks::ConditionalSendFuture;
21use rig_core::effect::Answered;
22use rig_core::effect::{EffectId, EffectKind};
23use rig_core::error::ErrorReport;
24use rig_core::serve::{
25    Dispatch, DispatchContext, ErasedHandler, RecordBuffer, Recorder, Reply, catch_panics,
26};
27use rig_core::tool::{PublishedContext, ToolContext, ToolResultContext};
28
29/// What a handler's run left: its answer, and what a tool published beside
30/// it, if anything.
31pub(crate) type Ran = (Result<Answered, ErrorReport>, Option<ToolResultContext>);
32
33/// The effect ids, and the records of the effects running.
34#[derive(Resource)]
35pub struct Effects {
36    /// The last id taken.
37    last: AtomicU64,
38    /// The records of effect entities' dispatches, timed, until they settle;
39    /// without stream items.
40    buffer: RecordBuffer,
41}
42
43impl Default for Effects {
44    fn default() -> Self {
45        Self {
46            last: AtomicU64::new(0),
47            buffer: RecordBuffer::new().timed(),
48        }
49    }
50}
51
52impl Effects {
53    /// Effects with ids after `last`, such as the highest id of the effect
54    /// log a plugin continues, so ids keep increasing across restarts.
55    pub fn continue_after(last: Option<EffectId>) -> Self {
56        Self {
57            last: AtomicU64::new(last.map_or(0, EffectId::as_u64)),
58            ..Self::default()
59        }
60    }
61
62    /// Runs `kind` through `handler` as the effect `id` of an entity, from
63    /// `origin`, recorded in the kernel's buffer; `consume` turns the reply
64    /// into the outcome, such as `Reply::into_outcome`. A panic in the
65    /// handler or its stream becomes an internal error, recorded as its
66    /// outcome. What a tool publishes into its context comes back beside it.
67    pub(crate) fn run<F>(
68        &self,
69        id: EffectId,
70        origin: DispatchContext,
71        handler: ErasedHandler,
72        kind: EffectKind,
73        consume: impl FnOnce(Reply) -> F + Send + 'static,
74    ) -> impl ConditionalSendFuture<Output = Ran> + 'static
75    where
76        F: ConditionalSendFuture<Output = Result<Answered, ErrorReport>> + 'static,
77    {
78        let recorder: Arc<dyn Recorder + Send + Sync> = Arc::new(self.buffer.clone());
79        recorder.begin(id, handler.key().clone(), kind.clone(), origin);
80        let published = PublishedContext::new();
81        let dispatch = Dispatch::new(id, kind.streams())
82            .with_scope(published.clone())
83            .recorded_by(recorder.clone());
84        let work = async move { consume(handler.handle(kind, dispatch).await).await };
85        let work = catch_panics(Some(recorder), id, work);
86        async move {
87            let outcome = work.await;
88            let facts = published.take().map(ToolContext::into_result_context);
89            (
90                outcome,
91                facts.filter(|facts| *facts != ToolResultContext::default()),
92            )
93        }
94    }
95
96    /// The records of effect entities' dispatches.
97    pub(crate) fn buffer(&self) -> &RecordBuffer {
98        &self.buffer
99    }
100
101    /// Takes the next effect id.
102    pub(crate) fn next_id(&self) -> EffectId {
103        EffectId::from_raw(self.last.fetch_add(1, Ordering::Relaxed) + 1)
104    }
105}
106
107/// An effect handler as a component or resource keeps it. rig-core's
108/// [`ErasedHandler`] is `Send` and `Sync` except on browser wasm, where its
109/// provider clients hold JavaScript values, and Bevy asks for both on
110/// every target.
111#[derive(Clone)]
112pub struct Handler(pub ErasedHandler);
113
114// SAFETY: browser wasm without the `atomics` target feature runs one
115// thread, so a `Handler` is never sent to or shared with another one.
116#[cfg(all(
117    target_arch = "wasm32",
118    target_os = "unknown",
119    not(target_feature = "atomics")
120))]
121unsafe impl Send for Handler {}
122// SAFETY: as for `Send`: there is no other thread.
123#[cfg(all(
124    target_arch = "wasm32",
125    target_os = "unknown",
126    not(target_feature = "atomics")
127))]
128unsafe impl Sync for Handler {}
rig_ecs/effects.rs

rig_ecs/
effects.rs
