1use std::time::{Instant, SystemTime, UNIX_EPOCH};
10
11use serde::{Deserialize, Serialize};
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct TraceSpan {
19 pub name: String,
21 pub start_time_ms: u64,
23 pub end_time_ms: u64,
25 pub duration_us: u64,
27 pub status: TraceStatus,
29 pub attributes: Vec<(String, String)>,
31 pub children: Vec<TraceSpan>,
33}
34
35#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
37pub enum TraceStatus {
38 Ok,
39 Error,
40}
41
42pub(crate) struct SpanBuilder {
44 name: String,
45 start_wall: u64,
46 start_instant: Instant,
47 attributes: Vec<(String, String)>,
48 children: Vec<TraceSpan>,
49}
50
51impl SpanBuilder {
52 pub fn new(name: &str) -> Self {
53 Self {
54 name: name.to_string(),
55 start_wall: now_ms(),
56 start_instant: Instant::now(),
57 attributes: Vec::new(),
58 children: Vec::new(),
59 }
60 }
61
62 pub fn attr(mut self, key: &str, value: impl ToString) -> Self {
63 self.attributes.push((key.to_string(), value.to_string()));
64 self
65 }
66
67 pub fn set_attr(&mut self, key: &str, value: impl ToString) {
68 self.attributes.push((key.to_string(), value.to_string()));
69 }
70
71 pub fn add_child(&mut self, child: TraceSpan) {
72 self.children.push(child);
73 }
74
75 pub fn finish(self, status: TraceStatus) -> TraceSpan {
76 let elapsed = self.start_instant.elapsed();
77 TraceSpan {
78 name: self.name,
79 start_time_ms: self.start_wall,
80 end_time_ms: self.start_wall + elapsed.as_millis() as u64,
81 duration_us: elapsed.as_micros() as u64,
82 status,
83 attributes: self.attributes,
84 children: self.children,
85 }
86 }
87
88 pub fn finish_ok(self) -> TraceSpan {
89 self.finish(TraceStatus::Ok)
90 }
91
92 pub fn finish_error(self, error: &str) -> TraceSpan {
93 self.attr("zapcode.error", error).finish(TraceStatus::Error)
94 }
95}
96
97fn now_ms() -> u64 {
98 SystemTime::now()
99 .duration_since(UNIX_EPOCH)
100 .unwrap_or_default()
101 .as_millis() as u64
102}
103
104#[derive(Debug, Clone, Serialize, Deserialize)]
106pub struct ExecutionTrace {
107 pub root: TraceSpan,
108}
109
110impl ExecutionTrace {
111 pub fn print(&self) {
113 print_span(&self.root, 0, true);
114 }
115
116 pub fn to_string_pretty(&self) -> String {
118 let mut buf = String::new();
119 format_span(&self.root, 0, true, &mut buf);
120 buf
121 }
122}
123
124fn format_duration(us: u64) -> String {
125 if us < 1000 {
126 format!("{}µs", us)
127 } else if us < 1_000_000 {
128 format!("{:.1}ms", us as f64 / 1000.0)
129 } else {
130 format!("{:.2}s", us as f64 / 1_000_000.0)
131 }
132}
133
134fn format_span(span: &TraceSpan, depth: usize, is_last: bool, buf: &mut String) {
135 let icon = match span.status {
136 TraceStatus::Ok => "✓",
137 TraceStatus::Error => "✗",
138 };
139 let duration = format_duration(span.duration_us);
140
141 let prefix = if depth == 0 {
143 String::new()
144 } else {
145 let connector = if is_last { "└─ " } else { "├─ " };
146 let indent = "│ ".repeat(depth - 1);
147 format!("{}{}", indent, connector)
148 };
149
150 buf.push_str(&format!("{}{} {} ({})", prefix, icon, span.name, duration));
151
152 for (k, v) in &span.attributes {
154 if k == "zapcode.error" {
155 buf.push_str(&format!(" error=\"{}\"", v));
156 } else if k == "zapcode.tool.name" {
157 buf.push_str(&format!(" {}", v));
158 } else if k == "zapcode.tool.args" {
159 buf.push_str(&format!("({})", v));
160 } else if k == "zapcode.tool.result" {
161 let display = if v.len() > 60 { &v[..57] } else { v };
162 buf.push_str(&format!(" → {}", display));
163 if v.len() > 60 {
164 buf.push_str("...");
165 }
166 }
167 }
168 buf.push('\n');
169
170 for (i, child) in span.children.iter().enumerate() {
171 format_span(child, depth + 1, i == span.children.len() - 1, buf);
172 }
173}
174
175fn print_span(span: &TraceSpan, depth: usize, is_last: bool) {
176 let mut buf = String::new();
177 format_span(span, depth, is_last, &mut buf);
178 print!("{}", buf);
179}