laser_sdk/agent/
state.rs

1use crate::context::{Checkpoint, ContextAssembler, ContextMessage};
2use crate::error::LaserError;
3use crate::laser::Laser;
4use crate::provenance::AgentTopic;
5use crate::snapshot::SnapshotStore;
6use crate::types::ConversationId;
7use iggy::prelude::*;
8use laser_wire::snapshot::{FoldSnapshot, SnapshotOffset};
9use laser_wire::validate::Validate;
10use std::collections::{BTreeMap, BTreeSet};
11
12/// How much of a conversation's log a state fold replays. Explicit at every
13/// call (the bounded-reads law): a long-lived conversation is a long
14/// partition, and a full-partition walk must be the caller writing the word,
15/// never a silent default. Every bound except [`Last`](Self::Last) folds its
16/// whole range, read in bounded chunks.
17pub enum ReplayBound {
18    /// Fold only messages at or after these per-partition offsets, the
19    /// incremental form (a persisted cursor, a snapshot's resume offsets).
20    /// Each topic has its own partition offsets.
21    FromOffsets(BTreeMap<String, BTreeMap<u32, u64>>),
22    /// Fold only the last `n` messages found in the newest
23    /// [`CONTEXT_READ_WINDOW`](crate::context::CONTEXT_READ_WINDOW) records of
24    /// each partition, the window a context read covers. On a busy shared
25    /// partition a quiet conversation can have fewer than `n` there.
26    Last(usize),
27    /// Fold the whole partition from offset zero. Correct for a short
28    /// conversation and for a first snapshot build, expensive everywhere else.
29    Full,
30    /// Fold only what was appended after a [`Checkpoint`], per topic and
31    /// partition, up to the tail.
32    FromCheckpoint(Checkpoint),
33    /// Fold the history up to a [`Checkpoint`], per topic and partition, and
34    /// stop there.
35    At(Checkpoint),
36}
37
38/// Rebuilds in-memory state by folding a conversation's logged events (event sourcing).
39pub struct ConversationState;
40
41impl ConversationState {
42    /// Replay `topics` for `conversation` under the explicit `bound` and fold
43    /// every message through `fold`, starting from `init`.
44    pub async fn load<S, F>(
45        laser: &Laser,
46        conversation: ConversationId,
47        topics: Vec<AgentTopic<'static>>,
48        bound: ReplayBound,
49        init: S,
50        fold: F,
51    ) -> Result<S, LaserError>
52    where
53        F: FnMut(S, &ContextMessage) -> S,
54    {
55        let assembler = ContextAssembler::builder()
56            .conversation_id(conversation)
57            .topics(topics);
58        let history = match bound {
59            ReplayBound::FromOffsets(offsets) => {
60                assembler
61                    .policy(Box::new(crate::context::LastN(usize::MAX)))
62                    .from_checkpoint(Checkpoint::from_topic_offsets(offsets))
63                    .build()
64                    .replay(laser)
65                    .await?
66            }
67            ReplayBound::Last(n) => {
68                assembler
69                    .policy(Box::new(crate::context::LastN(n)))
70                    .build()
71                    .assemble(laser)
72                    .await?
73            }
74            ReplayBound::Full => {
75                assembler
76                    .policy(Box::new(crate::context::LastN(usize::MAX)))
77                    .build()
78                    .replay(laser)
79                    .await?
80            }
81            ReplayBound::FromCheckpoint(checkpoint) => {
82                assembler
83                    .policy(Box::new(crate::context::LastN(usize::MAX)))
84                    .from_checkpoint(checkpoint)
85                    .build()
86                    .replay(laser)
87                    .await?
88            }
89            ReplayBound::At(checkpoint) => {
90                assembler
91                    .policy(Box::new(crate::context::LastN(usize::MAX)))
92                    .to_checkpoint(checkpoint)
93                    .build()
94                    .replay(laser)
95                    .await?
96            }
97        };
98        Ok(history.iter().fold(init, fold))
99    }
100
101    /// Replay through a [`SnapshotStore`]: seed the fold with the newest
102    /// snapshot's state (decoded as JSON into `S`) and replay only from
103    /// `as_of + 1` per partition, so a thousand-record conversation
104    /// snapshotted at nine hundred folds one hundred records, not a thousand.
105    /// A conversation with no snapshot folds fully from `init`, the honest
106    /// first build.
107    pub async fn load_with<Store, S, F>(
108        laser: &Laser,
109        store: &Store,
110        conversation: ConversationId,
111        topics: Vec<AgentTopic<'static>>,
112        init: S,
113        fold: F,
114    ) -> Result<S, LaserError>
115    where
116        Store: SnapshotStore + Sync,
117        S: serde::de::DeserializeOwned,
118        F: FnMut(S, &ContextMessage) -> S,
119    {
120        let (seed, bound) = match store.latest(conversation.into()).await? {
121            Some(snapshot) => {
122                if snapshot.conversation != conversation.into() {
123                    return Err(LaserError::Invalid(
124                        "snapshot conversation does not match the fold".to_owned(),
125                    ));
126                }
127                let checkpoint = checkpoint_from_snapshot(laser, &snapshot, &topics).await?;
128                let state: S = serde_json::from_slice(&snapshot.state).map_err(|error| {
129                    LaserError::Codec(format!("decode snapshot state: {error}"))
130                })?;
131                (state, ReplayBound::FromCheckpoint(checkpoint))
132            }
133            None => (init, ReplayBound::Full),
134        };
135        Self::load(laser, conversation, topics, bound, seed, fold).await
136    }
137}
138
139/// The per-partition offsets a fold resumes from after `snapshot`: one past
140/// each partition's last folded offset.
141pub fn resume_offsets(snapshot: &FoldSnapshot) -> Vec<SnapshotOffset> {
142    snapshot
143        .as_of
144        .iter()
145        .map(|entry| {
146            SnapshotOffset::new(
147                entry.topic_id,
148                entry.topic_created_at_micros,
149                entry.partition_id,
150                entry.offset.saturating_add(1),
151            )
152        })
153        .collect()
154}
155
156/// A snapshot of `state` folded up to `checkpoint` for `conversation` under
157/// the fold named `fold`. It records the stream and every checkpointed
158/// topic by id and creation time, so a resume refuses a recreated source.
159/// Empty partitions are left out.
160pub async fn snapshot_from_checkpoint(
161    laser: &Laser,
162    conversation: ConversationId,
163    fold: &str,
164    checkpoint: &Checkpoint,
165    state: Vec<u8>,
166) -> Result<FoldSnapshot, LaserError> {
167    let stream_name = laser.stream_required()?;
168    let stream = Identifier::named(stream_name)?;
169    let client = laser.client();
170    let details = client
171        .get_stream(&stream)
172        .await?
173        .ok_or_else(|| LaserError::Invalid(format!("stream `{stream_name}` does not exist")))?;
174    let mut as_of = Vec::new();
175    for (topic, partitions) in checkpoint.topics() {
176        let topic_details = client
177            .get_topic(&stream, &Identifier::named(topic)?)
178            .await?
179            .ok_or_else(|| LaserError::Invalid(format!("topic `{topic}` does not exist")))?;
180        for (&partition, &next) in partitions {
181            if next > 0 {
182                as_of.push(SnapshotOffset::new(
183                    topic_details.id,
184                    topic_details.created_at.as_micros(),
185                    partition,
186                    next - 1,
187                ));
188            }
189        }
190    }
191    as_of.sort_by_key(|entry| {
192        (
193            entry.topic_id,
194            entry.topic_created_at_micros,
195            entry.partition_id,
196        )
197    });
198    let snapshot = FoldSnapshot {
199        stream: stream_name.to_owned(),
200        stream_id: details.id,
201        stream_created_at_micros: details.created_at.as_micros(),
202        conversation: conversation.into(),
203        fold: fold.to_owned(),
204        as_of,
205        state,
206    };
207    snapshot
208        .validate()
209        .map_err(|error| LaserError::Invalid(error.to_string()))?;
210    Ok(snapshot)
211}
212
213pub async fn checkpoint_from_snapshot(
214    laser: &Laser,
215    snapshot: &FoldSnapshot,
216    topics: &[AgentTopic<'static>],
217) -> Result<Checkpoint, LaserError> {
218    snapshot
219        .validate()
220        .map_err(|error| LaserError::Invalid(error.to_string()))?;
221    let stream_name = laser.stream_required()?;
222    if snapshot.stream != stream_name {
223        return Err(LaserError::Invalid(
224            "snapshot stream does not match the client".to_owned(),
225        ));
226    }
227    let stream = Identifier::named(stream_name)?;
228    let client = laser.client();
229    let current_stream = client
230        .get_stream(&stream)
231        .await?
232        .ok_or_else(|| LaserError::Invalid("snapshot stream no longer exists".to_owned()))?;
233    if snapshot.stream_id != current_stream.id
234        || snapshot.stream_created_at_micros != current_stream.created_at.as_micros()
235    {
236        return Err(LaserError::Invalid(
237            "snapshot stream generation changed".to_owned(),
238        ));
239    }
240    let mut per_topic = BTreeMap::new();
241    let mut current_sources = BTreeSet::new();
242    for topic in topics {
243        let name = topic.topic_string();
244        let details = client
245            .get_topic(&stream, &topic.as_identifier())
246            .await?
247            .ok_or_else(|| {
248                LaserError::Invalid(format!("snapshot source topic {name} is missing"))
249            })?;
250        let generation = details.created_at.as_micros();
251        current_sources.insert((details.id, generation));
252        let offsets = (0..crate::poll::bounded_partitions(details.partitions_count))
253            .map(|partition| {
254                (
255                    partition,
256                    snapshot.resume_offset(details.id, generation, partition),
257                )
258            })
259            .collect();
260        per_topic.insert(name, offsets);
261    }
262    if snapshot
263        .as_of
264        .iter()
265        .any(|entry| !current_sources.contains(&(entry.topic_id, entry.topic_created_at_micros)))
266    {
267        return Err(LaserError::Invalid(
268            "snapshot source topic generation changed".to_owned(),
269        ));
270    }
271    Ok(Checkpoint::from_topic_offsets(per_topic))
272}
273
274#[cfg(test)]
275mod tests {
276    use super::*;
277    use crate::snapshot::FoldSnapshot;
278    use laser_wire::snapshot::SnapshotOffset;
279
280    #[test]
281    fn given_a_snapshot_when_resuming_then_should_start_one_past_each_folded_offset() {
282        let snapshot = FoldSnapshot {
283            stream: "agents".to_owned(),
284            stream_id: 0,
285            stream_created_at_micros: 100,
286            conversation: laser_wire::agent::ConversationId::from_u128(1),
287            fold: "planner".to_owned(),
288            as_of: vec![
289                SnapshotOffset::new(2, 20, 0, 899),
290                SnapshotOffset::new(2, 20, 2, 41),
291            ],
292            state: b"{}".to_vec(),
293        };
294        let offsets = resume_offsets(&snapshot);
295        assert_eq!(
296            offsets,
297            vec![
298                SnapshotOffset::new(2, 20, 0, 900),
299                SnapshotOffset::new(2, 20, 2, 42)
300            ]
301        );
302    }
303}
laser_sdk/agent/state.rs

laser_sdk/agent/
state.rs
