Skip to main content

rig_ecs/
checkpoint.rs

1//! Save and restore reflected execution graphs with host-supplied handlers.
2//!
3//! Registered components, binary assets, and execution counters are persisted;
4//! live handlers and provider launch settings are not. Restoration validates
5//! original saved contracts before aliasing and reissues unfinished effects with
6//! saved dispatch IDs. Relationship targets are rebuilt in checkpoint order.
7//!
8//! ```
9//! use rig_ecs::{RigPlugin, checkpoint::{save_world, load_world, RestoreMode}};
10//! let mut app = bevy_app::App::new();
11//! app.add_plugins(RigPlugin::default());
12//! let checkpoint = save_world(app.world_mut())?;
13//! load_world(&checkpoint, app.world_mut(), RestoreMode::Strict, [])?;
14//! # Ok::<(), rig_ecs::checkpoint::CheckpointError>(())
15//! ```
16
17mod restore;
18pub use restore::{RestoreMode, load_world};
19
20use std::{any::TypeId, collections::HashMap};
21
22use bevy_ecs::{
23    entity_disabling::Disabled,
24    prelude::*,
25    reflect::{AppTypeRegistry, ReflectComponent},
26    resource::IsResource,
27};
28use bevy_reflect::{
29    PartialReflect, TypePath, TypeRegistration, TypeRegistry,
30    serde::{
31        ReflectDeserializerProcessor, ReflectSerializer, ReflectSerializerProcessor,
32        TypedReflectDeserializer,
33    },
34};
35use rig_core::effect::{EffectFamily, HandlerKey};
36use serde::{Deserialize, Serialize, de::DeserializeSeed};
37
38use crate::{
39    agent::{
40        self, Utterance,
41        content::{
42            binary::{BinaryAssets, BinaryError},
43            parts::*,
44        },
45    },
46    bus::{
47        self, Bound, EffectOutcome, HandlerIndex, IdCounter, InFlight, Issued, PendingEffect,
48        Reserved, Seq, SeqCounter, Streamed,
49    },
50};
51
52/// One entity of a [`Checkpoint`]: its components by type path.
53pub type CheckpointEntity = serde_json::Map<String, serde_json::Value>;
54
55/// An entity's reflected components, by type path.
56type Reflected<'a> = Vec<(&'a str, Box<dyn PartialReflect>)>;
57
58/// The [`Checkpoint`] envelope format this crate writes and reads.
59pub const CHECKPOINT_FORMAT: u32 = 2;
60
61/// Reflected execution state with binary assets and counters.
62/// Deserialization rejects unknown envelope fields and formats other than
63/// [`CHECKPOINT_FORMAT`].
64#[derive(Debug, Clone, Serialize, Deserialize)]
65#[serde(deny_unknown_fields)]
66pub struct Checkpoint {
67    /// The envelope format ([`CHECKPOINT_FORMAT`]).
68    #[serde(deserialize_with = "checkpoint_format")]
69    pub format: u32,
70    /// Every entity with a reflected component, parents before children,
71    /// siblings in `Children` order.
72    pub entities: Vec<CheckpointEntity>,
73    /// The world's counters: the next run number (the id and sequence
74    /// counters bump themselves from what is loaded).
75    pub counters: Counters,
76    /// The binary store: every retained payload, once per content hash.
77    pub binaries: Vec<crate::agent::content::binary::BinaryRecord>,
78}
79
80impl Default for Checkpoint {
81    fn default() -> Self {
82        Self {
83            format: CHECKPOINT_FORMAT,
84            entities: Vec::new(),
85            counters: Counters::default(),
86            binaries: Vec::new(),
87        }
88    }
89}
90
91/// Deserialize the envelope's `format`, refusing any other than
92/// [`CHECKPOINT_FORMAT`] by name.
93fn checkpoint_format<'de, D: serde::Deserializer<'de>>(deserializer: D) -> Result<u32, D::Error> {
94    let format = u32::deserialize(deserializer)?;
95    if format == CHECKPOINT_FORMAT {
96        Ok(format)
97    } else {
98        Err(serde::de::Error::custom(format!(
99            "load refused: the checkpoint is format {format}, this rig reads format {CHECKPOINT_FORMAT}"
100        )))
101    }
102}
103
104/// The world's counters, so nothing loaded collides with what comes after.
105#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
106pub struct Counters {
107    /// [`agent::RunCounter`].
108    pub next_run: u64,
109    /// [`IdCounter`].
110    pub next_id: u64,
111}
112
113/// What [`load_world`] spawned, by checkpoint index; an index the load
114/// merged into an existing handler entity names that entity.
115#[derive(Debug, Clone, Default)]
116pub struct Loaded {
117    /// The entities, by [`Checkpoint::entities`] index.
118    pub entities: Vec<Entity>,
119}
120
121impl Loaded {
122    /// The loaded entities that carry `C`.
123    pub fn with<C: Component>(&self, world: &World) -> Vec<Entity> {
124        self.entities
125            .iter()
126            .copied()
127            .filter(|entity| world.get::<C>(*entity).is_some())
128            .collect()
129    }
130}
131
132/// Why a checkpoint cannot be saved, parsed, validated, or restored.
133#[derive(Debug, thiserror::Error)]
134#[non_exhaustive]
135pub enum CheckpointError {
136    /// The world lacks a resource that [`RigPlugin`](crate::RigPlugin) installs.
137    #[error("the world has no {0}: install RigPlugin first")]
138    NotInstalled(&'static str),
139    /// The envelope is a format other than [`CHECKPOINT_FORMAT`].
140    #[error("the checkpoint is format {found}, this rig reads format {CHECKPOINT_FORMAT}")]
141    UnsupportedFormat {
142        /// The envelope's format.
143        found: u32,
144    },
145    /// The checkpoint carries a component this format no longer persists.
146    #[error(
147        "the checkpoint contains removed `{path}`: migrate provider launch settings to the host"
148    )]
149    RemovedComponent {
150        /// The removed component's type path.
151        path: String,
152    },
153    /// The checkpoint envelope does not serialize to or parse from JSON.
154    #[error("checkpoint JSON: {0}")]
155    Json(#[from] serde_json::Error),
156    /// A live component does not serialize through reflection.
157    #[error("`{path}` does not serialize: {source}")]
158    Serialize {
159        /// The component's type path.
160        path: String,
161        /// The serializer's error.
162        #[source]
163        source: serde_json::Error,
164    },
165    /// A saved component does not deserialize as its registered type.
166    #[error("checkpoint entity {entity}: `{path}`: {source}")]
167    Deserialize {
168        /// The entity's index in [`Checkpoint::entities`].
169        entity: usize,
170        /// The component's type path.
171        path: String,
172        /// The deserializer's error.
173        #[source]
174        source: serde_json::Error,
175    },
176    /// A saved type path is not a component registered in the destination.
177    #[error("checkpoint entity {entity}: `{path}` is not a registered component")]
178    Unregistered {
179        /// The entity's index in [`Checkpoint::entities`].
180        entity: usize,
181        /// The unknown type path.
182        path: String,
183    },
184    /// A saved handler's binding key differs from its descriptor's key.
185    #[error("saved handler `{key}` has a different descriptor key")]
186    HandlerKeyMismatch {
187        /// The binding key.
188        key: HandlerKey,
189    },
190    /// Two saved handlers share a key.
191    #[error("duplicate saved handler `{key}`")]
192    DuplicateSavedHandler {
193        /// The repeated key.
194        key: HandlerKey,
195    },
196    /// The host supplied two handlers for one key.
197    #[error("duplicate supplied handler `{key}`")]
198    DuplicateSuppliedHandler {
199        /// The repeated key.
200        key: HandlerKey,
201    },
202    /// The host supplied a handler the checkpoint does not require.
203    #[error("supplied handler `{key}` is not required by this checkpoint")]
204    UnrequiredHandler {
205        /// The supplied key.
206        key: HandlerKey,
207    },
208    /// A required handler is neither supplied nor installed in the destination.
209    #[error("no implementation supplied for `{key}`")]
210    MissingHandler {
211        /// The required key.
212        key: HandlerKey,
213    },
214    /// An unfinished effect names a handler the checkpoint did not save.
215    #[error("unfinished effect requires missing saved handler `{key}`")]
216    MissingSavedHandler {
217        /// The effect's handler key.
218        key: HandlerKey,
219    },
220    /// A handler serves a different effect family than the saved one.
221    #[error("handler `{key}` serves {found}; the checkpoint's serves {saved}")]
222    FamilyChanged {
223        /// The handler key.
224        key: HandlerKey,
225        /// The saved family.
226        saved: EffectFamily,
227        /// The destination's family.
228        found: EffectFamily,
229    },
230    /// Under [`RestoreMode::Strict`], a handler's descriptor differs from the
231    /// saved one. [`RestoreMode::Replace`] accepts the change.
232    #[error("handler `{key}` differs from its original saved descriptor")]
233    DescriptorChanged {
234        /// The handler key.
235        key: HandlerKey,
236    },
237    /// An unfinished stream already delivered progress and has no provider cursor.
238    #[error("an unfinished stream with delivered progress cannot resume")]
239    UnresumableStream,
240    /// A saved utterance does not convert to a message.
241    #[error("invalid saved content: {0}")]
242    Content(#[from] ContentError),
243    /// The saved binary store is invalid or conflicts with the destination's.
244    #[error("invalid saved binary store: {0}")]
245    Binary(#[from] BinaryError),
246    /// The saved graph violates a structural invariant.
247    #[error("invalid checkpoint graph: {0}")]
248    InvalidGraph(&'static str),
249}
250
251/// The relationship targets: rebuilt by their sources' hooks, never saved.
252fn is_relationship_target(id: TypeId) -> bool {
253    [
254        TypeId::of::<bevy_ecs::hierarchy::Children>(),
255        TypeId::of::<bus::Serves>(),
256        TypeId::of::<agent::ModelOf>(),
257        TypeId::of::<agent::RememberedBy>(),
258        TypeId::of::<agent::RetrievedBy>(),
259        TypeId::of::<agent::RoutedTo>(),
260        TypeId::of::<agent::Grants>(),
261        TypeId::of::<agent::ContextOf>(),
262        TypeId::of::<agent::AttachedTo>(),
263        TypeId::of::<agent::Runs>(),
264        TypeId::of::<agent::AdvertisedOn>(),
265        TypeId::of::<agent::content::parts::EditedBy>(),
266        TypeId::of::<agent::checkpoint::AssistantForTurns>(),
267        TypeId::of::<agent::checkpoint::ResultsForTurns>(),
268    ]
269    .contains(&id)
270}
271
272/// The world's entities in checkpoint order: roots by spawn index (an
273/// `Entity`'s own `Ord` is over its opaque bits, which invert the index),
274/// then each root's subtree depth-first in `Children` order.
275fn ordered_entities(world: &mut World) -> Vec<Entity> {
276    // Disabled runs must survive checkpoints even though ordinary queries omit them.
277    let mut roots: Vec<Entity> = world
278        .query_filtered::<Entity, (Without<ChildOf>, Without<IsResource>, Allow<Disabled>)>()
279        .iter(world)
280        .collect();
281    roots.sort_by_key(|entity| (entity.index().index(), entity.generation().to_bits()));
282    let mut order = Vec::with_capacity(roots.len());
283    let mut stack: Vec<Entity> = roots.into_iter().rev().collect();
284    while let Some(entity) = stack.pop() {
285        order.push(entity);
286        if let Some(children) = world.get::<Children>(entity) {
287            stack.extend(children.iter().rev());
288        }
289    }
290    order
291}
292
293/// Save registered reflected components, binary assets, and execution counters.
294/// Returns an error if the type registry is absent or reflection serialization fails.
295#[must_use = "saving a checkpoint does not remove anything from the world"]
296pub fn save_world(world: &mut World) -> Result<Checkpoint, CheckpointError> {
297    let registry = world
298        .get_resource::<AppTypeRegistry>()
299        .ok_or(CheckpointError::NotInstalled("type registry"))?
300        .clone();
301    let registry = registry.read();
302    let mut components: Vec<(&str, &ReflectComponent)> = registry
303        .iter_with_data::<ReflectComponent>()
304        .filter(|(registration, _)| !is_relationship_target(registration.type_id()))
305        .map(|(registration, data)| (registration.type_info().type_path(), data))
306        .collect();
307    components.sort_by_key(|(path, _)| *path);
308    let order = ordered_entities(world);
309    let mut rows: Vec<(Entity, Reflected<'_>)> = Vec::new();
310    for entity in order {
311        let entity_ref = world.entity(entity);
312        let reflected: Reflected<'_> = components
313            .iter()
314            .filter_map(|(path, data)| data.reflect(entity_ref).map(|c| (*path, c.to_dynamic())))
315            .collect();
316        if !reflected.is_empty() {
317            rows.push((entity, reflected));
318        }
319    }
320    let index: HashMap<Entity, usize> = rows
321        .iter()
322        .enumerate()
323        .map(|(index, (entity, _))| (*entity, index))
324        .collect();
325    let indexed = Indexed { index };
326    let mut entities = Vec::with_capacity(rows.len());
327    for (_, components) in &rows {
328        let mut object = CheckpointEntity::new();
329        for (path, component) in components {
330            object.insert(
331                (*path).to_owned(),
332                reflect_to_json(component.as_ref(), &registry, &indexed)?,
333            );
334        }
335        entities.push(object);
336    }
337    Ok(Checkpoint {
338        format: CHECKPOINT_FORMAT,
339        entities,
340        counters: Counters {
341            next_run: world.get_resource::<agent::RunCounter>().map_or(0, |c| c.0),
342            next_id: world.get_resource::<IdCounter>().map_or(0, |c| c.0),
343        },
344        binaries: world
345            .get_resource::<BinaryAssets>()
346            .map(BinaryAssets::snapshot)
347            .unwrap_or_default(),
348    })
349}
350
351fn reflect_to_json(
352    value: &dyn PartialReflect,
353    registry: &TypeRegistry,
354    indexed: &Indexed,
355) -> Result<serde_json::Value, CheckpointError> {
356    let path = value
357        .get_represented_type_info()
358        .map(|info| info.type_path().to_owned())
359        .unwrap_or_default();
360    let json = serde_json::to_value(ReflectSerializer::with_processor(value, registry, indexed))
361        .map_err(|source| CheckpointError::Serialize {
362            path: path.clone(),
363            source,
364        })?;
365    // `ReflectSerializer` wraps the value in a one-key map by type path.
366    Ok(match json {
367        serde_json::Value::Object(mut map) if map.len() == 1 => {
368            map.shift_remove(&path).unwrap_or(serde_json::Value::Null)
369        }
370        other => other,
371    })
372}
373
374/// Serializes every `Entity` as the index the checkpoint gives it.
375struct Indexed {
376    index: HashMap<Entity, usize>,
377}
378
379impl ReflectSerializerProcessor for Indexed {
380    fn try_serialize<S>(
381        &self,
382        value: &dyn PartialReflect,
383        _registry: &TypeRegistry,
384        serializer: S,
385    ) -> Result<Result<S::Ok, S>, S::Error>
386    where
387        S: serde::Serializer,
388    {
389        match value.try_downcast_ref::<Entity>() {
390            Some(entity) => match self.index.get(entity) {
391                Some(index) => serde::Serialize::serialize(index, serializer).map(Ok),
392                None => serializer.serialize_none().map(Ok),
393            },
394            None => Ok(Err(serializer)),
395        }
396    }
397}
398
399/// Deserializes every `Entity` from its checkpoint index.
400struct Remapped<'a> {
401    entities: &'a [Entity],
402}
403
404impl ReflectDeserializerProcessor for Remapped<'_> {
405    fn try_deserialize<'de, D>(
406        &mut self,
407        registration: &TypeRegistration,
408        _registry: &TypeRegistry,
409        deserializer: D,
410    ) -> Result<Result<Box<dyn PartialReflect>, D>, D::Error>
411    where
412        D: serde::Deserializer<'de>,
413    {
414        if registration.type_id() != TypeId::of::<Entity>() {
415            return Ok(Err(deserializer));
416        }
417        let index: usize = serde::Deserialize::deserialize(deserializer)?;
418        let entity = self.entities.get(index).copied().ok_or_else(|| {
419            serde::de::Error::custom(format!("entity {index} is not in the checkpoint"))
420        })?;
421        Ok(Ok(Box::new(entity)))
422    }
423}
424
425/// Reuse destination entities for matching dispatch keys. Implementation
426/// selection and original-descriptor validation belong to restoration preflight.
427fn aliases(
428    checkpoint: &Checkpoint,
429    world: &World,
430) -> Result<HashMap<usize, Entity>, CheckpointError> {
431    let Some(index) = world.get_resource::<HandlerIndex>() else {
432        return Ok(HashMap::new());
433    };
434    let mut aliases = HashMap::new();
435    for (row, entity) in checkpoint.entities.iter().enumerate() {
436        let Some(bound) = entity.get(Bound::type_path()) else {
437            continue;
438        };
439        let bound = saved_bound(row, bound)?;
440        if let Some(existing) = index.entity(&bound.key) {
441            if let Some(served) = world.get::<Bound>(existing)
442                && served.family() != bound.family()
443            {
444                return Err(CheckpointError::FamilyChanged {
445                    saved: bound.family(),
446                    found: served.family(),
447                    key: bound.key,
448                });
449            }
450            aliases.insert(row, existing);
451        }
452    }
453    Ok(aliases)
454}
455
456/// Parse a saved [`Bound`] component of checkpoint entity `row`.
457fn saved_bound(row: usize, value: &serde_json::Value) -> Result<Bound, CheckpointError> {
458    serde_json::from_value(value.clone()).map_err(|source| CheckpointError::Deserialize {
459        entity: row,
460        path: Bound::type_path().to_owned(),
461        source,
462    })
463}
464
465/// Spawn `checkpoint` into `world` (no validation, no merging of the
466/// binary store): the mechanics of a load.
467fn spawn_into(
468    checkpoint: &Checkpoint,
469    world: &mut World,
470    aliases: &HashMap<usize, Entity>,
471) -> Result<Vec<Entity>, CheckpointError> {
472    let registry = world
473        .get_resource::<AppTypeRegistry>()
474        .ok_or(CheckpointError::NotInstalled("type registry"))?
475        .clone();
476    let registry = registry.read();
477    let entities: Vec<Entity> = (0..checkpoint.entities.len())
478        .map(|row| match aliases.get(&row) {
479            Some(existing) => *existing,
480            None => world.spawn_empty().id(),
481        })
482        .collect();
483    let ordered: Vec<Entity> = (0..checkpoint.entities.len())
484        .map(|row| entities.get(row).copied().unwrap_or(Entity::PLACEHOLDER))
485        .collect();
486    let entities = ordered;
487    let mut remap = Remapped {
488        entities: &entities,
489    };
490    // Offset saved sequences so loaded effects follow existing ones without
491    // changing their relative order or gaps.
492    let base = world.resource::<SeqCounter>().0;
493    let mut effects: Vec<(u64, Entity)> = Vec::new();
494    for (row, components) in checkpoint.entities.iter().enumerate() {
495        let Some(&entity) = entities.get(row) else {
496            continue;
497        };
498        let merged = aliases.contains_key(&row);
499        for (path, value) in components {
500            let unregistered = || CheckpointError::Unregistered {
501                entity: row,
502                path: path.clone(),
503            };
504            let registration = registry.get_with_type_path(path).ok_or_else(unregistered)?;
505            if is_relationship_target(registration.type_id()) {
506                continue;
507            }
508            if merged && registration.type_id() == TypeId::of::<Bound>() {
509                continue;
510            }
511            let component = registration
512                .data::<ReflectComponent>()
513                .ok_or_else(unregistered)?;
514            let malformed = |source| CheckpointError::Deserialize {
515                entity: row,
516                path: path.clone(),
517                source,
518            };
519            let reflected =
520                TypedReflectDeserializer::with_processor(registration, &registry, &mut remap)
521                    .deserialize(value)
522                    .map_err(malformed)?;
523            if registration.type_id() == TypeId::of::<Seq>() {
524                let seq: Seq = serde_json::from_value(value.clone()).map_err(malformed)?;
525                effects.push((seq.0, entity));
526                continue;
527            }
528            component.insert(&mut world.entity_mut(entity), reflected.as_ref(), &registry);
529        }
530    }
531    // Insertion hooks assign fresh sequences, which must not replace saved ordering.
532    let mut next = base;
533    for (saved, entity) in effects {
534        let seq = base.saturating_add(saved);
535        next = next.max(seq.saturating_add(1));
536        world.entity_mut(entity).insert(Seq(seq));
537    }
538    let mut counter = world.resource_mut::<SeqCounter>();
539    counter.0 = counter.0.max(next);
540    // Saved IDs let handlers recognize operations retried after restoration.
541    let taken: Vec<(Entity, Issued)> = world
542        .query_filtered::<(Entity, &Issued), (With<InFlight>, Without<EffectOutcome>)>()
543        .iter(world)
544        .filter(|(entity, _)| entities.contains(entity))
545        .map(|(entity, issued)| (entity, *issued))
546        .collect();
547    for (entity, Issued(id)) in taken {
548        world
549            .entity_mut(entity)
550            .remove::<(InFlight, Issued)>()
551            .insert(Reserved(id));
552    }
553    // Never reuse an ID or run number already allocated by the destination.
554    let mut ids = world.resource_mut::<IdCounter>();
555    ids.0 = ids.0.max(checkpoint.counters.next_id);
556    let mut runs = world.get_resource_or_init::<agent::RunCounter>();
557    runs.0 = runs.0.max(checkpoint.counters.next_run);
558    Ok(entities)
559}
560
561/// Every invariant a loaded graph must hold, checked in the scratch world.
562fn validate(world: &mut World, entities: &[Entity]) -> Result<(), CheckpointError> {
563    for &entity in entities {
564        // Resuming requires a saved handler contract, but advertised families
565        // cannot restrict world handlers that accept arbitrary effects.
566        if let Some(pending) = world.get::<PendingEffect>(entity)
567            && world.get::<EffectOutcome>(entity).is_none()
568        {
569            world
570                .resource::<HandlerIndex>()
571                .entity(&pending.key)
572                .and_then(|handler| world.get::<Bound>(handler))
573                .ok_or_else(|| CheckpointError::MissingSavedHandler {
574                    key: pending.key.clone(),
575                })?;
576        }
577        if world
578            .get::<Streamed>(entity)
579            .is_some_and(|streamed| !streamed.events.is_empty() || !streamed.errors.is_empty())
580            && world.get::<EffectOutcome>(entity).is_none()
581        {
582            return Err(CheckpointError::UnresumableStream);
583        }
584        if world.get::<Utterance>(entity).is_some() {
585            read_message(world, entity)?;
586        }
587        if world.get::<ContentPart>(entity).is_some() {
588            let parent = world
589                .get::<ChildOf>(entity)
590                .ok_or(CheckpointError::InvalidGraph("content part has no parent"))?
591                .parent();
592            if world.get::<Utterance>(parent).is_none()
593                && !matches!(
594                    world.get::<ContentPart>(parent),
595                    Some(ContentPart::ToolResult { .. })
596                )
597            {
598                return Err(CheckpointError::InvalidGraph(
599                    "content part has an invalid parent",
600                ));
601            }
602        }
603        if world.get::<ToolResultStatus>(entity).is_some()
604            && !matches!(
605                world.get::<ContentPart>(entity),
606                Some(ContentPart::ToolResult { .. })
607            )
608        {
609            return Err(CheckpointError::InvalidGraph(
610                "tool result status is not on a tool result part",
611            ));
612        }
613        if world.get::<agent::Run>(entity).is_none() {
614            if world.get::<agent::Ready>(entity).is_some() {
615                return Err(CheckpointError::InvalidGraph("ready is not on a run"));
616            }
617            if world.get::<agent::Prompt>(entity).is_some() {
618                return Err(CheckpointError::InvalidGraph("prompt is not on a run"));
619            }
620        }
621        if world.get::<RequestPartEdit>(entity).is_some() {
622            let parent = world
623                .get::<ChildOf>(entity)
624                .ok_or(CheckpointError::InvalidGraph("request edit has no turn"))?
625                .parent();
626            let target = world
627                .get::<EditTarget>(entity)
628                .ok_or(CheckpointError::InvalidGraph("request edit has no target"))?
629                .0;
630            if world.get::<agent::Turn>(parent).is_none()
631                || world.get::<ContentPart>(target).is_none()
632            {
633                return Err(CheckpointError::InvalidGraph(
634                    "invalid request edit relationship",
635                ));
636            }
637        }
638        if world.get::<PendingEffect>(entity).is_some()
639            && world
640                .get::<bus::HoldOwners>(entity)
641                .is_some_and(|owners| owners.owners().any(|owner| owner.name == "rig-ecs/batch"))
642            && world.get::<crate::systems::BatchHeld>(entity).is_none()
643        {
644            return Err(CheckpointError::InvalidGraph(
645                "batch owner is missing its runtime hold marker",
646            ));
647        }
648        if world.get::<crate::systems::BatchHeld>(entity).is_some()
649            && (world.get::<bus::Held>(entity).is_none()
650                || world.get::<agent::ToolCallSlot>(entity).is_none()
651                || !world.get::<bus::HoldOwners>(entity).is_some_and(|owners| {
652                    owners.owners().any(|owner| owner.name == "rig-ecs/batch")
653                }))
654        {
655            return Err(CheckpointError::InvalidGraph(
656                "batch hold is missing its barrier, owner or tool slot",
657            ));
658        }
659        crate::agent::checkpoint::validate(world, entity)?;
660    }
661    wire_expansion(world)?;
662    Ok(())
663}
664
665/// Every binary reference in the graph resolves in the store.
666fn wire_expansion(world: &mut World) -> Result<(), CheckpointError> {
667    let ids: Vec<Entity> = world
668        .query_filtered::<Entity, With<Utterance>>()
669        .iter(world)
670        .collect();
671    for entity in ids {
672        read_message(world, entity)?;
673    }
674    Ok(())
675}
676
677/// Preflight the original reflected graph and binary store without mutating
678/// the destination. Handler completeness is validated separately.
679fn validated_state(
680    checkpoint: &Checkpoint,
681    world: &World,
682) -> Result<(BinaryAssets, HashMap<usize, Entity>), CheckpointError> {
683    if checkpoint.format != CHECKPOINT_FORMAT {
684        return Err(CheckpointError::UnsupportedFormat {
685            found: checkpoint.format,
686        });
687    }
688    // These rejected wire identifiers must stay fixed across Rust module renames.
689    for entity in &checkpoint.entities {
690        if let Some(path) = entity.keys().find(|path| {
691            matches!(
692                path.as_str(),
693                "rig_ecs::bus::binding::ProviderBinding" | "rig_ecs::bus::binding::CredentialRef"
694            )
695        }) {
696            return Err(CheckpointError::RemovedComponent { path: path.clone() });
697        }
698    }
699    if !world.contains_resource::<SeqCounter>() || !world.contains_resource::<IdCounter>() {
700        return Err(CheckpointError::NotInstalled("execution counters"));
701    }
702    let registry = world
703        .get_resource::<AppTypeRegistry>()
704        .ok_or(CheckpointError::NotInstalled("type registry"))?
705        .clone();
706    let empty = BinaryAssets::default();
707    let assets = world
708        .get_resource::<BinaryAssets>()
709        .unwrap_or(&empty)
710        .merged(&checkpoint.binaries)?;
711    let aliases = aliases(checkpoint, world)?;
712    // Validate the original graph, including every saved descriptor. Aliasing
713    // in the destination must not hide malformed or inconsistent saved data.
714    let mut scratch = World::new();
715    scratch.insert_resource(registry);
716    scratch.insert_resource(assets);
717    scratch.init_resource::<SeqCounter>();
718    scratch.init_resource::<IdCounter>();
719    scratch.init_resource::<HandlerIndex>();
720    let scratch_entities = spawn_into(checkpoint, &mut scratch, &HashMap::new())?;
721    validate(&mut scratch, &scratch_entities)?;
722    let assets = scratch
723        .remove_resource::<BinaryAssets>()
724        .ok_or(CheckpointError::InvalidGraph(
725            "validated binary store missing",
726        ))?;
727    Ok((assets, aliases))
728}
729
730fn load_state(
731    checkpoint: &Checkpoint,
732    world: &mut World,
733    assets: BinaryAssets,
734    aliases: HashMap<usize, Entity>,
735) -> Result<Loaded, CheckpointError> {
736    world.insert_resource(assets);
737    let entities = spawn_into(checkpoint, world, &aliases)?;
738    Ok(Loaded { entities })
739}
740
741impl Checkpoint {
742    /// Validate saved execution data without installing state or constructing
743    /// implementations. Hosts can call this before side-effectful assembly.
744    /// Returns an error for invalid saved contracts, graphs, or binaries, or missing
745    /// destination registry/counters. [`load_world`] checks handler compatibility
746    /// and completeness.
747    pub fn validate(&self, world: &World) -> Result<(), CheckpointError> {
748        self.requirements()?;
749        validated_state(self, world).map(|_| ())
750    }
751
752    /// Serialize the checkpoint to JSON.
753    pub fn to_json(&self) -> Result<String, CheckpointError> {
754        Ok(serde_json::to_string(self)?)
755    }
756
757    /// Parse checkpoint JSON. Returns [`CheckpointError::UnsupportedFormat`]
758    /// for another envelope format and [`CheckpointError::Json`] for invalid
759    /// data or unknown fields.
760    pub fn from_json(json: &str) -> Result<Self, CheckpointError> {
761        #[derive(Deserialize)]
762        struct Envelope {
763            format: u32,
764        }
765        let Envelope { format } = serde_json::from_str(json)?;
766        if format != CHECKPOINT_FORMAT {
767            return Err(CheckpointError::UnsupportedFormat { found: format });
768        }
769        Ok(serde_json::from_str(json)?)
770    }
771}
772
773/// Register every component and resource of the bus and the graph, and every
774/// remote wrapper, with the world's [`AppTypeRegistry`] (created if absent).
775pub fn register_types(world: &mut World) {
776    crate::reflect::install_reflect(world);
777}