rig_ecs/agent/
checkpoint.rs

1//! Durable tool-turn commits and host-owned holds at their checkpoint boundary.
2//!
3//! A hold is armed before the desired tool turn. It only prevents advancement
4//! after a real batch commit, never dispatch of the batch being awaited. Owners
5//! release their own named hold; all remaining owners must release before the
6//! run continues. A checkpoint preserves holds, but does not persist host tools
7//! or external side effects.
8//!
9//! ```
10//! use rig_ecs::agent::checkpoint::ToolTurnHolds;
11//! let holds = ToolTurnHolds::default();
12//! assert!(!holds.blocks(1));
13//! ```
14
15use bevy_reflect::Reflect;
16use std::collections::BTreeMap;
17
18use bevy_ecs::prelude::*;
19use serde::{Deserialize, Serialize};
20
21use super::{Failed, Run, Settled};
22
23/// A completed tool batch on its existing turn entity. The number is the run's
24/// model-turn cursor, including earlier non-tool turns, not a batch count.
25#[derive(Component, Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Reflect)]
26#[reflect(Component)]
27pub struct ToolTurnCommit {
28    /// Model turn whose complete results became history.
29    pub turn: usize,
30}
31
32/// The assistant utterance committed for this tool-bearing turn.
33#[derive(Component, Debug, Clone, Copy, PartialEq, Eq, Reflect)]
34#[relationship(relationship_target = AssistantForTurns)]
35#[reflect(Component)]
36pub struct TurnAssistant(pub Entity);
37
38/// Turns whose assistant utterance is this entity.
39#[derive(Component, Debug, Default, Reflect)]
40#[relationship_target(relationship = TurnAssistant)]
41#[reflect(Component)]
42pub struct AssistantForTurns(Vec<Entity>);
43
44/// The single ordered user utterance containing the committed tool results.
45#[derive(Component, Debug, Clone, Copy, PartialEq, Eq, Reflect)]
46#[relationship(relationship_target = ResultsForTurns)]
47#[reflect(Component)]
48pub struct TurnResults(pub Entity);
49
50/// Turns whose tool-result utterance is this entity.
51#[derive(Component, Debug, Default, Reflect)]
52#[relationship_target(relationship = TurnResults)]
53#[reflect(Component)]
54pub struct ResultsForTurns(Vec<Entity>);
55
56/// Armed checkpoint holds on a run, indexed by host owner name. Each value is
57/// the earliest committed model-turn number at which that owner blocks advance.
58#[derive(Component, Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, Reflect)]
59#[reflect(Component)]
60pub struct ToolTurnHolds(BTreeMap<String, usize>);
61
62impl ToolTurnHolds {
63    /// Owners and their minimum committed model-turn numbers.
64    pub fn owners(&self) -> impl Iterator<Item = (&str, usize)> {
65        self.0.iter().map(|(owner, turn)| (owner.as_str(), *turn))
66    }
67
68    /// Whether any owner is holding this committed turn.
69    pub fn blocks(&self, committed_turn: usize) -> bool {
70        self.0.values().any(|turn| *turn <= committed_turn)
71    }
72}
73
74/// A live batch commit notification, emitted after deferred graph writes and
75/// phase changes are visible. A checkpoint load does not emit this event. Inspect
76/// ToolTurnCommit for durable state; do not treat component insertion on load
77/// as a notification that tools ran again. A terminal observer that removes
78/// the owning graph before delivery suppresses this live event.
79#[derive(Event, Debug, Clone, Copy)]
80pub struct ToolTurnCommitted {
81    /// Run owning the committed turn.
82    pub run: Entity,
83    /// Existing turn entity carrying ToolTurnCommit and utterance links.
84    pub turn: Entity,
85}
86
87/// An invalid host request to arm or release a checkpoint hold.
88#[derive(Debug, thiserror::Error, PartialEq, Eq)]
89pub enum TurnHoldError {
90    /// The entity does not name a live run.
91    #[error("checkpoint hold needs a live run")]
92    NotLiveRun,
93    /// An owner must have a nonempty stable name.
94    #[error("checkpoint hold owner must not be empty")]
95    EmptyOwner,
96    /// Model turns are numbered from one.
97    #[error("checkpoint turn must be at least one")]
98    ZeroTurn,
99}
100
101/// Arm a hold for this owner. Re-arming is idempotent and cannot move an
102/// existing owner's boundary later; release that hold explicitly to replace it.
103/// Returns true when a hold was added or moved earlier. Hosts must arm before
104/// advancement; this cannot retract a request that has already dispatched.
105/// Returns an error for a non-live run, empty owner, or zero turn.
106pub fn hold_after_tool_turn(
107    world: &mut World,
108    run: Entity,
109    owner: impl Into<String>,
110    turn: usize,
111) -> Result<bool, TurnHoldError> {
112    if world.get::<Run>(run).is_none()
113        || world.get::<Failed>(run).is_some()
114        || world.get::<Settled>(run).is_some()
115    {
116        return Err(TurnHoldError::NotLiveRun);
117    }
118    let owner = owner.into();
119    if owner.is_empty() {
120        return Err(TurnHoldError::EmptyOwner);
121    }
122    if turn == 0 {
123        return Err(TurnHoldError::ZeroTurn);
124    }
125    let mut holds = world.get::<ToolTurnHolds>(run).cloned().unwrap_or_default();
126    let previous = holds.0.get(&owner).copied();
127    if previous.is_some_and(|previous| previous <= turn) {
128        return Ok(false);
129    }
130    holds.0.insert(owner, turn);
131    world.entity_mut(run).insert(holds);
132    Ok(true)
133}
134
135/// Release only the named owner's hold. Also permitted on a terminal run for
136/// host cleanup. Unknown owners are an idempotent no-op. The next schedule
137/// update continues a nonterminal run once no hold blocks its committed turn.
138/// Returns whether an entry was removed, or an error for an empty owner or
139/// entity without [`Run`].
140pub fn release_tool_turn_hold(
141    world: &mut World,
142    run: Entity,
143    owner: &str,
144) -> Result<bool, TurnHoldError> {
145    if world.get::<Run>(run).is_none() {
146        return Err(TurnHoldError::NotLiveRun);
147    }
148    if owner.is_empty() {
149        return Err(TurnHoldError::EmptyOwner);
150    }
151    let Some(mut holds) = world.get_mut::<ToolTurnHolds>(run) else {
152        return Ok(false);
153    };
154    let removed = holds.0.remove(owner).is_some();
155    if holds.0.is_empty() {
156        world.entity_mut(run).remove::<ToolTurnHolds>();
157    }
158    Ok(removed)
159}
160
161/// The invariants a loaded tool-turn checkpoint holds on `entity`: a hold
162/// on a run with named owners and non-zero turns; a commit or its links on
163/// a turn of a run, naming utterances of that run in their roles; a commit
164/// with both links, no batch out, and a turn the run's cursor has reached.
165pub(crate) fn validate(
166    world: &World,
167    entity: Entity,
168) -> Result<(), crate::checkpoint::CheckpointError> {
169    use super::{Batch, Cursor, Role, Turn, Utterance};
170    let refused = crate::checkpoint::CheckpointError::InvalidGraph;
171    if let Some(holds) = world.get::<ToolTurnHolds>(entity)
172        && (world.get::<Run>(entity).is_none()
173            || holds
174                .owners()
175                .any(|(owner, turn)| owner.is_empty() || turn == 0))
176    {
177        return Err(refused("invalid tool-turn hold"));
178    }
179    let assistant = world.get::<TurnAssistant>(entity);
180    let results = world.get::<TurnResults>(entity);
181    let commit = world.get::<ToolTurnCommit>(entity);
182    if assistant.is_none() && results.is_none() && commit.is_none() {
183        return Ok(());
184    }
185    let run = world
186        .get::<ChildOf>(entity)
187        .map(ChildOf::parent)
188        .filter(|run| world.get::<Run>(*run).is_some())
189        .ok_or(refused("tool-turn link has no run"))?;
190    if world.get::<Turn>(entity).is_none() {
191        return Err(refused("tool-turn link is not on a turn"));
192    }
193    for (target, role) in [
194        (assistant.map(|link| link.0), Role::Assistant),
195        (results.map(|link| link.0), Role::User),
196    ] {
197        if let Some(target) = target
198            && (world.get::<ChildOf>(target).map(ChildOf::parent) != Some(run)
199                || world.get::<Utterance>(target).is_none()
200                || world.get::<Role>(target) != Some(&role))
201        {
202            return Err(refused("tool-turn utterance has invalid ownership or role"));
203        }
204    }
205    if let Some(commit) = commit {
206        if commit.turn == 0
207            || assistant.is_none()
208            || results.is_none()
209            || world.get::<Batch>(entity).is_some()
210            || world
211                .get::<Cursor>(run)
212                .is_none_or(|cursor| commit.turn > cursor.turn)
213        {
214            return Err(refused("invalid committed tool turn"));
215        }
216    } else if results.is_some() {
217        return Err(refused("tool results have no commit"));
218    }
219    Ok(())
220}
rig_ecs/agent/checkpoint.rs

rig_ecs/agent/
checkpoint.rs
