1mod 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
52pub type CheckpointEntity = serde_json::Map<String, serde_json::Value>;
54
55type Reflected<'a> = Vec<(&'a str, Box<dyn PartialReflect>)>;
57
58pub const CHECKPOINT_FORMAT: u32 = 2;
60
61#[derive(Debug, Clone, Serialize, Deserialize)]
65#[serde(deny_unknown_fields)]
66pub struct Checkpoint {
67 #[serde(deserialize_with = "checkpoint_format")]
69 pub format: u32,
70 pub entities: Vec<CheckpointEntity>,
73 pub counters: Counters,
76 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
91fn 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#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
106pub struct Counters {
107 pub next_run: u64,
109 pub next_id: u64,
111}
112
113#[derive(Debug, Clone, Default)]
116pub struct Loaded {
117 pub entities: Vec<Entity>,
119}
120
121impl Loaded {
122 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#[derive(Debug, thiserror::Error)]
134#[non_exhaustive]
135pub enum CheckpointError {
136 #[error("the world has no {0}: install RigPlugin first")]
138 NotInstalled(&'static str),
139 #[error("the checkpoint is format {found}, this rig reads format {CHECKPOINT_FORMAT}")]
141 UnsupportedFormat {
142 found: u32,
144 },
145 #[error(
147 "the checkpoint contains removed `{path}`: migrate provider launch settings to the host"
148 )]
149 RemovedComponent {
150 path: String,
152 },
153 #[error("checkpoint JSON: {0}")]
155 Json(#[from] serde_json::Error),
156 #[error("`{path}` does not serialize: {source}")]
158 Serialize {
159 path: String,
161 #[source]
163 source: serde_json::Error,
164 },
165 #[error("checkpoint entity {entity}: `{path}`: {source}")]
167 Deserialize {
168 entity: usize,
170 path: String,
172 #[source]
174 source: serde_json::Error,
175 },
176 #[error("checkpoint entity {entity}: `{path}` is not a registered component")]
178 Unregistered {
179 entity: usize,
181 path: String,
183 },
184 #[error("saved handler `{key}` has a different descriptor key")]
186 HandlerKeyMismatch {
187 key: HandlerKey,
189 },
190 #[error("duplicate saved handler `{key}`")]
192 DuplicateSavedHandler {
193 key: HandlerKey,
195 },
196 #[error("duplicate supplied handler `{key}`")]
198 DuplicateSuppliedHandler {
199 key: HandlerKey,
201 },
202 #[error("supplied handler `{key}` is not required by this checkpoint")]
204 UnrequiredHandler {
205 key: HandlerKey,
207 },
208 #[error("no implementation supplied for `{key}`")]
210 MissingHandler {
211 key: HandlerKey,
213 },
214 #[error("unfinished effect requires missing saved handler `{key}`")]
216 MissingSavedHandler {
217 key: HandlerKey,
219 },
220 #[error("handler `{key}` serves {found}; the checkpoint's serves {saved}")]
222 FamilyChanged {
223 key: HandlerKey,
225 saved: EffectFamily,
227 found: EffectFamily,
229 },
230 #[error("handler `{key}` differs from its original saved descriptor")]
233 DescriptorChanged {
234 key: HandlerKey,
236 },
237 #[error("an unfinished stream with delivered progress cannot resume")]
239 UnresumableStream,
240 #[error("invalid saved content: {0}")]
242 Content(#[from] ContentError),
243 #[error("invalid saved binary store: {0}")]
245 Binary(#[from] BinaryError),
246 #[error("invalid checkpoint graph: {0}")]
248 InvalidGraph(&'static str),
249}
250
251fn 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
272fn ordered_entities(world: &mut World) -> Vec<Entity> {
276 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#[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(), ®istry, &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 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
374struct 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
399struct 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
425fn 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
456fn 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
465fn 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 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, ®istry, &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(), ®istry);
529 }
530 }
531 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 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 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
561fn validate(world: &mut World, entities: &[Entity]) -> Result<(), CheckpointError> {
563 for &entity in entities {
564 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
665fn 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
677fn 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 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 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 pub fn validate(&self, world: &World) -> Result<(), CheckpointError> {
748 self.requirements()?;
749 validated_state(self, world).map(|_| ())
750 }
751
752 pub fn to_json(&self) -> Result<String, CheckpointError> {
754 Ok(serde_json::to_string(self)?)
755 }
756
757 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
773pub fn register_types(world: &mut World) {
776 crate::reflect::install_reflect(world);
777}