kaizen/extensions/
atif.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2
3use crate::core::event::{Event, EventKind, EventSource, SessionRecord};
4use crate::store::Store;
5use anyhow::{Context, Result};
6use serde::{Deserialize, Serialize};
7use serde_json::{Value, json};
8use std::path::Path;
9
10#[derive(Clone, Debug, Serialize, Deserialize)]
11pub struct AtifDocument {
12    pub schema_version: String,
13    pub trajectory_id: String,
14    pub session: SessionRecord,
15    pub steps: Vec<AtifStep>,
16    pub metadata: Value,
17}
18
19#[derive(Clone, Debug, Serialize, Deserialize)]
20pub struct AtifStep {
21    pub seq: u64,
22    pub ts_ms: u64,
23    pub kind: String,
24    pub source: String,
25    pub tool: Option<String>,
26    pub tokens: AtifTokens,
27    pub cost_usd_e6: Option<i64>,
28    pub payload: Value,
29}
30
31#[derive(Clone, Debug, Default, Serialize, Deserialize)]
32pub struct AtifTokens {
33    pub input: Option<u32>,
34    pub output: Option<u32>,
35    pub reasoning: Option<u32>,
36    pub cache_read: Option<u32>,
37    pub cache_create: Option<u32>,
38}
39
40pub fn export_session(store: &Store, session_id: &str) -> Result<AtifDocument> {
41    let session = store
42        .get_session(session_id)?
43        .with_context(|| format!("session not found: {session_id}"))?;
44    let steps = store
45        .list_events_for_session(session_id)?
46        .into_iter()
47        .map(step_from_event)
48        .collect();
49    Ok(AtifDocument {
50        schema_version: "kaizen.atif.v1".into(),
51        trajectory_id: session_id.to_string(),
52        session,
53        steps,
54        metadata: json!({"producer":"kaizen"}),
55    })
56}
57
58pub fn import_file(store: &Store, path: &Path, workspace: &str) -> Result<AtifDocument> {
59    let text = std::fs::read_to_string(path)?;
60    let mut doc: AtifDocument = serde_json::from_str(&text)?;
61    import_document(store, &mut doc, workspace)?;
62    Ok(doc)
63}
64
65pub fn import_document(store: &Store, doc: &mut AtifDocument, workspace: &str) -> Result<()> {
66    doc.session.workspace = workspace.to_string();
67    store.upsert_session(&doc.session)?;
68    doc.steps
69        .iter()
70        .map(|step| event_from_step(&doc.session.id, step))
71        .try_for_each(|event| store.append_event(&event?))
72}
73
74fn step_from_event(event: Event) -> AtifStep {
75    let tokens = tokens(&event);
76    AtifStep {
77        seq: event.seq,
78        ts_ms: event.ts_ms,
79        kind: format!("{:?}", event.kind),
80        source: format!("{:?}", event.source),
81        tool: event.tool,
82        tokens,
83        cost_usd_e6: event.cost_usd_e6,
84        payload: event.payload,
85    }
86}
87
88fn tokens(event: &Event) -> AtifTokens {
89    AtifTokens {
90        input: event.tokens_in,
91        output: event.tokens_out,
92        reasoning: event.reasoning_tokens,
93        cache_read: event.cache_read_tokens,
94        cache_create: event.cache_creation_tokens,
95    }
96}
97
98fn event_from_step(session_id: &str, step: &AtifStep) -> Result<Event> {
99    Ok(Event {
100        session_id: session_id.to_string(),
101        seq: step.seq,
102        ts_ms: step.ts_ms,
103        ts_exact: true,
104        kind: kind(&step.kind),
105        source: source(&step.source),
106        tool: step.tool.clone(),
107        tool_call_id: None,
108        tokens_in: step.tokens.input,
109        tokens_out: step.tokens.output,
110        reasoning_tokens: step.tokens.reasoning,
111        cost_usd_e6: step.cost_usd_e6,
112        stop_reason: None,
113        latency_ms: None,
114        ttft_ms: None,
115        retry_count: None,
116        context_used_tokens: None,
117        context_max_tokens: None,
118        cache_creation_tokens: step.tokens.cache_create,
119        cache_read_tokens: step.tokens.cache_read,
120        system_prompt_tokens: None,
121        payload: step.payload.clone(),
122    })
123}
124
125fn kind(value: &str) -> EventKind {
126    match value {
127        "ToolCall" => EventKind::ToolCall,
128        "ToolResult" => EventKind::ToolResult,
129        "Error" => EventKind::Error,
130        "Cost" => EventKind::Cost,
131        "Hook" => EventKind::Hook,
132        "Lifecycle" => EventKind::Lifecycle,
133        _ => EventKind::Message,
134    }
135}
136
137fn source(value: &str) -> EventSource {
138    match value {
139        "Hook" => EventSource::Hook,
140        "Proxy" => EventSource::Proxy,
141        _ => EventSource::Tail,
142    }
143}
kaizen/extensions/atif.rs

kaizen/extensions/
atif.rs
