rig_ecs/bus/
effect.rs

1//! Components for effect intent, ordering, identity, execution, and outcomes.
2//!
3//! ```
4//! use rig_ecs::bus::Held;
5//! let mut world = bevy_ecs::world::World::new();
6//! let effect = world.spawn(Held).id();
7//! assert!(world.get::<Held>(effect).is_some());
8//! ```
9
10use bevy_reflect::Reflect;
11use std::sync::Arc;
12
13use bevy_ecs::{lifecycle::HookContext, prelude::*, world::DeferredWorld};
14use bevy_tasks::Task;
15use rig_core::{
16    effect::{CustomEffect, EffectId, EffectKind, Family, HandlerKey, Key, Outcome},
17    error::ErrorReport,
18    streaming::{StreamEvents, Transcript},
19    tool::{PublishedContext, ToolContext},
20};
21use serde::{Deserialize, Serialize};
22
23/// A dispatch intent containing its handler key and request.
24/// Streaming completions accumulate in [`Streamed`]; other requests are unary.
25/// Adding this component requires the bus's [`SeqCounter`] resource and stamps
26/// [`Seq`] in spawn order, replacing any supplied sequence value.
27#[derive(Component, Debug, Clone, Serialize, Deserialize, Reflect)]
28#[require(Seq)]
29#[component(on_add = stamp_seq)]
30#[reflect(Component)]
31pub struct PendingEffect {
32    /// The handler key the effect is routed to.
33    #[reflect(remote = crate::bus::reflect::HandlerKeyReflect)]
34    pub key: HandlerKey,
35    /// The effect.
36    #[reflect(remote = crate::bus::reflect::EffectKindReflect)]
37    pub kind: EffectKind,
38}
39
40impl PendingEffect {
41    /// A pending effect for `key`.
42    pub fn new(key: impl Into<HandlerKey>, kind: EffectKind) -> Self {
43        Self {
44            key: key.into(),
45            kind,
46        }
47    }
48
49    /// Wrap a typed request for `key`, returning any family conversion error.
50    pub fn typed<F: Family>(key: &Key<F>, request: F::Request) -> Result<Self, ErrorReport> {
51        Ok(Self {
52            key: key.raw().clone(),
53            kind: F::wrap(request)?,
54        })
55    }
56
57    /// Build a custom request for `key` with `E::KIND`, returning a request error
58    /// if its payload cannot be serialized.
59    pub fn custom<E: CustomEffect>(
60        key: impl Into<HandlerKey>,
61        effect: &E,
62    ) -> Result<Self, ErrorReport> {
63        let payload = serde_json::to_value(effect).map_err(|error| {
64            ErrorReport::new(
65                rig_core::error::ErrorKind::Request,
66                format!("the custom effect `{}` did not serialize: {error}", E::KIND),
67            )
68        })?;
69        Ok(Self {
70            key: key.into(),
71            kind: EffectKind::Custom {
72                kind: Arc::from(E::KIND),
73                payload,
74            },
75        })
76    }
77
78    /// Whether this effect is dispatched as a stream.
79    pub const fn is_stream(&self) -> bool {
80        matches!(self.kind, EffectKind::Completion { stream: true, .. })
81    }
82}
83
84/// Global dispatch order, stamped from [`SeqCounter`] whenever [`PendingEffect`]
85/// is added, replacing manually supplied values. Scene loading preserves relative
86/// order by spawning effects in saved order.
87#[derive(
88    Component,
89    Debug,
90    Clone,
91    Copy,
92    Default,
93    PartialEq,
94    Eq,
95    PartialOrd,
96    Ord,
97    Hash,
98    Serialize,
99    Deserialize,
100    Reflect,
101)]
102#[reflect(Component)]
103pub struct Seq(pub u64);
104
105/// The world's one dispatch-order counter (see [`Seq`]).
106#[derive(Resource, Debug, Default, Reflect)]
107#[reflect(Resource)]
108pub struct SeqCounter(pub u64);
109
110fn stamp_seq(mut world: DeferredWorld<'_>, context: HookContext) {
111    let next = {
112        let mut counter = world.resource_mut::<SeqCounter>();
113        let next = counter.0;
114        counter.0 += 1;
115        next
116    };
117    if let Some(mut seq) = world.get_mut::<Seq>(context.entity) {
118        seq.0 = next;
119    } else {
120        world.commands().entity(context.entity).insert(Seq(next));
121    }
122}
123
124/// The world's one effect-id counter: `Dispatch` mints ids from it, strictly
125/// increasing, unless the entity carries a [`Reserved`] id. Every
126/// `Reserved` or `Issued` inserted anywhere (a scene load, a log load, a
127/// host's own) bumps it past that id, so a minted id never collides with
128/// a saved one. `u64::MAX` is the exhausted counter sentinel: fresh dispatch
129/// then returns a request error without minting or recording an effect. The
130/// maximum allocatable ID is `u64::MAX - 1`; invalid direct component insertion
131/// saturates the counter rather than wrapping it.
132#[derive(Resource, Debug, Default, Reflect)]
133#[reflect(Resource)]
134pub struct IdCounter(pub u64);
135
136/// An id the effect must be dispatched under: a scene's saved id, a
137/// replayed record's. Consumed by `Dispatch`, which bumps [`IdCounter`]
138/// past it so a minted id never collides with a reserved one.
139#[derive(Component, Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Reflect)]
140#[component(on_insert = bump_ids_past_reserved)]
141#[reflect(Component)]
142pub struct Reserved(#[reflect(remote = crate::bus::reflect::EffectIdReflect)] pub EffectId);
143
144fn bump_ids_past_reserved(mut world: DeferredWorld<'_>, context: HookContext) {
145    if let Some(Reserved(id)) = world.get::<Reserved>(context.entity).copied() {
146        let mut counter = world.resource_mut::<IdCounter>();
147        counter.0 = counter.0.max(id.as_u64().saturating_add(1));
148    }
149}
150
151fn bump_ids_past_issued(mut world: DeferredWorld<'_>, context: HookContext) {
152    if let Some(Issued(id)) = world.get::<Issued>(context.entity).copied() {
153        let mut counter = world.resource_mut::<IdCounter>();
154        counter.0 = counter.0.max(id.as_u64().saturating_add(1));
155    }
156}
157
158/// The id the effect was dispatched under. Inserted by `Dispatch` and kept
159/// for the entity's life: what a child's record names as its `parent`, what
160/// a scene saves.
161#[derive(Component, Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Reflect)]
162#[component(on_insert = bump_ids_past_issued)]
163#[reflect(Component)]
164pub struct Issued(#[reflect(remote = crate::bus::reflect::EffectIdReflect)] pub EffectId);
165
166/// A pending effect a user system is still deciding about: `Dispatch`
167/// leaves it alone until the marker is removed (approve), the effect is
168/// denied (`EffectOutcome(Err(..))` inserted) or the entity despawned. The
169/// world-side spelling of a layer that suspends in `before`.
170/// Policies with independent ownership use [`super::acquire_hold`] and
171/// [`super::release_hold`]; removing this marker directly bypasses all owners
172/// (the batch's own marker follows: a call the batch held and a host
173/// approved is dispatched and counted as active).
174#[derive(Component, Debug, Clone, Copy, Default, Serialize, Deserialize, Reflect)]
175#[reflect(Component)]
176pub struct Held;
177
178/// The effect was taken: a handler is serving it. Present from `Dispatch`
179/// until `settle` closes the record; for a stream, until the returned
180/// stream reaches EOF. Carries the key it occupies so serial serving is a
181/// query over this component. Never serialized: a scene stores intent.
182#[derive(Component, Debug, Clone, Reflect)]
183#[reflect(Component)]
184pub struct InFlight {
185    /// The key the effect occupies.
186    #[reflect(remote = crate::bus::reflect::HandlerKeyReflect)]
187    pub key: HandlerKey,
188}
189
190/// The marker of an initial handler task, held for the entity in the
191/// world's [`Executions`] table.
192#[derive(Component)]
193pub struct Serving;
194
195/// Delivery receiver and fold for a stream driven by an effect-owned worker
196/// (held in [`Executions`]). Dropping the component drops the receiver, and
197/// the worker returns at its next send.
198#[derive(Component)]
199pub struct Streaming {
200    /// Bounded worker delivery, exclusively consumed by Collect.
201    pub events: futures::channel::mpsc::Receiver<
202        Result<rig_core::streaming::Relayed, rig_core::error::ErrorReport>,
203    >,
204    /// The shared core fold of delivered events.
205    pub fold: rig_core::serve::StreamTap,
206    /// Items consumed from this receiver, across collection passes.
207    pub delivered: usize,
208}
209
210/// Drive one owned stream on the pool with bounded delivery to Collect;
211/// every delivered item raises `wake`. Returns the receiver half and the
212/// worker task.
213fn spawn_stream_worker(
214    mut stream: StreamEvents,
215    capacity: usize,
216    wake: super::plugin::Wake,
217) -> (
218    futures::channel::mpsc::Receiver<
219        Result<rig_core::streaming::Relayed, rig_core::error::ErrorReport>,
220    >,
221    Task<()>,
222) {
223    use futures::{SinkExt, StreamExt};
224    let (mut sender, events) = futures::channel::mpsc::channel(capacity.max(1));
225    let task = bevy_tasks::IoTaskPool::get().spawn(async move {
226        while let Some(item) = stream.next().await {
227            if sender.send(item).await.is_err() {
228                return;
229            }
230            wake.signal();
231        }
232        wake.signal();
233    });
234    (events, task)
235}
236
237/// The tasks behind [`Serving`] and [`Streaming`], rows in the non-send
238/// [`Executions`] table: a task is `!Send` on wasm and cannot be a
239/// component, and one storage keeps one code path on both targets.
240#[derive(bevy_ecs::system::SystemParam)]
241pub struct Tasks<'w> {
242    executions: NonSendMut<'w, Executions>,
243}
244
245impl Tasks<'_> {
246    /// The [`Serving`] component for `task`.
247    pub fn serving(&mut self, entity: Entity, task: Task<rig_core::serve::Reply>) -> Serving {
248        self.executions.tasks.insert(entity, task);
249        Serving
250    }
251
252    /// Whether the serving task finished: its reply, once.
253    pub fn poll(&mut self, entity: Entity) -> Option<rig_core::serve::Reply> {
254        let task = self.executions.tasks.get_mut(&entity)?;
255        let reply = bevy_tasks::futures::check_ready(task)?;
256        self.executions.tasks.remove(&entity);
257        Some(reply)
258    }
259
260    /// A [`Streaming`] component driving `stream` on the pool.
261    pub fn streaming(
262        &mut self,
263        entity: Entity,
264        stream: StreamEvents,
265        capacity: usize,
266        wake: super::plugin::Wake,
267    ) -> Streaming {
268        let (events, task) = spawn_stream_worker(stream, capacity, wake);
269        self.executions.streams.insert(entity, task);
270        Streaming {
271            events,
272            fold: rig_core::serve::StreamTap::new(),
273            delivered: 0,
274        }
275    }
276
277    /// The stream worker of `entity` is done with: dropped now.
278    pub fn forget_stream(&mut self, entity: Entity) {
279        self.executions.streams.remove(&entity);
280    }
281
282    /// The tasks of `world`, for an exclusive system. `None` when the bus
283    /// is not installed (the table is missing).
284    pub fn with<T>(world: &mut World, f: impl FnOnce(&mut Tasks<'_>) -> T) -> Option<T> {
285        let mut state = bevy_ecs::system::SystemState::<Tasks>::new(world);
286        let mut tasks = state.get_mut(world).ok()?;
287        Some(f(&mut tasks))
288    }
289}
290
291/// The tasks of in-flight effects, indexed by their effect entity and
292/// dropped by [`drop_execution`] when the effect leaves flight. Non-send
293/// because a task is `!Send` on wasm.
294#[derive(Default)]
295pub struct Executions {
296    /// Initial tasks, indexed by their in-flight effect entity.
297    pub tasks: std::collections::HashMap<Entity, Task<rig_core::serve::Reply>>,
298    /// Stream workers, cancelled when the effect leaves flight.
299    pub streams: std::collections::HashMap<Entity, Task<()>>,
300}
301
302/// Remove owned execution immediately when an effect leaves flight.
303pub fn drop_execution(removed: On<Remove, InFlight>, mut executions: NonSendMut<Executions>) {
304    let entity = removed.event().entity;
305    executions.tasks.remove(&entity);
306    executions.streams.remove(&entity);
307}
308
309/// Accumulated stream events, error positions, text, and the first folded outcome.
310/// [`EffectOutcome`] is published at EOF, retaining the serial key through stream
311/// completion. Observe `Changed<Streamed>` for newly collected data.
312///
313/// Live collection allows 64 queue checks per effect and 4,096 per pass.
314/// Exact replay grouping requires delivery metadata and kept items; event bytes
315/// alone establish order, and folded recordings provide only the final answer.
316/// Checkpoints restore completed streams but reject unfinished streams with
317/// delivered progress because no continuation cursor is saved.
318#[derive(Component, Debug, Default, Clone, Serialize, Deserialize, Reflect)]
319#[reflect(Component)]
320pub struct Streamed {
321    /// Every error item and its zero-based position among all stream items,
322    /// including errors after the first terminal outcome. Live observation is
323    /// independent of recorder event retention; `outcome` remains the first fold.
324    #[serde(default, skip_serializing_if = "Vec::is_empty")]
325    #[reflect(remote = crate::bus::reflect::StreamErrorsReflect)]
326    pub errors: Vec<(usize, ErrorReport)>,
327    /// Every event, in order.
328    #[reflect(remote = crate::bus::reflect::StreamEventsReflect)]
329    pub events: Transcript,
330    /// The text deltas concatenated.
331    pub text: String,
332    /// Who the reply is from, as the relay stated it before its first item.
333    #[serde(default, skip_serializing_if = "Option::is_none")]
334    #[reflect(remote = crate::agent::content::reflect::OriginReflect)]
335    pub origin: Option<rig_core::message::Origin>,
336    /// The fold's outcome at the terminal record, or the error that ended
337    /// the stream.
338    #[reflect(remote = crate::bus::reflect::StreamedOutcomeReflect)]
339    pub outcome: Option<Result<Outcome, ErrorReport>>,
340}
341
342/// The answer. Inserted by `Collect` when a handler's task or stream
343/// finished, by a `Gate` system that denies, or by a `Judge` system that
344/// replaces. Serde, so a scene keeps answered effects answered.
345#[derive(Component, Debug, Clone, Serialize, Deserialize, Reflect)]
346#[reflect(Component)]
347pub struct EffectOutcome(
348    #[reflect(remote = crate::bus::reflect::OutcomeReflect)] pub Result<Outcome, ErrorReport>,
349);
350
351/// An open world's submitted answer. Insert with [`WorldOutcome::new`];
352/// `Collect` publishes it as [`EffectOutcome`] in submission order. Typed
353/// [`Answer<E>`] uses the same path. Submission after Collect is visible next
354/// pass, so policy observes the same boundaries live and during replay.
355/// Like a ready task, this inbox is transient; save after collection to retain
356/// its answer in a scene. Applications observe `EffectOutcome`, not this inbox.
357#[derive(Component, Debug)]
358#[component(on_add = stamp_world_outcome)]
359pub struct WorldOutcome {
360    /// The submitted answer, published unchanged by the collector.
361    pub outcome: Result<Outcome, ErrorReport>,
362    order: u64,
363}
364
365impl WorldOutcome {
366    /// Submit one answer for an in-flight world-served effect.
367    pub fn new(outcome: Result<Outcome, ErrorReport>) -> Self {
368        Self { outcome, order: 0 }
369    }
370
371    /// Submission order within this world, stamped on component insertion.
372    pub fn order(&self) -> u64 {
373        self.order
374    }
375}
376
377/// Monotonic submission order for the world's answer inbox.
378#[derive(Resource, Default)]
379pub struct WorldOutcomeCounter(pub u64);
380
381fn stamp_world_outcome(mut world: DeferredWorld<'_>, context: HookContext) {
382    let order = {
383        let mut counter = world.resource_mut::<WorldOutcomeCounter>();
384        let order = counter.0;
385        counter.0 += 1;
386        order
387    };
388    if let Some(mut outcome) = world.get_mut::<WorldOutcome>(context.entity) {
389        outcome.order = order;
390    }
391}
392
393impl EffectOutcome {
394    /// Return the family's typed answer, propagating a recorded or conversion error.
395    pub fn typed<F: Family>(&self) -> Result<F::Answer, ErrorReport> {
396        F::unwrap(self.0.clone()?)
397    }
398
399    /// Deserialize a custom answer as `E::Answer`, returning recorded errors or a
400    /// response error for a different outcome family or invalid payload.
401    pub fn custom<E: CustomEffect>(&self) -> Result<E::Answer, ErrorReport> {
402        match self.0.clone()? {
403            Outcome::Custom { payload: value } => serde_json::from_value(value).map_err(|error| {
404                ErrorReport::new(
405                    rig_core::error::ErrorKind::Response,
406                    format!("the answer to `{}` did not deserialize: {error}", E::KIND),
407                )
408            }),
409            other => Err(ErrorReport::new(
410                rig_core::error::ErrorKind::Response,
411                format!(
412                    "`{}` was answered with a {} outcome",
413                    E::KIND,
414                    other.family()
415                ),
416            )),
417        }
418    }
419}
420
421/// Serializable inbound tool context attached beside the effect payload.
422/// Dispatch supplies it to the handler; absent inputs use an empty context.
423#[derive(Component, Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, Reflect)]
424#[reflect(Component)]
425pub struct ToolInputs(#[reflect(remote = crate::bus::reflect::ToolContextReflect)] pub ToolContext);
426
427/// What the tool published into its context: read off the dispatch's
428/// [`PublishedContext`] when the outcome lands (`Collect`), or inserted by
429/// the system that answers an open tool key. Data, beside the outcome.
430#[derive(Component, Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, Reflect)]
431#[reflect(Component)]
432pub struct ToolOutputs(
433    #[reflect(remote = crate::bus::reflect::ToolContextReflect)] pub ToolContext,
434);
435
436/// The slot a task-served tool call publishes into, shared with its dispatch context
437/// for the length of the call; `Collect` reads it into [`ToolOutputs`].
438/// Never serialized: in-flight state.
439#[derive(Component, Clone)]
440pub struct Publishing(pub Arc<PublishedContext>);
441
442/// A stable serializable program scope, not a runtime handle.
443/// Dispatch records the nearest scope along `ChildOf`, checking the effect first.
444#[derive(Component, Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Reflect)]
445#[reflect(Component)]
446pub struct Scope(pub String);
447
448/// A [`CustomEffect`] a system can serve: its payload and its answer live
449/// in components, so both are `Send + Sync` on every target (plain data
450/// is). Blanket-implemented.
451pub trait WorldEffect: CustomEffect + Send + Sync {
452    /// The answer, as a component field: `E::Answer`.
453    type Reply: Serialize + serde::de::DeserializeOwned + Send + Sync + 'static;
454}
455
456impl<E> WorldEffect for E
457where
458    E: CustomEffect + Send + Sync,
459    E::Answer: Send + Sync,
460{
461    type Reply = E::Answer;
462}
463
464/// A custom effect a system answers: what a [`WorldHandler`](super::WorldHandler)
465/// dispatch lands as on the effect entity, deserialized. A user system with
466/// any `World` access reads it and inserts [`Answer<E>`].
467#[derive(Component, Debug, Clone)]
468pub struct Asked<E: WorldEffect>(pub E);
469
470/// A system's answer to an [`Asked<E>`]: inserting it resolves the effect
471/// (the plugin serializes it into the [`EffectOutcome`]).
472#[derive(Component, Debug, Clone)]
473pub struct Answer<E: WorldEffect>(pub E::Reply);
474
475/// A typed key as a component: the typed view over a handler entity's key,
476/// placed wherever a system wants it. `Send + Sync` on every target.
477#[derive(Component, Debug, Clone, Serialize, Deserialize)]
478pub struct Typed<F: Family>(pub Key<F>);
479
480impl<F: Family> Typed<F> {
481    /// Build a pending effect for this key, returning any family conversion error.
482    pub fn pending(&self, request: F::Request) -> Result<PendingEffect, ErrorReport> {
483        PendingEffect::typed(&self.0, request)
484    }
485
486    /// The key.
487    pub fn key(&self) -> &HandlerKey {
488        self.0.raw()
489    }
490}
491
492// Every component a system holds is `Send + Sync` on every target: the
493// tasks, owned streams, and erased handlers live elsewhere.
494const _: () = {
495    const fn assert_send_sync<T: Send + Sync + 'static>() {}
496    assert_send_sync::<PendingEffect>();
497    assert_send_sync::<Seq>();
498    assert_send_sync::<Issued>();
499    assert_send_sync::<Reserved>();
500    assert_send_sync::<Held>();
501    assert_send_sync::<InFlight>();
502    assert_send_sync::<Streamed>();
503    assert_send_sync::<EffectOutcome>();
504    assert_send_sync::<Scope>();
505    assert_send_sync::<ToolInputs>();
506    assert_send_sync::<ToolOutputs>();
507    assert_send_sync::<Publishing>();
508    assert_send_sync::<Serving>();
509    assert_send_sync::<Streaming>();
510    assert_send_sync::<Typed<rig_core::effect::family::Completion>>();
511};
rig_ecs/bus/effect.rs

rig_ecs/bus/
effect.rs
