1use std::fs::{File, OpenOptions};
2use std::io::{BufWriter, Write};
3use std::path::Path;
4
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7use std::time::{SystemTime, UNIX_EPOCH};
8
9pub const TRACE_SCHEMA_VERSION: u32 = 1;
10const DEFAULT_MAX_STRING_CHARS: usize = 16 * 1024;
11
12#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
13#[serde(default)]
14pub struct TraceCorrelation {
15 pub message_id: Option<String>,
16 pub tool_call_id: Option<String>,
17 pub parent_event_id: Option<String>,
18 pub verification_gate_id: Option<String>,
19 pub evidence_id: Option<String>,
20}
21
22#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
23#[serde(default)]
24pub struct TraceRedaction {
25 pub contains_redactions: bool,
26 pub truncated_fields: Vec<String>,
27 pub content_hash: Option<String>,
28}
29
30#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
31#[serde(default)]
32pub struct TraceEvent {
33 pub schema_version: u32,
34 pub sequence: u64,
35 pub timestamp_ms: i128,
36 pub run_id: String,
37 pub workflow_id: Option<String>,
38 pub session_id: Option<String>,
39 pub turn: Option<u32>,
40 pub kind: String,
41 pub correlation: TraceCorrelation,
42 pub redaction: TraceRedaction,
43 pub payload: Value,
44}
45
46impl TraceEvent {
47 pub fn new(run_id: impl Into<String>, kind: impl Into<String>, payload: Value) -> Self {
48 Self {
49 schema_version: TRACE_SCHEMA_VERSION,
50 sequence: 0,
51 timestamp_ms: SystemTime::now()
52 .duration_since(UNIX_EPOCH)
53 .map(|duration| duration.as_millis() as i128)
54 .unwrap_or_default(),
55 run_id: run_id.into(),
56 workflow_id: None,
57 session_id: None,
58 turn: None,
59 kind: kind.into(),
60 correlation: TraceCorrelation::default(),
61 redaction: TraceRedaction::default(),
62 payload,
63 }
64 }
65
66 pub fn with_turn(mut self, turn: u32) -> Self {
67 self.turn = Some(turn);
68 self
69 }
70
71 pub fn with_tool_call_id(mut self, tool_call_id: impl Into<String>) -> Self {
72 self.correlation.tool_call_id = Some(tool_call_id.into());
73 self
74 }
75
76 pub fn with_workflow_id(mut self, workflow_id: impl Into<String>) -> Self {
77 self.workflow_id = Some(workflow_id.into());
78 self
79 }
80}
81
82impl Default for TraceEvent {
83 fn default() -> Self {
84 Self::new(String::new(), String::new(), Value::Null)
85 }
86}
87
88#[derive(Debug, Clone)]
89pub struct TraceWriterOptions {
90 pub max_string_chars: usize,
91}
92
93impl Default for TraceWriterOptions {
94 fn default() -> Self {
95 Self {
96 max_string_chars: DEFAULT_MAX_STRING_CHARS,
97 }
98 }
99}
100
101pub struct TraceWriter {
102 writer: BufWriter<File>,
103 next_sequence: u64,
104 options: TraceWriterOptions,
105}
106
107impl TraceWriter {
108 pub fn create(path: impl AsRef<Path>) -> std::io::Result<Self> {
109 Self::create_with_options(path, TraceWriterOptions::default())
110 }
111
112 pub fn create_with_options(
113 path: impl AsRef<Path>,
114 options: TraceWriterOptions,
115 ) -> std::io::Result<Self> {
116 if let Some(parent) = path.as_ref().parent() {
117 std::fs::create_dir_all(parent)?;
118 }
119 let file = OpenOptions::new()
120 .create(true)
121 .truncate(true)
122 .write(true)
123 .open(path)?;
124 Ok(Self {
125 writer: BufWriter::new(file),
126 next_sequence: 0,
127 options,
128 })
129 }
130
131 pub fn append(path: impl AsRef<Path>) -> std::io::Result<Self> {
132 if let Some(parent) = path.as_ref().parent() {
133 std::fs::create_dir_all(parent)?;
134 }
135 let file = OpenOptions::new().create(true).append(true).open(path)?;
136 Ok(Self {
137 writer: BufWriter::new(file),
138 next_sequence: 0,
139 options: TraceWriterOptions::default(),
140 })
141 }
142
143 pub fn write_event(&mut self, mut event: TraceEvent) -> std::io::Result<u64> {
144 event.sequence = self.next_sequence;
145 self.next_sequence += 1;
146 truncate_event_strings(&mut event, self.options.max_string_chars);
147 serde_json::to_writer(&mut self.writer, &event)?;
148 self.writer.write_all(b"\n")?;
149 Ok(event.sequence)
150 }
151
152 pub fn flush(&mut self) -> std::io::Result<()> {
153 self.writer.flush()
154 }
155}
156
157fn truncate_event_strings(event: &mut TraceEvent, max_chars: usize) {
158 truncate_value_strings(
159 &mut event.payload,
160 max_chars,
161 "payload",
162 &mut event.redaction,
163 );
164}
165
166fn truncate_value_strings(
167 value: &mut Value,
168 max_chars: usize,
169 path: &str,
170 redaction: &mut TraceRedaction,
171) {
172 match value {
173 Value::String(text) => {
174 if text.chars().count() > max_chars {
175 let truncated = text.chars().take(max_chars).collect::<String>();
176 *text = format!("{truncated}…[truncated]");
177 redaction.contains_redactions = true;
178 redaction.truncated_fields.push(path.to_string());
179 }
180 }
181 Value::Array(items) => {
182 for (index, item) in items.iter_mut().enumerate() {
183 truncate_value_strings(item, max_chars, &format!("{path}[{index}]"), redaction);
184 }
185 }
186 Value::Object(map) => {
187 for (key, value) in map.iter_mut() {
188 truncate_value_strings(value, max_chars, &format!("{path}.{key}"), redaction);
189 }
190 }
191 Value::Null | Value::Bool(_) | Value::Number(_) => {}
192 }
193}
194
195#[cfg(test)]
196mod tests {
197 use super::*;
198 use serde_json::json;
199
200 #[test]
201 fn trace_jsonl_writes_ordered_roundtrippable_events() {
202 let temp = tempfile::TempDir::new().unwrap();
203 let path = temp.path().join("trace.jsonl");
204 let mut writer = TraceWriter::create(&path).unwrap();
205 writer
206 .write_event(TraceEvent::new(
207 "run-1",
208 "agent.start",
209 json!({"model": "test"}),
210 ))
211 .unwrap();
212 writer
213 .write_event(TraceEvent::new("run-1", "turn.start", json!({"index": 1})))
214 .unwrap();
215 writer.flush().unwrap();
216
217 let contents = std::fs::read_to_string(path).unwrap();
218 let events = contents
219 .lines()
220 .map(|line| serde_json::from_str::<TraceEvent>(line).unwrap())
221 .collect::<Vec<_>>();
222
223 assert_eq!(events.len(), 2);
224 assert_eq!(events[0].sequence, 0);
225 assert_eq!(events[1].sequence, 1);
226 assert_eq!(events[0].schema_version, TRACE_SCHEMA_VERSION);
227 assert_eq!(events[0].kind, "agent.start");
228 }
229
230 #[test]
231 fn trace_jsonl_truncates_large_strings_and_marks_redaction() {
232 let temp = tempfile::TempDir::new().unwrap();
233 let path = temp.path().join("trace.jsonl");
234 let mut writer = TraceWriter::create_with_options(
235 &path,
236 TraceWriterOptions {
237 max_string_chars: 8,
238 },
239 )
240 .unwrap();
241 writer
242 .write_event(TraceEvent::new(
243 "run-1",
244 "tool.output.delta",
245 json!({"text": "abcdefghijklmnopqrstuvwxyz"}),
246 ))
247 .unwrap();
248 writer.flush().unwrap();
249
250 let contents = std::fs::read_to_string(path).unwrap();
251 let event: TraceEvent = serde_json::from_str(contents.lines().next().unwrap()).unwrap();
252 assert!(event.redaction.contains_redactions);
253 assert_eq!(event.redaction.truncated_fields, vec!["payload.text"]);
254 assert_eq!(event.payload["text"], "abcdefgh…[truncated]");
255 }
256}