feat(otel): add request params, response metadata, tool call parity, and agent identification (#11261)

Co-authored-by: Jack Amadeo <jackamadeo@squareup.com>
This commit is contained in:
Steven Pousty
2026-08-21 01:46:22 +00:00
committed by GitHub
parent 1f6c7524e1
commit f45ccd46f7
14 changed files with 540 additions and 63 deletions
@@ -12,6 +12,10 @@ pub struct ProviderUsage {
pub cost: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cost_source: Option<CostSource>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub finish_reasons: Option<Vec<String>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub response_id: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
@@ -48,6 +52,8 @@ impl ProviderUsage {
stats: None,
cost: None,
cost_source: None,
finish_reasons: None,
response_id: None,
}
}
@@ -61,6 +67,16 @@ impl ProviderUsage {
self.cost_source = Some(source);
self
}
pub fn with_finish_reasons(mut self, reasons: Vec<String>) -> Self {
self.finish_reasons = Some(reasons);
self
}
pub fn with_response_id(mut self, id: String) -> Self {
self.response_id = Some(id);
self
}
}
/// `input_tokens` is the total input including cache read/write tokens;
@@ -209,6 +225,46 @@ mod tests {
assert_eq!(usage.cache_write_input_tokens, Some(1000));
}
#[test]
fn provider_usage_finish_reasons_and_response_id_roundtrip() -> Result<()> {
let usage = ProviderUsage::new("gpt-4".to_string(), Usage::new(Some(10), Some(20), None))
.with_finish_reasons(vec!["stop".to_string()])
.with_response_id("resp-abc".to_string());
let json = serde_json::to_value(&usage)?;
assert_eq!(json["finish_reasons"], json!(["stop"]));
assert_eq!(json["response_id"], json!("resp-abc"));
let roundtripped: ProviderUsage = serde_json::from_value(json)?;
assert_eq!(
roundtripped.finish_reasons.as_deref(),
Some(&["stop".to_string()][..])
);
assert_eq!(roundtripped.response_id.as_deref(), Some("resp-abc"));
Ok(())
}
#[test]
fn provider_usage_omits_none_fields_in_json() -> Result<()> {
let usage = ProviderUsage::new("gpt-4".to_string(), Usage::new(Some(5), Some(10), None));
let json = serde_json::to_value(&usage)?;
assert!(!json.as_object().unwrap().contains_key("finish_reasons"));
assert!(!json.as_object().unwrap().contains_key("response_id"));
Ok(())
}
#[test]
fn provider_usage_deserializes_without_new_fields() -> Result<()> {
let json = json!({
"model": "gpt-4",
"usage": {"input_tokens": 5, "output_tokens": 10, "total_tokens": 15}
});
let usage: ProviderUsage = serde_json::from_value(json)?;
assert!(usage.finish_reasons.is_none());
assert!(usage.response_id.is_none());
Ok(())
}
#[test]
fn test_usage_addition_includes_cached_tokens() {
let usage_a =
@@ -1095,7 +1095,9 @@ where
let category = str_field("category");
// The refusal delta carries the request's usage;
// flush it so refused turns are still accounted.
if let Some(usage) = final_usage.take() {
if let Some(mut usage) = final_usage.take() {
usage.finish_reasons = Some(vec![STOP_REASON_REFUSAL.to_string()]);
usage.response_id = message_id.clone();
yield (None, Some(usage));
}
Err(ProviderError::Refusal { details, category })?;
@@ -1165,12 +1167,18 @@ where
if stop_reason.as_deref() == Some("max_tokens") {
let mut message = Message::assistant();
message.id = message_id;
message.id = message_id.clone();
message.metadata.output_token_limit_reached = true;
yield (Some(message), None);
}
if let Some(usage) = final_usage {
if let Some(mut usage) = final_usage {
if let Some(reason) = stop_reason {
usage.finish_reasons = Some(vec![reason]);
}
if let Some(id) = message_id {
usage.response_id = Some(id);
}
yield (None, Some(usage));
}
}
@@ -2441,6 +2449,11 @@ mod tests {
assert_eq!(usage.usage.output_tokens, Some(25));
assert_eq!(usage.usage.cache_read_input_tokens, Some(5000));
assert_eq!(usage.usage.cache_write_input_tokens, Some(10000));
assert_eq!(
usage.finish_reasons.as_deref(),
Some(&["end_turn".to_string()][..])
);
assert_eq!(usage.response_id.as_deref(), Some("msg_1"));
}
#[tokio::test]
@@ -2550,6 +2563,8 @@ mod tests {
.expect("a refused request should still yield its usage");
assert_eq!(usage.usage.input_tokens, Some(10));
assert_eq!(usage.usage.output_tokens, Some(5));
assert_eq!(usage.finish_reasons, Some(vec!["refusal".to_string()]));
assert_eq!(usage.response_id.as_deref(), Some("msg_1"));
let (details, category) = expect_refusal(results);
assert_eq!(details, "This request violates the usage policy.");
@@ -457,6 +457,8 @@ where
let mut last_signature: Option<String> = None;
let stream_id = Uuid::new_v4().to_string();
let mut incomplete_data: Option<String> = None;
let mut last_finish_reason: Option<String> = None;
let mut last_response_id: Option<String> = None;
while let Some(line_result) = stream.next().await {
let line = line_result?;
@@ -533,10 +535,22 @@ where
}
}
let parts = chunk
if let Some(response_id) = chunk.get("responseId").and_then(|v| v.as_str()) {
last_response_id = Some(response_id.to_string());
}
let candidate = chunk
.get("candidates")
.and_then(|v| v.as_array())
.and_then(|c| c.first())
.and_then(|c| c.first());
if let Some(reason) = candidate
.and_then(|c| c.get("finishReason"))
.and_then(|v| v.as_str())
{
last_finish_reason = Some(reason.to_string());
}
let parts = candidate
.and_then(|c| c.get("content"))
.and_then(|c| c.get("parts"))
.and_then(|p| p.as_array());
@@ -555,7 +569,11 @@ where
}
}
if let Some(usage) = final_usage {
if let Some(mut usage) = final_usage {
if let Some(reason) = last_finish_reason {
usage.finish_reasons = Some(vec![reason]);
}
usage.response_id = last_response_id;
yield (None, Some(usage));
}
}
@@ -1450,6 +1468,30 @@ mod tests {
);
}
#[tokio::test]
async fn test_streaming_response_metadata() {
use futures::StreamExt;
let lines = vec![Ok(
r#"data: {"candidates":[{"content":{"role":"model","parts":[{"text":"done"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":1,"totalTokenCount":3},"modelVersion":"gemini-test","responseId":"response-123"}"#
.to_string(),
)];
let stream = Box::pin(futures::stream::iter(lines));
let mut message_stream = std::pin::pin!(response_to_streaming_message(stream));
let mut final_usage = None;
while let Some(result) = message_stream.next().await {
let (_, usage) = result.unwrap();
if usage.is_some() {
final_usage = usage;
}
}
let usage = final_usage.unwrap();
assert_eq!(usage.finish_reasons, Some(vec!["STOP".to_string()]));
assert_eq!(usage.response_id.as_deref(), Some("response-123"));
}
#[tokio::test]
async fn test_streaming_function_call() {
use futures::StreamExt;
@@ -170,13 +170,12 @@ where
let mut accumulated_text = String::new();
let mut xml_detected = false;
let mut last_usage: Option<ProviderUsage> = None;
let mut buffered_usage: Option<ProviderUsage> = None;
while let Some(result) = base_stream.next().await {
let (message_opt, usage) = result?;
if usage.is_some() {
last_usage = usage.clone();
buffered_usage = usage.clone();
}
if let Some(message) = message_opt {
@@ -194,7 +193,7 @@ where
}
yield (Some(message), usage);
} else {
} else if usage.is_some() && !xml_detected {
yield (None, usage);
}
}
@@ -215,7 +214,7 @@ where
contents,
);
yield (Some(msg), last_usage);
yield (Some(msg), buffered_usage);
} else {
let msg = Message::new(
Role::Assistant,
@@ -224,7 +223,7 @@ where
)
.with_generated_id();
yield (Some(msg), last_usage);
yield (Some(msg), buffered_usage);
}
}
}
@@ -389,7 +388,9 @@ data: [DONE]"#;
.next()
.await
.expect("expected invalid XML fallback message")?;
assert!(usage.is_none());
let usage = usage.expect("expected buffered response metadata");
assert_eq!(usage.response_id.as_deref(), Some("ollama-source-id"));
assert_eq!(usage.finish_reasons, Some(vec!["stop".to_string()]));
let message = message.expect("expected invalid XML fallback message");
assert_eq!(message.role, Role::Assistant);
assert_eq!(message.content.len(), 1);
+123 -15
View File
@@ -679,6 +679,25 @@ pub fn format_tools(tools: &[Tool]) -> anyhow::Result<Vec<Value>> {
Ok(result)
}
pub fn record_response_metadata(usage: &mut ProviderUsage, response: &Value) {
usage.response_id = response
.get("id")
.and_then(Value::as_str)
.map(str::to_string);
let finish_reasons = response
.get("choices")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(|choice| choice.get("finish_reason").and_then(Value::as_str))
.map(str::to_string)
.collect::<Vec<_>>();
if !finish_reasons.is_empty() {
usage.finish_reasons = Some(finish_reasons);
}
}
/// Convert OpenAI's API response to internal Message format
pub fn response_to_message(response: &Value) -> anyhow::Result<Message> {
let output_token_limit_reached = response
@@ -1220,8 +1239,10 @@ where
let mut pending_inline_thinking = String::new();
let mut last_seen_model: Option<String> = None;
let mut last_response_id: Option<String> = None;
let mut last_finish_reason: Option<String> = None;
let mut output_token_limit_reached = false;
let mut output_token_limit_metadata_emitted = false;
let mut usage_emitted = false;
'outer: while let Some(response) = stream.next().await {
let response_str = response?;
@@ -1260,14 +1281,22 @@ where
}
}
if let Some(reason) = chunk.choices.first().and_then(|c| c.finish_reason.clone()) {
last_finish_reason = Some(reason);
}
let mut usage = extract_usage_with_output_tokens(&chunk, last_seen_model.as_deref());
output_token_limit_reached |= chunk
.choices
.first()
.and_then(|choice| choice.finish_reason.as_deref())
== Some("length");
if let Some(u) = usage.as_mut() {
if let Some(reason) = &last_finish_reason {
u.finish_reasons = Some(vec![reason.clone()]);
}
if let Some(id) = &last_response_id {
u.response_id = Some(id.clone());
}
}
output_token_limit_reached |= last_finish_reason.as_deref() == Some("length");
if chunk.choices.is_empty() {
usage_emitted |= usage.is_some();
yield (None, usage)
} else if chunk.choices[0].delta.tool_calls.as_ref().is_some_and(|tc| !tc.is_empty()) {
let mut tool_call_data: ToolCallData = HashMap::new();
@@ -1308,8 +1337,17 @@ where
if let Some(id) = &tool_chunk.id {
last_response_id = Some(id.clone());
}
if let Some(reason) = tool_chunk.choices.first().and_then(|c| c.finish_reason.clone()) {
last_finish_reason = Some(reason);
}
if let Some(chunk_usage) = extract_usage_with_output_tokens(&tool_chunk, last_seen_model.as_deref()) {
if let Some(mut chunk_usage) = extract_usage_with_output_tokens(&tool_chunk, last_seen_model.as_deref()) {
if let Some(reason) = &last_finish_reason {
chunk_usage.finish_reasons = Some(vec![reason.clone()]);
}
if let Some(id) = &last_response_id {
chunk_usage.response_id = Some(id.clone());
}
usage = Some(chunk_usage);
}
@@ -1492,6 +1530,7 @@ where
msg.metadata.output_token_limit_reached = output_token_limit_reached;
output_token_limit_metadata_emitted |= output_token_limit_reached;
usage_emitted |= usage.is_some();
yield (
Some(msg),
usage,
@@ -1534,18 +1573,19 @@ where
msg = msg.with_id(id);
}
yield (
Some(msg),
if chunk.choices[0].finish_reason.is_some() {
usage
} else {
None
},
)
let final_usage = if chunk.choices[0].finish_reason.is_some() {
usage
} else {
None
};
usage_emitted |= final_usage.is_some();
yield (Some(msg), final_usage)
} else if usage.is_some() {
usage_emitted = true;
yield (None, usage)
}
} else if usage.is_some() {
usage_emitted = true;
yield (None, usage)
}
}
@@ -1585,7 +1625,17 @@ where
}
if output_token_limit_reached && !output_token_limit_metadata_emitted {
yield (Some(output_token_limit_marker(last_response_id)), None)
yield (Some(output_token_limit_marker(last_response_id.clone())), None)
}
if !usage_emitted && (last_response_id.is_some() || last_finish_reason.is_some()) {
let mut usage = ProviderUsage::new(
last_seen_model.unwrap_or_else(|| "unknown".to_string()),
Usage::default(),
);
usage.response_id = last_response_id;
usage.finish_reasons = last_finish_reason.map(|reason| vec![reason]);
yield (None, Some(usage))
}
}
}
@@ -2355,6 +2405,26 @@ mod tests {
Ok(())
}
#[test]
fn test_record_response_metadata() {
let response = json!({
"id": "chatcmpl-123",
"choices": [
{"finish_reason": "stop"},
{"finish_reason": "tool_calls"}
]
});
let mut usage = ProviderUsage::new("test-model".to_string(), Usage::default());
record_response_metadata(&mut usage, &response);
assert_eq!(usage.response_id.as_deref(), Some("chatcmpl-123"));
assert_eq!(
usage.finish_reasons,
Some(vec!["stop".to_string(), "tool_calls".to_string()])
);
}
#[test]
fn test_response_to_message_marks_length_finish_reason() -> anyhow::Result<()> {
let response = json!({
@@ -3222,6 +3292,26 @@ data: [DONE]
Ok(())
}
#[tokio::test]
async fn test_streaming_metadata_without_usage() -> anyhow::Result<()> {
let response_lines = r#"
data: {"id":"chatcmpl-no-usage","model":"test-model","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]}
data: {"id":"chatcmpl-no-usage","model":"test-model","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}
data: [DONE]
"#;
let result = run_streaming_test(response_lines).await?;
assert_eq!(result.usage_count, 1);
let usage = result.usage.unwrap();
assert_eq!(usage.model, "test-model");
assert_eq!(usage.usage, Usage::default());
assert_eq!(usage.finish_reasons, Some(vec!["stop".to_string()]));
assert_eq!(usage.response_id.as_deref(), Some("chatcmpl-no-usage"));
Ok(())
}
#[tokio::test]
async fn test_openrouter_streaming_usage_yielded_once() -> anyhow::Result<()> {
let response_lines = r#"
@@ -3239,6 +3329,15 @@ data: [DONE]
assert!(result.has_text_content, "Expected text content in response");
assert_usage_yielded_once(&result, 7007, 49, 7056);
let usage = result.usage.as_ref().unwrap();
assert_eq!(
usage.finish_reasons.as_deref(),
Some(&["stop".to_string()][..])
);
assert_eq!(
usage.response_id.as_deref(),
Some("gen-1768896871-9HgAQqS1Z72C6gApaidi")
);
Ok(())
}
@@ -3262,6 +3361,15 @@ data: [DONE]
result.usage.as_ref().map(|usage| usage.model.as_str()),
Some("gpt-5.2-1106-preview")
);
let usage = result.usage.as_ref().unwrap();
assert_eq!(
usage.finish_reasons.as_deref(),
Some(&["tool_calls".to_string()][..])
);
assert_eq!(
usage.response_id.as_deref(),
Some("chatcmpl-Bk9Ye6Y0t9E7bC3DOMxCpW8eJkTKU")
);
Ok(())
}
@@ -1098,7 +1098,10 @@ where
Usage::default,
ResponseUsage::to_usage,
);
final_usage = Some(ProviderUsage::new(model.clone(), usage));
let mut pu = ProviderUsage::new(model.clone(), usage);
pu.finish_reasons = Some(vec![response.status.clone()]);
pu.response_id = Some(response.id.clone());
final_usage = Some(pu);
// For complete output, use the response output items
if !response.output.is_empty() {
@@ -1114,7 +1117,14 @@ where
Usage::default,
ResponseUsage::to_usage,
);
final_usage = Some(ProviderUsage::new(model.clone(), usage));
let mut pu = ProviderUsage::new(model.clone(), usage);
pu.finish_reasons = Some(vec![response
.incomplete_details
.as_ref()
.and_then(|details| details.reason.clone())
.unwrap_or_else(|| response.status.clone())]);
pu.response_id = Some(response.id.clone());
final_usage = Some(pu);
response_id = Some(response.id.clone());
output_token_limit_reached = response_reached_output_token_limit(
&response.status,
@@ -1383,6 +1393,10 @@ mod tests {
assert_eq!(usage.usage.input_tokens, Some(10));
assert_eq!(usage.usage.output_tokens, Some(5));
assert_eq!(usage.usage.total_tokens, Some(15));
assert_eq!(
usage.finish_reasons.as_deref(),
Some(&["max_output_tokens".to_string()][..])
);
Ok(())
}
+8 -1
View File
@@ -9,7 +9,7 @@ use crate::errors::ProviderError;
use crate::formats::openai::is_openai_responses_model;
use crate::formats::openai::{
create_request_with_options, get_cost, get_usage, is_reserved_request_param_key,
response_to_message, OpenAiFormatOptions,
record_response_metadata, response_to_message, OpenAiFormatOptions,
};
use crate::formats::openai_responses::{
create_responses_request_for_model, get_responses_usage, responses_api_to_message,
@@ -338,6 +338,12 @@ impl OpenAiProvider {
let usage_data = get_responses_usage(&parsed);
let usage_json = json.get("usage").unwrap_or(&serde_json::Value::Null);
let mut usage = ProviderUsage::new(model_config.model_name.clone(), usage_data);
usage.response_id = Some(parsed.id.clone());
let finish_reason = json
.pointer("/incomplete_details/reason")
.and_then(serde_json::Value::as_str)
.unwrap_or(&parsed.status);
usage.finish_reasons = Some(vec![finish_reason.to_string()]);
if let Some(cost) = get_cost(usage_json) {
usage = usage.with_cost(cost, CostSource::ProviderReported);
}
@@ -822,6 +828,7 @@ impl Provider for OpenAiProvider {
let usage_json = json.get("usage").unwrap_or(&serde_json::Value::Null);
let usage_data = get_usage(usage_json);
let mut usage = ProviderUsage::new(model_config.model_name.clone(), usage_data);
record_response_metadata(&mut usage, &json);
if let Some(cost) = get_cost(usage_json) {
usage = usage.with_cost(cost, CostSource::ProviderReported);
}
@@ -19,7 +19,8 @@ use crate::conversation::message::Message;
use crate::errors::ProviderError;
use crate::formats::openai::{
create_request, create_request_for_model_with_options, get_cost, get_usage,
response_to_message, response_to_streaming_message, OpenAiFormatOptions,
record_response_metadata, response_to_message, response_to_streaming_message,
OpenAiFormatOptions,
};
use crate::formats::openai_responses::responses_api_to_streaming_message;
use crate::model::ModelConfig;
@@ -134,6 +135,7 @@ impl OpenAiCompatibleProvider {
let usage_json = json.get("usage").unwrap_or(&Value::Null);
let usage_data = get_usage(usage_json);
let mut usage = ProviderUsage::new(model_config.model_name.clone(), usage_data);
record_response_metadata(&mut usage, &json);
if let Some(cost) = get_cost(usage_json) {
usage = usage.with_cost(cost, CostSource::ProviderReported);
}
+21 -18
View File
@@ -719,12 +719,8 @@ impl Agent {
if capture_message_content {
let output = gen_ai_telemetry::tool_result_json(&processed_result);
span.record("output", output.as_str());
if let Some(result) =
gen_ai_telemetry::successful_tool_result_json(&processed_result)
{
span.record("gen_ai.tool.call.result", result.as_str());
}
}
gen_ai_telemetry::record_tool_result(&span, &processed_result);
let event = match &processed_result {
Ok(call_result) if call_result.is_error != Some(true) => {
crate::hooks::HookEvent::PostToolUse
@@ -1189,17 +1185,7 @@ impl Agent {
"arguments": tool_call.arguments,
});
tracing::Span::current().record("input", tracing::field::display(&input_summary));
if gen_ai_telemetry::capture_message_content() {
let arguments = tool_call
.arguments
.as_ref()
.map(|arguments| Value::Object(arguments.clone()))
.unwrap_or_else(|| Value::Object(serde_json::Map::new()));
tracing::Span::current().record(
"gen_ai.tool.call.arguments",
tracing::field::display(arguments),
);
}
gen_ai_telemetry::record_tool_arguments(&tracing::Span::current(), &tool_call);
self.prompt_manager
.lock()
@@ -1899,6 +1885,7 @@ impl Agent {
trace_output = tracing::field::Empty,
session.id = %session_config.id,
gen_ai.operation.name = "invoke_agent",
gen_ai.agent.name = tracing::field::Empty,
gen_ai.input.messages = tracing::field::Empty,
gen_ai.output.messages = tracing::field::Empty,
gen_ai.usage.input_tokens = tracing::field::Empty,
@@ -1911,13 +1898,18 @@ impl Agent {
session_config: SessionConfig,
cancel_token: Option<CancellationToken>,
) -> Result<BoxStream<'_, Result<AgentEvent>>> {
let reply_span = tracing::Span::current();
let events = self
.reply_impl(user_message, session_config, cancel_token)
.await?;
// This is the single live-event identity boundary. Callers that intentionally stream
// multiple events for one logical message must assign their shared ID before this point.
Ok(Box::pin(events.map_ok(ensure_message_event_id)))
Ok(Box::pin(
events
.map_ok(ensure_message_event_id)
.instrument(reply_span),
))
}
async fn reply_impl(
@@ -1991,6 +1983,8 @@ impl Agent {
let session = session_manager
.get_session(&session_config.id, true)
.await?;
tracing::Span::current()
.record("gen_ai.agent.name", gen_ai_telemetry::agent_name(&session));
let is_first_agent_turn = session
.conversation
.as_ref()
@@ -2162,6 +2156,7 @@ impl Agent {
let conversation_to_compact = conversation.clone();
let reply_span = tracing::Span::current();
reply_span.record("gen_ai.agent.name", gen_ai_telemetry::agent_name(&session));
Ok(Box::pin(async_stream::try_stream! {
for event in command_preamble {
@@ -2233,7 +2228,8 @@ impl Agent {
}
};
let mut reply_stream = self.reply_internal(final_conversation, session_config, session, cancel_token, reply_span.clone()).await?;
let parent_span = tracing::Span::current();
let mut reply_stream = self.reply_internal(final_conversation, session_config, session, cancel_token, parent_span.clone()).await?;
while let Some(event) = reply_stream.next().await {
yield event?;
}
@@ -2337,14 +2333,21 @@ impl Agent {
session.host = %crate::session_context::session_host(),
session.agent_type = "goose",
gen_ai.operation.name = "invoke_agent",
gen_ai.agent.name = tracing::field::Empty,
gen_ai.conversation.id = %session_config.id,
gen_ai.request.model = %model_config.model_name,
gen_ai.request.temperature = tracing::field::Empty,
gen_ai.request.max_tokens = tracing::field::Empty,
gen_ai.provider.name = %provider_name,
gen_ai.input.messages = tracing::field::Empty,
gen_ai.output.messages = tracing::field::Empty,
gen_ai.response.finish_reasons = tracing::field::Empty,
gen_ai.response.id = tracing::field::Empty,
gen_ai.usage.input_tokens = tracing::field::Empty,
gen_ai.usage.output_tokens = tracing::field::Empty,
);
gen_ai_telemetry::record_request_params(&reply_stream_span, &model_config);
reply_stream_span.record("gen_ai.agent.name", gen_ai_telemetry::agent_name(&session));
if gen_ai_telemetry::capture_message_content() {
if let Some(last_user_msg) = conversation
.messages()
+210
View File
@@ -1,5 +1,7 @@
use crate::conversation::message::{Message, MessageContent, ToolResult};
use crate::session::Session;
use goose_providers::conversation::token_usage::{ProviderUsage, Usage};
use goose_providers::model::ModelConfig;
use rmcp::model::{CallToolRequestParams, CallToolResult, Role};
use serde_json::{json, Value};
use tracing::Span;
@@ -66,6 +68,51 @@ pub(super) fn record_usage(span: &Span, usage: &Usage) {
pub(super) fn record_provider_usage(span: &Span, usage: &ProviderUsage) {
span.record("gen_ai.response.model", usage.model.as_str());
record_usage(span, &usage.usage);
if let Some(reasons) = &usage.finish_reasons {
let reasons_json = serde_json::to_string(reasons).unwrap_or_default();
span.record("gen_ai.response.finish_reasons", reasons_json.as_str());
}
if let Some(id) = &usage.response_id {
span.record("gen_ai.response.id", id.as_str());
}
}
pub(super) fn record_request_params(span: &Span, model_config: &ModelConfig) {
if let Some(temperature) = model_config.temperature {
span.record("gen_ai.request.temperature", temperature as f64);
}
if let Some(max_tokens) = model_config.max_tokens {
span.record("gen_ai.request.max_tokens", max_tokens as i64);
}
}
pub(super) fn record_tool_arguments(span: &Span, tool_call: &CallToolRequestParams) {
if capture_message_content() {
let arguments = tool_call
.arguments
.as_ref()
.map(|arguments| Value::Object(arguments.clone()))
.unwrap_or_else(|| Value::Object(serde_json::Map::new()));
span.record(
"gen_ai.tool.call.arguments",
tracing::field::display(arguments),
);
}
}
pub(super) fn record_tool_result(span: &Span, result: &ToolResult<CallToolResult>) {
if capture_message_content() {
if let Some(result_json) = successful_tool_result_json(result) {
span.record("gen_ai.tool.call.result", result_json.as_str());
}
}
}
pub(super) fn agent_name(session: &Session) -> &str {
session
.recipe
.as_ref()
.map_or("goose", |recipe| recipe.title.as_str())
}
pub(super) fn tool_result_json(result: &ToolResult<CallToolResult>) -> String {
@@ -325,6 +372,15 @@ mod tests {
use goose_test_support::otel::clear_otel_env;
use rmcp::{model::CallToolRequestParams, object};
fn test_recipe(title: &str) -> crate::recipe::Recipe {
serde_json::from_value(serde_json::json!({
"title": title,
"description": "test recipe",
"instructions": "do stuff",
}))
.unwrap()
}
#[test]
fn content_capture_requires_explicit_opt_in() {
let _env = clear_otel_env(&[]);
@@ -367,4 +423,158 @@ mod tests {
let value: Value = serde_json::from_str(&output_message_json(&message)).unwrap();
assert_eq!(value[0]["finish_reason"], "tool_call");
}
#[test]
fn record_request_params_records_temperature_and_max_tokens() {
let capture = test_support::SpanFieldCapture::new("test_span");
let _guard = capture.clone().set_default();
let config = ModelConfig::new("test-model")
.with_temperature(Some(0.5))
.with_max_tokens(Some(4096));
let span = tracing::info_span!(
"test_span",
"gen_ai.request.temperature" = tracing::field::Empty,
"gen_ai.request.max_tokens" = tracing::field::Empty,
);
record_request_params(&span, &config);
let fields = capture.fields();
assert_eq!(fields["gen_ai.request.temperature"], 0.5);
assert_eq!(fields["gen_ai.request.max_tokens"], 4096);
}
#[test]
fn record_request_params_skips_none_values() {
let capture = test_support::SpanFieldCapture::new("test_span");
let _guard = capture.clone().set_default();
let config = ModelConfig::new("test-model");
let span = tracing::info_span!(
"test_span",
"gen_ai.request.temperature" = tracing::field::Empty,
"gen_ai.request.max_tokens" = tracing::field::Empty,
);
record_request_params(&span, &config);
let fields = capture.fields();
assert!(!fields.contains_key("gen_ai.request.temperature"));
assert!(!fields.contains_key("gen_ai.request.max_tokens"));
}
#[test]
fn record_tool_arguments_gated_by_content_env() {
let _env = clear_otel_env(&[]);
let capture = test_support::SpanFieldCapture::new("test_span");
let _guard = capture.clone().set_default();
let tool_call =
CallToolRequestParams::new("my_tool").with_arguments(object!({ "key": "value" }));
let span = tracing::info_span!(
"test_span",
"gen_ai.tool.call.arguments" = tracing::field::Empty,
);
record_tool_arguments(&span, &tool_call);
let fields = capture.fields();
assert!(!fields.contains_key("gen_ai.tool.call.arguments"));
drop(_env);
let _env = clear_otel_env(&[(CAPTURE_MESSAGE_CONTENT_ENV, "true")]);
let capture2 = test_support::SpanFieldCapture::new("test_span2");
let _guard2 = capture2.clone().set_default();
let span2 = tracing::info_span!(
"test_span2",
"gen_ai.tool.call.arguments" = tracing::field::Empty,
);
record_tool_arguments(&span2, &tool_call);
let fields2 = capture2.fields();
assert!(fields2.contains_key("gen_ai.tool.call.arguments"));
let args: Value =
serde_json::from_str(fields2["gen_ai.tool.call.arguments"].as_str().unwrap()).unwrap();
assert_eq!(args["key"], "value");
}
#[test]
fn record_tool_result_only_on_success() {
let _env = clear_otel_env(&[(CAPTURE_MESSAGE_CONTENT_ENV, "true")]);
let capture = test_support::SpanFieldCapture::new("test_span");
let _guard = capture.clone().set_default();
let success_result: ToolResult<CallToolResult> = Ok(CallToolResult::success(vec![
rmcp::model::ContentBlock::text("ok"),
]));
let span = tracing::info_span!(
"test_span",
"gen_ai.tool.call.result" = tracing::field::Empty,
);
record_tool_result(&span, &success_result);
let fields = capture.fields();
assert!(fields.contains_key("gen_ai.tool.call.result"));
let capture2 = test_support::SpanFieldCapture::new("test_span2");
let _guard2 = capture2.clone().set_default();
let error_result: ToolResult<CallToolResult> = Err(rmcp::model::ErrorData::new(
rmcp::model::ErrorCode::INTERNAL_ERROR,
"failed".to_string(),
None,
));
let span2 = tracing::info_span!(
"test_span2",
"gen_ai.tool.call.result" = tracing::field::Empty,
);
record_tool_result(&span2, &error_result);
let fields2 = capture2.fields();
assert!(!fields2.contains_key("gen_ai.tool.call.result"));
}
#[test]
fn record_provider_usage_includes_finish_reasons_and_response_id() {
let capture = test_support::SpanFieldCapture::new("test_span");
let _guard = capture.clone().set_default();
let usage = ProviderUsage::new(
"test-model".to_string(),
Usage::new(Some(10), Some(20), None),
)
.with_finish_reasons(vec!["stop".to_string()])
.with_response_id("resp-123".to_string());
let span = tracing::info_span!(
"test_span",
"gen_ai.response.model" = tracing::field::Empty,
"gen_ai.response.finish_reasons" = tracing::field::Empty,
"gen_ai.response.id" = tracing::field::Empty,
"gen_ai.usage.input_tokens" = tracing::field::Empty,
"gen_ai.usage.output_tokens" = tracing::field::Empty,
);
record_provider_usage(&span, &usage);
let fields = capture.fields();
assert_eq!(fields["gen_ai.response.model"], "test-model");
assert_eq!(fields["gen_ai.response.finish_reasons"], "[\"stop\"]");
assert_eq!(fields["gen_ai.response.id"], "resp-123");
assert_eq!(fields["gen_ai.usage.input_tokens"], 10);
assert_eq!(fields["gen_ai.usage.output_tokens"], 20);
}
#[test]
fn agent_name_returns_recipe_title_when_present() {
let session = Session {
recipe: Some(test_recipe("My Recipe")),
..Default::default()
};
assert_eq!(agent_name(&session), "My Recipe");
}
#[test]
fn agent_name_returns_goose_default() {
let session = Session::default();
assert_eq!(agent_name(&session), "goose");
}
}
+5
View File
@@ -315,7 +315,11 @@ pub(crate) fn prepare_tools_for_provider(
gen_ai.provider.name = %provider.get_name(),
gen_ai.request.model = %model_config.model_name,
gen_ai.request.stream = true,
gen_ai.request.temperature = tracing::field::Empty,
gen_ai.request.max_tokens = tracing::field::Empty,
gen_ai.response.model = tracing::field::Empty,
gen_ai.response.finish_reasons = tracing::field::Empty,
gen_ai.response.id = tracing::field::Empty,
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,
@@ -350,6 +354,7 @@ pub(crate) async fn stream_response_from_provider(
filtered_messages
};
let span = tracing::Span::current();
gen_ai_telemetry::record_request_params(&span, &model_config);
let capture_message_content = gen_ai_telemetry::capture_message_content();
if capture_message_content {
let input_messages =
@@ -86,29 +86,29 @@ pub(super) fn chat_span(
session_id: &str,
purpose: &'static str,
) -> tracing::Span {
tracing::info_span!(
let span = tracing::info_span!(
target: "goose::state_machine",
"chat",
"gen_ai.operation.name" = "chat",
"gen_ai.provider.name" = %provider.get_name(),
"gen_ai.request.model" = %model_config.model_name,
"gen_ai.request.temperature" = tracing::field::Empty,
"gen_ai.request.max_tokens" = tracing::field::Empty,
"gen_ai.response.model" = tracing::field::Empty,
"gen_ai.response.finish_reasons" = tracing::field::Empty,
"gen_ai.response.id" = tracing::field::Empty,
"gen_ai.usage.input_tokens" = tracing::field::Empty,
"gen_ai.usage.output_tokens" = tracing::field::Empty,
"goose.chat.purpose" = purpose,
"error.type" = tracing::field::Empty,
session.id = %session_id,
)
);
super::super::gen_ai_telemetry::record_request_params(&span, model_config);
span
}
pub(super) fn record_chat_usage(span: &tracing::Span, usage: &ProviderUsage) {
span.record("gen_ai.response.model", usage.model.as_str());
if let Some(tokens) = usage.usage.input_tokens {
span.record("gen_ai.usage.input_tokens", tokens);
}
if let Some(tokens) = usage.usage.output_tokens {
span.record("gen_ai.usage.output_tokens", tokens);
}
super::super::gen_ai_telemetry::record_provider_usage(span, usage);
}
pub struct InferenceRunner<'a> {
@@ -575,10 +575,16 @@ impl Inference<Session, GooseEffect> for InferenceRunner<'_> {
return yielded_with(usage_effects);
}
let has_recorded_usage = usage_effects
.iter()
.any(|effect| matches!(effect, GooseEffect::RecordUsage(_)));
if !has_recorded_usage {
let has_recorded_tokens = usage_effects.iter().any(|effect| {
matches!(
effect,
GooseEffect::RecordUsage(usage)
if usage.usage.input_tokens.is_some()
|| usage.usage.output_tokens.is_some()
|| usage.usage.total_tokens.is_some()
)
});
if !has_recorded_tokens {
let mut usage = ProviderUsage::new(
self.model_config.model_name.clone(),
goose_providers::conversation::token_usage::Usage::default(),
@@ -75,6 +75,8 @@ pub(super) fn tool_span(tool_name: &str, tool_call_id: &str, session_id: &str) -
"gen_ai.operation.name" = "execute_tool",
"gen_ai.tool.name" = %tool_name,
"gen_ai.tool.call.id" = %tool_call_id,
"gen_ai.tool.call.arguments" = tracing::field::Empty,
"gen_ai.tool.call.result" = tracing::field::Empty,
"error.type" = tracing::field::Empty,
session.id = %session_id,
)
@@ -257,6 +259,7 @@ pub(super) fn with_post_tool_hooks(
let future = async move {
let result =
crate::agents::large_response_handler::process_tool_response(result.result.await);
crate::agents::gen_ai_telemetry::record_tool_result(&tracing::Span::current(), &result);
match &result {
Ok(result) if result.is_error == Some(true) => {
tracing::Span::current().record("error.type", "tool_error");
@@ -340,6 +343,7 @@ impl<'a> ToolExecutionOperation<'a> {
session: &Session,
) -> std::result::Result<ToolCallResult, ErrorData> {
let span = tool_span(&tool_call.name, &request_id, &session.id);
crate::agents::gen_ai_telemetry::record_tool_arguments(&span, &tool_call);
let result_span = span.clone();
async {
@@ -158,6 +158,10 @@ pub(crate) async fn run(
emit: &Emitter,
) -> Result<Session> {
let entry_session = runtime.load(session_id).await?;
tracing::Span::current().record(
"gen_ai.agent.name",
crate::agents::gen_ai_telemetry::agent_name(&entry_session),
);
if let Some(input) = entry_session
.conversation()
.and_then(|conversation| {