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 pub fn reopen_one_shot_for_resend(&self) {
324 if let Some(gate) = &self.one_shot_gate {
325 gate.store(false, Ordering::SeqCst);
326 }
327 }
328
329 fn pass_one_shot_gate(&self) -> bool {
332 match &self.one_shot_gate {
333 None => true,
334 Some(gate) => !gate.swap(true, Ordering::SeqCst),
335 }
336 }
337
338 pub fn cache_snapshots(&self) -> &Arc<super::diagnostics::CacheSnapshotStore> {
340 &self.cache_snapshots
341 }
342
343 pub fn capture_request_content(&self, request_id: &TraceId, body_bytes: &[u8]) {
347 if let Some(capture) = &self.content_capture {
348 capture.capture(request_id, body_bytes);
349 }
350 }
351
352 pub fn capture_unsupported(&self, reason: &str) {
360 if let Some(capture) = &self.content_capture {
361 if capture.mark_unsupported() {
362 self.note_degraded();
363 tracing::warn!(
364 reason,
365 "trace content capture unsupported on this provider path; \
366 no capture bundle was written (see /trace status degraded counter)"
367 );
368 }
369 }
370 }
371
372 pub fn with_key_path(mut self, path: PathBuf) -> Self {
374 self.key_path = Some(Arc::new(path));
375 self
376 }
377
378 pub fn enabled(&self) -> bool {
379 self.sink.enabled()
383 && self
384 .one_shot_gate
385 .as_ref()
386 .map(|gate| !gate.load(Ordering::SeqCst))
387 .unwrap_or(true)
388 }
389
390 pub fn degraded_records(&self) -> u64 {
392 self.degraded.load(Ordering::Relaxed)
393 }
394
395 pub async fn digest_key(&self) -> Option<Arc<TraceDigestKey>> {
400 self.key
401 .get_or_init(|| async {
402 let path = self.key_path.clone();
403 let loaded = tokio::task::spawn_blocking(move || match path {
404 Some(p) => load_or_create_digest_key_at(&p),
405 None => load_or_create_digest_key(),
406 })
407 .await;
408 match loaded {
409 Ok(Ok(key)) => Some(Arc::new(key)),
410 Ok(Err(err)) => {
411 self.note_key_failure(&err.to_string());
412 None
413 }
414 Err(join_err) => {
415 self.note_key_failure(&join_err.to_string());
416 None
417 }
418 }
419 })
420 .await
421 .clone()
422 }
423
424 fn note_key_failure(&self, reason: &str) {
425 self.degraded.fetch_add(1, Ordering::Relaxed);
426 if !self.key_warned.swap(true, Ordering::Relaxed) {
427 tracing::warn!(
430 reason,
431 "trace digest key unavailable — traces degrade to digest-free records"
432 );
433 }
434 }
435
436 fn note_degraded(&self) {
437 self.degraded.fetch_add(1, Ordering::Relaxed);
438 }
439
440 fn next_request_seq(&self) -> u64 {
441 self.request_seq.fetch_add(1, Ordering::Relaxed)
442 }
443}
444
445fn unix_ms_now() -> u64 {
446 SystemTime::now()
447 .duration_since(UNIX_EPOCH)
448 .map(|d| d.as_millis() as u64)
449 .unwrap_or(0)
450}
451
452#[derive(Debug, Clone, Copy)]
457pub struct AttemptClock {
458 start: Instant,
459 send_start_unix_ms: u64,
460 headers_ms: Option<u64>,
461 first_byte_ms: Option<u64>,
462 first_model_event_ms: Option<u64>,
463 stream_end_ms: Option<u64>,
464}
465
466impl AttemptClock {
467 pub fn start() -> Self {
469 Self {
470 start: Instant::now(),
471 send_start_unix_ms: unix_ms_now(),
472 headers_ms: None,
473 first_byte_ms: None,
474 first_model_event_ms: None,
475 stream_end_ms: None,
476 }
477 }
478
479 fn elapsed_ms(&self) -> u64 {
480 self.start.elapsed().as_millis() as u64
481 }
482
483 pub fn mark_headers(&mut self) {
485 if self.headers_ms.is_none() {
486 self.headers_ms = Some(self.elapsed_ms());
487 }
488 }
489
490 pub fn mark_first_byte(&mut self) {
492 if self.first_byte_ms.is_none() {
493 self.first_byte_ms = Some(self.elapsed_ms());
494 }
495 }
496
497 pub fn set_first_model_event_after_headers(&mut self, offset_ms: u64) {
500 if self.first_model_event_ms.is_none() {
501 if let Some(headers) = self.headers_ms {
502 self.first_model_event_ms = Some(headers.saturating_add(offset_ms));
503 }
504 }
505 }
506
507 pub fn mark_first_model_event(&mut self) {
513 if self.first_model_event_ms.is_none() {
514 self.first_model_event_ms = Some(self.elapsed_ms());
515 }
516 }
517
518 pub fn set_first_model_event_offset(&mut self, offset_ms: u64) {
524 if self.first_model_event_ms.is_none() {
525 self.first_model_event_ms = Some(offset_ms);
526 }
527 }
528
529 pub fn mark_stream_end(&mut self) {
531 if self.stream_end_ms.is_none() {
532 self.stream_end_ms = Some(self.elapsed_ms());
533 }
534 }
535
536 fn stages(&self) -> TimingStages {
537 TimingStages {
538 send_start_unix_ms: Some(self.send_start_unix_ms),
539 headers_ms: self.headers_ms,
540 first_byte_ms: self.first_byte_ms,
541 first_model_event_ms: self.first_model_event_ms,
542 stream_end_ms: self.stream_end_ms,
543 }
544 }
545}
546
547#[derive(Debug, Clone, Default)]
553pub struct RequestStructure {
554 pub anatomy: RequestAnatomy,
555 pub wire: Option<WireMeta>,
557 pub system_segments: Vec<SystemSegmentMeta>,
558 pub messages: Vec<MessageMeta>,
559 pub tools: Vec<ToolMeta>,
560 pub cache: CacheMeta,
561 pub translation: Vec<super::types::TranslationLoss>,
564}
565
566pub fn wire_meta_from_sent_bytes(key: &TraceDigestKey, sent_bytes: &[u8]) -> WireMeta {
570 WireMeta {
571 byte_len: sent_bytes.len() as u64,
572 digest: keyed_digest(key, DigestDomain::Wire, sent_bytes),
573 }
574}
575
576#[derive(Debug, Clone, PartialEq, Eq)]
577pub struct RequestCorrelation {
578 pub session_id: TraceId,
579 pub turn_id: TraceId,
580 pub request_id: TraceId,
581}
582
583pub struct RequestTracer {
588 ctx: TraceContext,
589 session_id: TraceId,
590 turn_id: TraceId,
591 request_id: TraceId,
592 model: QualifiedModelId,
593 transport: TransportKind,
594 endpoint: EndpointMeta,
595 structure: RequestStructure,
596 retries: Vec<RetryMeta>,
597 attempt: u32,
598}
599
600impl RequestTracer {
601 pub fn begin(
605 ctx: &TraceContext,
606 correlation: Option<RequestCorrelation>,
607 model: QualifiedModelId,
608 transport: TransportKind,
609 endpoint: EndpointMeta,
610 structure: RequestStructure,
611 ) -> Option<Self> {
612 if !ctx.enabled() {
613 return None;
614 }
615 let correlation = match correlation.or_else(|| ctx.reserve_request_correlation()) {
616 Some(correlation) => correlation,
617 None => {
618 ctx.note_degraded();
619 return None;
620 }
621 };
622 let request_id = correlation.request_id;
623 let turn_id = correlation.turn_id;
624 let session_id = correlation.session_id;
625 if !ctx.pass_one_shot_gate() {
632 return None;
633 }
634 Some(Self {
635 ctx: ctx.clone(),
636 session_id,
637 turn_id,
638 request_id,
639 model,
640 transport,
641 endpoint,
642 structure,
643 retries: Vec::new(),
644 attempt: 1,
645 })
646 }
647
648 pub fn request_id(&self) -> &TraceId {
650 &self.request_id
651 }
652
653 pub fn correlation_ids(&self) -> (&TraceId, &TraceId, &TraceId) {
654 (&self.session_id, &self.turn_id, &self.request_id)
655 }
656
657 #[allow(clippy::too_many_arguments)]
661 pub fn attempt_failed(
662 &mut self,
663 clock: AttemptClock,
664 class: RetryClass,
665 delay: Duration,
666 http_status: Option<u16>,
667 provider_request_id: Option<TraceId>,
668 code: &str,
669 ) {
670 let outcome = TransportOutcome {
671 timings: clock.stages(),
672 retries: self.retries.clone(),
673 provider_request_id,
674 http_status,
675 stop_reason: None,
676 usage: None,
677 terminal: TurnOutcome::ProviderFailed {
678 code: code.to_string(),
679 correlation_id: self.request_id.as_str().to_string(),
680 },
681 };
682 self.emit_record(outcome);
683 self.structure.cache.server_diagnosis = None;
686 self.retries.push(RetryMeta {
687 attempt: self.attempt,
688 class,
689 delay_ms: delay.as_millis() as u64,
690 });
691 self.attempt += 1;
692 }
693
694 #[allow(clippy::too_many_arguments)]
696 pub fn finish(
697 self,
698 clock: AttemptClock,
699 http_status: Option<u16>,
700 provider_request_id: Option<TraceId>,
701 stop_reason: Option<StopReason>,
702 usage: Option<UsageMeta>,
703 terminal: TurnOutcome,
704 ) {
705 let outcome = TransportOutcome {
706 timings: clock.stages(),
707 retries: self.retries.clone(),
708 provider_request_id,
709 http_status,
710 stop_reason,
711 usage,
712 terminal,
713 };
714 self.emit_record(outcome);
715 }
716
717 pub fn failed_terminal(&self, code: &str) -> TurnOutcome {
720 TurnOutcome::ProviderFailed {
721 code: code.to_string(),
722 correlation_id: self.request_id.as_str().to_string(),
723 }
724 }
725
726 pub fn set_server_diagnosis(&mut self, diagnosis: Option<super::ServerCacheDiagnosis>) {
730 self.structure.cache.server_diagnosis = diagnosis;
731 }
732
733 fn emit_record(&self, outcome: TransportOutcome) {
734 let record = RequestTrace {
735 schema: TraceSchemaVersion,
736 session_id: self.session_id.clone(),
737 turn_id: self.turn_id.clone(),
738 request_id: self.request_id.clone(),
739 execution_events: self
740 .ctx
741 .execution_events
742 .lock()
743 .unwrap_or_else(std::sync::PoisonError::into_inner)
744 .get(self.request_id.as_str())
745 .cloned()
746 .unwrap_or_default(),
747 attempt: self.attempt,
748 model: self.model.clone(),
749 transport: self.transport,
750 endpoint: self.endpoint.clone(),
751 anatomy: self.structure.anatomy,
752 wire: self.structure.wire.clone(),
753 system_segments: self.structure.system_segments.clone(),
754 messages: self.structure.messages.clone(),
755 tools: self.structure.tools.clone(),
756 cache: self.structure.cache.clone(),
757 translation_losses: self.structure.translation.clone(),
758 outcome,
759 };
760 self.ctx.sink.emit(record);
761 }
762}