use crate::OneOrMany;
use crate::completion::{AssistantContent, GetTokenUsage, Message};
use crate::message::{
DocumentSourceKind, Image, MimeType, Reasoning, ReasoningContent, ToolResult,
ToolResultContent, UserContent,
};
use base64::Engine;
use serde::Serialize;
use std::collections::HashSet;
use std::sync::{LazyLock, Mutex};
use tracing::callsite::Identifier;
#[doc(hidden)]
pub use tracing as __tracing;
pub use tracing::field::Empty;
#[doc(hidden)]
#[macro_export]
macro_rules! __rig_canonical_completion_span {
(
target: $target:literal,
$(parent: $parent:expr,)?
name: $name:literal,
// Both blocks are spliced verbatim into `info_span!`: the header block
// must end with a trailing comma, the extras block must begin with one.
// Violating either surfaces as an `info_span!` parse error at the call
// site, not here.
{ $($header:tt)* }
{ $($extra:tt)* }
) => {
$crate::telemetry::__tracing::info_span!(
target: $target,
$(parent: $parent,)?
$name,
$($header)*
gen_ai.response.id = $crate::telemetry::__tracing::field::Empty,
gen_ai.response.model = $crate::telemetry::__tracing::field::Empty,
gen_ai.usage.input_tokens = $crate::telemetry::__tracing::field::Empty,
gen_ai.usage.output_tokens = $crate::telemetry::__tracing::field::Empty,
gen_ai.usage.cache_read.input_tokens = $crate::telemetry::__tracing::field::Empty,
gen_ai.usage.cache_creation.input_tokens = $crate::telemetry::__tracing::field::Empty,
gen_ai.usage.tool_use_prompt_tokens = $crate::telemetry::__tracing::field::Empty,
gen_ai.usage.reasoning_tokens = $crate::telemetry::__tracing::field::Empty,
gen_ai.input.messages = $crate::telemetry::__tracing::field::Empty,
gen_ai.output.messages = $crate::telemetry::__tracing::field::Empty
$($extra)*
)
};
}
macro_rules! new_completion_span {
($name:literal, $provider:expr, $request_model:expr, $operation:expr, $system:expr) => {
$crate::__rig_canonical_completion_span!(
target: "rig::completions",
name: $name,
{
gen_ai.operation.name = $operation,
gen_ai.provider.name = $provider,
gen_ai.request.model = $request_model,
gen_ai.system_instructions = $system,
}
{}
)
};
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum CompletionOperation {
Chat,
ChatStreaming,
GenerateContent,
Interactions,
InteractionsStreaming,
}
impl CompletionOperation {
fn as_str(self) -> &'static str {
match self {
Self::Chat => "chat",
Self::ChatStreaming => "chat_streaming",
Self::GenerateContent => "generate_content",
Self::Interactions => "interactions",
Self::InteractionsStreaming => "interactions_streaming",
}
}
}
pub const COMPLETION_PARENT_MARKER_FIELD: &str = "rig.completion_parent";
pub const COMPLETION_PARENT_REQUIRED_FIELDS: &[&str] = &[
"gen_ai.operation.name",
"gen_ai.provider.name",
"gen_ai.request.model",
"gen_ai.system_instructions",
"gen_ai.response.id",
"gen_ai.response.model",
"gen_ai.usage.input_tokens",
"gen_ai.usage.output_tokens",
"gen_ai.usage.cache_read.input_tokens",
"gen_ai.usage.cache_creation.input_tokens",
"gen_ai.usage.tool_use_prompt_tokens",
"gen_ai.usage.reasoning_tokens",
"gen_ai.input.messages",
"gen_ai.output.messages",
];
#[macro_export]
macro_rules! completion_parent_span {
(
target: $target:literal,
parent: $parent:expr,
name: $name:literal,
operation: $operation:expr,
system_instructions: $system:expr
$(, $($extra:tt)*)?
) => {
$crate::__rig_canonical_completion_span!(
target: $target,
parent: $parent,
name: $name,
{
rig.completion_parent = true,
gen_ai.operation.name = $operation,
gen_ai.system_instructions = $system,
gen_ai.provider.name = $crate::telemetry::__tracing::field::Empty,
gen_ai.request.model = $crate::telemetry::__tracing::field::Empty,
}
{ $(, $($extra)*)? }
)
};
(
target: $target:literal,
name: $name:literal,
operation: $operation:expr,
system_instructions: $system:expr
$(, $($extra:tt)*)?
) => {
$crate::completion_parent_span!(
target: $target,
parent: $crate::telemetry::__tracing::Span::current(),
name: $name,
operation: $operation,
system_instructions: $system
$(, $($extra)*)?
)
};
}
pub use crate::completion_parent_span;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CompletionParentVerdict {
Adopt,
RejectMissingFields,
NotAParent,
}
fn missing_required_fields(metadata: &tracing::Metadata<'_>) -> Vec<&'static str> {
let fields = metadata.fields();
COMPLETION_PARENT_REQUIRED_FIELDS
.iter()
.copied()
.filter(|name| fields.field(name).is_none())
.collect()
}
fn classify_completion_parent(metadata: &tracing::Metadata<'_>) -> CompletionParentVerdict {
let fields = metadata.fields();
if fields.field(COMPLETION_PARENT_MARKER_FIELD).is_none() {
return CompletionParentVerdict::NotAParent;
}
if COMPLETION_PARENT_REQUIRED_FIELDS
.iter()
.all(|name| fields.field(name).is_some())
{
CompletionParentVerdict::Adopt
} else {
CompletionParentVerdict::RejectMissingFields
}
}
static NEAR_MISS_WARNED: LazyLock<Mutex<HashSet<Identifier>>> =
LazyLock::new(|| Mutex::new(HashSet::new()));
#[cfg(test)]
fn reset_near_miss_warnings() {
NEAR_MISS_WARNED
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.clear();
}
fn warn_once_on_completion_parent_verdict(
verdict: CompletionParentVerdict,
metadata: &tracing::Metadata<'_>,
) {
match verdict {
CompletionParentVerdict::Adopt | CompletionParentVerdict::NotAParent => {}
CompletionParentVerdict::RejectMissingFields => {
let first_sighting = {
let mut warned = NEAR_MISS_WARNED
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
warned.insert(metadata.callsite())
};
if !first_sighting {
return;
}
tracing::warn!(
marker = COMPLETION_PARENT_MARKER_FIELD,
missing_fields = ?missing_required_fields(metadata),
"completion-parent span declares the marker but not every required field \
and is not adopted; provider telemetry lands on a fresh child span \
instead — declare the span with \
`rig_core::telemetry::completion_parent_span!`"
);
}
}
}
pub struct CompletionSpanBuilder<'a> {
provider: &'a str,
request_model: &'a str,
operation: CompletionOperation,
system_instructions: Option<String>,
}
impl<'a> CompletionSpanBuilder<'a> {
pub fn new(provider: &'a str, request_model: &'a str, operation: CompletionOperation) -> Self {
Self {
provider,
request_model,
operation,
system_instructions: None,
}
}
pub fn system_instructions(
mut self,
system_instructions: Option<&'a str>,
record_content: bool,
) -> Self {
self.system_instructions = system_instructions_json(system_instructions, record_content);
self
}
pub fn build(self) -> tracing::Span {
let current = tracing::Span::current();
if let Some(metadata) = current.metadata() {
let verdict = classify_completion_parent(metadata);
warn_once_on_completion_parent_verdict(verdict, metadata);
if verdict == CompletionParentVerdict::Adopt {
current.record("gen_ai.operation.name", self.operation.as_str());
current.record("gen_ai.provider.name", self.provider);
current.record("gen_ai.request.model", self.request_model);
if let Some(system_instructions) = self.system_instructions.as_deref() {
current.record("gen_ai.system_instructions", system_instructions);
}
return current;
}
}
let operation = self.operation.as_str();
let system_instructions = self.system_instructions.as_deref();
match self.operation {
CompletionOperation::Chat => new_completion_span!(
"chat",
self.provider,
self.request_model,
operation,
system_instructions
),
CompletionOperation::ChatStreaming => new_completion_span!(
"chat_streaming",
self.provider,
self.request_model,
operation,
system_instructions
),
CompletionOperation::GenerateContent => new_completion_span!(
"generate_content",
self.provider,
self.request_model,
operation,
system_instructions
),
CompletionOperation::Interactions => new_completion_span!(
"interactions",
self.provider,
self.request_model,
operation,
system_instructions
),
CompletionOperation::InteractionsStreaming => new_completion_span!(
"interactions_streaming",
self.provider,
self.request_model,
operation,
system_instructions
),
}
}
}
#[derive(Serialize)]
struct TelemetryChatMessage {
role: &'static str,
parts: Vec<TelemetryPart>,
}
#[derive(Serialize)]
struct TelemetryOutputMessage {
role: &'static str,
parts: Vec<TelemetryPart>,
finish_reason: &'static str,
}
#[derive(Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum TelemetryPart {
Text {
content: String,
},
ToolCall {
#[serde(skip_serializing_if = "Option::is_none")]
id: Option<String>,
name: String,
arguments: serde_json::Value,
},
ToolCallResponse {
#[serde(skip_serializing_if = "Option::is_none")]
id: Option<String>,
response: serde_json::Value,
},
Reasoning {
content: String,
},
Uri {
#[serde(skip_serializing_if = "Option::is_none")]
mime_type: Option<String>,
modality: &'static str,
uri: String,
},
File {
#[serde(skip_serializing_if = "Option::is_none")]
mime_type: Option<String>,
modality: &'static str,
file_id: String,
},
Blob {
#[serde(skip_serializing_if = "Option::is_none")]
mime_type: Option<String>,
modality: &'static str,
content: String,
},
}
fn media_part<T>(
data: &DocumentSourceKind,
media_type: Option<&T>,
modality: &'static str,
) -> Option<TelemetryPart>
where
T: MimeType,
{
let mime_type = media_type.map(|media_type| media_type.to_mime_type().to_string());
match data {
DocumentSourceKind::Url(uri) => Some(TelemetryPart::Uri {
mime_type,
modality,
uri: uri.clone(),
}),
DocumentSourceKind::FileId(file_id) => Some(TelemetryPart::File {
mime_type,
modality,
file_id: file_id.clone(),
}),
DocumentSourceKind::Base64(content) => Some(TelemetryPart::Blob {
mime_type,
modality,
content: content.clone(),
}),
DocumentSourceKind::Raw(content) => Some(TelemetryPart::Blob {
mime_type,
modality,
content: base64::engine::general_purpose::STANDARD.encode(content),
}),
DocumentSourceKind::String(content) => Some(TelemetryPart::Text {
content: content.clone(),
}),
DocumentSourceKind::Unknown => None,
}
}
fn image_part(image: &Image) -> Option<TelemetryPart> {
media_part(&image.data, image.media_type.as_ref(), "image")
}
fn reasoning_parts(reasoning: &Reasoning) -> Vec<TelemetryPart> {
reasoning
.content
.iter()
.map(|content| {
let content = match content {
ReasoningContent::Text { text, .. } | ReasoningContent::Summary(text) => text,
ReasoningContent::Encrypted(content) => content,
ReasoningContent::Redacted { data } => data,
};
TelemetryPart::Reasoning {
content: content.clone(),
}
})
.collect()
}
fn tool_result_response(result: &ToolResult) -> serde_json::Value {
let mut content = result
.content
.iter()
.filter_map(|content| match content {
ToolResultContent::Text(text) => Some(serde_json::Value::String(text.text.clone())),
ToolResultContent::Json { value } => Some(value.clone()),
ToolResultContent::Image(image) => {
image_part(image).and_then(|part| serde_json::to_value(part).ok())
}
})
.collect::<Vec<_>>();
if content.len() == 1 {
content.pop().unwrap_or(serde_json::Value::Null)
} else {
serde_json::Value::Array(content)
}
}
fn user_parts(content: &OneOrMany<UserContent>) -> Vec<TelemetryPart> {
content
.iter()
.filter_map(|content| match content {
UserContent::Text(text) => Some(TelemetryPart::Text {
content: text.text.clone(),
}),
UserContent::ToolResult(result) => Some(TelemetryPart::ToolCallResponse {
id: Some(result.id.clone()),
response: tool_result_response(result),
}),
UserContent::Image(image) => image_part(image),
UserContent::Audio(audio) => {
media_part(&audio.data, audio.media_type.as_ref(), "audio")
}
UserContent::Video(video) => {
media_part(&video.data, video.media_type.as_ref(), "video")
}
UserContent::Document(document) => {
media_part(&document.data, document.media_type.as_ref(), "document")
}
})
.collect()
}
fn assistant_parts(content: &OneOrMany<AssistantContent>) -> Vec<TelemetryPart> {
content
.iter()
.flat_map(|content| match content {
AssistantContent::Text(text) => vec![TelemetryPart::Text {
content: text.text.clone(),
}],
AssistantContent::ToolCall(tool_call) => vec![TelemetryPart::ToolCall {
id: Some(tool_call.id.clone()),
name: tool_call.function.name.clone(),
arguments: tool_call.function.arguments.clone(),
}],
AssistantContent::Reasoning(reasoning) => reasoning_parts(reasoning),
AssistantContent::Image(image) => image_part(image).into_iter().collect(),
})
.collect()
}
fn input_messages(messages: &[Message]) -> Vec<TelemetryChatMessage> {
messages
.iter()
.map(|message| match message {
Message::System { content } => TelemetryChatMessage {
role: "system",
parts: vec![TelemetryPart::Text {
content: content.clone(),
}],
},
Message::User { content } => TelemetryChatMessage {
role: "user",
parts: user_parts(content),
},
Message::Assistant { content, .. } => TelemetryChatMessage {
role: "assistant",
parts: assistant_parts(content),
},
})
.collect()
}
fn output_messages(content: &OneOrMany<AssistantContent>) -> Vec<TelemetryOutputMessage> {
let finish_reason = if content
.iter()
.any(|content| matches!(content, AssistantContent::ToolCall(_)))
{
"tool_call"
} else {
"unknown"
};
vec![TelemetryOutputMessage {
role: "assistant",
parts: assistant_parts(content),
finish_reason,
}]
}
pub fn system_instructions_json(instructions: Option<&str>, enabled: bool) -> Option<String> {
if !enabled {
return None;
}
instructions.and_then(|instructions| {
serde_json::to_string(&vec![TelemetryPart::Text {
content: instructions.to_string(),
}])
.ok()
})
}
pub fn record_model_input(span: &tracing::Span, messages: &[Message], enabled: bool) {
if !enabled || span.is_disabled() {
return;
}
if let Ok(messages) = serde_json::to_string(&input_messages(messages)) {
span.record("gen_ai.input.messages", messages);
}
}
pub fn record_model_output(
span: &tracing::Span,
content: &OneOrMany<AssistantContent>,
enabled: bool,
) {
if !enabled || span.is_disabled() {
return;
}
let messages = output_messages(content);
if let Ok(messages) = serde_json::to_string(&messages) {
span.record("gen_ai.output.messages", messages);
}
}
pub trait ProviderResponseExt {
type OutputMessage: Serialize;
type Usage: Serialize;
fn get_response_id(&self) -> Option<String>;
fn get_response_model_name(&self) -> Option<String>;
fn get_output_messages(&self) -> Vec<Self::OutputMessage>;
fn get_text_response(&self) -> Option<String>;
fn get_usage(&self) -> Option<Self::Usage>;
}
pub trait SpanCombinator {
fn record_token_usage<U>(&self, usage: &U)
where
U: GetTokenUsage;
fn record_response_metadata<R>(&self, response: &R)
where
R: ProviderResponseExt;
}
impl SpanCombinator for tracing::Span {
fn record_token_usage<U>(&self, usage: &U)
where
U: GetTokenUsage,
{
if self.is_disabled() {
return;
}
let usage = usage.token_usage();
if usage.has_values() {
self.record("gen_ai.usage.input_tokens", usage.input_tokens);
self.record("gen_ai.usage.output_tokens", usage.output_tokens);
self.record(
"gen_ai.usage.cache_read.input_tokens",
usage.cached_input_tokens,
);
self.record(
"gen_ai.usage.cache_creation.input_tokens",
usage.cache_creation_input_tokens,
);
self.record(
"gen_ai.usage.tool_use_prompt_tokens",
usage.tool_use_prompt_tokens,
);
self.record("gen_ai.usage.reasoning_tokens", usage.reasoning_tokens);
}
}
fn record_response_metadata<R>(&self, response: &R)
where
R: ProviderResponseExt,
{
if self.is_disabled() {
return;
}
if let Some(id) = response.get_response_id() {
self.record("gen_ai.response.id", id);
}
if let Some(model_name) = response.get_response_model_name() {
self.record("gen_ai.response.model", model_name);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::completion::{AssistantContent, GetTokenUsage, Message, Usage};
use serde_json::json;
use std::sync::{Arc, Mutex};
use tracing::field::{Field, Visit};
use tracing::{Id, Subscriber};
use tracing_subscriber::layer::{Context, SubscriberExt};
use tracing_subscriber::{Layer, Registry, registry::LookupSpan};
#[derive(Clone)]
struct TestUsage(Usage);
impl GetTokenUsage for TestUsage {
fn token_usage(&self) -> Usage {
self.0
}
}
#[test]
fn content_attributes_follow_gen_ai_semantic_convention_json_shapes() {
assert_eq!(
system_instructions_json(Some("follow policy"), true).as_deref(),
Some(r#"[{"type":"text","content":"follow policy"}]"#)
);
assert_eq!(system_instructions_json(Some("secret"), false), None);
let input = input_messages(&[
Message::system("follow policy"),
Message::user("hello"),
Message::tool_result("call_1", "sunny"),
]);
assert_eq!(
serde_json::to_value(input).expect("semantic-convention input DTOs serialize"),
json!([
{
"role": "system",
"parts": [{"type": "text", "content": "follow policy"}]
},
{
"role": "user",
"parts": [{"type": "text", "content": "hello"}]
},
{
"role": "user",
"parts": [{
"type": "tool_call_response",
"id": "call_1",
"response": "sunny"
}]
}
])
);
let output = OneOrMany::one(AssistantContent::tool_call(
"call_1",
"weather",
json!({"city": "Paris"}),
));
assert_eq!(
serde_json::to_value(output_messages(&output))
.expect("semantic-convention output DTOs serialize"),
json!([{
"role": "assistant",
"parts": [{
"type": "tool_call",
"id": "call_1",
"name": "weather",
"arguments": {"city": "Paris"}
}],
"finish_reason": "tool_call"
}])
);
let text_output = OneOrMany::one(AssistantContent::text("done"));
assert_eq!(
serde_json::to_value(output_messages(&text_output))
.expect("semantic-convention text output DTOs serialize"),
json!([{
"role": "assistant",
"parts": [{"type": "text", "content": "done"}],
"finish_reason": "unknown"
}])
);
}
#[derive(Clone, Default)]
struct CapturedFields(Arc<Mutex<Vec<(String, u64)>>>);
impl CapturedFields {
fn push(&self, name: &str, value: u64) {
if let Ok(mut fields) = self.0.lock() {
fields.push((name.to_string(), value));
}
}
fn contains(&self, name: &str, value: u64) -> bool {
self.0.lock().is_ok_and(|fields| {
fields
.iter()
.any(|field| field == &(name.to_string(), value))
})
}
}
struct FieldCaptureLayer {
fields: CapturedFields,
}
impl<S> Layer<S> for FieldCaptureLayer
where
S: Subscriber,
S: for<'lookup> LookupSpan<'lookup>,
{
fn on_record(&self, _span: &Id, values: &tracing::span::Record<'_>, _ctx: Context<'_, S>) {
values.record(&mut FieldCaptureVisitor {
fields: self.fields.clone(),
});
}
}
struct FieldCaptureVisitor {
fields: CapturedFields,
}
impl Visit for FieldCaptureVisitor {
fn record_u64(&mut self, field: &Field, value: u64) {
self.fields.push(field.name(), value);
}
fn record_debug(&mut self, _field: &Field, _value: &dyn std::fmt::Debug) {}
}
#[derive(Clone, Default)]
struct CapturedWarnings(Arc<Mutex<Vec<String>>>);
impl CapturedWarnings {
fn push(&self, rendered: String) {
if let Ok(mut events) = self.0.lock() {
events.push(rendered);
}
}
fn take(&self) -> Vec<String> {
self.0
.lock()
.map(|mut events| std::mem::take(&mut *events))
.unwrap_or_default()
}
}
struct WarningCaptureLayer {
warnings: CapturedWarnings,
}
impl<S> Layer<S> for WarningCaptureLayer
where
S: Subscriber,
S: for<'lookup> LookupSpan<'lookup>,
{
fn on_event(&self, event: &tracing::Event<'_>, _ctx: Context<'_, S>) {
if *event.metadata().level() != tracing::Level::WARN {
return;
}
let mut visitor = WarningCaptureVisitor::default();
event.record(&mut visitor);
self.warnings.push(visitor.rendered);
}
}
#[derive(Default)]
struct WarningCaptureVisitor {
rendered: String,
}
impl Visit for WarningCaptureVisitor {
fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
use std::fmt::Write;
let _ = write!(&mut self.rendered, " {}={value:?}", field.name());
}
fn record_str(&mut self, field: &Field, value: &str) {
use std::fmt::Write;
let _ = write!(&mut self.rendered, " {}={value}", field.name());
}
}
#[derive(Clone, Default)]
struct CapturedSpan(Arc<Mutex<Option<CapturedSpanData>>>);
struct CapturedSpanData {
name: String,
target: String,
parent_name: Option<String>,
fields: Vec<String>,
initial_values: Vec<(String, String)>,
recorded_values: Vec<(String, String)>,
}
struct SpanCaptureLayer {
span: CapturedSpan,
}
#[derive(Default)]
struct StringFieldVisitor {
values: Vec<(String, String)>,
}
impl Visit for StringFieldVisitor {
fn record_str(&mut self, field: &Field, value: &str) {
self.values
.push((field.name().to_owned(), value.to_owned()));
}
fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
self.values
.push((field.name().to_owned(), format!("{value:?}")));
}
}
impl<S> Layer<S> for SpanCaptureLayer
where
S: Subscriber,
S: for<'lookup> LookupSpan<'lookup>,
{
fn on_new_span(
&self,
attrs: &tracing::span::Attributes<'_>,
_id: &Id,
ctx: Context<'_, S>,
) {
if let Ok(mut captured) = self.span.0.lock() {
let mut visitor = StringFieldVisitor::default();
attrs.record(&mut visitor);
let parent_name = if let Some(parent) = attrs.parent() {
ctx.span(parent)
.map(|span| span.metadata().name().to_owned())
} else if attrs.is_contextual() {
ctx.lookup_current()
.map(|span| span.metadata().name().to_owned())
} else {
None
};
*captured = Some(CapturedSpanData {
name: attrs.metadata().name().to_owned(),
target: attrs.metadata().target().to_owned(),
parent_name,
fields: attrs
.metadata()
.fields()
.iter()
.map(|field| field.name().to_owned())
.collect(),
initial_values: visitor.values,
recorded_values: Vec::new(),
});
}
}
fn on_record(&self, _span: &Id, values: &tracing::span::Record<'_>, _ctx: Context<'_, S>) {
if let Ok(mut captured) = self.span.0.lock()
&& let Some(captured) = captured.as_mut()
{
let mut visitor = StringFieldVisitor::default();
values.record(&mut visitor);
captured.recorded_values.extend(visitor.values);
}
}
}
fn contains_string(values: &[(String, String)], field: &str, value: &str) -> bool {
values
.iter()
.any(|candidate| candidate == &(field.to_owned(), value.to_owned()))
}
#[test]
fn completion_span_uses_canonical_names_fields_and_initial_attributes() {
let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
for (operation, expected_name) in [
(CompletionOperation::Chat, "chat"),
(CompletionOperation::ChatStreaming, "chat_streaming"),
(CompletionOperation::GenerateContent, "generate_content"),
(CompletionOperation::Interactions, "interactions"),
(
CompletionOperation::InteractionsStreaming,
"interactions_streaming",
),
] {
let captured = CapturedSpan::default();
let subscriber = Registry::default().with(SpanCaptureLayer {
span: captured.clone(),
});
tracing::subscriber::with_default(subscriber, || {
let span = CompletionSpanBuilder::new("openai", "gpt-5", operation)
.system_instructions(Some("system prompt"), true)
.build();
assert!(!span.is_disabled());
});
let Ok(captured) = captured.0.lock() else {
panic!("captured span lock poisoned");
};
let Some(span) = captured.as_ref() else {
panic!("completion span was not created");
};
assert_eq!(span.name, expected_name);
assert_eq!(span.target, "rig::completions");
assert_eq!(span.parent_name, None);
for (field, value) in [
("gen_ai.operation.name", expected_name),
("gen_ai.provider.name", "openai"),
("gen_ai.request.model", "gpt-5"),
(
"gen_ai.system_instructions",
r#"[{"type":"text","content":"system prompt"}]"#,
),
] {
assert!(
contains_string(&span.initial_values, field, value),
"missing initial {field}={value}"
);
}
assert!(span.recorded_values.is_empty());
assert!(
!span
.initial_values
.iter()
.any(|(field, _)| field == "gen_ai.response.model")
);
for field in COMPLETION_PARENT_REQUIRED_FIELDS {
assert!(
span.fields.iter().any(|candidate| candidate == field),
"missing {field}"
);
}
}
}
#[test]
fn completion_parent_span_macro_honours_its_parent_argument() {
fn captured_parent(captured: &CapturedSpan) -> Option<String> {
let Ok(captured) = captured.0.lock() else {
panic!("captured span lock poisoned");
};
let Some(span) = captured.as_ref() else {
panic!("completion-parent span was not captured");
};
assert_eq!(span.target, "third_party_runtime");
span.parent_name.clone()
}
let captured = CapturedSpan::default();
let subscriber = Registry::default().with(SpanCaptureLayer {
span: captured.clone(),
});
let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
tracing::subscriber::with_default(subscriber, || {
let ambient = tracing::info_span!(target: "application", "ambient");
ambient.in_scope(|| {
let _default_arm = completion_parent_span!(
target: "third_party_runtime",
name: "chat",
operation: Empty,
system_instructions: Option::<&str>::None,
);
});
assert_eq!(captured_parent(&captured).as_deref(), Some("ambient"));
ambient.in_scope(|| {
let _explicit_arm = completion_parent_span!(
target: "third_party_runtime",
parent: None,
name: "chat",
operation: Empty,
system_instructions: Option::<&str>::None,
);
});
assert_eq!(captured_parent(&captured), None);
});
}
#[test]
fn unrelated_ambient_span_is_parent_not_adopted() {
let captured = CapturedSpan::default();
let subscriber = Registry::default().with(SpanCaptureLayer {
span: captured.clone(),
});
let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
tracing::subscriber::with_default(subscriber, || {
let ambient = tracing::info_span!(target: "application", "ambient");
let _guard = ambient.enter();
let span =
CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
assert_ne!(span.id(), ambient.id());
});
let Ok(captured) = captured.0.lock() else {
panic!("captured span lock poisoned");
};
let Some(span) = captured.as_ref() else {
panic!("completion span was not captured");
};
assert_eq!(span.target, "rig::completions");
assert_eq!(span.parent_name.as_deref(), Some("ambient"));
}
#[test]
fn marker_span_missing_required_fields_is_not_adopted() {
let captured = CapturedSpan::default();
let subscriber = Registry::default().with(SpanCaptureLayer {
span: captured.clone(),
});
let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
tracing::subscriber::with_default(subscriber, || {
let partial_marker = tracing::info_span!(
target: "third_party_runtime",
"chat",
rig.completion_parent = true,
gen_ai.operation.name = tracing::field::Empty,
gen_ai.provider.name = tracing::field::Empty,
gen_ai.request.model = tracing::field::Empty,
);
let Some(metadata) = partial_marker.metadata() else {
panic!("partial marker span was disabled");
};
assert!(
metadata
.fields()
.field(COMPLETION_PARENT_MARKER_FIELD)
.is_some(),
"hand-written marker literal is stale; update it to {COMPLETION_PARENT_MARKER_FIELD}"
);
let _guard = partial_marker.enter();
let span =
CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
assert_ne!(span.id(), partial_marker.id());
});
let Ok(captured) = captured.0.lock() else {
panic!("captured span lock poisoned");
};
let Some(span) = captured.as_ref() else {
panic!("completion span was not captured");
};
assert_eq!(span.target, "rig::completions");
assert_eq!(span.parent_name.as_deref(), Some("chat"));
for (field, value) in [
("gen_ai.operation.name", "chat"),
("gen_ai.provider.name", "openai"),
("gen_ai.request.model", "gpt-5"),
] {
assert!(contains_string(&span.initial_values, field, value));
}
}
#[test]
fn completion_parent_span_macro_matches_the_contract_exactly() {
use std::collections::HashSet;
let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
tracing::subscriber::with_default(Registry::default(), || {
let span = completion_parent_span!(
target: "contract_test",
name: "chat",
operation: "chat",
system_instructions: Option::<&str>::None,
);
let Some(metadata) = span.metadata() else {
panic!("contract span was disabled");
};
let declared: HashSet<&str> =
metadata.fields().iter().map(|field| field.name()).collect();
let expected: HashSet<&str> = COMPLETION_PARENT_REQUIRED_FIELDS
.iter()
.copied()
.chain([COMPLETION_PARENT_MARKER_FIELD])
.collect();
assert_eq!(declared, expected);
assert_eq!(metadata.fields().len(), expected.len());
assert_eq!(
classify_completion_parent(metadata),
CompletionParentVerdict::Adopt
);
let span = completion_parent_span!(
target: "contract_test",
parent: None,
name: "chat",
operation: "chat",
system_instructions: Option::<&str>::None,
);
let Some(metadata) = span.metadata() else {
panic!("contract span with explicit parent was disabled");
};
let declared: HashSet<&str> =
metadata.fields().iter().map(|field| field.name()).collect();
assert_eq!(declared, expected);
assert_eq!(metadata.fields().len(), expected.len());
assert_eq!(
classify_completion_parent(metadata),
CompletionParentVerdict::Adopt
);
let span = completion_parent_span!(
target: "contract_test",
name: "chat",
operation: "chat",
system_instructions: Option::<&str>::None,
gen_ai.agent.name = "assistant",
);
let Some(metadata) = span.metadata() else {
panic!("contract span with extras was disabled");
};
let declared: HashSet<&str> =
metadata.fields().iter().map(|field| field.name()).collect();
let expected: HashSet<&str> =
expected.into_iter().chain(["gen_ai.agent.name"]).collect();
assert_eq!(declared, expected);
assert_eq!(metadata.fields().len(), expected.len());
assert_eq!(
classify_completion_parent(metadata),
CompletionParentVerdict::Adopt
);
});
}
#[test]
fn canonical_completion_span_declares_exactly_the_required_fields() {
use std::collections::HashSet;
let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
tracing::subscriber::with_default(Registry::default(), || {
let span =
CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
let Some(metadata) = span.metadata() else {
panic!("completion span was disabled");
};
let declared: HashSet<&str> =
metadata.fields().iter().map(|field| field.name()).collect();
let expected: HashSet<&str> =
COMPLETION_PARENT_REQUIRED_FIELDS.iter().copied().collect();
assert_eq!(declared, expected);
assert_eq!(metadata.fields().len(), expected.len());
});
}
#[test]
fn completion_parent_required_fields_are_pinned() {
assert_eq!(COMPLETION_PARENT_MARKER_FIELD, "rig.completion_parent");
assert_eq!(
COMPLETION_PARENT_REQUIRED_FIELDS,
&[
"gen_ai.operation.name",
"gen_ai.provider.name",
"gen_ai.request.model",
"gen_ai.system_instructions",
"gen_ai.response.id",
"gen_ai.response.model",
"gen_ai.usage.input_tokens",
"gen_ai.usage.output_tokens",
"gen_ai.usage.cache_read.input_tokens",
"gen_ai.usage.cache_creation.input_tokens",
"gen_ai.usage.tool_use_prompt_tokens",
"gen_ai.usage.reasoning_tokens",
"gen_ai.input.messages",
"gen_ai.output.messages",
]
);
}
#[test]
fn classify_completion_parent_covers_the_decision_table() {
let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
tracing::subscriber::with_default(Registry::default(), || {
let verdict = |span: &tracing::Span| {
let Some(metadata) = span.metadata() else {
panic!("classifier fixture span was disabled");
};
classify_completion_parent(metadata)
};
let conforming = completion_parent_span!(
target: "classifier_test",
name: "chat",
operation: tracing::field::Empty,
system_instructions: tracing::field::Empty,
);
assert_eq!(verdict(&conforming), CompletionParentVerdict::Adopt);
let partial = tracing::info_span!(
target: "classifier_test",
"chat",
rig.completion_parent = true,
gen_ai.operation.name = tracing::field::Empty,
);
assert_eq!(
verdict(&partial),
CompletionParentVerdict::RejectMissingFields
);
let Some(partial_metadata) = partial.metadata() else {
panic!("classifier fixture span was disabled");
};
assert_eq!(
missing_required_fields(partial_metadata),
COMPLETION_PARENT_REQUIRED_FIELDS
.iter()
.copied()
.filter(|name| *name != "gen_ai.operation.name")
.collect::<Vec<_>>()
);
let ambient = tracing::info_span!(target: "application", "ambient");
assert_eq!(verdict(&ambient), CompletionParentVerdict::NotAParent);
let lookalike = tracing::info_span!(
target: "application",
"ambient",
rig.completion_parent.id = "abc",
rig.completion_parent_id = "abc",
);
assert_eq!(verdict(&lookalike), CompletionParentVerdict::NotAParent);
});
}
#[test]
fn near_miss_completion_parent_warns_once_per_callsite() {
let warnings = CapturedWarnings::default();
let subscriber = Registry::default().with(WarningCaptureLayer {
warnings: warnings.clone(),
});
let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
reset_near_miss_warnings();
tracing::subscriber::with_default(subscriber, || {
tracing::callsite::rebuild_interest_cache();
let near_miss = tracing::info_span!(
target: "third_party_runtime",
"chat",
rig.completion_parent = true,
gen_ai.operation.name = tracing::field::Empty,
);
let _guard = near_miss.enter();
CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
});
let captured = warnings.take();
assert_eq!(
captured.len(),
1,
"a near-miss callsite warns exactly once, got: {captured:?}"
);
let Some(message) = captured.first() else {
panic!("near miss did not warn");
};
assert!(
message.contains("gen_ai.provider.name"),
"warning must name the missing fields, got: {message}"
);
assert!(
message.contains("completion_parent_span!"),
"warning must point at the supported fix, got: {message}"
);
}
#[test]
fn distinct_near_miss_callsites_each_warn() {
let warnings = CapturedWarnings::default();
let subscriber = Registry::default().with(WarningCaptureLayer {
warnings: warnings.clone(),
});
let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
reset_near_miss_warnings();
tracing::subscriber::with_default(subscriber, || {
tracing::callsite::rebuild_interest_cache();
let first = tracing::info_span!(
target: "runtime_a",
"chat",
rig.completion_parent = true,
gen_ai.operation.name = tracing::field::Empty,
);
first.in_scope(|| {
CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
});
let second = tracing::info_span!(
target: "runtime_b",
"chat",
rig.completion_parent = true,
gen_ai.operation.name = tracing::field::Empty,
);
second.in_scope(|| {
CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
});
});
let captured = warnings.take();
assert_eq!(
captured.len(),
2,
"each offending callsite warns once, got: {captured:?}"
);
}
#[test]
fn conforming_completion_parent_never_warns() {
let warnings = CapturedWarnings::default();
let subscriber = Registry::default().with(WarningCaptureLayer {
warnings: warnings.clone(),
});
let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
reset_near_miss_warnings();
tracing::subscriber::with_default(subscriber, || {
tracing::callsite::rebuild_interest_cache();
let near_miss = tracing::info_span!(
target: "third_party_runtime",
"chat",
rig.completion_parent = true,
gen_ai.operation.name = tracing::field::Empty,
);
near_miss.in_scope(|| {
CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
});
assert_eq!(
warnings.take().len(),
1,
"control: a near miss must warn, or this test cannot detect silence"
);
let conforming = completion_parent_span!(
target: "third_party_runtime",
name: "chat",
operation: Empty,
system_instructions: Option::<&str>::None,
);
let _guard = conforming.enter();
CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
drop(_guard);
let ambient = tracing::info_span!(target: "application", "ambient");
let _ambient_guard = ambient.enter();
CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
});
let captured = warnings.take();
assert!(
captured.is_empty(),
"adoption and non-participation are both silent, got: {captured:?}"
);
}
#[test]
fn agent_chat_span_is_adopted_and_enriched() {
let captured = CapturedSpan::default();
let subscriber = Registry::default().with(SpanCaptureLayer {
span: captured.clone(),
});
let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
tracing::subscriber::with_default(subscriber, || {
let completion_parent = completion_parent_span!(
target: "rig::agent_chat",
name: "chat_streaming",
operation: tracing::field::Empty,
system_instructions: tracing::field::Empty,
);
let _guard = completion_parent.enter();
let span = CompletionSpanBuilder::new(
"anthropic",
"claude-sonnet",
CompletionOperation::ChatStreaming,
)
.system_instructions(Some("provider system"), true)
.build();
assert_eq!(span.id(), completion_parent.id());
});
let Ok(captured) = captured.0.lock() else {
panic!("captured span lock poisoned");
};
let Some(span) = captured.as_ref() else {
panic!("completion-parent span was not captured");
};
assert_eq!(span.target, "rig::agent_chat");
for (field, value) in [
("gen_ai.operation.name", "chat_streaming"),
("gen_ai.provider.name", "anthropic"),
("gen_ai.request.model", "claude-sonnet"),
(
"gen_ai.system_instructions",
r#"[{"type":"text","content":"provider system"}]"#,
),
] {
assert!(contains_string(&span.recorded_values, field, value));
}
}
#[test]
fn neutral_completion_parent_span_is_adopted_and_enriched() {
let captured = CapturedSpan::default();
let subscriber = Registry::default().with(SpanCaptureLayer {
span: captured.clone(),
});
let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
tracing::subscriber::with_default(subscriber, || {
let completion_parent = completion_parent_span!(
target: "test_runtime",
name: "chat",
operation: tracing::field::Empty,
system_instructions: tracing::field::Empty,
);
let _guard = completion_parent.enter();
let span = CompletionSpanBuilder::new(
"neutral-provider",
"neutral-model",
CompletionOperation::Chat,
)
.build();
assert_eq!(span.id(), completion_parent.id());
});
let Ok(captured) = captured.0.lock() else {
panic!("captured span lock poisoned");
};
let Some(span) = captured.as_ref() else {
panic!("neutral completion-parent span was not captured");
};
assert_eq!(span.target, "test_runtime");
for (field, value) in [
("gen_ai.operation.name", "chat"),
("gen_ai.provider.name", "neutral-provider"),
("gen_ai.request.model", "neutral-model"),
] {
assert!(contains_string(&span.recorded_values, field, value));
}
}
#[test]
fn absent_provider_system_does_not_overwrite_agent_instructions() {
let captured = CapturedSpan::default();
let subscriber = Registry::default().with(SpanCaptureLayer {
span: captured.clone(),
});
let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
tracing::subscriber::with_default(subscriber, || {
let completion_parent = completion_parent_span!(
target: "test_runtime",
name: "chat",
operation: tracing::field::Empty,
system_instructions: "effective agent instructions",
);
let _guard = completion_parent.enter();
CompletionSpanBuilder::new("openai", "gpt-5", CompletionOperation::Chat).build();
});
let Ok(captured) = captured.0.lock() else {
panic!("captured span lock poisoned");
};
let Some(span) = captured.as_ref() else {
panic!("completion-parent span was not captured");
};
assert!(contains_string(
&span.initial_values,
"gen_ai.system_instructions",
"effective agent instructions"
));
assert!(
!span
.recorded_values
.iter()
.any(|(field, _)| field == "gen_ai.system_instructions")
);
}
#[test]
fn record_token_usage_records_tool_use_prompt_tokens() {
let fields = CapturedFields::default();
let subscriber = Registry::default().with(FieldCaptureLayer {
fields: fields.clone(),
});
let usage = TestUsage(Usage {
input_tokens: 1,
output_tokens: 2,
total_tokens: 15,
cached_input_tokens: 3,
cache_creation_input_tokens: 4,
tool_use_prompt_tokens: 12,
reasoning_tokens: 5,
});
let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
tracing::subscriber::with_default(subscriber, || {
let span = tracing::info_span!(
"usage_recording",
gen_ai.usage.input_tokens = tracing::field::Empty,
gen_ai.usage.output_tokens = tracing::field::Empty,
gen_ai.usage.cache_read.input_tokens = tracing::field::Empty,
gen_ai.usage.cache_creation.input_tokens = tracing::field::Empty,
gen_ai.usage.tool_use_prompt_tokens = tracing::field::Empty,
gen_ai.usage.reasoning_tokens = tracing::field::Empty,
);
span.record_token_usage(&usage);
});
assert!(fields.contains("gen_ai.usage.tool_use_prompt_tokens", 12));
}
}