use chrono::{DateTime, Utc};
use cortex_core::{EventId, Trace, TraceId, TraceStatus, SCHEMA_VERSION};
use thiserror::Error;
#[derive(Debug, Error, Clone, PartialEq, Eq)]
pub enum TraceError {
#[error("trace {trace_id} is closed; cannot attach event {event_id}")]
Closed {
trace_id: TraceId,
event_id: EventId,
},
#[error("trace {trace_id} is already closed")]
AlreadyClosed {
trace_id: TraceId,
},
}
#[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,
event_ids: Vec<EventId>,
}
impl TraceAssembler {
#[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(),
}
}
#[must_use]
pub fn id(&self) -> TraceId {
self.id
}
#[must_use]
pub fn status(&self) -> TraceStatus {
self.status
}
#[must_use]
pub fn len(&self) -> usize {
self.event_ids.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.event_ids.is_empty()
}
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)
}
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())
}
#[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()
}
#[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"));
}
assert_eq!(
ordinals,
(0u64..16).collect::<Vec<_>>(),
"ordinals must be dense"
);
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());
let e2 = EventId::new();
let err = a.attach(e2).unwrap_err();
assert!(matches!(err, TraceError::Closed { .. }));
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());
assert_eq!(a.status(), TraceStatus::Open);
}
}