1use std::collections::{HashMap, VecDeque};
19use std::path::{Path, PathBuf};
20use std::sync::atomic::{AtomicU64, Ordering};
21use std::time::{SystemTime, UNIX_EPOCH};
22
23use rand::RngExt as _;
24use serde::{Deserialize, Serialize};
25
26use crate::redact::scrub_content;
27
28static SPAN_COUNTER: AtomicU64 = AtomicU64::new(0);
31
32#[must_use]
33fn new_trace_id() -> [u8; 16] {
34 rand::rng().random()
35}
36
37#[must_use]
38fn new_span_id() -> [u8; 8] {
39 let mut id: [u8; 8] = rand::rng().random();
40 id[0] ^= (SPAN_COUNTER.fetch_add(1, Ordering::Relaxed) & 0xFF) as u8;
42 id
43}
44
45#[must_use]
46fn hex16(b: &[u8; 16]) -> String {
47 use std::fmt::Write as _;
48 let mut s = String::with_capacity(32);
49 for x in b {
50 let _ = write!(s, "{x:02x}");
51 }
52 s
53}
54
55#[must_use]
56fn hex8(b: [u8; 8]) -> String {
57 use std::fmt::Write as _;
58 let mut s = String::with_capacity(16);
59 for x in b {
60 let _ = write!(s, "{x:02x}");
61 }
62 s
63}
64
65#[must_use]
66fn now_unix_nanos() -> u64 {
67 SystemTime::now()
68 .duration_since(UNIX_EPOCH)
69 .map_or(0, |d| u64::try_from(d.as_nanos()).unwrap_or(u64::MAX))
70}
71
72#[non_exhaustive]
75#[derive(Debug, Clone, Serialize, Deserialize)]
77pub enum SpanStatus {
78 Ok,
79 Error { message: String },
80 Unset,
81}
82
83#[derive(Debug, Clone)]
85pub struct SpanData {
86 pub trace_id: [u8; 16],
87 pub span_id: [u8; 8],
88 pub parent_span_id: Option<[u8; 8]>,
89 pub name: String,
90 pub start_time_unix_nanos: u64,
91 pub end_time_unix_nanos: u64,
92 pub attributes: Vec<(String, String)>,
93 pub status: SpanStatus,
94}
95
96pub struct SpanGuard {
100 pub span_id: [u8; 8],
101 pub parent_span_id: [u8; 8],
102 pub name: String,
103 pub start_time_unix_nanos: u64,
104}
105
106pub struct LlmAttributes {
108 pub model: String,
109 pub prompt_tokens: u64,
110 pub completion_tokens: u64,
111 pub latency_ms: u64,
112 pub streaming: bool,
113 pub cache_hit: bool,
114}
115
116pub struct ToolAttributes {
118 pub latency_ms: u64,
119 pub is_error: bool,
120 pub error_kind: Option<String>,
121}
122
123pub struct MemorySearchAttributes {
125 pub query_preview: String,
126 pub result_count: usize,
127 pub latency_ms: u64,
128}
129
130#[derive(Debug)]
137pub struct TraceEvent {
138 pub trace_id: [u8; 16],
139 pub spans: Vec<SpanData>,
140}
141
142struct IterationEntry {
146 guard: SpanGuard,
147 user_msg_preview: String,
148}
149
150const DEFAULT_MAX_SPANS: usize = 10_000;
158
159pub struct TracingCollector {
160 trace_id: [u8; 16],
161 session_span_id: [u8; 8],
162 session_start: u64,
163 service_name: String,
164 trace_metadata: HashMap<String, String>,
167 output_dir: PathBuf,
168 active_iterations: HashMap<usize, IterationEntry>,
170 completed_spans: VecDeque<SpanData>,
171 max_spans: usize,
173 redact: bool,
175 flushed: bool,
177 trace_tx: Option<tokio::sync::mpsc::UnboundedSender<TraceEvent>>,
180}
181
182impl TracingCollector {
183 pub fn new(
192 output_dir: &Path,
193 service_name: impl Into<String>,
194 trace_metadata: HashMap<String, String>,
195 redact: bool,
196 trace_tx: Option<tokio::sync::mpsc::UnboundedSender<TraceEvent>>,
197 ) -> std::io::Result<Self> {
198 std::fs::create_dir_all(output_dir)?;
199 Ok(Self {
200 trace_id: new_trace_id(),
201 session_span_id: new_span_id(),
202 session_start: now_unix_nanos(),
203 service_name: service_name.into(),
204 trace_metadata,
205 output_dir: output_dir.to_owned(),
206 active_iterations: HashMap::new(),
207 completed_spans: VecDeque::new(),
208 max_spans: DEFAULT_MAX_SPANS,
209 redact,
210 flushed: false,
211 trace_tx,
212 })
213 }
214
215 fn maybe_redact<'a>(&self, text: &'a str) -> std::borrow::Cow<'a, str> {
216 if self.redact {
217 scrub_content(text)
218 } else {
219 std::borrow::Cow::Borrowed(text)
220 }
221 }
222
223 fn push_span(&mut self, span: SpanData) {
225 if self.completed_spans.len() >= self.max_spans {
226 tracing::warn!(
227 max_spans = self.max_spans,
228 "trace span cap reached, dropping oldest span"
229 );
230 self.completed_spans.pop_front();
231 }
232 self.completed_spans.push_back(span);
233 }
234
235 pub fn begin_iteration(&mut self, index: usize, user_msg_preview: &str) {
239 let preview = self
240 .maybe_redact(user_msg_preview)
241 .chars()
242 .take(100)
243 .collect::<String>();
244 let entry = IterationEntry {
245 guard: SpanGuard {
246 span_id: new_span_id(),
247 parent_span_id: self.session_span_id,
248 name: format!("iteration.{index}"),
249 start_time_unix_nanos: now_unix_nanos(),
250 },
251 user_msg_preview: preview,
252 };
253 self.active_iterations.insert(index, entry);
254 }
255
256 pub fn end_iteration(&mut self, index: usize, status: SpanStatus) {
258 let end_time = now_unix_nanos();
259 if let Some(entry) = self.active_iterations.remove(&index) {
260 let span = SpanData {
261 trace_id: self.trace_id,
262 span_id: entry.guard.span_id,
263 parent_span_id: Some(entry.guard.parent_span_id),
264 name: entry.guard.name,
265 start_time_unix_nanos: entry.guard.start_time_unix_nanos,
266 end_time_unix_nanos: end_time,
267 attributes: vec![(
268 "zeph.iteration.user_message_preview".to_owned(),
269 entry.user_msg_preview,
270 )],
271 status,
272 };
273 self.push_span(span);
274 } else {
275 tracing::warn!(index, "end_iteration without matching begin_iteration");
276 }
277 }
278
279 #[must_use]
283 pub fn begin_llm_request(&self, iteration_span_id: [u8; 8]) -> SpanGuard {
284 SpanGuard {
285 span_id: new_span_id(),
286 parent_span_id: iteration_span_id,
287 name: "llm.request".to_owned(),
288 start_time_unix_nanos: now_unix_nanos(),
289 }
290 }
291
292 pub fn end_llm_request(&mut self, guard: SpanGuard, attrs: &LlmAttributes) {
294 let end_time = now_unix_nanos();
295 let model_clean = self.maybe_redact(&attrs.model).into_owned();
296 self.push_span(SpanData {
297 trace_id: self.trace_id,
298 span_id: guard.span_id,
299 parent_span_id: Some(guard.parent_span_id),
300 name: guard.name,
301 start_time_unix_nanos: guard.start_time_unix_nanos,
302 end_time_unix_nanos: end_time,
303 attributes: vec![
304 ("zeph.llm.model".to_owned(), model_clean),
305 (
306 "zeph.llm.prompt_tokens".to_owned(),
307 attrs.prompt_tokens.to_string(),
308 ),
309 (
310 "zeph.llm.completion_tokens".to_owned(),
311 attrs.completion_tokens.to_string(),
312 ),
313 (
314 "zeph.llm.latency_ms".to_owned(),
315 attrs.latency_ms.to_string(),
316 ),
317 ("zeph.llm.streaming".to_owned(), attrs.streaming.to_string()),
318 ("zeph.llm.cache_hit".to_owned(), attrs.cache_hit.to_string()),
319 ],
320 status: SpanStatus::Ok,
321 });
322 }
323
324 #[must_use]
328 pub fn begin_tool_call(&self, tool_name: &str, iteration_span_id: [u8; 8]) -> SpanGuard {
329 self.begin_tool_call_at(tool_name, iteration_span_id, &std::time::Instant::now())
330 }
331
332 #[must_use]
338 pub fn begin_tool_call_at(
339 &self,
340 tool_name: &str,
341 iteration_span_id: [u8; 8],
342 started_at: &std::time::Instant,
343 ) -> SpanGuard {
344 let elapsed_nanos = u64::try_from(started_at.elapsed().as_nanos()).unwrap_or(u64::MAX);
345 let start_time_unix_nanos = now_unix_nanos().saturating_sub(elapsed_nanos);
346 SpanGuard {
347 span_id: new_span_id(),
348 parent_span_id: iteration_span_id,
349 name: format!("tool.{}", sanitize_name(tool_name)),
350 start_time_unix_nanos,
351 }
352 }
353
354 pub fn end_tool_call(&mut self, guard: SpanGuard, tool_name: &str, attrs: ToolAttributes) {
356 let end_time = now_unix_nanos();
357 let tool_clean = sanitize_name(tool_name);
358 let mut attributes = vec![
359 ("zeph.tool.name".to_owned(), tool_clean),
360 (
361 "zeph.tool.latency_ms".to_owned(),
362 attrs.latency_ms.to_string(),
363 ),
364 ("zeph.tool.is_error".to_owned(), attrs.is_error.to_string()),
365 ];
366 if let Some(kind) = attrs.error_kind {
367 let kind_clean = self.maybe_redact(&kind).into_owned();
369 attributes.push(("zeph.tool.error_kind".to_owned(), kind_clean));
370 }
371 let status = if attrs.is_error {
372 SpanStatus::Error {
373 message: "tool call failed".to_owned(),
374 }
375 } else {
376 SpanStatus::Ok
377 };
378 self.push_span(SpanData {
379 trace_id: self.trace_id,
380 span_id: guard.span_id,
381 parent_span_id: Some(guard.parent_span_id),
382 name: guard.name,
383 start_time_unix_nanos: guard.start_time_unix_nanos,
384 end_time_unix_nanos: end_time,
385 attributes,
386 status,
387 });
388 }
389
390 #[must_use]
394 pub fn begin_memory_search(&self, parent_span_id: [u8; 8]) -> SpanGuard {
395 SpanGuard {
396 span_id: new_span_id(),
397 parent_span_id,
398 name: "memory.search".to_owned(),
399 start_time_unix_nanos: now_unix_nanos(),
400 }
401 }
402
403 pub fn end_memory_search(&mut self, guard: SpanGuard, attrs: &MemorySearchAttributes) {
405 let end_time = now_unix_nanos();
406 let query_clean = self
407 .maybe_redact(&attrs.query_preview)
408 .chars()
409 .take(100)
410 .collect::<String>();
411 self.push_span(SpanData {
412 trace_id: self.trace_id,
413 span_id: guard.span_id,
414 parent_span_id: Some(guard.parent_span_id),
415 name: guard.name,
416 start_time_unix_nanos: guard.start_time_unix_nanos,
417 end_time_unix_nanos: end_time,
418 attributes: vec![
419 ("zeph.memory.query_preview".to_owned(), query_clean),
420 (
421 "zeph.memory.result_count".to_owned(),
422 attrs.result_count.to_string(),
423 ),
424 (
425 "zeph.memory.latency_ms".to_owned(),
426 attrs.latency_ms.to_string(),
427 ),
428 ],
429 status: SpanStatus::Ok,
430 });
431 }
432
433 #[must_use]
437 pub fn trace_json_path(&self) -> PathBuf {
438 self.output_dir.join("trace.json")
439 }
440
441 #[must_use]
443 pub fn current_iteration_span_id(&self, index: usize) -> Option<[u8; 8]> {
444 self.active_iterations.get(&index).map(|e| e.guard.span_id)
445 }
446
447 #[must_use]
449 pub fn session_span_id(&self) -> [u8; 8] {
450 self.session_span_id
451 }
452
453 #[must_use]
455 pub fn trace_id(&self) -> [u8; 16] {
456 self.trace_id
457 }
458
459 pub fn finish(&mut self) -> Option<tokio::task::JoinHandle<()>> {
474 if self.flushed {
475 return None;
476 }
477 self.flushed = true;
478
479 let open_keys: Vec<usize> = self.active_iterations.keys().copied().collect();
481 let end_time = now_unix_nanos();
482 for index in open_keys {
483 if let Some(entry) = self.active_iterations.remove(&index) {
484 self.push_span(SpanData {
485 trace_id: self.trace_id,
486 span_id: entry.guard.span_id,
487 parent_span_id: Some(entry.guard.parent_span_id),
488 name: entry.guard.name,
489 start_time_unix_nanos: entry.guard.start_time_unix_nanos,
490 end_time_unix_nanos: end_time,
491 attributes: vec![(
492 "zeph.iteration.user_message_preview".to_owned(),
493 entry.user_msg_preview,
494 )],
495 status: SpanStatus::Unset,
496 });
497 }
498 }
499
500 let session_span = SpanData {
501 trace_id: self.trace_id,
502 span_id: self.session_span_id,
503 parent_span_id: None,
504 name: "session".to_owned(),
505 start_time_unix_nanos: self.session_start,
506 end_time_unix_nanos: end_time,
507 attributes: vec![
508 ("service.name".to_owned(), self.service_name.clone()),
509 ("zeph.session.trace_id".to_owned(), hex16(&self.trace_id)),
510 ],
511 status: SpanStatus::Ok,
512 };
513
514 let mut all_spans = vec![session_span];
515 all_spans.extend(self.completed_spans.drain(..));
516
517 let json = serialize_otlp_json(&all_spans, &self.service_name, &self.trace_metadata);
518 let path = self.output_dir.join("trace.json");
519 let handle = write_trace_file(&path, json.as_bytes());
520
521 if let Some(ref tx) = self.trace_tx {
523 let event = TraceEvent {
524 trace_id: self.trace_id,
525 spans: all_spans,
526 };
527 if tx.send(event).is_err() {
528 tracing::debug!("OTLP trace channel closed, skipping export");
529 }
530 }
531
532 handle
533 }
534}
535
536impl Drop for TracingCollector {
541 fn drop(&mut self) {
542 let _ = self.finish();
543 }
544}
545
546fn span_status_code(status: &SpanStatus) -> u8 {
549 match status {
550 SpanStatus::Unset => 0,
551 SpanStatus::Ok => 1,
552 SpanStatus::Error { .. } => 2,
553 }
554}
555
556#[must_use]
563pub fn serialize_otlp_json<S: std::hash::BuildHasher>(
564 spans: &[SpanData],
565 service_name: &str,
566 trace_metadata: &HashMap<String, String, S>,
567) -> String {
568 let otlp_spans: Vec<serde_json::Value> = spans
569 .iter()
570 .map(|s| {
571 let attrs: Vec<serde_json::Value> = s
572 .attributes
573 .iter()
574 .map(|(k, v)| {
575 serde_json::json!({
576 "key": k,
577 "value": { "stringValue": v }
578 })
579 })
580 .collect();
581
582 let mut obj = serde_json::json!({
583 "traceId": hex16(&s.trace_id),
584 "spanId": hex8(s.span_id),
585 "name": s.name,
586 "startTimeUnixNano": s.start_time_unix_nanos.to_string(),
588 "endTimeUnixNano": s.end_time_unix_nanos.to_string(),
589 "attributes": attrs,
590 "status": {
591 "code": span_status_code(&s.status)
592 }
593 });
594
595 if let Some(parent) = s.parent_span_id {
596 obj["parentSpanId"] = serde_json::json!(hex8(parent));
597 }
598
599 if let SpanStatus::Error { message } = &s.status {
600 obj["status"]["message"] = serde_json::json!(message);
601 }
602
603 obj
604 })
605 .collect();
606
607 let mut resource_attrs = vec![serde_json::json!({
608 "key": "service.name",
609 "value": { "stringValue": service_name }
610 })];
611 for (k, v) in trace_metadata {
612 if k == "service.name" {
613 continue;
614 }
615 resource_attrs.push(serde_json::json!({
616 "key": k,
617 "value": { "stringValue": v }
618 }));
619 }
620
621 let payload = serde_json::json!({
622 "resourceSpans": [{
623 "resource": {
624 "attributes": resource_attrs
625 },
626 "scopeSpans": [{
627 "scope": {
628 "name": "zeph",
629 "version": env!("CARGO_PKG_VERSION")
630 },
631 "spans": otlp_spans
632 }]
633 }]
634 });
635
636 serde_json::to_string_pretty(&payload)
637 .unwrap_or_else(|e| format!("{{\"error\": \"serialization failed: {e}\"}}"))
638}
639
640fn write_trace_file(path: &Path, data: &[u8]) -> Option<tokio::task::JoinHandle<()>> {
656 let path = path.to_owned();
657 let data = data.to_vec();
658 let write_and_log = move || match zeph_common::fs_secure::write_private(&path, &data) {
659 Ok(()) => tracing::info!(path = %path.display(), "OTel trace written"),
660 Err(e) => tracing::warn!(path = %path.display(), error = %e, "trace.json write failed"),
661 };
662 if tokio::runtime::Handle::try_current().is_ok() {
663 Some(tokio::task::spawn_blocking(write_and_log))
664 } else {
665 write_and_log();
666 None
667 }
668}
669
670fn sanitize_name(name: &str) -> String {
671 name.chars()
672 .map(|c| {
673 if c.is_alphanumeric() || c == '-' || c == '_' || c == '.' {
674 c
675 } else {
676 '_'
677 }
678 })
679 .collect()
680}
681
682#[cfg(test)]
685mod tests {
686 use super::*;
687 use tempfile::tempdir;
688
689 fn make_collector(dir: &Path) -> TracingCollector {
690 TracingCollector::new(dir, "zeph-test", HashMap::new(), false, None).unwrap()
691 }
692
693 #[test]
694 fn span_id_generation_is_unique() {
695 let ids: Vec<[u8; 8]> = (0..100).map(|_| new_span_id()).collect();
696 let unique: std::collections::HashSet<[u8; 8]> = ids.into_iter().collect();
697 assert_eq!(unique.len(), 100);
698 }
699
700 #[test]
701 fn hex_lengths_correct() {
702 assert_eq!(hex16(&new_trace_id()).len(), 32);
703 assert_eq!(hex8(new_span_id()).len(), 16);
704 }
705
706 #[test]
707 fn collector_creates_output_dir() {
708 let tmp = tempdir().unwrap();
709 let sub = tmp.path().join("traces");
710 make_collector(&sub);
711 assert!(sub.exists());
712 }
713
714 #[test]
715 fn finish_writes_trace_json() {
716 let tmp = tempdir().unwrap();
717 let mut c = make_collector(tmp.path());
718 c.begin_iteration(0, "hello world");
719 c.end_iteration(0, SpanStatus::Ok);
720 c.finish();
721
722 let path = tmp.path().join("trace.json");
723 assert!(path.exists(), "trace.json must be written");
724
725 let v: serde_json::Value =
726 serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
727 assert!(v["resourceSpans"].is_array());
728 let spans = v["resourceSpans"][0]["scopeSpans"][0]["spans"]
730 .as_array()
731 .unwrap();
732 assert_eq!(spans.len(), 2);
733 }
734
735 #[test]
736 fn span_hierarchy_parent_child_correct() {
737 let tmp = tempdir().unwrap();
738 let mut c = make_collector(tmp.path());
739 c.begin_iteration(0, "test");
740 let iter_id = c.current_iteration_span_id(0).unwrap();
741 let guard = c.begin_llm_request(iter_id);
742 c.end_llm_request(
743 guard,
744 &LlmAttributes {
745 model: "test-model".to_owned(),
746 prompt_tokens: 100,
747 completion_tokens: 50,
748 latency_ms: 200,
749 streaming: false,
750 cache_hit: false,
751 },
752 );
753 c.end_iteration(0, SpanStatus::Ok);
754 c.finish();
755
756 let content = std::fs::read_to_string(tmp.path().join("trace.json")).unwrap();
757 let v: serde_json::Value = serde_json::from_str(&content).unwrap();
758 let spans = v["resourceSpans"][0]["scopeSpans"][0]["spans"]
759 .as_array()
760 .unwrap();
761 assert_eq!(spans.len(), 3, "session + iteration + llm = 3 spans");
762
763 let llm_span = spans
764 .iter()
765 .find(|s| s["name"] == "llm.request")
766 .expect("llm.request span missing");
767 let iter_span = spans
768 .iter()
769 .find(|s| s["name"] == "iteration.0")
770 .expect("iteration.0 span missing");
771
772 assert_eq!(
773 llm_span["parentSpanId"], iter_span["spanId"],
774 "llm span parent must be iteration span"
775 );
776 }
777
778 #[test]
779 fn redaction_applied_to_text_attributes() {
780 let tmp = tempdir().unwrap();
781 let mut c = TracingCollector::new(tmp.path(), "test", HashMap::new(), true, None).unwrap();
782 let iter_id = c.session_span_id();
783 let guard = c.begin_memory_search(iter_id);
784 c.end_memory_search(
785 guard,
786 &MemorySearchAttributes {
787 query_preview: "search sk-secretkey123 here".to_owned(),
788 result_count: 3,
789 latency_ms: 10,
790 },
791 );
792 c.finish();
793
794 let content = std::fs::read_to_string(tmp.path().join("trace.json")).unwrap();
795 assert!(
796 !content.contains("sk-secretkey123"),
797 "raw secret must be redacted from trace"
798 );
799 }
800
801 #[test]
802 fn otlp_json_format_spec_compliant() {
803 let trace_id = [0xAB_u8; 16];
804 let span_id = [0xCD_u8; 8];
805 let spans = vec![SpanData {
806 trace_id,
807 span_id,
808 parent_span_id: None,
809 name: "session".to_owned(),
810 start_time_unix_nanos: 1_000_000,
811 end_time_unix_nanos: 2_000_000,
812 attributes: vec![("service.name".to_owned(), "zeph".to_owned())],
813 status: SpanStatus::Ok,
814 }];
815
816 let json = serialize_otlp_json(&spans, "zeph", &HashMap::new());
817 let v: serde_json::Value = serde_json::from_str(&json).unwrap();
818
819 let span = &v["resourceSpans"][0]["scopeSpans"][0]["spans"][0];
820 assert_eq!(
821 span["traceId"],
822 "abababababababababababababababababab"[..32]
823 );
824 assert_eq!(span["spanId"], "cdcdcdcdcdcdcdcd");
825 assert_eq!(span["name"], "session");
826 assert!(span["startTimeUnixNano"].is_string());
828 assert_eq!(span["status"]["code"], 1_u64);
829 }
830
831 #[test]
832 fn drop_flushes_trace() {
833 let tmp = tempdir().unwrap();
834 {
835 let mut c = make_collector(tmp.path());
836 c.begin_iteration(0, "hello");
837 }
839 assert!(
840 tmp.path().join("trace.json").exists(),
841 "Drop must flush trace.json"
842 );
843 }
844
845 #[test]
846 fn trace_metadata_included_in_resource_attributes() {
847 let tmp = tempdir().unwrap();
848 let mut meta = HashMap::new();
849 meta.insert("deployment.environment".to_owned(), "staging".to_owned());
850 meta.insert("service.name".to_owned(), "should-be-skipped".to_owned());
851 let mut c = TracingCollector::new(tmp.path(), "zeph-test", meta, false, None).unwrap();
852 c.finish();
853
854 let content = std::fs::read_to_string(tmp.path().join("trace.json")).unwrap();
855 let v: serde_json::Value = serde_json::from_str(&content).unwrap();
856 let attrs = v["resourceSpans"][0]["resource"]["attributes"]
857 .as_array()
858 .unwrap();
859
860 let svc = attrs
862 .iter()
863 .find(|a| a["key"] == "service.name")
864 .expect("service.name must be in resource attributes");
865 assert_eq!(svc["value"]["stringValue"], "zeph-test");
866
867 let env_attr = attrs.iter().find(|a| a["key"] == "deployment.environment");
869 assert!(
870 env_attr.is_some(),
871 "deployment.environment must be in resource attributes"
872 );
873 assert_eq!(env_attr.unwrap()["value"]["stringValue"], "staging");
874 }
875
876 #[test]
877 fn finish_is_idempotent() {
878 let tmp = tempdir().unwrap();
879 let mut c = make_collector(tmp.path());
880 c.finish();
881 c.finish();
882 assert!(tmp.path().join("trace.json").exists());
883 }
884
885 #[tokio::test]
891 async fn finish_dispatches_via_spawn_blocking_under_active_runtime() {
892 let tmp = tempdir().unwrap();
893 let mut c = make_collector(tmp.path());
894 c.begin_iteration(0, "hello");
895 c.end_iteration(0, SpanStatus::Ok);
896
897 let handle = c.finish();
898 assert!(
899 handle.is_some(),
900 "finish() must dispatch the write via spawn_blocking (Some) when a Tokio runtime \
901 is active, not silently take the synchronous fallback path (None)"
902 );
903 handle
904 .unwrap()
905 .await
906 .expect("spawn_blocking write task must not panic");
907
908 let path = tmp.path().join("trace.json");
909 assert!(
910 path.exists(),
911 "trace.json must exist once the spawn_blocking handle has been awaited"
912 );
913 }
914
915 #[test]
916 fn concurrent_iterations_tracked_independently() {
917 let tmp = tempdir().unwrap();
918 let mut c = make_collector(tmp.path());
919 c.begin_iteration(0, "first");
920 c.begin_iteration(1, "second");
921 let id0 = c.current_iteration_span_id(0).unwrap();
922 let id1 = c.current_iteration_span_id(1).unwrap();
923 assert_ne!(
924 id0, id1,
925 "concurrent iterations must have distinct span IDs"
926 );
927 c.end_iteration(0, SpanStatus::Ok);
928 c.end_iteration(1, SpanStatus::Ok);
929 c.finish();
930
931 let v: serde_json::Value =
932 serde_json::from_str(&std::fs::read_to_string(tmp.path().join("trace.json")).unwrap())
933 .unwrap();
934 let spans = v["resourceSpans"][0]["scopeSpans"][0]["spans"]
935 .as_array()
936 .unwrap();
937 assert_eq!(spans.len(), 3);
939 }
940
941 #[test]
942 fn trace_format_skips_legacy_numbered_files() {
943 use crate::debug_dump::{DebugDumper, DumpFormat, RequestDebugDump};
944
945 let tmp = tempdir().unwrap();
946 let d = DebugDumper::new(tmp.path(), DumpFormat::Trace).unwrap();
947 let session_dir = d.dir().to_owned();
948 let id = d.dump_request(&RequestDebugDump {
949 model_name: "test",
950 messages: &[],
951 tools: &[],
952 provider_request: serde_json::json!({}),
953 memcot_state: None,
954 });
955 d.dump_response(id, "resp");
956 d.dump_tool_output("shell", "output");
957
958 let files: Vec<_> = std::fs::read_dir(&session_dir)
960 .unwrap()
961 .filter_map(std::result::Result::ok)
962 .filter(|e| {
963 e.file_name()
964 .to_string_lossy()
965 .chars()
966 .next()
967 .is_some_and(|c| c.is_ascii_digit())
968 })
969 .collect();
970 assert!(
971 files.is_empty(),
972 "no legacy numbered files in Trace format session dir"
973 );
974
975 assert!(session_dir.is_dir(), "session subdir must exist");
978 }
979
980 #[test]
981 fn tool_call_span_emitted() {
982 let tmp = tempdir().unwrap();
983 let mut c = make_collector(tmp.path());
984 c.begin_iteration(0, "test");
985 let iter_id = c.current_iteration_span_id(0).unwrap();
986 let guard = c.begin_tool_call("shell", iter_id);
987 c.end_tool_call(
988 guard,
989 "shell",
990 ToolAttributes {
991 latency_ms: 50,
992 is_error: false,
993 error_kind: None,
994 },
995 );
996 c.end_iteration(0, SpanStatus::Ok);
997 c.finish();
998
999 let content = std::fs::read_to_string(c.trace_json_path()).unwrap();
1000 let v: serde_json::Value = serde_json::from_str(&content).unwrap();
1001 let spans = v["resourceSpans"][0]["scopeSpans"][0]["spans"]
1002 .as_array()
1003 .unwrap();
1004 assert!(
1005 spans.iter().any(|s| s["name"] == "tool.shell"),
1006 "tool.shell span must be emitted"
1007 );
1008 }
1009
1010 #[test]
1011 fn tool_call_error_span_emitted() {
1012 let tmp = tempdir().unwrap();
1013 let mut c = make_collector(tmp.path());
1014 c.begin_iteration(0, "test");
1015 let iter_id = c.current_iteration_span_id(0).unwrap();
1016 let guard = c.begin_tool_call("shell", iter_id);
1017 c.end_tool_call(
1018 guard,
1019 "shell",
1020 ToolAttributes {
1021 latency_ms: 10,
1022 is_error: true,
1023 error_kind: Some("permission denied".to_owned()),
1024 },
1025 );
1026 c.end_iteration(0, SpanStatus::Ok);
1027 c.finish();
1028
1029 let content = std::fs::read_to_string(c.trace_json_path()).unwrap();
1030 let v: serde_json::Value = serde_json::from_str(&content).unwrap();
1031 let spans = v["resourceSpans"][0]["scopeSpans"][0]["spans"]
1032 .as_array()
1033 .unwrap();
1034 let tool_span = spans
1035 .iter()
1036 .find(|s| s["name"] == "tool.shell")
1037 .expect("tool.shell span missing");
1038 assert_eq!(
1039 tool_span["status"]["code"], 2_u64,
1040 "error span must have status code 2"
1041 );
1042 }
1043
1044 #[test]
1045 fn begin_tool_call_at_timestamps_precede_end_time() {
1046 let tmp = tempdir().unwrap();
1047 let mut c = make_collector(tmp.path());
1048 c.begin_iteration(0, "test");
1049 let iter_id = c.current_iteration_span_id(0).unwrap();
1050
1051 let started_at = std::time::Instant::now();
1053 std::thread::sleep(std::time::Duration::from_millis(2));
1054 let guard = c.begin_tool_call_at("shell", iter_id, &started_at);
1055 let span_start = guard.start_time_unix_nanos;
1056 c.end_tool_call(
1057 guard,
1058 "shell",
1059 ToolAttributes {
1060 latency_ms: 2,
1061 is_error: false,
1062 error_kind: None,
1063 },
1064 );
1065 c.end_iteration(0, SpanStatus::Ok);
1066 c.finish();
1067
1068 let content = std::fs::read_to_string(c.trace_json_path()).unwrap();
1069 let v: serde_json::Value = serde_json::from_str(&content).unwrap();
1070 let spans = v["resourceSpans"][0]["scopeSpans"][0]["spans"]
1071 .as_array()
1072 .unwrap();
1073 let tool_span = spans
1074 .iter()
1075 .find(|s| s["name"] == "tool.shell")
1076 .expect("tool.shell span missing");
1077 let recorded_start: u64 = tool_span["startTimeUnixNano"]
1078 .as_str()
1079 .unwrap()
1080 .parse()
1081 .unwrap();
1082 let recorded_end: u64 = tool_span["endTimeUnixNano"]
1083 .as_str()
1084 .unwrap()
1085 .parse()
1086 .unwrap();
1087 assert!(
1089 recorded_start < recorded_end,
1090 "start ({recorded_start}) must precede end ({recorded_end})"
1091 );
1092 assert_eq!(
1094 span_start, recorded_start,
1095 "guard start must match serialized start"
1096 );
1097 }
1098
1099 #[test]
1100 fn session_to_iteration_parent_span_id() {
1101 let tmp = tempdir().unwrap();
1102 let mut c = make_collector(tmp.path());
1103 let session_id = c.session_span_id();
1104 c.begin_iteration(0, "test");
1105 c.end_iteration(0, SpanStatus::Ok);
1106 c.finish();
1107
1108 let content = std::fs::read_to_string(c.trace_json_path()).unwrap();
1109 let v: serde_json::Value = serde_json::from_str(&content).unwrap();
1110 let spans = v["resourceSpans"][0]["scopeSpans"][0]["spans"]
1111 .as_array()
1112 .unwrap();
1113 let iter_span = spans
1114 .iter()
1115 .find(|s| s["name"] == "iteration.0")
1116 .expect("iteration.0 span missing");
1117 assert_eq!(
1118 iter_span["parentSpanId"],
1119 serde_json::json!(hex8(session_id)),
1120 "iteration span parent must be session span"
1121 );
1122 }
1123
1124 #[tokio::test]
1125 async fn json_and_raw_formats_still_write_files() {
1126 use crate::debug_dump::{DebugDumper, DumpFormat, RequestDebugDump};
1127
1128 let tmp = tempdir().unwrap();
1129 for fmt in [DumpFormat::Json, DumpFormat::Raw] {
1130 let d = DebugDumper::new(tmp.path(), fmt).unwrap();
1131 let id = d.dump_request(&RequestDebugDump {
1132 model_name: "test-model",
1133 messages: &[],
1134 tools: &[],
1135 provider_request: serde_json::json!({"model": "test-model", "max_tokens": 100}),
1136 memcot_state: None,
1137 });
1138 d.dump_response(id, "hello");
1139 let session_dir = std::fs::read_dir(tmp.path())
1140 .unwrap()
1141 .filter_map(std::result::Result::ok)
1142 .find(|e| e.file_type().is_ok_and(|t| t.is_dir()))
1143 .unwrap()
1144 .path();
1145 let request_path = session_dir.join("0000-request.json");
1148 let mut written = false;
1149 for _ in 0..200 {
1150 if request_path.exists() {
1151 written = true;
1152 break;
1153 }
1154 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1155 }
1156 assert!(written, "request file must exist for format {fmt:?}");
1157 }
1158 }
1159}