cortex_ledger/

trace.rs

//! Trace assembly: open, attach event, close. Emits dense, monotonic ordinals.
//!
//! A `TraceAssembler` is the in-memory bookkeeping for one open trace. It
//! tracks the trace id, the open status, and the next ordinal to assign.
//! Attach returns the ordinal an event should be persisted with (in
//! `trace_events.ordinal`, BUILD_SPEC §10).
//!
//! ## Ordinal invariants
//!
//! - The first attached event gets ordinal `0`.
//! - Ordinals are **dense** (no gaps) and **strictly monotonic** within a
//!   trace.
//! - Closing a trace freezes its ordinals; further attaches return
//!   `Err(TraceError::Closed)`.
//!
//! These match the `UNIQUE (trace_id, ordinal)` constraint in
//! `trace_events`.
//!
//! ## What this module does NOT do
//!
//! - Persist anything. The caller composes a `TraceAssembler` with a
//!   `JsonlLog` (and later, a `cortex-store` repository) to actually
//!   write rows.
//! - Compute hashes. Hash chaining is `crate::hash`'s job; the assembler
//!   is purely about ordinal bookkeeping and trace lifecycle.

use chrono::{DateTime, Utc};
use cortex_core::{EventId, Trace, TraceId, TraceStatus, SCHEMA_VERSION};
use thiserror::Error;

/// Errors raised by [`TraceAssembler`].
#[derive(Debug, Error, Clone, PartialEq, Eq)]
pub enum TraceError {
    /// Attempted to attach an event to a closed (sealed) trace.
    #[error("trace {trace_id} is closed; cannot attach event {event_id}")]
    Closed {
        /// Trace identifier.
        trace_id: TraceId,
        /// Event identifier the caller tried to attach.
        event_id: EventId,
    },
    /// Attempted to close a trace that was already closed.
    #[error("trace {trace_id} is already closed")]
    AlreadyClosed {
        /// Trace identifier.
        trace_id: TraceId,
    },
}

/// In-memory assembler for one trace.
///
/// Construct via [`TraceAssembler::open`]. Attach events via
/// [`TraceAssembler::attach`]. Seal via [`TraceAssembler::close`].
#[derive(Debug, Clone)]
pub struct TraceAssembler {
    id: TraceId,
    schema_version: u16,
    opened_at: DateTime<Utc>,
    closed_at: Option<DateTime<Utc>>,
    trace_type: String,
    status: TraceStatus,
    /// Ordered event ids in the trace. Position is the ordinal.
    event_ids: Vec<EventId>,
}

impl TraceAssembler {
    /// Open a new trace with the given id, type tag, and open timestamp.
    #[must_use]
    pub fn open(id: TraceId, trace_type: impl Into<String>, opened_at: DateTime<Utc>) -> Self {
        Self {
            id,
            schema_version: SCHEMA_VERSION,
            opened_at,
            closed_at: None,
            trace_type: trace_type.into(),
            status: TraceStatus::Open,
            event_ids: Vec::new(),
        }
    }

    /// Trace identifier.
    #[must_use]
    pub fn id(&self) -> TraceId {
        self.id
    }

    /// Current status.
    #[must_use]
    pub fn status(&self) -> TraceStatus {
        self.status
    }

    /// Number of events attached so far.
    #[must_use]
    pub fn len(&self) -> usize {
        self.event_ids.len()
    }

