Skip to main content

zeph_core/debug_dump/
trace.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! OpenTelemetry-compatible trace collector for debug sessions.
5//!
6//! Collects span data during an agent session and serializes to OTLP JSON format
7//! at session end. All text-bearing attributes are redacted via `crate::redact::scrub_content`
8//! before storage (C-01).
9//!
10//! Design notes:
11//! - Uses explicit `begin_X` / `end_X` methods with owned `SpanGuard` — safe
12//!   across async `.await` boundaries because no borrow to `TracingCollector` is held (C-02).
13//! - A `HashMap<usize, IterationEntry>` tracks concurrent iterations (I-03).
14//! - `Drop` on `TracingCollector` flushes partial traces on error/panic paths (C-04).
15//! - When the `otel` feature is enabled an `mpsc` channel forwards completed spans to
16//!   the OTLP exporter in `tracing_init.rs` (C-05).
17
18use 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
28// ─── Span ID generation ───────────────────────────────────────────────────────
29
30static 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    // XOR low byte with counter to guarantee distinct IDs even under high concurrency.
41    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// ─── Public types ─────────────────────────────────────────────────────────────
73
74#[non_exhaustive]
75/// Span status code (matches OTLP `StatusCode` values).
76#[derive(Debug, Clone, Serialize, Deserialize)]
77pub enum SpanStatus {
78    Ok,
79    Error { message: String },
80    Unset,
81}
82
83/// A completed span ready for OTLP serialization.
84#[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
96/// Owned guard returned by `begin_*` methods. Pass back to `end_*` to close the span.
97///
98/// Does NOT hold a reference to `TracingCollector` — safe across async `.await` boundaries (C-02).
99pub 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
106/// Attributes for a completed LLM request span.
107pub 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
116/// Attributes for a completed tool call span.
117pub struct ToolAttributes {
118    pub latency_ms: u64,
119    pub is_error: bool,
120    pub error_kind: Option<String>,
121}
122
123/// Attributes for a completed memory search span.
124pub struct MemorySearchAttributes {
125    pub query_preview: String,
126    pub result_count: usize,
127    pub latency_ms: u64,
128}
129
130// ─── Bridge event for OTLP export (C-05) ─────────────────────────────────────
131
132/// Event sent over the mpsc channel to the OTLP exporter.
133///
134/// The root crate wires the sender when the `otel` feature is enabled.
135/// `zeph-core` compiles this unconditionally so the struct is always available.
136#[derive(Debug)]
137pub struct TraceEvent {
138    pub trace_id: [u8; 16],
139    pub spans: Vec<SpanData>,
140}
141
142// ─── Internal ─────────────────────────────────────────────────────────────────
143
144/// Internal entry for a still-open iteration.
145struct IterationEntry {
146    guard: SpanGuard,
147    user_msg_preview: String,
148}
149
150// ─── TracingCollector ─────────────────────────────────────────────────────────
151
152/// Collects OTel-compatible spans for a single agent session.
153///
154/// All methods take `&mut self`. The agent loop is single-threaded within a session.
155/// For concurrent iteration support (I-03), a `HashMap<usize, IterationEntry>` is used.
156/// Default cap on collected spans per session (SEC-02).
157const 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    /// User-defined resource attributes from `telemetry.trace_metadata`, included in
165    /// `resourceSpans[].resource.attributes` when serializing to OTLP JSON.
166    trace_metadata: HashMap<String, String>,
167    output_dir: PathBuf,
168    /// Active (open) iterations keyed by iteration index (I-03).
169    active_iterations: HashMap<usize, IterationEntry>,
170    completed_spans: VecDeque<SpanData>,
171    /// Hard cap on `completed_spans` length. Oldest span dropped when exceeded (SEC-02).
172    max_spans: usize,
173    /// Whether to redact text attributes. Defaults to `true` (C-01).
174    redact: bool,
175    /// Guards against double-write on explicit `finish()` followed by `Drop`.
176    flushed: bool,
177    /// Optional channel to forward completed spans to the OTLP exporter.
178    /// Wired by the root crate when the `otel` feature is enabled (C-05).
179    trace_tx: Option<tokio::sync::mpsc::UnboundedSender<TraceEvent>>,
180}
181
182impl TracingCollector {
183    /// Create a new collector.
184    ///
185    /// `trace_metadata` is included as additional resource attributes in the OTLP JSON output.
186    /// The reserved key `service.name` in `trace_metadata` is silently skipped.
187    ///
188    /// # Errors
189    ///
190    /// Returns an error if `output_dir` cannot be created.
191    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    /// Append a span, dropping the oldest when `max_spans` is exceeded (SEC-02).
224    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    // ── Iteration spans ───────────────────────────────────────────────────────
236
237    /// Open an iteration span. Call at the start of `process_user_message`.
238    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    /// Close an iteration span.
257    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    // ── LLM request spans ─────────────────────────────────────────────────────
280
281    /// Open an LLM request span. Returns an owned `SpanGuard` safe to hold across `.await`.
282    #[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    /// Close an LLM request span.
293    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    // ── Tool call spans ───────────────────────────────────────────────────────
325
326    /// Open a tool call span, recording the start time as now.
327    #[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    /// Open a tool call span with a pre-recorded start time.
333    ///
334    /// Use this variant when the tool has already executed (post-hoc assembly pattern) and
335    /// `started_at` was captured *before* the call. The Unix start timestamp is back-computed
336    /// from `started_at.elapsed()` so the span is correctly positioned on the timeline.
337    #[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    /// Close a tool call span.
355    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            // IMP-04: apply redaction to error messages (may contain secret data).
368            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    // ── Memory search spans ───────────────────────────────────────────────────
391
392    /// Open a memory search span.
393    #[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    /// Close a memory search span.
404    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    // ── Accessors ─────────────────────────────────────────────────────────────
434
435    /// Return the path to the `trace.json` file that will be written on `finish()`.
436    #[must_use]
437    pub fn trace_json_path(&self) -> PathBuf {
438        self.output_dir.join("trace.json")
439    }
440
441    /// Return the span ID of the currently active iteration, if any.
442    #[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    /// Return the session root span ID (fallback parent when no iteration is active).
448    #[must_use]
449    pub fn session_span_id(&self) -> [u8; 8] {
450        self.session_span_id
451    }
452
453    /// Return the trace ID for this session.
454    #[must_use]
455    pub fn trace_id(&self) -> [u8; 16] {
456        self.trace_id
457    }
458
459    // ── Flush ─────────────────────────────────────────────────────────────────
460
461    /// Finalize the session span and write `trace.json`.
462    ///
463    /// Safe to call multiple times — subsequent calls after the first are no-ops (returning
464    /// `None`). Also sends spans over the `OTel` channel when the `otel` feature is enabled
465    /// (C-05).
466    ///
467    /// Returns the `trace.json` write's `JoinHandle` when it was dispatched to the blocking
468    /// pool (a Tokio runtime was active), so a caller for whom this is the *only* remaining
469    /// write of the session — nothing runs after it — can `.await` it to guarantee the file
470    /// actually lands before the process/runtime tears down (#6107). Callers that don't need
471    /// that guarantee (a mid-session format switch, or `Drop`, which cannot `.await` anyway)
472    /// are free to drop the handle and keep the write fire-and-forget.
473    pub fn finish(&mut self) -> Option<tokio::task::JoinHandle<()>> {
474        if self.flushed {
475            return None;
476        }
477        self.flushed = true;
478
479        // Close any still-open iteration spans with `Unset` status (partial trace on error/cancel).
480        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        // C-05: forward spans to OTLP exporter when the root crate has wired the channel.
522        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
536// C-04: Drop flushes partial traces on error/panic/cancellation. Fire-and-forget: `Drop::drop`
537// cannot `.await` the returned handle, but this is a best-effort partial-trace capture (the
538// normal-exit path calls `finish()` explicitly and awaits it — see `agent/mod.rs`), not the
539// session's only copy of the trace.
540impl Drop for TracingCollector {
541    fn drop(&mut self) {
542        let _ = self.finish();
543    }
544}
545
546// ─── OTLP JSON serialization (I-04) ───────────────────────────────────────────
547
548fn 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/// Serialize spans to OTLP JSON Protobuf encoding.
557///
558/// `trace_metadata` entries are appended as additional resource attributes after `service.name`.
559/// The reserved key `service.name` in `trace_metadata` is skipped to avoid duplication.
560///
561/// Format: <https://opentelemetry.io/docs/specs/otlp/#json-protobuf-encoding>
562#[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                // OTLP JSON spec requires int64 fields as strings.
587                "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
640// ─── Helpers ──────────────────────────────────────────────────────────────────
641
642/// Write `data` to `path` with mode 0o600 on Unix (SEC-01) via `zeph_common::fs_secure`.
643///
644/// Offloaded to the blocking thread pool when a Tokio runtime is available — `finish()` is
645/// reachable from the async agent hot path (#6107) and must never block it on synchronous
646/// file I/O — returning the `JoinHandle` so a caller for whom this write is session-critical
647/// (nothing else will write this session's trace) can await completion instead of racing
648/// process/runtime teardown. Callers that don't need that guarantee may drop the handle,
649/// mirroring the fire-and-forget `DebugDumper::write` pattern from #6101.
650///
651/// Falls back to an inline synchronous write (returning `None`) when no runtime is present,
652/// since `finish()` is also reachable from `Drop` (which cannot `.await`) and from plain
653/// non-async unit tests; this mirrors the runtime-detection idiom in
654/// `zeph_common::task_supervisor::TaskSupervisor::new`.
655fn 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// ─── Tests ────────────────────────────────────────────────────────────────────
683
684#[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        // session + iteration = 2 spans.
729        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        // int64 must be serialized as string per OTLP JSON spec.
827        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            // Drop without explicit finish.
838        }
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        // service.name must appear with the collector's service_name, not the skipped value.
861        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        // Custom metadata attribute must be present.
868        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    // #6107: every other test in this module runs as a plain `#[test]` (no Tokio runtime), so
886    // they only ever exercise `write_trace_file`'s synchronous fallback branch — the actual
887    // `spawn_blocking` dispatch added for #6107 was previously untested. This test runs inside
888    // a real runtime so `finish()` takes that branch (asserted via `Some`), then awaits the
889    // returned handle to confirm the write actually completes.
890    #[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        // session + 2 iterations = 3.
938        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        // No legacy numbered files should be written in Trace format.
959        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        // trace.json is written into the session subdir by TracingCollector when wired.
976        // Here we only verify the session dir itself exists (TracingCollector is not wired in this test).
977        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        // Simulate post-hoc assembly: capture start before "execution", then call begin_tool_call_at.
1052        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        // The span start must be earlier than the end.
1088        assert!(
1089            recorded_start < recorded_end,
1090            "start ({recorded_start}) must precede end ({recorded_end})"
1091        );
1092        // The guard's start_time matches what was serialized.
1093        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            // write() now offloads to spawn_blocking and is fire-and-forget (#6029), so poll
1146            // rather than asserting the file exists immediately.
1147            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}