Skip to main content

kaizen/interchange/
atif.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2//! Pure ATIF mapping types. No storage or CLI wiring.
3
4mod error;
5mod types;
6
7use super::jsonl::JsonlEvent;
8pub use error::AtifImportError;
9use std::collections::BTreeMap;
10pub use types::{
11    AtifDocument, AtifEvent, AtifSession, InterchangeEvent, InterchangeSession, InterchangeTrace,
12};
13
14pub const ATIF_FORMAT: &str = "atif";
15pub const ATIF_VERSION: u16 = 1;
16
17pub fn export_atif(trace: &InterchangeTrace) -> AtifDocument {
18    AtifDocument {
19        format: ATIF_FORMAT.into(),
20        version: ATIF_VERSION,
21        session: atif_session(&trace.session),
22        events: trace.events.iter().map(atif_event).collect(),
23    }
24}
25
26pub fn import_atif(doc: &AtifDocument) -> Result<InterchangeTrace, AtifImportError> {
27    validate_document(doc)?;
28    Ok(InterchangeTrace {
29        session: interchange_session(&doc.session),
30        events: doc
31            .events
32            .iter()
33            .map(|event| interchange_event(&doc.session.id, event))
34            .collect(),
35    })
36}
37
38pub fn trace_from_jsonl(session: InterchangeSession, events: Vec<JsonlEvent>) -> InterchangeTrace {
39    InterchangeTrace {
40        session,
41        events: events.into_iter().map(InterchangeEvent::from).collect(),
42    }
43}
44
45impl From<JsonlEvent> for InterchangeEvent {
46    fn from(event: JsonlEvent) -> Self {
47        Self {
48            session_id: event.session_id,
49            seq: event.seq,
50            ts_ms: event.ts_ms,
51            kind: event.kind,
52            source: event.source,
53            tool: event.tool,
54            tool_call_id: event.tool_call_id,
55            payload: event.payload,
56            attributes: BTreeMap::new(),
57        }
58    }
59}
60
61fn validate_document(doc: &AtifDocument) -> Result<(), AtifImportError> {
62    validate_header(doc)?;
63    doc.events
64        .iter()
65        .find(|event| !event.id.starts_with(&format!("{}:", doc.session.id)))
66        .map_or(Ok(()), |event| Err(mismatch(&event.id, &doc.session.id)))
67}
68
69fn validate_header(doc: &AtifDocument) -> Result<(), AtifImportError> {
70    if doc.format != ATIF_FORMAT {
71        return Err(AtifImportError::UnsupportedFormat(doc.format.clone()));
72    }
73    (doc.version == ATIF_VERSION)
74        .then_some(())
75        .ok_or(AtifImportError::UnsupportedVersion(doc.version))
76}
77
78fn atif_session(session: &InterchangeSession) -> AtifSession {
79    AtifSession {
80        id: session.id.clone(),
81        agent: session.agent.clone(),
82        model: session.model.clone(),
83        workspace: session.workspace.clone(),
84        started_at_ms: session.started_at_ms,
85        ended_at_ms: session.ended_at_ms,
86        attributes: session.attributes.clone(),
87    }
88}
89
90fn interchange_session(session: &AtifSession) -> InterchangeSession {
91    InterchangeSession {
92        id: session.id.clone(),
93        agent: session.agent.clone(),
94        model: session.model.clone(),
95        workspace: session.workspace.clone(),
96        started_at_ms: session.started_at_ms,
97        ended_at_ms: session.ended_at_ms,
98        attributes: session.attributes.clone(),
99    }
100}
101
102fn atif_event(event: &InterchangeEvent) -> AtifEvent {
103    AtifEvent {
104        id: event_id(&event.session_id, event.seq),
105        sequence: event.seq,
106        timestamp_ms: event.ts_ms,
107        event_type: event.kind.clone(),
108        source: event.source.clone(),
109        tool: event.tool.clone(),
110        tool_call_id: event.tool_call_id.clone(),
111        payload: event.payload.clone(),
112        attributes: event.attributes.clone(),
113    }
114}
115
116fn interchange_event(session_id: &str, event: &AtifEvent) -> InterchangeEvent {
117    InterchangeEvent {
118        session_id: session_id.into(),
119        seq: event.sequence,
120        ts_ms: event.timestamp_ms,
121        kind: event.event_type.clone(),
122        source: event.source.clone(),
123        tool: event.tool.clone(),
124        tool_call_id: event.tool_call_id.clone(),
125        payload: event.payload.clone(),
126        attributes: event.attributes.clone(),
127    }
128}
129
130fn event_id(session_id: &str, seq: u64) -> String {
131    format!("{session_id}:{seq}")
132}
133
134fn mismatch(event_id: &str, session_id: &str) -> AtifImportError {
135    AtifImportError::SessionMismatch {
136        event_id: event_id.into(),
137        session_id: session_id.into(),
138    }
139}