1use super::key::{
31 keyed_digest, load_or_create_digest_key, load_or_create_digest_key_at, DigestDomain,
32 TraceDigestKey,
33};
34use super::types::{
35 CacheMeta, EndpointMeta, MessageMeta, RequestAnatomy, RequestTrace, RetryClass, RetryMeta,
36 StopReason, SystemSegmentMeta, TimingStages, ToolMeta, TraceId, TraceSchemaVersion,
37 TransportKind, TransportOutcome, UsageMeta, WireMeta,
38};
39use agent_core::prompt::QualifiedModelId;
40use agent_core::TurnOutcome;
41use std::path::PathBuf;
42use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
43use std::sync::Arc;
44use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
45
46pub trait TraceSink: Send + Sync + std::fmt::Debug {
51 fn emit(&self, record: RequestTrace);
52 fn enrich_execution_events(
57 &self,
58 _request_id: &TraceId,
59 _events: Vec<super::types::ToolExecutionEvent>,
60 ) {
61 }
62 fn snapshot_for_request(&self, _request_id: &TraceId) -> Option<RequestTrace> {
65 None
66 }
67 fn enabled(&self) -> bool {
69 true
70 }
71}
72
73#[derive(Debug, Default)]
78pub struct NoopTraceSink;
79
80impl TraceSink for NoopTraceSink {
81 fn emit(&self, _record: RequestTrace) {}
82 fn enabled(&self) -> bool {
83 false
84 }
85}
86
87#[derive(Debug, Default)]
89pub struct CollectingTraceSink {
90 records: std::sync::Mutex<Vec<RequestTrace>>,
91}
92
93impl CollectingTraceSink {
94 pub fn new() -> Arc<Self> {
95 Arc::new(Self::default())
96 }
97
98 pub fn records(&self) -> Vec<RequestTrace> {
100 self.records.lock().expect("trace sink poisoned").clone()
101 }
102}
103
104impl TraceSink for CollectingTraceSink {
105 fn emit(&self, record: RequestTrace) {
106 self.records
107 .lock()
108 .expect("trace sink poisoned")
109 .push(record);
110 }
111
112 fn enrich_execution_events(
113 &self,
114 request_id: &TraceId,
115 events: Vec<super::types::ToolExecutionEvent>,
116 ) {
117 if let Some(record) = self
118 .records
119 .lock()
120 .expect("trace sink poisoned")
121 .iter_mut()
122 .rev()
123 .find(|record| &record.request_id == request_id)
124 {
125 record.execution_events = events;
126 }
127 }
128
129 fn snapshot_for_request(&self, request_id: &TraceId) -> Option<RequestTrace> {
130 self.records
131 .lock()
132 .expect("trace sink poisoned")
133 .iter()
134 .rev()
135 .find(|record| &record.request_id == request_id)
136 .cloned()
137 }
138}
139
140#[derive(Clone)]
145pub struct TraceContext {
146 sink: Arc<dyn TraceSink>,
147 session_id: TraceId,
148 key_path: Option<Arc<PathBuf>>,
150 key: Arc<tokio::sync::OnceCell<Option<Arc<TraceDigestKey>>>>,
153 key_warned: Arc<AtomicBool>,
154 degraded: Arc<AtomicU64>,
157 request_seq: Arc<AtomicU64>,
158 cache_snapshots: Arc<super::diagnostics::CacheSnapshotStore>,
161 content_capture: Option<Arc<super::controls::ContentCapture>>,
164 one_shot_gate: Option<Arc<AtomicBool>>,
172 execution_events: Arc<
176 std::sync::Mutex<std::collections::BTreeMap<String, Vec<super::types::ToolExecutionEvent>>>,
177 >,
178}
179
180impl std::fmt::Debug for TraceContext {
181 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
182 f.debug_struct("TraceContext")
183 .field("enabled", &self.sink.enabled())
184 .field("session_id", &self.session_id)
185 .finish_non_exhaustive()
186 }
187}
188
189impl Default for TraceContext {
190 fn default() -> Self {
191 Self::disabled()
192 }
193}
194
195impl TraceContext {
196 pub fn emit_execution_enriched(&self, request_id: &TraceId) {
200 self.sink
201 .enrich_execution_events(request_id, self.execution_events(request_id));
202 }
203
204 pub fn session_id(&self) -> &TraceId {
206 &self.session_id
207 }
208
209 pub fn record_execution_event(&self, event: super::types::ToolExecutionEvent) {
212 self.execution_events
213 .lock()
214 .unwrap_or_else(std::sync::PoisonError::into_inner)
215 .entry(event.request_id.as_str().to_string())
216 .or_default()
217 .push(event);
218 }
219
220 pub fn execution_events(&self, request_id: &TraceId) -> Vec<super::types::ToolExecutionEvent> {
221 self.execution_events
222 .lock()
223 .unwrap_or_else(std::sync::PoisonError::into_inner)
224 .get(request_id.as_str())
225 .cloned()
226 .unwrap_or_default()
227 }
228
229 pub fn latest_request_correlation(&self) -> Option<(TraceId, TraceId, TraceId)> {
231 let seq = self.request_seq.load(Ordering::Relaxed);
232 if seq == 0 {
233 return None;
234 }
235 let pid = std::process::id();
236 Some((
237 self.session_id.clone(),
238 TraceId::new(format!("turn-{pid}-{seq}")).ok()?,
239 TraceId::new(format!("req-{pid}-{seq}")).ok()?,
240 ))
241 }
242
243 pub fn reserve_request_correlation(&self) -> Option<RequestCorrelation> {
244 if !self.enabled() {
245 return None;
246 }
247 let seq = self.next_request_seq();
248 let pid = std::process::id();
249 Some(RequestCorrelation {
250 session_id: self.session_id.clone(),
251 turn_id: TraceId::new(format!("turn-{pid}-{seq}")).ok()?,
252 request_id: TraceId::new(format!("req-{pid}-{seq}")).ok()?,
253 })
254 }
255
256 pub fn disabled() -> Self {
258 Self::with_sink(Arc::new(NoopTraceSink))
259 }
260
261 pub fn with_sink(sink: Arc<dyn TraceSink>) -> Self {
263 let session_id = TraceId::new(format!("session-{}-{}", std::process::id(), unix_ms_now()))
264 .expect("generated session id is valid");
265 Self {
266 sink,
267 session_id,
268 key_path: None,
269 key: Arc::new(tokio::sync::OnceCell::new()),
270 key_warned: Arc::new(AtomicBool::new(false)),
271 degraded: Arc::new(AtomicU64::new(0)),
272 request_seq: Arc::new(AtomicU64::new(0)),
273 cache_snapshots: Arc::new(super::diagnostics::CacheSnapshotStore::new()),
274 content_capture: None,
275 one_shot_gate: None,
276 execution_events: Arc::new(std::sync::Mutex::new(std::collections::BTreeMap::new())),
277 }
278 }
279
280 pub fn fork_with_sink(&self, sink: Arc<dyn TraceSink>) -> Self {
287 Self {
288 sink,
289 session_id: self.session_id.clone(),
290 key_path: self.key_path.clone(),
291 key: Arc::clone(&self.key),
292 key_warned: Arc::clone(&self.key_warned),
293 degraded: Arc::clone(&self.degraded),
294 request_seq: Arc::clone(&self.request_seq),
295 cache_snapshots: Arc::clone(&self.cache_snapshots),
296 content_capture: None,
297 one_shot_gate: None,
298 execution_events: Arc::clone(&self.execution_events),
299 }
300 }
301
302 pub fn with_content_capture(mut self, capture: Arc<super::controls::ContentCapture>) -> Self {
304 self.content_capture = Some(capture);
305 self
306 }
307
308 pub fn with_one_shot_request_gate(mut self) -> Self {
313 self.one_shot_gate = Some(Arc::new(AtomicBool::new(false)));
314 self
315 }
316
317 fn pass_one_shot_gate(&self) -> bool {
320 match &self.one_shot_gate {
321 None => true,
322 Some(gate) => !gate.swap(true, Ordering::SeqCst),
323 }
324 }
325
326 pub fn cache_snapshots(&self) -> &Arc<super::diagnostics::CacheSnapshotStore> {
328 &self.cache_snapshots
329 }
330
331 pub fn capture_request_content(&self, request_id: &TraceId, body_bytes: &[u8]) {
335 if let Some(capture) = &self.content_capture {
336 capture.capture(request_id, body_bytes);
337 }
338 }
339
340 pub fn capture_unsupported(&self, reason: &str) {
348 if let Some(capture) = &self.content_capture {
349 if capture.mark_unsupported() {
350 self.note_degraded();
351 tracing::warn!(
352 reason,
353 "trace content capture unsupported on this provider path; \
354 no capture bundle was written (see /trace status degraded counter)"
355 );
356 }
357 }
358 }
359
360 pub fn with_key_path(mut self, path: PathBuf) -> Self {
362 self.key_path = Some(Arc::new(path));
363 self
364 }
365
366 pub fn enabled(&self) -> bool {
367 self.sink.enabled()
371 && self
372 .one_shot_gate
373 .as_ref()
374 .map(|gate| !gate.load(Ordering::SeqCst))
375 .unwrap_or(true)
376 }
377
378 pub fn degraded_records(&self) -> u64 {
380 self.degraded.load(Ordering::Relaxed)
381 }
382
383 pub async fn digest_key(&self) -> Option<Arc<TraceDigestKey>> {
388 self.key
389 .get_or_init(|| async {
390 let path = self.key_path.clone();
391 let loaded = tokio::task::spawn_blocking(move || match path {
392 Some(p) => load_or_create_digest_key_at(&p),
393 None => load_or_create_digest_key(),
394 })
395 .await;
396 match loaded {
397 Ok(Ok(key)) => Some(Arc::new(key)),
398 Ok(Err(err)) => {
399 self.note_key_failure(&err.to_string());
400 None
401 }
402 Err(join_err) => {
403 self.note_key_failure(&join_err.to_string());
404 None
405 }
406 }
407 })
408 .await
409 .clone()
410 }
411
412 fn note_key_failure(&self, reason: &str) {
413 self.degraded.fetch_add(1, Ordering::Relaxed);
414 if !self.key_warned.swap(true, Ordering::Relaxed) {
415 tracing::warn!(
418 reason,
419 "trace digest key unavailable — traces degrade to digest-free records"
420 );
421 }
422 }
423
424 fn note_degraded(&self) {
425 self.degraded.fetch_add(1, Ordering::Relaxed);
426 }
427
428 fn next_request_seq(&self) -> u64 {
429 self.request_seq.fetch_add(1, Ordering::Relaxed)
430 }
431}
432
433fn unix_ms_now() -> u64 {
434 SystemTime::now()
435 .duration_since(UNIX_EPOCH)
436 .map(|d| d.as_millis() as u64)
437 .unwrap_or(0)
438}
439
440#[derive(Debug, Clone, Copy)]
445pub struct AttemptClock {
446 start: Instant,
447 send_start_unix_ms: u64,
448 headers_ms: Option<u64>,
449 first_byte_ms: Option<u64>,
450 first_model_event_ms: Option<u64>,
451 stream_end_ms: Option<u64>,
452}
453
454impl AttemptClock {
455 pub fn start() -> Self {
457 Self {
458 start: Instant::now(),
459 send_start_unix_ms: unix_ms_now(),
460 headers_ms: None,
461 first_byte_ms: None,
462 first_model_event_ms: None,
463 stream_end_ms: None,
464 }
465 }
466
467 fn elapsed_ms(&self) -> u64 {
468 self.start.elapsed().as_millis() as u64
469 }
470
471 pub fn mark_headers(&mut self) {
473 if self.headers_ms.is_none() {
474 self.headers_ms = Some(self.elapsed_ms());
475 }
476 }
477
478 pub fn mark_first_byte(&mut self) {
480 if self.first_byte_ms.is_none() {
481 self.first_byte_ms = Some(self.elapsed_ms());
482 }
483 }
484
485 pub fn set_first_model_event_after_headers(&mut self, offset_ms: u64) {
488 if self.first_model_event_ms.is_none() {
489 if let Some(headers) = self.headers_ms {
490 self.first_model_event_ms = Some(headers.saturating_add(offset_ms));
491 }
492 }
493 }
494
495 pub fn mark_first_model_event(&mut self) {
501 if self.first_model_event_ms.is_none() {
502 self.first_model_event_ms = Some(self.elapsed_ms());
503 }
504 }
505
506 pub fn set_first_model_event_offset(&mut self, offset_ms: u64) {
512 if self.first_model_event_ms.is_none() {
513 self.first_model_event_ms = Some(offset_ms);
514 }
515 }
516
517 pub fn mark_stream_end(&mut self) {
519 if self.stream_end_ms.is_none() {
520 self.stream_end_ms = Some(self.elapsed_ms());
521 }
522 }
523
524 fn stages(&self) -> TimingStages {
525 TimingStages {
526 send_start_unix_ms: Some(self.send_start_unix_ms),
527 headers_ms: self.headers_ms,
528 first_byte_ms: self.first_byte_ms,
529 first_model_event_ms: self.first_model_event_ms,
530 stream_end_ms: self.stream_end_ms,
531 }
532 }
533}
534
535#[derive(Debug, Clone, Default)]
541pub struct RequestStructure {
542 pub anatomy: RequestAnatomy,
543 pub wire: Option<WireMeta>,
545 pub system_segments: Vec<SystemSegmentMeta>,
546 pub messages: Vec<MessageMeta>,
547 pub tools: Vec<ToolMeta>,
548 pub cache: CacheMeta,
549 pub translation: Vec<super::types::TranslationLoss>,
552}
553
554pub fn wire_meta_from_sent_bytes(key: &TraceDigestKey, sent_bytes: &[u8]) -> WireMeta {
558 WireMeta {
559 byte_len: sent_bytes.len() as u64,
560 digest: keyed_digest(key, DigestDomain::Wire, sent_bytes),
561 }
562}
563
564#[derive(Debug, Clone, PartialEq, Eq)]
565pub struct RequestCorrelation {
566 pub session_id: TraceId,
567 pub turn_id: TraceId,
568 pub request_id: TraceId,
569}
570
571pub struct RequestTracer {
576 ctx: TraceContext,
577 session_id: TraceId,
578 turn_id: TraceId,
579 request_id: TraceId,
580 model: QualifiedModelId,
581 transport: TransportKind,
582 endpoint: EndpointMeta,
583 structure: RequestStructure,
584 retries: Vec<RetryMeta>,
585 attempt: u32,
586}
587
588impl RequestTracer {
589 pub fn begin(
593 ctx: &TraceContext,
594 correlation: Option<RequestCorrelation>,
595 model: QualifiedModelId,
596 transport: TransportKind,
597 endpoint: EndpointMeta,
598 structure: RequestStructure,
599 ) -> Option<Self> {
600 if !ctx.enabled() {
601 return None;
602 }
603 let correlation = match correlation.or_else(|| ctx.reserve_request_correlation()) {
604 Some(correlation) => correlation,
605 None => {
606 ctx.note_degraded();
607 return None;
608 }
609 };
610 let request_id = correlation.request_id;
611 let turn_id = correlation.turn_id;
612 let session_id = correlation.session_id;
613 if !ctx.pass_one_shot_gate() {
620 return None;
621 }
622 Some(Self {
623 ctx: ctx.clone(),
624 session_id,
625 turn_id,
626 request_id,
627 model,
628 transport,
629 endpoint,
630 structure,
631 retries: Vec::new(),
632 attempt: 1,
633 })
634 }
635
636 pub fn request_id(&self) -> &TraceId {
638 &self.request_id
639 }
640
641 pub fn correlation_ids(&self) -> (&TraceId, &TraceId, &TraceId) {
642 (&self.session_id, &self.turn_id, &self.request_id)
643 }
644
645 #[allow(clippy::too_many_arguments)]
649 pub fn attempt_failed(
650 &mut self,
651 clock: AttemptClock,
652 class: RetryClass,
653 delay: Duration,
654 http_status: Option<u16>,
655 provider_request_id: Option<TraceId>,
656 code: &str,
657 ) {
658 let outcome = TransportOutcome {
659 timings: clock.stages(),
660 retries: self.retries.clone(),
661 provider_request_id,
662 http_status,
663 stop_reason: None,
664 usage: None,
665 terminal: TurnOutcome::ProviderFailed {
666 code: code.to_string(),
667 correlation_id: self.request_id.as_str().to_string(),
668 },
669 };
670 self.emit_record(outcome);
671 self.retries.push(RetryMeta {
672 attempt: self.attempt,
673 class,
674 delay_ms: delay.as_millis() as u64,
675 });
676 self.attempt += 1;
677 }
678
679 #[allow(clippy::too_many_arguments)]
681 pub fn finish(
682 self,
683 clock: AttemptClock,
684 http_status: Option<u16>,
685 provider_request_id: Option<TraceId>,
686 stop_reason: Option<StopReason>,
687 usage: Option<UsageMeta>,
688 terminal: TurnOutcome,
689 ) {
690 let outcome = TransportOutcome {
691 timings: clock.stages(),
692 retries: self.retries.clone(),
693 provider_request_id,
694 http_status,
695 stop_reason,
696 usage,
697 terminal,
698 };
699 self.emit_record(outcome);
700 }
701
702 pub fn failed_terminal(&self, code: &str) -> TurnOutcome {
705 TurnOutcome::ProviderFailed {
706 code: code.to_string(),
707 correlation_id: self.request_id.as_str().to_string(),
708 }
709 }
710
711 fn emit_record(&self, outcome: TransportOutcome) {
712 let record = RequestTrace {
713 schema: TraceSchemaVersion,
714 session_id: self.session_id.clone(),
715 turn_id: self.turn_id.clone(),
716 request_id: self.request_id.clone(),
717 execution_events: self
718 .ctx
719 .execution_events
720 .lock()
721 .unwrap_or_else(std::sync::PoisonError::into_inner)
722 .get(self.request_id.as_str())
723 .cloned()
724 .unwrap_or_default(),
725 attempt: self.attempt,
726 model: self.model.clone(),
727 transport: self.transport,
728 endpoint: self.endpoint.clone(),
729 anatomy: self.structure.anatomy,
730 wire: self.structure.wire.clone(),
731 system_segments: self.structure.system_segments.clone(),
732 messages: self.structure.messages.clone(),
733 tools: self.structure.tools.clone(),
734 cache: self.structure.cache.clone(),
735 translation_losses: self.structure.translation.clone(),
736 outcome,
737 };
738 self.ctx.sink.emit(record);
739 }
740}