1use std::fs::{File, OpenOptions};
2use std::io::Write;
3use std::path::Path;
4use std::sync::Mutex;
5use std::time::{SystemTime, UNIX_EPOCH};
6
7use serde::{Deserialize, Serialize};
8use serde_json::json;
9
10pub trait TraceSink: Send + Sync {
11 fn on_span_start(&self, span: &Span);
12 fn on_span_end(&self, span: &Span);
13
14 fn flush(&self) -> Result<(), String> {
15 Ok(())
16 }
17}
18
19#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
20pub struct Span {
21 pub trace_id: String,
22 pub span_id: String,
23 pub name: String,
24 pub run_id: String,
25 pub agent_name: Option<String>,
26}
27
28impl Span {
29 pub fn new(
30 run_id: impl Into<String>,
31 name: impl Into<String>,
32 agent_name: Option<String>,
33 ) -> Self {
34 let run_id = run_id.into();
35 let name = name.into();
36 Self {
37 trace_id: run_id.clone(),
38 span_id: format!("{}_{}", name, timestamp_millis()),
39 name,
40 run_id,
41 agent_name,
42 }
43 }
44}
45
46pub struct JsonlTraceExporter {
47 file: Mutex<File>,
48}
49
50impl JsonlTraceExporter {
51 pub fn new(path: impl AsRef<Path>) -> Result<Self, String> {
52 let file = OpenOptions::new()
53 .create(true)
54 .append(true)
55 .open(path)
56 .map_err(|error| error.to_string())?;
57 Ok(Self {
58 file: Mutex::new(file),
59 })
60 }
61
62 fn write_event(&self, event: &str, span: &Span) {
63 if let Ok(mut file) = self.file.lock() {
64 let _ = writeln!(
65 file,
66 "{}",
67 json!({
68 "event": event,
69 "timestamp_ms": timestamp_millis(),
70 "span": span,
71 })
72 );
73 }
74 }
75}
76
77impl TraceSink for JsonlTraceExporter {
78 fn on_span_start(&self, span: &Span) {
79 self.write_event("span_start", span);
80 }
81
82 fn on_span_end(&self, span: &Span) {
83 self.write_event("span_end", span);
84 }
85
86 fn flush(&self) -> Result<(), String> {
87 self.file
88 .lock()
89 .map_err(|_| "trace exporter lock poisoned".to_string())?
90 .flush()
91 .map_err(|error| error.to_string())
92 }
93}
94
95fn timestamp_millis() -> u128 {
96 SystemTime::now()
97 .duration_since(UNIX_EPOCH)
98 .map(|duration| duration.as_millis())
99 .unwrap_or_default()
100}