Skip to main content

rig_core/telemetry/
mod.rs

1//! This module primarily concerns being able to orchestrate telemetry across a given pipeline or workflow.
2//! This includes tracing, being able to send traces to an OpenTelemetry collector, setting up your
3//! agents with the correct tracing style so you can emit the right traces for platforms like Langfuse,
4//! and more.
5
6use crate::OneOrMany;
7use crate::completion::{AssistantContent, GetTokenUsage, Message};
8use crate::message::{
9    DocumentSourceKind, Image, MimeType, Reasoning, ReasoningContent, ToolResult,
10    ToolResultContent, UserContent,
11};
12use base64::Engine;
13use serde::Serialize;
14use std::collections::HashSet;
15use std::sync::{LazyLock, Mutex};
16use tracing::callsite::Identifier;
17
18/// Macro implementation dependency; public because exported macro expansions
19/// must be able to resolve it from downstream crates.
20#[doc(hidden)]
21pub use tracing as __tracing;
22
23/// Marks a span field as declared but not yet valued.
24///
25/// Re-exported so a runtime can declare a contract field through
26/// [`completion_parent_span!`](crate::completion_parent_span) without taking a
27/// direct `tracing` dependency — the crate's own `__tracing` path is hidden,
28/// and `Option::<&str>::None` is the only other portable spelling. `rig-core`
29/// already exposes `tracing` types across this module's public API
30/// ([`CompletionSpanBuilder::build`] returns a [`tracing::Span`]), so this adds
31/// no new semver surface.
32pub use tracing::field::Empty;
33
34/// Implementation detail of [`new_completion_span!`] and
35/// [`completion_parent_span!`](crate::completion_parent_span): declares a span
36/// with the caller's header fields followed by the canonical completion
37/// telemetry fields recorded over a completion's lifetime.
38///
39/// `tracing` bakes a span's field set into static metadata, so the canonical
40/// list can only be single-sourced by a macro that owns the whole
41/// `info_span!` invocation — a `const` list can never be spliced in. This is
42/// the one copy; [`COMPLETION_PARENT_REQUIRED_FIELDS`] is the checklist form
43/// of the same contract and a test asserts the two agree exactly.
44#[doc(hidden)]
45#[macro_export]
46macro_rules! __rig_canonical_completion_span {
47    (
48        target: $target:literal,
49        $(parent: $parent:expr,)?
50        name: $name:literal,
51        // Both blocks are spliced verbatim into `info_span!`: the header block
52        // must end with a trailing comma, the extras block must begin with one.
53        // Violating either surfaces as an `info_span!` parse error at the call
54        // site, not here.
55        { $($header:tt)* }
56        { $($extra:tt)* }
57    ) => {
58        $crate::telemetry::__tracing::info_span!(
59            target: $target,
60            $(parent: $parent,)?
61            $name,
62            $($header)*
63            gen_ai.response.id = $crate::telemetry::__tracing::field::Empty,
64            gen_ai.response.model = $crate::telemetry::__tracing::field::Empty,
65            gen_ai.usage.input_tokens = $crate::telemetry::__tracing::field::Empty,
66            gen_ai.usage.output_tokens = $crate::telemetry::__tracing::field::Empty,
67            gen_ai.usage.cache_read.input_tokens = $crate::telemetry::__tracing::field::Empty,
68            gen_ai.usage.cache_creation.input_tokens = $crate::telemetry::__tracing::field::Empty,
69            gen_ai.usage.tool_use_prompt_tokens = $crate::telemetry::__tracing::field::Empty,
70            gen_ai.usage.reasoning_tokens = $crate::telemetry::__tracing::field::Empty,
71            gen_ai.input.messages = $crate::telemetry::__tracing::field::Empty,
72            gen_ai.output.messages = $crate::telemetry::__tracing::field::Empty
73            $($extra)*
74        )
75    };
76}
77
78macro_rules! new_completion_span {
79    ($name:literal, $provider:expr, $request_model:expr, $operation:expr, $system:expr) => {
80        $crate::__rig_canonical_completion_span!(
81            target: "rig::completions",
82            name: $name,
83            {
84                gen_ai.operation.name = $operation,
85                gen_ai.provider.name = $provider,
86                gen_ai.request.model = $request_model,
87                gen_ai.system_instructions = $system,
88            }
89            {}
90        )
91    };
92}
93
94/// A supported GenAI completion operation and its canonical span name.
95#[derive(Debug, Clone, Copy, PartialEq, Eq)]
96#[non_exhaustive]
97pub enum CompletionOperation {
98    /// A chat completion.
99    Chat,
100    /// A streaming chat completion.
101    ChatStreaming,
102    /// A Gemini generate-content request.
103    GenerateContent,
104    /// A Gemini Interactions API request.
105    Interactions,
106    /// A streaming Gemini Interactions API request.
107    InteractionsStreaming,
108}
109
110impl CompletionOperation {
111    fn as_str(self) -> &'static str {
112        match self {
113            Self::Chat => "chat",
114            Self::ChatStreaming => "chat_streaming",
115            Self::GenerateContent => "generate_content",
116            Self::Interactions => "interactions",
117            Self::InteractionsStreaming => "interactions_streaming",
118        }
119    }
120}
121
122/// Core-owned marker field for a runtime span that may absorb provider completion fields.
123///
124/// Runtimes declare this field on their completion-parent spans without having
125/// to share a tracing target. The field's presence (its value is ignored) marks
126/// the span as a *candidate* for adoption, but adoption also requires the span
127/// to statically declare every field in [`COMPLETION_PARENT_REQUIRED_FIELDS`].
128///
129/// This second requirement exists because [`tracing::Span::record`] silently
130/// no-ops for any field absent from a span's static metadata: a span carrying
131/// only the marker would be adopted and then drop every recorded completion
132/// field, losing telemetry with no error. A span missing any required field is
133/// therefore *not* adopted — [`CompletionSpanBuilder::build`] creates a fresh
134/// `rig::completions` child span instead, so telemetry is never silently lost.
135///
136/// The declarative source of the contract is the
137/// [`completion_parent_span!`](crate::completion_parent_span) macro; declaring
138/// a span through it is what guarantees the marker and the field set stay in
139/// agreement. Changing [`COMPLETION_PARENT_REQUIRED_FIELDS`] therefore needs no
140/// marker change to degrade safely: a hand-written span built against an older
141/// field list simply fails the gate above and gets a fresh child span, with a
142/// warning (once per offending callsite) naming what it is missing.
143pub const COMPLETION_PARENT_MARKER_FIELD: &str = "rig.completion_parent";
144
145/// Fields a completion-parent span must statically declare (as
146/// [`tracing::field::Empty`] or a value) to be adopted by
147/// [`CompletionSpanBuilder::build`], alongside [`COMPLETION_PARENT_MARKER_FIELD`].
148///
149/// These mirror the canonical fields the builder's own `rig::completions` span
150/// declares, so an adopted span can absorb the request, response, usage, and
151/// content telemetry recorded over a completion's lifetime without dropping any.
152///
153/// This constant is the *checklist* form of the contract used by the adoption
154/// gate; the *declarative* form is the
155/// [`completion_parent_span!`](crate::completion_parent_span) macro, which is
156/// how a runtime declares a conforming span. A test asserts the two forms
157/// agree exactly, so neither can drift from the other.
158pub const COMPLETION_PARENT_REQUIRED_FIELDS: &[&str] = &[
159    "gen_ai.operation.name",
160    "gen_ai.provider.name",
161    "gen_ai.request.model",
162    "gen_ai.system_instructions",
163    "gen_ai.response.id",
164    "gen_ai.response.model",
165    "gen_ai.usage.input_tokens",
166    "gen_ai.usage.output_tokens",
167    "gen_ai.usage.cache_read.input_tokens",
168    "gen_ai.usage.cache_creation.input_tokens",
169    "gen_ai.usage.tool_use_prompt_tokens",
170    "gen_ai.usage.reasoning_tokens",
171    "gen_ai.input.messages",
172    "gen_ai.output.messages",
173];
174
175/// Declare a completion-parent span conforming to the adoption contract.
176///
177/// This macro is the single declarative source of the contract: it declares
178/// [`COMPLETION_PARENT_MARKER_FIELD`] and every field in
179/// [`COMPLETION_PARENT_REQUIRED_FIELDS`], so a span it builds is always
180/// adoptable by [`CompletionSpanBuilder::build`] and can absorb every field
181/// recorded over the completion's lifetime. Runtimes should declare their
182/// completion-parent spans through it rather than hand-writing the marker and
183/// field list — [`tracing::Span::record`] silently no-ops on undeclared
184/// fields, so a hand-written span that omits one field loses that telemetry
185/// with no error (and a span that omits enough to fail the adoption gate is
186/// not adopted at all). A span declared through this macro tracks the contract
187/// automatically; a hand-written one has to be maintained by hand.
188///
189/// By default the span is explicitly parented on
190/// [`tracing::Span::current()`]; pass `parent: <expr>` between `target` and
191/// `name` to override it with any parent expression [`tracing::info_span!`]
192/// accepts (including `None`). `operation` and `system_instructions` are
193/// declared with the given values; the provider records
194/// `gen_ai.provider.name` and `gen_ai.request.model` at adoption time.
195/// Additional runtime-specific fields may be appended after the named
196/// arguments. Extra fields must not repeat the marker or a required contract
197/// field: a duplicate compiles, but the span then declares two same-named
198/// fields and [`tracing::Span::record`] targets only the first.
199///
200/// Invoking this macro does not require a direct `tracing` dependency. To
201/// declare a field with no value yet, use [`Empty`] (re-exported here for
202/// exactly this reason) or `Option::<&str>::None`; [`tracing::field::Empty`]
203/// itself is the same type but only nameable by a crate that depends on
204/// `tracing` directly.
205///
206/// # Examples
207///
208/// ```
209/// use rig_core::telemetry::completion_parent_span;
210///
211/// let span = completion_parent_span!(
212///     target: "my_runtime",
213///     name: "chat",
214///     operation: "chat",
215///     system_instructions: Option::<&str>::None,
216///     gen_ai.agent.name = "assistant",
217/// );
218/// ```
219#[macro_export]
220macro_rules! completion_parent_span {
221    (
222        target: $target:literal,
223        parent: $parent:expr,
224        name: $name:literal,
225        operation: $operation:expr,
226        system_instructions: $system:expr
227        $(, $($extra:tt)*)?
228    ) => {
229        $crate::__rig_canonical_completion_span!(
230            target: $target,
231            parent: $parent,
232            name: $name,
233            {
234                rig.completion_parent = true,
235                gen_ai.operation.name = $operation,
236                gen_ai.system_instructions = $system,
237                gen_ai.provider.name = $crate::telemetry::__tracing::field::Empty,
238                gen_ai.request.model = $crate::telemetry::__tracing::field::Empty,
239            }
240            { $(, $($extra)*)? }
241        )
242    };
243    // Default arm: delegates to the explicit-parent arm so the two cannot
244    // drift in the fields they declare.
245    (
246        target: $target:literal,
247        name: $name:literal,
248        operation: $operation:expr,
249        system_instructions: $system:expr
250        $(, $($extra:tt)*)?
251    ) => {
252        $crate::completion_parent_span!(
253            target: $target,
254            parent: $crate::telemetry::__tracing::Span::current(),
255            name: $name,
256            operation: $operation,
257            system_instructions: $system
258            $(, $($extra)*)?
259        )
260    };
261}
262
263// `#[macro_export]` places the macro at the crate root; re-export it here so
264// it is also reachable at its documented home alongside the contract
265// constants it implements.
266pub use crate::completion_parent_span;
267
268/// What [`CompletionSpanBuilder::build`] decided about the span that is
269/// current when it runs.
270///
271/// This is the whole adoption decision table in one place, kept free of
272/// side effects so every row can be asserted directly. Emitting the
273/// diagnostics for it is [`warn_once_on_completion_parent_verdict`]'s job.
274#[derive(Debug, Clone, Copy, PartialEq, Eq)]
275enum CompletionParentVerdict {
276    /// The marker plus the complete contract — adopt and enrich.
277    Adopt,
278    /// The marker, but the span omits at least one field in
279    /// [`COMPLETION_PARENT_REQUIRED_FIELDS`]. The missing names are computed
280    /// only if a warning is emitted.
281    RejectMissingFields,
282    /// No marker at all: an ordinary ambient span that becomes the parent of a
283    /// fresh `rig::completions` child. Never warns.
284    NotAParent,
285}
286
287/// Fields in [`COMPLETION_PARENT_REQUIRED_FIELDS`] that `metadata` does not
288/// statically declare. Only called on the warning path — the adoption gate
289/// itself needs a yes/no, not the names, and allocating a `Vec` on every
290/// completion would be waste.
291fn missing_required_fields(metadata: &tracing::Metadata<'_>) -> Vec<&'static str> {
292    let fields = metadata.fields();
293    COMPLETION_PARENT_REQUIRED_FIELDS
294        .iter()
295        .copied()
296        .filter(|name| fields.field(name).is_none())
297        .collect()
298}
299
300/// Classify the span `metadata` as a completion parent. Pure: no logging, no
301/// global state — see [`CompletionParentVerdict`] for the decision table.
302fn classify_completion_parent(metadata: &tracing::Metadata<'_>) -> CompletionParentVerdict {
303    let fields = metadata.fields();
304    // Exact match, never a prefix: a runtime field that merely starts with the
305    // marker name (`rig.completion_parent.id`, say) is not the marker and must
306    // not make its span a rejected parent.
307    if fields.field(COMPLETION_PARENT_MARKER_FIELD).is_none() {
308        return CompletionParentVerdict::NotAParent;
309    }
310    if COMPLETION_PARENT_REQUIRED_FIELDS
311        .iter()
312        .all(|name| fields.field(name).is_some())
313    {
314        CompletionParentVerdict::Adopt
315    } else {
316        CompletionParentVerdict::RejectMissingFields
317    }
318}
319
320/// Near-miss parent callsites already reported.
321///
322/// Keyed per callsite rather than by one process-wide flag: two runtimes can
323/// each declare a near-miss parent, and a single flag would report whichever
324/// won the race and stay silent about the other — while the `missing_fields`
325/// list it printed describes only that one span. Callsites are static, so this
326/// set is bounded by the number of distinct near-miss spans in the program
327/// (normally zero).
328static NEAR_MISS_WARNED: LazyLock<Mutex<HashSet<Identifier>>> =
329    LazyLock::new(|| Mutex::new(HashSet::new()));
330
331/// Clear the per-callsite warn budget.
332///
333/// The budget is process-global, so without this the warning tests are coupled:
334/// whichever runs first consumes the budget for any callsite they share, and the
335/// other sees silence. `cargo nextest` hides that (one process per test) while
336/// `cargo test` exposes it, so the coupling would be green in CI and red
337/// locally — the worst orientation for a latent test bug. Resetting makes each
338/// test independent of callsite identity, ordering, and runner.
339#[cfg(test)]
340fn reset_near_miss_warnings() {
341    NEAR_MISS_WARNED
342        .lock()
343        .unwrap_or_else(|poisoned| poisoned.into_inner())
344        .clear();
345}
346
347/// Surface a verdict that a human should act on, once per offending callsite.
348///
349/// A rejected near-miss degrades safely — the fresh child span loses no
350/// telemetry — but silently, so the operator's only symptom would be a
351/// duplicated span layer in dashboards.
352///
353/// Neither a `Once` nor a poison-propagating lock: a subscriber panic must not
354/// turn a diagnostic into a hard failure on every subsequent completion, and a
355/// dedup set stays perfectly valid across an unrelated panic.
356fn warn_once_on_completion_parent_verdict(
357    verdict: CompletionParentVerdict,
358    metadata: &tracing::Metadata<'_>,
359) {
360    match verdict {
361        CompletionParentVerdict::Adopt | CompletionParentVerdict::NotAParent => {}
362        CompletionParentVerdict::RejectMissingFields => {
363            let first_sighting = {
364                let mut warned = NEAR_MISS_WARNED
365                    .lock()
366                    .unwrap_or_else(|poisoned| poisoned.into_inner());
367                warned.insert(metadata.callsite())
368            };
369            // The guard is released above, before `warn!`, and that scoping is
370            // load-bearing: `warn!` dispatches into arbitrary subscriber code,
371            // and a subscriber that transitively reaches `build` would
372            // deadlock re-entering this non-reentrant lock.
373            if !first_sighting {
374                return;
375            }
376            tracing::warn!(
377                marker = COMPLETION_PARENT_MARKER_FIELD,
378                missing_fields = ?missing_required_fields(metadata),
379                "completion-parent span declares the marker but not every required field \
380                 and is not adopted; provider telemetry lands on a fresh child span \
381                 instead — declare the span with \
382                 `rig_core::telemetry::completion_parent_span!`"
383            );
384        }
385    }
386}
387
388/// Builder for a canonical GenAI completion span.
389///
390/// Runtime spans declaring [`COMPLETION_PARENT_MARKER_FIELD`] and the
391/// [`COMPLETION_PARENT_REQUIRED_FIELDS`] are enriched and reused so one model
392/// turn has exactly one model span. Other ambient spans — and marker spans that
393/// omit a required field — remain parents of a newly created `rig::completions`
394/// span.
395pub struct CompletionSpanBuilder<'a> {
396    provider: &'a str,
397    request_model: &'a str,
398    operation: CompletionOperation,
399    system_instructions: Option<String>,
400}
401
402impl<'a> CompletionSpanBuilder<'a> {
403    /// Create a completion-span builder for a provider request.
404    pub fn new(provider: &'a str, request_model: &'a str, operation: CompletionOperation) -> Self {
405        Self {
406            provider,
407            request_model,
408            operation,
409            system_instructions: None,
410        }
411    }
412
413    /// Set the system instructions sent with the request when sensitive content
414    /// telemetry has been explicitly enabled.
415    pub fn system_instructions(
416        mut self,
417        system_instructions: Option<&'a str>,
418        record_content: bool,
419    ) -> Self {
420        self.system_instructions = system_instructions_json(system_instructions, record_content);
421        self
422    }
423
424    /// Build a canonical completion span or enrich Rig's current completion-parent span.
425    pub fn build(self) -> tracing::Span {
426        let current = tracing::Span::current();
427        if let Some(metadata) = current.metadata() {
428            let verdict = classify_completion_parent(metadata);
429            warn_once_on_completion_parent_verdict(verdict, metadata);
430            if verdict == CompletionParentVerdict::Adopt {
431                current.record("gen_ai.operation.name", self.operation.as_str());
432                current.record("gen_ai.provider.name", self.provider);
433                current.record("gen_ai.request.model", self.request_model);
434                if let Some(system_instructions) = self.system_instructions.as_deref() {
435                    current.record("gen_ai.system_instructions", system_instructions);
436                }
437                return current;
438            }
439        }
440
441        let operation = self.operation.as_str();
442        let system_instructions = self.system_instructions.as_deref();
443        match self.operation {
444            CompletionOperation::Chat => new_completion_span!(
445                "chat",
446                self.provider,
447                self.request_model,
448                operation,
449                system_instructions
450            ),
451            CompletionOperation::ChatStreaming => new_completion_span!(
452                "chat_streaming",
453                self.provider,
454                self.request_model,
455                operation,
456                system_instructions
457            ),
458            CompletionOperation::GenerateContent => new_completion_span!(
459                "generate_content",
460                self.provider,
461                self.request_model,
462                operation,
463                system_instructions
464            ),
465            CompletionOperation::Interactions => new_completion_span!(
466                "interactions",
467                self.provider,
468                self.request_model,
469                operation,
470                system_instructions
471            ),
472            CompletionOperation::InteractionsStreaming => new_completion_span!(
473                "interactions_streaming",
474                self.provider,
475                self.request_model,
476                operation,
477                system_instructions
478            ),
479        }
480    }
481}
482
483#[derive(Serialize)]
484struct TelemetryChatMessage {
485    role: &'static str,
486    parts: Vec<TelemetryPart>,
487}
488
489#[derive(Serialize)]
490struct TelemetryOutputMessage {
491    role: &'static str,
492    parts: Vec<TelemetryPart>,
493    finish_reason: &'static str,
494}
495
496#[derive(Serialize)]
497#[serde(tag = "type", rename_all = "snake_case")]
498enum TelemetryPart {
499    Text {
500        content: String,
501    },
502    ToolCall {
503        #[serde(skip_serializing_if = "Option::is_none")]
504        id: Option<String>,
505        name: String,
506        arguments: serde_json::Value,
507    },
508    ToolCallResponse {
509        #[serde(skip_serializing_if = "Option::is_none")]
510        id: Option<String>,
511        response: serde_json::Value,
512    },
513    Reasoning {
514        content: String,
515    },
516    Uri {
517        #[serde(skip_serializing_if = "Option::is_none")]
518        mime_type: Option<String>,
519        modality: &'static str,
520        uri: String,
521    },
522    File {
523        #[serde(skip_serializing_if = "Option::is_none")]
524        mime_type: Option<String>,
525        modality: &'static str,
526        file_id: String,
527    },
528    Blob {
529        #[serde(skip_serializing_if = "Option::is_none")]
530        mime_type: Option<String>,
531        modality: &'static str,
532        content: String,
533    },
534}
535
536fn media_part<T>(
537    data: &DocumentSourceKind,
538    media_type: Option<&T>,
539    modality: &'static str,
540) -> Option<TelemetryPart>
541where
542    T: MimeType,
543{
544    let mime_type = media_type.map(|media_type| media_type.to_mime_type().to_string());
545    match data {
546        DocumentSourceKind::Url(uri) => Some(TelemetryPart::Uri {
547            mime_type,
548            modality,
549            uri: uri.clone(),
550        }),
551        DocumentSourceKind::FileId(file_id) => Some(TelemetryPart::File {
552            mime_type,
553            modality,
554            file_id: file_id.clone(),
555        }),
556        DocumentSourceKind::Base64(content) => Some(TelemetryPart::Blob {
557            mime_type,
558            modality,
559            content: content.clone(),
560        }),
561        DocumentSourceKind::Raw(content) => Some(TelemetryPart::Blob {
562            mime_type,
563            modality,
564            content: base64::engine::general_purpose::STANDARD.encode(content),
565        }),
566        DocumentSourceKind::String(content) => Some(TelemetryPart::Text {
567            content: content.clone(),
568        }),
569        DocumentSourceKind::Unknown => None,
570    }
571}
572
573fn image_part(image: &Image) -> Option<TelemetryPart> {
574    media_part(&image.data, image.media_type.as_ref(), "image")
575}
576
577fn reasoning_parts(reasoning: &Reasoning) -> Vec<TelemetryPart> {
578    reasoning
579        .content
580        .iter()
581        .map(|content| {
582            let content = match content {
583                ReasoningContent::Text { text, .. } | ReasoningContent::Summary(text) => text,
584                ReasoningContent::Encrypted(content) => content,
585                ReasoningContent::Redacted { data } => data,
586            };
587            TelemetryPart::Reasoning {
588                content: content.clone(),
589            }
590        })
591        .collect()
592}
593
594fn tool_result_response(result: &ToolResult) -> serde_json::Value {
595    let mut content = result
596        .content
597        .iter()
598        .filter_map(|content| match content {
599            ToolResultContent::Text(text) => Some(serde_json::Value::String(text.text.clone())),
600            ToolResultContent::Json { value } => Some(value.clone()),
601            ToolResultContent::Image(image) => {
602                image_part(image).and_then(|part| serde_json::to_value(part).ok())
603            }
604        })
605        .collect::<Vec<_>>();
606
607    if content.len() == 1 {
608        content.pop().unwrap_or(serde_json::Value::Null)
609    } else {
610        serde_json::Value::Array(content)
611    }
612}
613
614fn user_parts(content: &OneOrMany<UserContent>) -> Vec<TelemetryPart> {
615    content
616        .iter()
617        .filter_map(|content| match content {
618            UserContent::Text(text) => Some(TelemetryPart::Text {
619                content: text.text.clone(),
620            }),
621            UserContent::ToolResult(result) => Some(TelemetryPart::ToolCallResponse {
622                id: Some(result.id.clone()),
623                response: tool_result_response(result),
624            }),
625            UserContent::Image(image) => image_part(image),
626            UserContent::Audio(audio) => {
627                media_part(&audio.data, audio.media_type.as_ref(), "audio")
628            }
629            UserContent::Video(video) => {
630                media_part(&video.data, video.media_type.as_ref(), "video")
631            }
632            UserContent::Document(document) => {
633                media_part(&document.data, document.media_type.as_ref(), "document")
634            }
635        })
636        .collect()
637}
638
639fn assistant_parts(content: &OneOrMany<AssistantContent>) -> Vec<TelemetryPart> {
640    content
641        .iter()
642        .flat_map(|content| match content {
643            AssistantContent::Text(text) => vec![TelemetryPart::Text {
644                content: text.text.clone(),
645            }],
646            AssistantContent::ToolCall(tool_call) => vec![TelemetryPart::ToolCall {
647                id: Some(tool_call.id.clone()),
648                name: tool_call.function.name.clone(),
649                arguments: tool_call.function.arguments.clone(),
650            }],
651            AssistantContent::Reasoning(reasoning) => reasoning_parts(reasoning),
652            AssistantContent::Image(image) => image_part(image).into_iter().collect(),
653        })
654        .collect()
655}
656
657fn input_messages(messages: &[Message]) -> Vec<TelemetryChatMessage> {
658    messages
659        .iter()
660        .map(|message| match message {
661            Message::System { content } => TelemetryChatMessage {
662                role: "system",
663                parts: vec![TelemetryPart::Text {
664                    content: content.clone(),
665                }],
666            },
667            Message::User { content } => TelemetryChatMessage {
668                role: "user",
669                parts: user_parts(content),
670            },
671            Message::Assistant { content, .. } => TelemetryChatMessage {
672                role: "assistant",
673                parts: assistant_parts(content),
674            },
675        })
676        .collect()
677}
678
679fn output_messages(content: &OneOrMany<AssistantContent>) -> Vec<TelemetryOutputMessage> {
680    let finish_reason = if content
681        .iter()
682        .any(|content| matches!(content, AssistantContent::ToolCall(_)))
683    {
684        "tool_call"
685    } else {
686        // Rig's normalized assistant content does not retain provider finish
687        // reasons such as length or content filtering. Avoid claiming a clean
688        // stop when the actual reason is unavailable.
689        "unknown"
690    };
691    vec![TelemetryOutputMessage {
692        role: "assistant",
693        parts: assistant_parts(content),
694        finish_reason,
695    }]
696}
697
698/// Serializes system instructions using the normalized GenAI telemetry shape.
699pub fn system_instructions_json(instructions: Option<&str>, enabled: bool) -> Option<String> {
700    if !enabled {
701        return None;
702    }
703
704    instructions.and_then(|instructions| {
705        serde_json::to_string(&vec![TelemetryPart::Text {
706            content: instructions.to_string(),
707        }])
708        .ok()
709    })
710}
711
712/// Records serialized model input messages on `gen_ai.input.messages` when
713/// content telemetry is explicitly enabled.
714///
715/// Message content can contain prompts, retrieved context, tool results, and
716/// other sensitive or high-cardinality data. Keep this disabled unless the
717/// caller has explicitly opted in for debugging/observability.
718pub fn record_model_input(span: &tracing::Span, messages: &[Message], enabled: bool) {
719    if !enabled || span.is_disabled() {
720        return;
721    }
722
723    if let Ok(messages) = serde_json::to_string(&input_messages(messages)) {
724        span.record("gen_ai.input.messages", messages);
725    }
726}
727
728/// Records serialized model output messages on `gen_ai.output.messages` when
729/// content telemetry is explicitly enabled.
730///
731/// Message content can contain model responses, tool calls, and other sensitive
732/// or high-cardinality data. Keep this disabled unless the caller has explicitly
733/// opted in for debugging/observability.
734pub fn record_model_output(
735    span: &tracing::Span,
736    content: &OneOrMany<AssistantContent>,
737    enabled: bool,
738) {
739    if !enabled || span.is_disabled() {
740        return;
741    }
742
743    let messages = output_messages(content);
744    if let Ok(messages) = serde_json::to_string(&messages) {
745        span.record("gen_ai.output.messages", messages);
746    }
747}
748
749/// Provider response metadata used to populate GenAI telemetry spans.
750pub trait ProviderResponseExt {
751    /// Provider-native output message type.
752    type OutputMessage: Serialize;
753    /// Provider-native usage type.
754    type Usage: Serialize;
755
756    /// Returns the provider response ID, if supplied.
757    fn get_response_id(&self) -> Option<String>;
758
759    /// Returns the provider response model name, if supplied.
760    fn get_response_model_name(&self) -> Option<String>;
761
762    /// Returns serialized output messages produced by the provider.
763    fn get_output_messages(&self) -> Vec<Self::OutputMessage>;
764
765    /// Returns the primary text response, when available.
766    fn get_text_response(&self) -> Option<String>;
767
768    /// Returns provider-native usage metrics, if supplied.
769    fn get_usage(&self) -> Option<Self::Usage>;
770}
771
772/// A trait designed specifically to be used with Spans for the purpose of recording telemetry.
773/// Implemented for [`tracing::Span`] to record GenAI semantic convention fields.
774pub trait SpanCombinator {
775    /// Record Rig-normalized token usage fields on the span.
776    fn record_token_usage<U>(&self, usage: &U)
777    where
778        U: GetTokenUsage;
779
780    /// Record provider response metadata such as response ID and model name.
781    fn record_response_metadata<R>(&self, response: &R)
782    where
783        R: ProviderResponseExt;
784}
785
786impl SpanCombinator for tracing::Span {
787    fn record_token_usage<U>(&self, usage: &U)
788    where
789        U: GetTokenUsage,
790    {
791        if self.is_disabled() {
792            return;
793        }
794
795        let usage = usage.token_usage();
796        // Zero-valued usage is the documented sentinel for missing provider
797        // usage metrics; leave the span fields unset.
798        if usage.has_values() {
799            self.record("gen_ai.usage.input_tokens", usage.input_tokens);
800            self.record("gen_ai.usage.output_tokens", usage.output_tokens);
801            self.record(
802                "gen_ai.usage.cache_read.input_tokens",
803                usage.cached_input_tokens,
804            );
805            self.record(
806                "gen_ai.usage.cache_creation.input_tokens",
807                usage.cache_creation_input_tokens,
808            );
809            self.record(
810                "gen_ai.usage.tool_use_prompt_tokens",
811                usage.tool_use_prompt_tokens,
812            );
813            self.record("gen_ai.usage.reasoning_tokens", usage.reasoning_tokens);
814        }
815    }
816
817    fn record_response_metadata<R>(&self, response: &R)
818    where
819        R: ProviderResponseExt,
820    {
821        if self.is_disabled() {
822            return;
823        }
824
825        if let Some(id) = response.get_response_id() {
826            self.record("gen_ai.response.id", id);
827        }
828
829        if let Some(model_name) = response.get_response_model_name() {
830            self.record("gen_ai.response.model", model_name);
831        }
832    }
833}
834
835#[cfg(test)]
836mod tests {
837    use super::*;
838    use crate::completion::{AssistantContent, GetTokenUsage, Message, Usage};
839    use serde_json::json;
840    use std::sync::{Arc, Mutex};
841    use tracing::field::{Field, Visit};
842    use tracing::{Id, Subscriber};
843    use tracing_subscriber::layer::{Context, SubscriberExt};
844    use tracing_subscriber::{Layer, Registry, registry::LookupSpan};
845
846    #[derive(Clone)]
847    struct TestUsage(Usage);
848
849    impl GetTokenUsage for TestUsage {
850        fn token_usage(&self) -> Usage {
851            self.0
852        }
853    }
854
855    #[test]
856    fn content_attributes_follow_gen_ai_semantic_convention_json_shapes() {
857        assert_eq!(
858            system_instructions_json(Some("follow policy"), true).as_deref(),
859            Some(r#"[{"type":"text","content":"follow policy"}]"#)
860        );
861        assert_eq!(system_instructions_json(Some("secret"), false), None);
862
863        let input = input_messages(&[
864            Message::system("follow policy"),
865            Message::user("hello"),
866            Message::tool_result("call_1", "sunny"),
867        ]);
868        assert_eq!(
869            serde_json::to_value(input).expect("semantic-convention input DTOs serialize"),
870            json!([
871                {
872                    "role": "system",
873                    "parts": [{"type": "text", "content": "follow policy"}]
874                },
875                {
876                    "role": "user",
877                    "parts": [{"type": "text", "content": "hello"}]
878                },
879                {
880                    "role": "user",
881                    "parts": [{
882                        "type": "tool_call_response",
883                        "id": "call_1",
884                        "response": "sunny"
885                    }]
886                }
887            ])
888        );
889
890        let output = OneOrMany::one(AssistantContent::tool_call(
891            "call_1",
892            "weather",
893            json!({"city": "Paris"}),
894        ));
895        assert_eq!(
896            serde_json::to_value(output_messages(&output))
897                .expect("semantic-convention output DTOs serialize"),
898            json!([{
899                "role": "assistant",
900                "parts": [{
901                    "type": "tool_call",
902                    "id": "call_1",
903                    "name": "weather",
904                    "arguments": {"city": "Paris"}
905                }],
906                "finish_reason": "tool_call"
907            }])
908        );
909
910        let text_output = OneOrMany::one(AssistantContent::text("done"));
911        assert_eq!(
912            serde_json::to_value(output_messages(&text_output))
913                .expect("semantic-convention text output DTOs serialize"),
914            json!([{
915                "role": "assistant",
916                "parts": [{"type": "text", "content": "done"}],
917                "finish_reason": "unknown"
918            }])
919        );
920    }
921
922    #[derive(Clone, Default)]
923    struct CapturedFields(Arc<Mutex<Vec<(String, u64)>>>);
924
925    impl CapturedFields {
926        fn push(&self, name: &str, value: u64) {
927            if let Ok(mut fields) = self.0.lock() {
928                fields.push((name.to_string(), value));
929            }
930        }
931
932        fn contains(&self, name: &str, value: u64) -> bool {
933            self.0.lock().is_ok_and(|fields| {
934                fields
935                    .iter()
936                    .any(|field| field == &(name.to_string(), value))
937            })
938        }
939    }
940
941    struct FieldCaptureLayer {
942        fields: CapturedFields,
943    }
944
945    impl<S> Layer<S> for FieldCaptureLayer
946    where
947        S: Subscriber,
948        S: for<'lookup> LookupSpan<'lookup>,
949    {
950        fn on_record(&self, _span: &Id, values: &tracing::span::Record<'_>, _ctx: Context<'_, S>) {
951            values.record(&mut FieldCaptureVisitor {
952                fields: self.fields.clone(),
953            });
954        }
955    }
956
957    struct FieldCaptureVisitor {
958        fields: CapturedFields,
959    }
960
961    impl Visit for FieldCaptureVisitor {
962        fn record_u64(&mut self, field: &Field, value: u64) {
963            self.fields.push(field.name(), value);
964        }
965
966        fn record_debug(&mut self, _field: &Field, _value: &dyn std::fmt::Debug) {}
967    }
968
969    /// WARN-level events, rendered as `field=value` pairs joined with the
970    /// event's message, in emission order.
971    #[derive(Clone, Default)]
972    struct CapturedWarnings(Arc<Mutex<Vec<String>>>);
973
974    impl CapturedWarnings {
975        fn push(&self, rendered: String) {
976            if let Ok(mut events) = self.0.lock() {
977                events.push(rendered);
978            }
979        }
980
981        /// Drains, so a test can assert on one phase and then assert that a
982        /// later phase added nothing. A cloning read would make the second
983        /// assertion see the first phase's events and quietly fail.
984        fn take(&self) -> Vec<String> {
985            self.0
986                .lock()
987                .map(|mut events| std::mem::take(&mut *events))
988                .unwrap_or_default()
989        }
990    }
991
992    struct WarningCaptureLayer {
993        warnings: CapturedWarnings,
994    }
995
996    impl<S> Layer<S> for WarningCaptureLayer
997    where
998        S: Subscriber,
999        S: for<'lookup> LookupSpan<'lookup>,
1000    {
1001        fn on_event(&self, event: &tracing::Event<'_>, _ctx: Context<'_, S>) {
1002            if *event.metadata().level() != tracing::Level::WARN {
1003                return;
1004            }
1005            let mut visitor = WarningCaptureVisitor::default();
1006            event.record(&mut visitor);
1007            self.warnings.push(visitor.rendered);
1008        }
1009    }
1010
1011    #[derive(Default)]
1012    struct WarningCaptureVisitor {
1013        rendered: String,
1014    }
1015
1016    impl Visit for WarningCaptureVisitor {
1017        fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
1018            use std::fmt::Write;
1019
1020            // `missing_fields = ?vec` arrives here; the message itself arrives
1021            // as the reserved `message` field. Both matter to the assertions,
1022            // so render every field rather than special-casing.
1023            let _ = write!(&mut self.rendered, " {}={value:?}", field.name());
1024        }
1025
1026        fn record_str(&mut self, field: &Field, value: &str) {
1027            use std::fmt::Write;
1028
1029            let _ = write!(&mut self.rendered, " {}={value}", field.name());
1030        }
1031    }
1032
1033    #[derive(Clone, Default)]
1034    struct CapturedSpan(Arc<Mutex<Option<CapturedSpanData>>>);
1035
1036    struct CapturedSpanData {
1037        name: String,
1038        target: String,
1039        parent_name: Option<String>,
1040        fields: Vec<String>,
1041        initial_values: Vec<(String, String)>,
1042        recorded_values: Vec<(String, String)>,
1043    }
1044
1045    struct SpanCaptureLayer {
1046        span: CapturedSpan,
1047    }
1048
1049    #[derive(Default)]
1050    struct StringFieldVisitor {
1051        values: Vec<(String, String)>,
1052    }
1053
1054    impl Visit for StringFieldVisitor {
1055        fn record_str(&mut self, field: &Field, value: &str) {
1056            self.values
1057                .push((field.name().to_owned(), value.to_owned()));
1058        }
1059
1060        fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
1061            self.values
1062                .push((field.name().to_owned(), format!("{value:?}")));
1063        }
1064    }
1065
1066    impl<S> Layer<S> for SpanCaptureLayer
1067    where
1068        S: Subscriber,
1069        S: for<'lookup> LookupSpan<'lookup>,
1070    {
1071        fn on_new_span(
1072            &self,
1073            attrs: &tracing::span::Attributes<'_>,
1074            _id: &Id,
1075            ctx: Context<'_, S>,
1076        ) {
1077            if let Ok(mut captured) = self.span.0.lock() {
1078                let mut visitor = StringFieldVisitor::default();
1079                attrs.record(&mut visitor);
1080                let parent_name = if let Some(parent) = attrs.parent() {
1081                    ctx.span(parent)
1082                        .map(|span| span.metadata().name().to_owned())
1083                } else if attrs.is_contextual() {
1084                    ctx.lookup_current()
1085                        .map(|span| span.metadata().name().to_owned())
1086                } else {
1087                    None
1088                };
1089                *captured = Some(CapturedSpanData {
1090                    name: attrs.metadata().name().to_owned(),
1091                    target: attrs.metadata().target().to_owned(),
1092                    parent_name,
1093                    fields: attrs
1094                        .metadata()
1095                        .fields()
1096                        .iter()
1097                        .map(|field| field.name().to_owned())
1098                        .collect(),
1099                    initial_values: visitor.values,
1100                    recorded_values: Vec::new(),
1101                });
1102            }
1103        }
1104
1105        fn on_record(&self, _span: &Id, values: &tracing::span::Record<'_>, _ctx: Context<'_, S>) {
1106            if let Ok(mut captured) = self.span.0.lock()
1107                && let Some(captured) = captured.as_mut()
1108            {
1109                let mut visitor = StringFieldVisitor::default();
1110                values.record(&mut visitor);
1111                captured.recorded_values.extend(visitor.values);
1112            }
1113        }
1114    }
1115
1116    fn contains_string(values: &[(String, String)], field: &str, value: &str) -> bool {
1117        values
1118            .iter()
1119            .any(|candidate| candidate == &(field.to_owned(), value.to_owned()))
1120    }
1121
1122    #[test]
1123    fn completion_span_uses_canonical_names_fields_and_initial_attributes() {
1124        let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
1125
1126        for (operation, expected_name) in [
1127            (CompletionOperation::Chat, "chat"),
1128            (CompletionOperation::ChatStreaming, "chat_streaming"),
1129            (CompletionOperation::GenerateContent, "generate_content"),
1130            (CompletionOperation::Interactions, "interactions"),
1131            (
1132                CompletionOperation::InteractionsStreaming,
1133                "interactions_streaming",
1134            ),
1135        ] {
1136            let captured = CapturedSpan::default();
1137            let subscriber = Registry::default().with(SpanCaptureLayer {
1138                span: captured.clone(),
1139            });
1140            tracing::subscriber::with_default(subscriber, || {
1141                let span = CompletionSpanBuilder::new("openai", "gpt-5", operation)
1142                    .system_instructions(Some("system prompt"), true)
1143                    .build();
1144                assert!(!span.is_disabled());
1145            });
1146
1147            let Ok(captured) = captured.0.lock() else {
1148                panic!("captured span lock poisoned");
1149            };
1150            let Some(span) = captured.as_ref() else {
1151                panic!("completion span was not created");
1152            };
1153            assert_eq!(span.name, expected_name);
1154            assert_eq!(span.target, "rig::completions");
1155            assert_eq!(span.parent_name, None);
1156            for (field, value) in [
1157                ("gen_ai.operation.name", expected_name),
1158                ("gen_ai.provider.name", "openai"),
1159                ("gen_ai.request.model", "gpt-5"),
1160                (
1161                    "gen_ai.system_instructions",
1162                    r#"[{"type":"text","content":"system prompt"}]"#,
1163                ),
1164            ] {
1165                assert!(
1166                    contains_string(&span.initial_values, field, value),
1167                    "missing initial {field}={value}"
1168                );
1169            }
1170            assert!(span.recorded_values.is_empty());
1171            assert!(
1172                !span
1173                    .initial_values
1174                    .iter()
1175                    .any(|(field, _)| field == "gen_ai.response.model")
1176            );
1177            for field in COMPLETION_PARENT_REQUIRED_FIELDS {
1178                assert!(
1179                    span.fields.iter().any(|candidate| candidate == field),
1180                    "missing {field}"
1181                );
1182            }
1183        }
1184    }
1185
1186    /// The default arm parents on the ambient span, and an explicit `parent:`
1187    /// overrides it.
1188    ///
1189    /// A regression to `parent: None` in the default arm would root every
1190    /// completion-parent span, detaching it from the surrounding trace. No
1191    /// field-set assertion in this module can see that — the fields are
1192    /// identical either way — while an operator sees completion spans floating
1193    /// as roots instead of nesting under the agent span.
1194    #[test]
1195    fn completion_parent_span_macro_honours_its_parent_argument() {
1196        /// `SpanCaptureLayer` has no target filter and keeps only the most
1197        /// recent span, so read it immediately after the span under test is
1198        /// created, and confirm the target before trusting the parent.
1199        fn captured_parent(captured: &CapturedSpan) -> Option<String> {
1200            let Ok(captured) = captured.0.lock() else {
1201                panic!("captured span lock poisoned");
1202            };
1203            let Some(span) = captured.as_ref() else {
1204                panic!("completion-parent span was not captured");
1205            };
1206            assert_eq!(span.target, "third_party_runtime");
1207            span.parent_name.clone()
1208        }
1209
1210        let captured = CapturedSpan::default();
1211        let subscriber = Registry::default().with(SpanCaptureLayer {
1212            span: captured.clone(),
1213        });
1214        let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
1215        tracing::subscriber::with_default(subscriber, || {
1216            let ambient = tracing::info_span!(target: "application", "ambient");
1217
1218            // Default arm: nests under whatever span is current.
1219            ambient.in_scope(|| {
1220                let _default_arm = completion_parent_span!(
1221                    target: "third_party_runtime",
1222                    name: "chat",
1223                    operation: Empty,
1224                    system_instructions: Option::<&str>::None,
1225                );
1226            });
1227            assert_eq!(captured_parent(&captured).as_deref(), Some("ambient"));
1228
1229            // Explicit arm: the caller's parent wins over the ambient span, so
1230            // this one is a root despite `ambient` being current.
1231            ambient.in_scope(|| {
1232                let _explicit_arm = completion_parent_span!(
1233                    target: "third_party_runtime",
1234                    parent: None,
1235                    name: "chat",
1236                    operation: Empty,
1237                    system_instructions: Option::<&str>::None,
1238                );
1239            });
1240            assert_eq!(captured_parent(&captured), None);
1241        });
1242    }
1243
1244    #[test]
1245    fn unrelated_ambient_span_is_parent_not_adopted() {
1246        let captured = CapturedSpan::default();
1247        let subscriber = Registry::default().with(SpanCaptureLayer {
1248            span: captured.clone(),
1249        });
1250        let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
1251        tracing::subscriber::with_default(subscriber, || {
1252            let ambient = tracing::info_span!(target: "application", "ambient");
1253            let _guard = ambient.enter();
1254            let span =
1255                CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
1256            assert_ne!(span.id(), ambient.id());
1257        });
1258
1259        let Ok(captured) = captured.0.lock() else {
1260            panic!("captured span lock poisoned");
1261        };
1262        let Some(span) = captured.as_ref() else {
1263            panic!("completion span was not captured");
1264        };
1265        assert_eq!(span.target, "rig::completions");
1266        assert_eq!(span.parent_name.as_deref(), Some("ambient"));
1267    }
1268
1269    #[test]
1270    fn marker_span_missing_required_fields_is_not_adopted() {
1271        // A span that carries the marker but omits required canonical fields
1272        // (here: everything past `gen_ai.request.model`) must NOT be adopted.
1273        // Adopting it would silently drop the response/usage/content telemetry
1274        // that `Span::record` no-ops on for undeclared fields. Instead the
1275        // builder creates a fresh `rig::completions` child so nothing is lost.
1276        let captured = CapturedSpan::default();
1277        let subscriber = Registry::default().with(SpanCaptureLayer {
1278            span: captured.clone(),
1279        });
1280        let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
1281        tracing::subscriber::with_default(subscriber, || {
1282            // Deliberately hand-written (not `completion_parent_span!`): the
1283            // point is a marker span that fails to declare required fields.
1284            let partial_marker = tracing::info_span!(
1285                target: "third_party_runtime",
1286                "chat",
1287                rig.completion_parent = true,
1288                gen_ai.operation.name = tracing::field::Empty,
1289                gen_ai.provider.name = tracing::field::Empty,
1290                gen_ai.request.model = tracing::field::Empty,
1291            );
1292            // Premise check: the hand-written marker literal above must still
1293            // match the constant — if the marker is ever renamed this fails
1294            // first, pointing at the stale literal, so the test cannot keep
1295            // passing for the wrong reason (no marker at all, rather than the
1296            // marker with fields missing).
1297            let Some(metadata) = partial_marker.metadata() else {
1298                panic!("partial marker span was disabled");
1299            };
1300            assert!(
1301                metadata
1302                    .fields()
1303                    .field(COMPLETION_PARENT_MARKER_FIELD)
1304                    .is_some(),
1305                "hand-written marker literal is stale; update it to {COMPLETION_PARENT_MARKER_FIELD}"
1306            );
1307            let _guard = partial_marker.enter();
1308            let span =
1309                CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
1310            assert_ne!(span.id(), partial_marker.id());
1311        });
1312
1313        let Ok(captured) = captured.0.lock() else {
1314            panic!("captured span lock poisoned");
1315        };
1316        let Some(span) = captured.as_ref() else {
1317            panic!("completion span was not captured");
1318        };
1319        // A canonical child span is created and parented under the marker span,
1320        // and it carries the completion fields the marker span could not absorb.
1321        assert_eq!(span.target, "rig::completions");
1322        assert_eq!(span.parent_name.as_deref(), Some("chat"));
1323        for (field, value) in [
1324            ("gen_ai.operation.name", "chat"),
1325            ("gen_ai.provider.name", "openai"),
1326            ("gen_ai.request.model", "gpt-5"),
1327        ] {
1328            assert!(contains_string(&span.initial_values, field, value));
1329        }
1330    }
1331
1332    #[test]
1333    fn completion_parent_span_macro_matches_the_contract_exactly() {
1334        use std::collections::HashSet;
1335
1336        let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
1337        tracing::subscriber::with_default(Registry::default(), || {
1338            let span = completion_parent_span!(
1339                target: "contract_test",
1340                name: "chat",
1341                operation: "chat",
1342                system_instructions: Option::<&str>::None,
1343            );
1344            let Some(metadata) = span.metadata() else {
1345                panic!("contract span was disabled");
1346            };
1347            let declared: HashSet<&str> =
1348                metadata.fields().iter().map(|field| field.name()).collect();
1349            let expected: HashSet<&str> = COMPLETION_PARENT_REQUIRED_FIELDS
1350                .iter()
1351                .copied()
1352                .chain([COMPLETION_PARENT_MARKER_FIELD])
1353                .collect();
1354            // Exact equality in both directions: a field added to the macro
1355            // but not the constant (or vice versa) is drift, not a superset.
1356            assert_eq!(declared, expected);
1357            // Duplicate field names collapse in a `HashSet`, so also pin the
1358            // count: set equality alone cannot catch a field declared twice.
1359            assert_eq!(metadata.fields().len(), expected.len());
1360            assert_eq!(
1361                classify_completion_parent(metadata),
1362                CompletionParentVerdict::Adopt
1363            );
1364
1365            // The explicit-`parent:` arm declares the identical field set.
1366            let span = completion_parent_span!(
1367                target: "contract_test",
1368                parent: None,
1369                name: "chat",
1370                operation: "chat",
1371                system_instructions: Option::<&str>::None,
1372            );
1373            let Some(metadata) = span.metadata() else {
1374                panic!("contract span with explicit parent was disabled");
1375            };
1376            let declared: HashSet<&str> =
1377                metadata.fields().iter().map(|field| field.name()).collect();
1378            assert_eq!(declared, expected);
1379            assert_eq!(metadata.fields().len(), expected.len());
1380            assert_eq!(
1381                classify_completion_parent(metadata),
1382                CompletionParentVerdict::Adopt
1383            );
1384
1385            // Runtime-specific extra fields are additive on top of the contract.
1386            let span = completion_parent_span!(
1387                target: "contract_test",
1388                name: "chat",
1389                operation: "chat",
1390                system_instructions: Option::<&str>::None,
1391                gen_ai.agent.name = "assistant",
1392            );
1393            let Some(metadata) = span.metadata() else {
1394                panic!("contract span with extras was disabled");
1395            };
1396            let declared: HashSet<&str> =
1397                metadata.fields().iter().map(|field| field.name()).collect();
1398            let expected: HashSet<&str> =
1399                expected.into_iter().chain(["gen_ai.agent.name"]).collect();
1400            assert_eq!(declared, expected);
1401            assert_eq!(metadata.fields().len(), expected.len());
1402            assert_eq!(
1403                classify_completion_parent(metadata),
1404                CompletionParentVerdict::Adopt
1405            );
1406        });
1407    }
1408
1409    #[test]
1410    fn canonical_completion_span_declares_exactly_the_required_fields() {
1411        use std::collections::HashSet;
1412
1413        let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
1414        tracing::subscriber::with_default(Registry::default(), || {
1415            let span =
1416                CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
1417            let Some(metadata) = span.metadata() else {
1418                panic!("completion span was disabled");
1419            };
1420            let declared: HashSet<&str> =
1421                metadata.fields().iter().map(|field| field.name()).collect();
1422            let expected: HashSet<&str> =
1423                COMPLETION_PARENT_REQUIRED_FIELDS.iter().copied().collect();
1424            // The adoption checklist and the span the builder itself creates
1425            // must be the same set, or an adopted parent could not absorb
1426            // every field the builder records.
1427            assert_eq!(declared, expected);
1428            // Duplicate field names collapse in a `HashSet`, so also pin the
1429            // count: set equality alone cannot catch a field declared twice.
1430            assert_eq!(metadata.fields().len(), expected.len());
1431        });
1432    }
1433
1434    #[test]
1435    fn completion_parent_required_fields_are_pinned() {
1436        // Changing this list is a contract change. An adopted parent declares
1437        // its fields statically, so a runtime whose span was hand-written
1438        // against the old list stops being adopted once the list moves — it
1439        // degrades gracefully (fresh child span, one-time warning naming what
1440        // is missing), but it does degrade. Confirm that is intended, note it
1441        // in the CHANGELOG, then update this snapshot.
1442        //
1443        // This is the only test that notices. Every other contract test
1444        // compares the three forms of the contract to each other, so a
1445        // *coherent* change — a field added to both this constant and the
1446        // macro — leaves them all agreeing, and green.
1447        assert_eq!(COMPLETION_PARENT_MARKER_FIELD, "rig.completion_parent");
1448        assert_eq!(
1449            COMPLETION_PARENT_REQUIRED_FIELDS,
1450            &[
1451                "gen_ai.operation.name",
1452                "gen_ai.provider.name",
1453                "gen_ai.request.model",
1454                "gen_ai.system_instructions",
1455                "gen_ai.response.id",
1456                "gen_ai.response.model",
1457                "gen_ai.usage.input_tokens",
1458                "gen_ai.usage.output_tokens",
1459                "gen_ai.usage.cache_read.input_tokens",
1460                "gen_ai.usage.cache_creation.input_tokens",
1461                "gen_ai.usage.tool_use_prompt_tokens",
1462                "gen_ai.usage.reasoning_tokens",
1463                "gen_ai.input.messages",
1464                "gen_ai.output.messages",
1465            ]
1466        );
1467    }
1468
1469    /// Every row of the adoption decision table, asserted against the pure
1470    /// classifier so no global warn-once state is involved.
1471    #[test]
1472    fn classify_completion_parent_covers_the_decision_table() {
1473        let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
1474        tracing::subscriber::with_default(Registry::default(), || {
1475            let verdict = |span: &tracing::Span| {
1476                let Some(metadata) = span.metadata() else {
1477                    panic!("classifier fixture span was disabled");
1478                };
1479                classify_completion_parent(metadata)
1480            };
1481
1482            // Marker + full contract.
1483            let conforming = completion_parent_span!(
1484                target: "classifier_test",
1485                name: "chat",
1486                operation: tracing::field::Empty,
1487                system_instructions: tracing::field::Empty,
1488            );
1489            assert_eq!(verdict(&conforming), CompletionParentVerdict::Adopt);
1490
1491            // Marker, incomplete contract.
1492            let partial = tracing::info_span!(
1493                target: "classifier_test",
1494                "chat",
1495                rig.completion_parent = true,
1496                gen_ai.operation.name = tracing::field::Empty,
1497            );
1498            assert_eq!(
1499                verdict(&partial),
1500                CompletionParentVerdict::RejectMissingFields
1501            );
1502            let Some(partial_metadata) = partial.metadata() else {
1503                panic!("classifier fixture span was disabled");
1504            };
1505            // The names only get computed on the warning path, so pin them here.
1506            assert_eq!(
1507                missing_required_fields(partial_metadata),
1508                COMPLETION_PARENT_REQUIRED_FIELDS
1509                    .iter()
1510                    .copied()
1511                    .filter(|name| *name != "gen_ai.operation.name")
1512                    .collect::<Vec<_>>()
1513            );
1514
1515            // An ordinary ambient span.
1516            let ambient = tracing::info_span!(target: "application", "ambient");
1517            assert_eq!(verdict(&ambient), CompletionParentVerdict::NotAParent);
1518
1519            // Marker detection is an exact field-name match, never a prefix: a
1520            // runtime field that merely starts with the marker name must not
1521            // make its span a rejected parent and warn at a runtime that never
1522            // opted in.
1523            let lookalike = tracing::info_span!(
1524                target: "application",
1525                "ambient",
1526                rig.completion_parent.id = "abc",
1527                rig.completion_parent_id = "abc",
1528            );
1529            assert_eq!(verdict(&lookalike), CompletionParentVerdict::NotAParent);
1530        });
1531    }
1532
1533    /// The near-miss diagnostic is the only thing that makes a rejected parent
1534    /// visible to an operator — otherwise the sole symptom is a duplicated span
1535    /// layer in dashboards — so its message, its `missing_fields` payload, and
1536    /// its once-per-callsite budget all need pinning.
1537    #[test]
1538    fn near_miss_completion_parent_warns_once_per_callsite() {
1539        let warnings = CapturedWarnings::default();
1540        let subscriber = Registry::default().with(WarningCaptureLayer {
1541            warnings: warnings.clone(),
1542        });
1543        let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
1544        // The warn budget is process-global; claim a clean one rather than
1545        // relying on this test's fixture span owning a callsite no other test
1546        // touches.
1547        reset_near_miss_warnings();
1548        tracing::subscriber::with_default(subscriber, || {
1549            // The `warn!` callsite lives in `warn_once_on_completion_parent_verdict`
1550            // and is shared with every other near-miss test, so its interest may
1551            // already be cached as `never` from a run under a different
1552            // subscriber. Same hazard `test_utils::scoped_tracing_subscriber_guard`
1553            // documents; same fix used in
1554            // `agent::prompt_request::streaming`'s scoped-subscriber tests.
1555            tracing::callsite::rebuild_interest_cache();
1556
1557            let near_miss = tracing::info_span!(
1558                target: "third_party_runtime",
1559                "chat",
1560                rig.completion_parent = true,
1561                gen_ai.operation.name = tracing::field::Empty,
1562            );
1563            let _guard = near_miss.enter();
1564            CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
1565            // Second completion under the *same* span callsite: the budget is
1566            // per callsite, so this one must stay silent.
1567            CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
1568        });
1569
1570        let captured = warnings.take();
1571        assert_eq!(
1572            captured.len(),
1573            1,
1574            "a near-miss callsite warns exactly once, got: {captured:?}"
1575        );
1576        let Some(message) = captured.first() else {
1577            panic!("near miss did not warn");
1578        };
1579        assert!(
1580            message.contains("gen_ai.provider.name"),
1581            "warning must name the missing fields, got: {message}"
1582        );
1583        assert!(
1584            message.contains("completion_parent_span!"),
1585            "warning must point at the supported fix, got: {message}"
1586        );
1587    }
1588
1589    /// The property that justifies keying the budget on the callsite rather
1590    /// than a single process-wide flag: two runtimes each declaring a broken
1591    /// parent are both reported. A global flag would report whichever ran first
1592    /// and stay silent about the other — and every other test in this module
1593    /// passes under that behaviour, so this is the only one that pins it.
1594    #[test]
1595    fn distinct_near_miss_callsites_each_warn() {
1596        let warnings = CapturedWarnings::default();
1597        let subscriber = Registry::default().with(WarningCaptureLayer {
1598            warnings: warnings.clone(),
1599        });
1600        let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
1601        reset_near_miss_warnings();
1602        tracing::subscriber::with_default(subscriber, || {
1603            tracing::callsite::rebuild_interest_cache();
1604
1605            // Two separate `info_span!` invocations, and they must stay
1606            // separate: the callsite *is* the dedup key, so extracting these
1607            // into a shared helper collapses them into one and this test would
1608            // assert 1, not 2. `reset_near_miss_warnings` cannot protect this
1609            // test the way it protects the others — distinct callsites are the
1610            // thing under test.
1611            let first = tracing::info_span!(
1612                target: "runtime_a",
1613                "chat",
1614                rig.completion_parent = true,
1615                gen_ai.operation.name = tracing::field::Empty,
1616            );
1617            first.in_scope(|| {
1618                CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
1619            });
1620
1621            let second = tracing::info_span!(
1622                target: "runtime_b",
1623                "chat",
1624                rig.completion_parent = true,
1625                gen_ai.operation.name = tracing::field::Empty,
1626            );
1627            second.in_scope(|| {
1628                CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
1629            });
1630        });
1631
1632        let captured = warnings.take();
1633        assert_eq!(
1634            captured.len(),
1635            2,
1636            "each offending callsite warns once, got: {captured:?}"
1637        );
1638    }
1639
1640    /// The happy path must stay quiet: a conforming parent is adopted silently,
1641    /// so the diagnostic above cannot become background noise on every
1642    /// completion.
1643    #[test]
1644    fn conforming_completion_parent_never_warns() {
1645        let warnings = CapturedWarnings::default();
1646        let subscriber = Registry::default().with(WarningCaptureLayer {
1647            warnings: warnings.clone(),
1648        });
1649        let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
1650        // Claim a clean budget: the control below must be able to warn even if
1651        // another test already reported this fixture's callsite.
1652        reset_near_miss_warnings();
1653        tracing::subscriber::with_default(subscriber, || {
1654            tracing::callsite::rebuild_interest_cache();
1655
1656            // Control. Asserting an absence proves nothing unless the warning
1657            // pipe is known live in *this* subscriber: without this, the
1658            // assertions below pass just as happily when the diagnostic has
1659            // been deleted, or when callsite interest is stuck at `never`.
1660            let near_miss = tracing::info_span!(
1661                target: "third_party_runtime",
1662                "chat",
1663                rig.completion_parent = true,
1664                gen_ai.operation.name = tracing::field::Empty,
1665            );
1666            near_miss.in_scope(|| {
1667                CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
1668            });
1669            assert_eq!(
1670                warnings.take().len(),
1671                1,
1672                "control: a near miss must warn, or this test cannot detect silence"
1673            );
1674
1675            let conforming = completion_parent_span!(
1676                target: "third_party_runtime",
1677                name: "chat",
1678                operation: Empty,
1679                system_instructions: Option::<&str>::None,
1680            );
1681            let _guard = conforming.enter();
1682            CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
1683
1684            // An ordinary ambient span is not a parent at all, and must not warn
1685            // either — a runtime that never opted in should never hear about
1686            // this contract.
1687            drop(_guard);
1688            let ambient = tracing::info_span!(target: "application", "ambient");
1689            let _ambient_guard = ambient.enter();
1690            CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
1691        });
1692
1693        // The control drained the buffer, so anything here was emitted by the
1694        // conforming or ambient span.
1695        let captured = warnings.take();
1696        assert!(
1697            captured.is_empty(),
1698            "adoption and non-participation are both silent, got: {captured:?}"
1699        );
1700    }
1701
1702    #[test]
1703    fn agent_chat_span_is_adopted_and_enriched() {
1704        let captured = CapturedSpan::default();
1705        let subscriber = Registry::default().with(SpanCaptureLayer {
1706            span: captured.clone(),
1707        });
1708        let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
1709        tracing::subscriber::with_default(subscriber, || {
1710            let completion_parent = completion_parent_span!(
1711                target: "rig::agent_chat",
1712                name: "chat_streaming",
1713                operation: tracing::field::Empty,
1714                system_instructions: tracing::field::Empty,
1715            );
1716            let _guard = completion_parent.enter();
1717            let span = CompletionSpanBuilder::new(
1718                "anthropic",
1719                "claude-sonnet",
1720                CompletionOperation::ChatStreaming,
1721            )
1722            .system_instructions(Some("provider system"), true)
1723            .build();
1724            assert_eq!(span.id(), completion_parent.id());
1725        });
1726
1727        let Ok(captured) = captured.0.lock() else {
1728            panic!("captured span lock poisoned");
1729        };
1730        let Some(span) = captured.as_ref() else {
1731            panic!("completion-parent span was not captured");
1732        };
1733        assert_eq!(span.target, "rig::agent_chat");
1734        for (field, value) in [
1735            ("gen_ai.operation.name", "chat_streaming"),
1736            ("gen_ai.provider.name", "anthropic"),
1737            ("gen_ai.request.model", "claude-sonnet"),
1738            (
1739                "gen_ai.system_instructions",
1740                r#"[{"type":"text","content":"provider system"}]"#,
1741            ),
1742        ] {
1743            assert!(contains_string(&span.recorded_values, field, value));
1744        }
1745    }
1746
1747    #[test]
1748    fn neutral_completion_parent_span_is_adopted_and_enriched() {
1749        let captured = CapturedSpan::default();
1750        let subscriber = Registry::default().with(SpanCaptureLayer {
1751            span: captured.clone(),
1752        });
1753        let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
1754        tracing::subscriber::with_default(subscriber, || {
1755            let completion_parent = completion_parent_span!(
1756                target: "test_runtime",
1757                name: "chat",
1758                operation: tracing::field::Empty,
1759                system_instructions: tracing::field::Empty,
1760            );
1761            let _guard = completion_parent.enter();
1762            let span = CompletionSpanBuilder::new(
1763                "neutral-provider",
1764                "neutral-model",
1765                CompletionOperation::Chat,
1766            )
1767            .build();
1768            assert_eq!(span.id(), completion_parent.id());
1769        });
1770
1771        let Ok(captured) = captured.0.lock() else {
1772            panic!("captured span lock poisoned");
1773        };
1774        let Some(span) = captured.as_ref() else {
1775            panic!("neutral completion-parent span was not captured");
1776        };
1777        assert_eq!(span.target, "test_runtime");
1778        for (field, value) in [
1779            ("gen_ai.operation.name", "chat"),
1780            ("gen_ai.provider.name", "neutral-provider"),
1781            ("gen_ai.request.model", "neutral-model"),
1782        ] {
1783            assert!(contains_string(&span.recorded_values, field, value));
1784        }
1785    }
1786
1787    #[test]
1788    fn absent_provider_system_does_not_overwrite_agent_instructions() {
1789        let captured = CapturedSpan::default();
1790        let subscriber = Registry::default().with(SpanCaptureLayer {
1791            span: captured.clone(),
1792        });
1793        let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
1794        tracing::subscriber::with_default(subscriber, || {
1795            let completion_parent = completion_parent_span!(
1796                target: "test_runtime",
1797                name: "chat",
1798                operation: tracing::field::Empty,
1799                system_instructions: "effective agent instructions",
1800            );
1801            let _guard = completion_parent.enter();
1802            CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
1803        });
1804
1805        let Ok(captured) = captured.0.lock() else {
1806            panic!("captured span lock poisoned");
1807        };
1808        let Some(span) = captured.as_ref() else {
1809            panic!("completion-parent span was not captured");
1810        };
1811        assert!(contains_string(
1812            &span.initial_values,
1813            "gen_ai.system_instructions",
1814            "effective agent instructions"
1815        ));
1816        assert!(
1817            !span
1818                .recorded_values
1819                .iter()
1820                .any(|(field, _)| field == "gen_ai.system_instructions")
1821        );
1822    }
1823
1824    #[test]
1825    fn record_token_usage_records_tool_use_prompt_tokens() {
1826        let fields = CapturedFields::default();
1827        let subscriber = Registry::default().with(FieldCaptureLayer {
1828            fields: fields.clone(),
1829        });
1830        let usage = TestUsage(Usage {
1831            input_tokens: 1,
1832            output_tokens: 2,
1833            total_tokens: 15,
1834            cached_input_tokens: 3,
1835            cache_creation_input_tokens: 4,
1836            tool_use_prompt_tokens: 12,
1837            reasoning_tokens: 5,
1838        });
1839
1840        // Scoped-subscriber tests must not run concurrently; see
1841        // `test_utils::scoped_tracing_subscriber_guard`.
1842        let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
1843        tracing::subscriber::with_default(subscriber, || {
1844            let span = tracing::info_span!(
1845                "usage_recording",
1846                gen_ai.usage.input_tokens = tracing::field::Empty,
1847                gen_ai.usage.output_tokens = tracing::field::Empty,
1848                gen_ai.usage.cache_read.input_tokens = tracing::field::Empty,
1849                gen_ai.usage.cache_creation.input_tokens = tracing::field::Empty,
1850                gen_ai.usage.tool_use_prompt_tokens = tracing::field::Empty,
1851                gen_ai.usage.reasoning_tokens = tracing::field::Empty,
1852            );
1853
1854            span.record_token_usage(&usage);
1855        });
1856
1857        assert!(fields.contains("gen_ai.usage.tool_use_prompt_tokens", 12));
1858    }
1859}