scouter_types/
span_capture.rs

1use crate::trace::{TraceSpanRecord, SCOUTER_EVAL_SCENARIO_ID_ATTR};
2use crate::TraceId as ScouterTraceId;
3use std::collections::{HashMap, HashSet};
4use std::sync::atomic::{AtomicBool, Ordering};
5use std::sync::RwLock;
6
7pub const CAPTURE_BUFFER_MAX: usize = 20_000;
8
9/// Whether local span capture is enabled.
10pub static CAPTURING: AtomicBool = AtomicBool::new(false);
11
12/// Global buffer of captured spans.
13pub static CAPTURE_BUFFER: RwLock<Vec<TraceSpanRecord>> = RwLock::new(Vec::new());
14
15/// Returns `true` if local span capture is currently enabled.
16pub fn is_capturing() -> bool {
17    CAPTURING.load(Ordering::Acquire)
18}
19
20/// Drain all captured spans from the buffer (takes ownership).
21pub fn drain_captured_spans() -> Vec<TraceSpanRecord> {
22    std::mem::take(&mut *CAPTURE_BUFFER.write().unwrap_or_else(|p| p.into_inner()))
23}
24
25/// Returns clones of spans matching the given trace_ids.
26/// Does NOT drain the buffer.
27pub fn get_captured_spans_by_trace_ids(
28    trace_ids: &HashSet<ScouterTraceId>,
29) -> Vec<TraceSpanRecord> {
30    let buf = CAPTURE_BUFFER.read().unwrap_or_else(|p| p.into_inner());
31    buf.iter()
32        .filter(|span| trace_ids.contains(&span.trace_id))
33        .cloned()
34        .collect()
35}
36
37/// Returns a clone of all captured spans without draining.
38pub fn get_all_captured_spans() -> Vec<TraceSpanRecord> {
39    CAPTURE_BUFFER
40        .read()
41        .unwrap_or_else(|p| p.into_inner())
42        .clone()
43}
44
45/// Two-pass buffer scan that groups all captured spans by `scouter.eval.scenario_id`.
46///
47/// **Pass 1**: Build a `trace_id → scenario_id` map from spans that carry the
48/// `scouter.eval.scenario_id` attribute (i.e. the orchestrator's wrapper spans).
49///
50/// **Pass 2**: Group every span whose `trace_id` appears in that map into the
51/// corresponding scenario bucket — this picks up child spans (e.g. LLM calls)
52/// that share the trace but don't carry the attribute directly.
53///
54/// Does NOT drain the buffer.
55pub fn get_spans_grouped_by_scenario_id(
56    scenario_ids: &HashSet<String>,
57) -> HashMap<String, Vec<TraceSpanRecord>> {
58    let buf = CAPTURE_BUFFER.read().unwrap_or_else(|p| p.into_inner());
59
60    // Pass 1: trace_id → scenario_id for spans that carry the attribute
61    let mut trace_to_scenario: HashMap<ScouterTraceId, String> = HashMap::new();
62    for span in buf.iter() {
63        for attr in &span.attributes {
64            if attr.key == SCOUTER_EVAL_SCENARIO_ID_ATTR {
65                if let Some(sid) = attr.value.as_str() {
66                    if scenario_ids.contains(sid) {
67                        trace_to_scenario.insert(span.trace_id, sid.to_string());
68                    }
69                }
70                break;
71            }
72        }
73    }
74
75    // Pass 2: group all spans (including children) by their scenario
76    let mut grouped: HashMap<String, Vec<TraceSpanRecord>> = HashMap::new();
77    for span in buf.iter() {
78        if let Some(sid) = trace_to_scenario.get(&span.trace_id) {
79            grouped.entry(sid.clone()).or_default().push(span.clone());
80        }
81    }
82    grouped
83}
84
85/// Returns the set of trace IDs for any captured span that carries a
86/// `scouter.eval.scenario_id` attribute whose value matches one of the
87/// provided scenario IDs.
88pub fn get_trace_ids_by_scenario_ids(scenario_ids: &HashSet<String>) -> HashSet<ScouterTraceId> {
89    let buf = CAPTURE_BUFFER.read().unwrap_or_else(|p| p.into_inner());
90    buf.iter()
91        .filter(|span| {
92            span.attributes.iter().any(|attr| {
93                attr.key == SCOUTER_EVAL_SCENARIO_ID_ATTR
94                    && attr
95                        .value
96                        .as_str()
97                        .map(|v| scenario_ids.contains(v))
98                        .unwrap_or(false)
99            })
100        })
101        .map(|span| span.trace_id)
102        .collect()
103}
scouter_types/span_capture.rs

scouter_types/
span_capture.rs
