aion_store/
observability.rs

1//! The durable observability (`O`) keyspace contract — NOI-5's durability spine.
2//!
3//! This module defines the persistence contract for the agent-observability
4//! transcript: an **append-only, per-`(workflow, activity, attempt)`** stream of
5//! [`aion_core::ActivityEvent`] records that survives kill-9 and failover, is
6//! replayable by `store_seq`, and is **never** part of the workflow replay log.
7//!
8//! # The `O` keyspace is NOT the `E`-stream (LOCKED)
9//!
10//! Workflow replay authority lives exclusively on the `E`-stream (the
11//! [`crate::WritableEventStore`] append path). An [`ActivityRecord`] is an
12//! observability record: the replay decoder never scans this keyspace and could
13//! not decode one of these records as an `Event` even if it did (different region
14//! tag, different schema). The byte-level disjointness is what makes "durable but
15//! non-replay-authoritative" a *guarantee*, not a hope — see
16//! `aion-store-haematite`'s `observability` module for the `O` (0x4F) region tag
17//! and the disjointness test.
18//!
19//! # Single-writer, server-allocated `store_seq`
20//!
21//! `store_seq` is **not** allocated by the store: it is a caller-supplied
22//! `expected_seq` under optimistic concurrency, exactly like the workflow-history
23//! append path. [`ObservabilityStore::append_activity_event`] returns the
24//! `SequenceConflict` the server's sequencer re-reads-head-and-retries on. The
25//! server is the *single writer* to this keyspace, so monotonicity is enforced by
26//! the server's read-head -> append(expected_seq) -> retry loop, not by any magic
27//! in the store. This mirrors [`StoreError::SequenceConflict`] on the workflow
28//! path and is why the store deliberately does not auto-allocate an id.
29
30use async_trait::async_trait;
31
32use aion_core::{ActivityEvent, ActivityId, WorkflowId};
33
34use crate::StoreError;
35
36/// The durable key of one observability stream: a `(workflow, activity, attempt)`
37/// triple. Every [`ActivityRecord`] for one running agent attempt shares this key
38/// and is ordered by `store_seq` within it.
39#[derive(Clone, Debug, PartialEq, Eq, Hash)]
40pub struct ActivityStreamKey {
41    /// The workflow the activity belongs to.
42    pub workflow_id: WorkflowId,
43    /// The activity within the workflow.
44    pub activity_id: ActivityId,
45    /// The attempt number — the third key axis (NOI-0). Two attempts of one
46    /// activity are DISTINCT streams; a within-attempt failover shares one stream
47    /// (so a dying + adopting worker's events dedupe), while a retry is a new
48    /// attempt and therefore a new stream.
49    pub attempt: u32,
50}
51
52impl ActivityStreamKey {
53    /// Build a stream key from its three components.
54    #[must_use]
55    pub const fn new(workflow_id: WorkflowId, activity_id: ActivityId, attempt: u32) -> Self {
56        Self {
57            workflow_id,
58            activity_id,
59            attempt,
60        }
61    }
62
63    /// The stream key an [`ActivityEvent`] belongs to.
64    #[must_use]
65    pub fn of(event: &ActivityEvent) -> Self {
66        Self {
67            workflow_id: event.workflow_id.clone(),
68            activity_id: event.activity_id.clone(),
69            attempt: event.attempt,
70        }
71    }
72}
73
74/// A durably persisted observability event: an [`ActivityEvent`] with its
75/// server-stamped `store_seq` guaranteed present.
76///
77/// The wire envelope carries `store_seq: Option<u64>` (`None` until persisted);
78/// once read back from the `O` keyspace the sequence is always present, so this
79/// record exposes it as a non-optional field alongside the event.
80#[derive(Clone, Debug, PartialEq)]
81pub struct ActivityRecord {
82    /// The monotonic, server-allocated sequence assigned at durable commit.
83    pub store_seq: u64,
84    /// The persisted event. Its `store_seq` field is populated to match
85    /// [`Self::store_seq`] so a record read back is self-describing.
86    pub event: ActivityEvent,
87}
88
89/// One retained transcript stream of a workflow: its key and its head
90/// (the number of durably retained records / the next `store_seq`).
91#[derive(Clone, Debug, PartialEq, Eq)]
92pub struct ActivityStreamSummary {
93    /// The stream's `(workflow, activity, attempt)` key.
94    pub key: ActivityStreamKey,
95    /// Next `store_seq` to be written == count of retained records.
96    pub head: u64,
97}
98
99/// Durable, append-only observability keyspace contract.
100///
101/// Implemented by the haematite backend for production and by
102/// [`InMemoryObservabilityStore`] for tests + conformance. The server is the
103/// single writer; every method keys on the `(workflow, activity, attempt)`
104/// triple, never on the workflow alone.
105#[async_trait]
106pub trait ObservabilityStore: Send + Sync + 'static {
107    /// Append `event` to its `(workflow, activity, attempt)` stream at
108    /// `expected_seq` (the current head the caller believes it holds).
109    ///
110    /// On success returns the newly assigned `store_seq` (which equals
111    /// `expected_seq`) — the caller advances its head to `store_seq + 1`. On a
112    /// stale expectation returns [`StoreError::SequenceConflict`] with the actual
113    /// head, leaving the stream unchanged, so the server's sequencer can re-read
114    /// the advanced head and retry. **Ephemeral events must never be passed
115    /// here** — they are WS-forward-only and are filtered out before this call.
116    ///
117    /// # Errors
118    /// [`StoreError::SequenceConflict`] on a stale `expected_seq`; otherwise a
119    /// backend or serialization error.
120    async fn append_activity_event(
121        &self,
122        expected_seq: u64,
123        event: &ActivityEvent,
124    ) -> Result<u64, StoreError>;
125
126    /// Read the current head (next `store_seq` to be written) for `key`.
127    ///
128    /// An unwritten stream reads head `0`. The server's sequencer seeds its
129    /// retry loop from this value.
130    ///
131    /// # Errors
132    /// A backend or serialization error.
133    async fn activity_head(&self, key: &ActivityStreamKey) -> Result<u64, StoreError>;
134
135    /// Read every record for `key` with `store_seq >= from_seq`, in order.
136    ///
137    /// This is the resume-by-`store_seq` primitive: a reconnecting transcript
138    /// client replays from its last-seen cursor without paying for the whole
139    /// stream. An unwritten stream (or a `from_seq` beyond the head) reads empty.
140    ///
141    /// # Errors
142    /// A backend or serialization error.
143    async fn read_activity_events_from(
144        &self,
145        key: &ActivityStreamKey,
146        from_seq: u64,
147    ) -> Result<Vec<ActivityRecord>, StoreError>;
148
149    /// Enumerate every retained transcript stream of `workflow_id`, ordered by
150    /// `(activity_id, attempt)` ascending. A workflow with no retained
151    /// transcript reads empty (old runs simply have none).
152    ///
153    /// # Errors
154    /// A backend or serialization error.
155    async fn list_activity_streams(
156        &self,
157        workflow_id: &WorkflowId,
158    ) -> Result<Vec<ActivityStreamSummary>, StoreError>;
159}
160
161/// An in-memory [`ObservabilityStore`] reference implementation for tests.
162///
163/// Enforces the SAME optimistic-concurrency contract the haematite backend does:
164/// an append with a stale `expected_seq` returns [`StoreError::SequenceConflict`]
165/// and writes nothing, so the server's retry loop can be exercised without a real
166/// database. A `std::sync::Mutex` serializes the read-compare-write so two racing
167/// appends on one stream cannot both win — the same single-shard-actor guarantee
168/// the haematite backend gives.
169#[derive(Debug, Default)]
170pub struct InMemoryObservabilityStore {
171    streams:
172        std::sync::Mutex<std::collections::HashMap<ActivityStreamKeyBytes, Vec<ActivityRecord>>>,
173}
174
175/// A hashable, owned encoding of [`ActivityStreamKey`] for the in-memory map.
176type ActivityStreamKeyBytes = (uuid::Uuid, u64, u32);
177
178fn key_bytes(key: &ActivityStreamKey) -> ActivityStreamKeyBytes {
179    (
180        key.workflow_id.as_uuid(),
181        key.activity_id.sequence_position(),
182        key.attempt,
183    )
184}
185
186/// The next `store_seq` for an in-memory stream = its record count.
187///
188/// A `Vec` length that does not fit in `u64` is unrepresentable on any supported
189/// target (a 64-bit `usize` maxes at `u64::MAX`), so the saturating conversion is
190/// exact in practice; it is written as a fallible convert to satisfy the
191/// deny-level pedantic cast lints without an `as` cast.
192fn stream_head(stream: &[ActivityRecord]) -> u64 {
193    u64::try_from(stream.len()).unwrap_or(u64::MAX)
194}
195
196#[async_trait]
197impl ObservabilityStore for InMemoryObservabilityStore {
198    async fn append_activity_event(
199        &self,
200        expected_seq: u64,
201        event: &ActivityEvent,
202    ) -> Result<u64, StoreError> {
203        let key = ActivityStreamKey::of(event);
204        let mut streams = self.streams.lock().map_err(|error| {
205            StoreError::Backend(format!("observability mutex poisoned: {error}"))
206        })?;
207        let stream = streams.entry(key_bytes(&key)).or_default();
208        let head = stream_head(stream);
209        if head != expected_seq {
210            return Err(StoreError::SequenceConflict {
211                expected: expected_seq,
212                found: head,
213            });
214        }
215        let mut event = event.clone();
216        event.store_seq = Some(head);
217        stream.push(ActivityRecord {
218            store_seq: head,
219            event,
220        });
221        Ok(head)
222    }
223
224    async fn activity_head(&self, key: &ActivityStreamKey) -> Result<u64, StoreError> {
225        let streams = self.streams.lock().map_err(|error| {
226            StoreError::Backend(format!("observability mutex poisoned: {error}"))
227        })?;
228        Ok(streams
229            .get(&key_bytes(key))
230            .map_or(0, |stream| stream_head(stream)))
231    }
232
233    async fn read_activity_events_from(
234        &self,
235        key: &ActivityStreamKey,
236        from_seq: u64,
237    ) -> Result<Vec<ActivityRecord>, StoreError> {
238        let streams = self.streams.lock().map_err(|error| {
239            StoreError::Backend(format!("observability mutex poisoned: {error}"))
240        })?;
241        Ok(streams
242            .get(&key_bytes(key))
243            .map_or_else(Vec::new, |stream| {
244                stream
245                    .iter()
246                    .filter(|record| record.store_seq >= from_seq)
247                    .cloned()
248                    .collect()
249            }))
250    }
251
252    async fn list_activity_streams(
253        &self,
254        workflow_id: &WorkflowId,
255    ) -> Result<Vec<ActivityStreamSummary>, StoreError> {
256        let streams = self.streams.lock().map_err(|error| {
257            StoreError::Backend(format!("observability mutex poisoned: {error}"))
258        })?;
259        let mut summaries: Vec<ActivityStreamSummary> = streams
260            .iter()
261            .filter(|((workflow, _activity, _attempt), _records)| {
262                *workflow == workflow_id.as_uuid()
263            })
264            .map(
265                |(&(workflow, activity_seq, attempt), records)| ActivityStreamSummary {
266                    key: ActivityStreamKey::new(
267                        WorkflowId::new(workflow),
268                        ActivityId::from_sequence_position(activity_seq),
269                        attempt,
270                    ),
271                    head: stream_head(records),
272                },
273            )
274            .collect();
275        summaries.sort_by_key(|summary| {
276            (
277                summary.key.activity_id.sequence_position(),
278                summary.key.attempt,
279            )
280        });
281        Ok(summaries)
282    }
283}
284
285#[cfg(test)]
286mod tests {
287    use super::*;
288    use aion_core::{ActivityEventKind, MessageRole};
289    use chrono::Utc;
290    use uuid::Uuid;
291
292    fn event(attempt: u32, worker_seq: u64, text: &str) -> ActivityEvent {
293        ActivityEvent {
294            workflow_id: WorkflowId::new(Uuid::from_u128(1)),
295            activity_id: ActivityId::from_sequence_position(3),
296            attempt,
297            agent_id: Uuid::from_u128(9),
298            agent_role: "orchestrator".to_owned(),
299            emitted_at: Utc::now(),
300            worker_seq,
301            store_seq: None,
302            ephemeral: false,
303            kind: ActivityEventKind::Message {
304                role: MessageRole::Assistant,
305                text: text.to_owned(),
306            },
307        }
308    }
309
310    #[tokio::test]
311    async fn append_assigns_contiguous_store_seq_from_zero() -> Result<(), StoreError> {
312        let store = InMemoryObservabilityStore::default();
313        let key = ActivityStreamKey::new(
314            WorkflowId::new(Uuid::from_u128(1)),
315            ActivityId::from_sequence_position(3),
316            0,
317        );
318        assert_eq!(store.activity_head(&key).await?, 0);
319        assert_eq!(store.append_activity_event(0, &event(0, 1, "a")).await?, 0);
320        assert_eq!(store.append_activity_event(1, &event(0, 2, "b")).await?, 1);
321        assert_eq!(store.activity_head(&key).await?, 2);
322        let records = store.read_activity_events_from(&key, 0).await?;
323        assert_eq!(records.len(), 2);
324        assert_eq!(records[0].store_seq, 0);
325        assert_eq!(records[0].event.store_seq, Some(0));
326        assert_eq!(records[1].store_seq, 1);
327        Ok(())
328    }
329
330    #[tokio::test]
331    async fn stale_expected_seq_conflicts_and_writes_nothing() -> Result<(), StoreError> {
332        let store = InMemoryObservabilityStore::default();
333        store.append_activity_event(0, &event(0, 1, "a")).await?;
334        // Re-appending at the already-consumed seq 0 conflicts against head 1.
335        let conflict = store.append_activity_event(0, &event(0, 2, "dup")).await;
336        assert_eq!(
337            conflict,
338            Err(StoreError::SequenceConflict {
339                expected: 0,
340                found: 1
341            })
342        );
343        let key = ActivityStreamKey::of(&event(0, 0, ""));
344        // Nothing partial was written: still exactly one record.
345        assert_eq!(store.read_activity_events_from(&key, 0).await?.len(), 1);
346        Ok(())
347    }
348
349    #[tokio::test]
350    async fn attempts_are_disjoint_streams() -> Result<(), StoreError> {
351        let store = InMemoryObservabilityStore::default();
352        store
353            .append_activity_event(0, &event(0, 1, "attempt-0"))
354            .await?;
355        // A different attempt is a fresh stream with its own head at 0.
356        store
357            .append_activity_event(0, &event(1, 1, "attempt-1"))
358            .await?;
359        let key0 = ActivityStreamKey::new(
360            WorkflowId::new(Uuid::from_u128(1)),
361            ActivityId::from_sequence_position(3),
362            0,
363        );
364        let key1 = ActivityStreamKey::new(
365            WorkflowId::new(Uuid::from_u128(1)),
366            ActivityId::from_sequence_position(3),
367            1,
368        );
369        assert_eq!(store.activity_head(&key0).await?, 1);
370        assert_eq!(store.activity_head(&key1).await?, 1);
371        Ok(())
372    }
373
374    /// Two activities x two attempts of wf-1 plus one stream of wf-2: listing
375    /// wf-1 yields exactly its three streams, ordered by `(activity, attempt)`
376    /// ascending, each with the correct head.
377    #[tokio::test]
378    async fn list_activity_streams_orders_by_activity_then_attempt() -> Result<(), StoreError> {
379        let store = InMemoryObservabilityStore::default();
380        let event_for = |activity_seq: u64, attempt: u32, workflow: u128| {
381            let mut event = event(attempt, 1, "x");
382            event.workflow_id = WorkflowId::new(Uuid::from_u128(workflow));
383            event.activity_id = ActivityId::from_sequence_position(activity_seq);
384            event
385        };
386        // wf-1: activity 3 attempt 0 (two records), activity 3 attempt 1 (one),
387        // activity 5 attempt 0 (one). Inserted deliberately out of order.
388        store.append_activity_event(0, &event_for(5, 0, 1)).await?;
389        store.append_activity_event(0, &event_for(3, 1, 1)).await?;
390        store.append_activity_event(0, &event_for(3, 0, 1)).await?;
391        store.append_activity_event(1, &event_for(3, 0, 1)).await?;
392        // wf-2: one stream that must not leak into wf-1's enumeration.
393        store.append_activity_event(0, &event_for(3, 0, 2)).await?;
394
395        let summaries = store
396            .list_activity_streams(&WorkflowId::new(Uuid::from_u128(1)))
397            .await?;
398        let listed: Vec<(u64, u32, u64)> = summaries
399            .iter()
400            .map(|summary| {
401                (
402                    summary.key.activity_id.sequence_position(),
403                    summary.key.attempt,
404                    summary.head,
405                )
406            })
407            .collect();
408        assert_eq!(listed, vec![(3, 0, 2), (3, 1, 1), (5, 0, 1)]);
409        Ok(())
410    }
411
412    #[tokio::test]
413    async fn list_activity_streams_is_empty_for_unknown_workflow() -> Result<(), StoreError> {
414        let store = InMemoryObservabilityStore::default();
415        store.append_activity_event(0, &event(0, 1, "a")).await?;
416        let summaries = store
417            .list_activity_streams(&WorkflowId::new(Uuid::from_u128(99)))
418            .await?;
419        assert!(summaries.is_empty(), "an unwritten workflow lists empty");
420        Ok(())
421    }
422
423    #[tokio::test]
424    async fn read_from_resumes_by_store_seq() -> Result<(), StoreError> {
425        let store = InMemoryObservabilityStore::default();
426        for seq in 0..5u64 {
427            store
428                .append_activity_event(seq, &event(0, seq, "x"))
429                .await?;
430        }
431        let key = ActivityStreamKey::of(&event(0, 0, ""));
432        let tail = store.read_activity_events_from(&key, 3).await?;
433        assert_eq!(tail.len(), 2);
434        assert_eq!(tail[0].store_seq, 3);
435        assert_eq!(tail[1].store_seq, 4);
436        Ok(())
437    }
438}
aion_store/observability.rs

aion_store/
observability.rs