    /// Whether this trace has zero events attached.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.event_ids.is_empty()
    }

    /// Attach an event to this trace and return its ordinal.
    ///
    /// Ordinals start at `0` and are strictly monotonic. Returns
    /// [`TraceError::Closed`] if the trace is sealed.
    pub fn attach(&mut self, event_id: EventId) -> Result<u64, TraceError> {
        if !matches!(self.status, TraceStatus::Open) {
            return Err(TraceError::Closed {
                trace_id: self.id,
                event_id,
            });
        }
        let ordinal = self.event_ids.len() as u64;
        self.event_ids.push(event_id);
        Ok(ordinal)
    }

    /// Close (seal) the trace.
    ///
    /// Returns the persistent [`Trace`] row that should be written to the
    /// store. Closing is idempotent on repeated calls *only* in the sense
    /// that a second call returns `AlreadyClosed`; callers must take the
    /// `Trace` from the first successful call.
    pub fn close(&mut self, closed_at: DateTime<Utc>) -> Result<Trace, TraceError> {
        if !matches!(self.status, TraceStatus::Open) {
            return Err(TraceError::AlreadyClosed { trace_id: self.id });
        }
        self.status = TraceStatus::Closed;
        self.closed_at = Some(closed_at);
        Ok(self.snapshot())
    }

    /// Snapshot the current state as a [`Trace`] row (without sealing).
    ///
    /// Useful for dashboards or mid-flight inspection. The returned `Trace`
    /// will reflect whatever `status` the assembler is currently in.
    #[must_use]
    pub fn snapshot(&self) -> Trace {
        Trace {
            id: self.id,
            schema_version: self.schema_version,
            opened_at: self.opened_at,
            closed_at: self.closed_at,
            event_ids: self.event_ids.clone(),
            trace_type: self.trace_type.clone(),
            status: self.status,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::TimeZone;

    fn now() -> DateTime<Utc> {
        Utc.with_ymd_and_hms(2026, 1, 1, 12, 0, 0).unwrap()
    }

    /// T-1.B.3 acceptance: ordinals are dense and monotonic.
    #[test]
    fn ordinals_are_dense_and_monotonic() {
        let tid = TraceId::new();
        let mut a = TraceAssembler::open(tid, "agent_run", now());

        let mut ordinals = Vec::new();
        for _ in 0..16 {
            let e = EventId::new();
            ordinals.push(a.attach(e).expect("open trace accepts attach"));
        }

        // Dense: ordinals are exactly 0..16.
        assert_eq!(
            ordinals,
            (0u64..16).collect::<Vec<_>>(),
            "ordinals must be dense"
        );
        // Monotonic strict: every next > prev.
        for w in ordinals.windows(2) {
            assert!(w[1] > w[0], "ordinals must be strictly monotonic");
        }
        assert_eq!(a.len(), 16);
        assert_eq!(a.status(), TraceStatus::Open);
    }

    #[test]
    fn close_seals_trace_and_blocks_further_attaches() {
        let tid = TraceId::new();
        let mut a = TraceAssembler::open(tid, "agent_run", now());
        let e1 = EventId::new();
        a.attach(e1).unwrap();

        let trace = a.close(now()).expect("close succeeds");
        assert_eq!(trace.status, TraceStatus::Closed);
        assert_eq!(trace.event_ids, vec![e1]);
        assert!(trace.closed_at.is_some());

        // Subsequent attach is rejected.
        let e2 = EventId::new();
        let err = a.attach(e2).unwrap_err();
        assert!(matches!(err, TraceError::Closed { .. }));

        // Subsequent close is rejected.
        let err = a.close(now()).unwrap_err();
        assert!(matches!(err, TraceError::AlreadyClosed { .. }));
    }

    #[test]
    fn empty_trace_closes_cleanly() {
        let tid = TraceId::new();
        let mut a = TraceAssembler::open(tid, "manual_session", now());
        assert!(a.is_empty());
        let trace = a.close(now()).unwrap();
        assert!(trace.event_ids.is_empty());
        assert_eq!(trace.status, TraceStatus::Closed);
    }

    #[test]
    fn snapshot_reflects_open_state_without_sealing() {
        let tid = TraceId::new();
        let mut a = TraceAssembler::open(tid, "agent_run", now());
        a.attach(EventId::new()).unwrap();
        let snap = a.snapshot();
        assert_eq!(snap.status, TraceStatus::Open);
        assert!(snap.closed_at.is_none());
        // Status on the assembler is unchanged.
        assert_eq!(a.status(), TraceStatus::Open);
    }
}