diff --git a/crates/goose-test-support/src/mcp.rs b/crates/goose-test-support/src/mcp.rs index 6938132850..4021db54ef 100644 --- a/crates/goose-test-support/src/mcp.rs +++ b/crates/goose-test-support/src/mcp.rs @@ -31,6 +31,8 @@ static NEXT_INSTANCE_ID: AtomicU64 = AtomicU64::new(1); const SESSION_ID_META_KEY: &str = "agent-session-id"; const WORKING_DIR_META_KEY: &str = "agent-working-dir"; +const TRACEPARENT_META_KEY: &str = "traceparent"; +const TRACESTATE_META_KEY: &str = "tracestate"; #[derive(Debug, serde::Deserialize, serde::Serialize)] #[serde(rename_all = "camelCase")] @@ -42,6 +44,8 @@ pub struct ContextReport { pub process_session_id: Option, pub request_session_id: Option, pub request_working_dir: Option, + pub request_traceparent: Option, + pub request_tracestate: Option, pub roots: Vec, } @@ -157,6 +161,8 @@ impl McpFixtureServer { process_session_id: std::env::var("AGENT_SESSION_ID").ok(), request_session_id: meta_value(SESSION_ID_META_KEY), request_working_dir: meta_value(WORKING_DIR_META_KEY), + request_traceparent: meta_value(TRACEPARENT_META_KEY), + request_tracestate: meta_value(TRACESTATE_META_KEY), roots, }; Ok(CallToolResult::success(vec![ContentBlock::text( diff --git a/crates/goose/src/acp/server.rs b/crates/goose/src/acp/server.rs index 836d0f0906..c0096bf3cb 100644 --- a/crates/goose/src/acp/server.rs +++ b/crates/goose/src/acp/server.rs @@ -2307,16 +2307,20 @@ impl GooseAcpAgent { } let user_message = Self::convert_acp_prompt_to_message(&args.prompt); + // The client's trace context must stay attached while the turn's spans and tool calls are polled. + #[cfg(feature = "otel")] + let trace_cx = crate::otel::trace_context::extract_from_meta(args.meta.as_ref()) + .unwrap_or_else(opentelemetry::Context::current); let session_config = SessionConfig { id: session_id.clone(), schedule_id: None, max_turns: None, }; - let stream = match agent - .reply(user_message, session_config, Some(cancel_token.clone())) - .await - { + let reply = agent.reply(user_message, session_config, Some(cancel_token.clone())); + #[cfg(feature = "otel")] + let reply = opentelemetry::context::FutureExt::with_context(reply, trace_cx.clone()); + let stream = match reply.await { Ok(stream) => stream, Err(error) => { self.clear_active_run(&session_id, &run_id).await; @@ -2325,16 +2329,17 @@ impl GooseAcpAgent { .data(format!("Error getting agent reply: {error}"))); } }; - let stream_result = self - .forward_agent_stream( - cx, - &args.session_id, - &session_id, - &agent, - &cancel_token, - stream, - ) - .await; + let forward = self.forward_agent_stream( + cx, + &args.session_id, + &session_id, + &agent, + &cancel_token, + stream, + ); + #[cfg(feature = "otel")] + let forward = opentelemetry::context::FutureExt::with_context(forward, trace_cx); + let stream_result = forward.await; self.clear_active_run(&session_id, &run_id).await; Self::send_active_run_update(cx, &args.session_id, None)?; let outcome = stream_result?; diff --git a/crates/goose/src/agents/mcp_client.rs b/crates/goose/src/agents/mcp_client.rs index 8d1cf84788..f01af02a22 100644 --- a/crates/goose/src/agents/mcp_client.rs +++ b/crates/goose/src/agents/mcp_client.rs @@ -758,12 +758,13 @@ impl McpClientTrait for McpClient { client.peer_info().map(|info| info.protocol_version.clone()) }; if protocol_version.as_ref() == Some(&ProtocolVersion::V_2026_07_28) { - let extensions = inject_session_context_into_extensions( - Extensions::new(), - Some(&ctx.session_id), - ctx.working_dir_str(), - ctx.tool_call_request_id.as_deref(), - ); + let extensions = + inject_trace_context_into_extensions(inject_session_context_into_extensions( + Extensions::new(), + Some(&ctx.session_id), + ctx.working_dir_str(), + ctx.tool_call_request_id.as_deref(), + )); if let Some(meta) = extensions.get::() { params.meta.get_or_insert_default().0 .0.extend( meta.0 @@ -935,6 +936,22 @@ fn inject_session_context_into_extensions( extensions } +#[cfg(feature = "otel")] +fn inject_trace_context_into_extensions(mut extensions: Extensions) -> Extensions { + let mut meta_map = extensions + .get::() + .map(|meta| meta.0.clone()) + .unwrap_or_default(); + crate::otel::trace_context::inject_current(&mut meta_map); + extensions.insert(MetaObject(meta_map)); + extensions +} + +#[cfg(not(feature = "otel"))] +fn inject_trace_context_into_extensions(extensions: Extensions) -> Extensions { + extensions +} + fn inject_session_context_into_request( request: ClientRequest, session_id: Option<&str>, @@ -970,12 +987,13 @@ fn inject_session_context_into_request( ClientRequest::ListToolsRequest(req) } ClientRequest::CallToolRequest(mut req) => { - req.extensions = inject_session_context_into_extensions( - req.extensions, - session_id, - working_dir, - tool_call_request_id, - ); + req.extensions = + inject_trace_context_into_extensions(inject_session_context_into_extensions( + req.extensions, + session_id, + working_dir, + tool_call_request_id, + )); ClientRequest::CallToolRequest(req) } ClientRequest::ListPromptsRequest(mut req) => { @@ -1174,6 +1192,59 @@ mod tests { assert_eq!(&mcp_meta.0, expected_meta.as_object().unwrap()); } + #[cfg(feature = "otel")] + #[test_case(list_resources_request, false; "list_resources")] + #[test_case(read_resource_request, false; "read_resource")] + #[test_case(list_tools_request, false; "list_tools")] + #[test_case(call_tool_request, true; "call_tool")] + #[test_case(list_prompts_request, false; "list_prompts")] + #[test_case(get_prompt_request, false; "get_prompt")] + fn test_request_injects_trace_context_only_on_call_tool( + request_builder: fn(Extensions) -> ClientRequest, + expect_trace_context: bool, + ) { + use opentelemetry::trace::{ + SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId, TraceState, + }; + + let _guard = opentelemetry::Context::new() + .with_remote_span_context(SpanContext::new( + TraceId::from_hex("4bf92f3577b34da6a3ce929d0e0e4736").unwrap(), + SpanId::from_hex("00f067aa0ba902b7").unwrap(), + TraceFlags::SAMPLED, + true, + "vendor=value".parse::().unwrap(), + )) + .attach(); + + let request = inject_session_context_into_request( + request_builder(Extensions::new()), + Some("test-session-id"), + None, + None, + ); + let meta = request_extensions(&request) + .and_then(|extensions| extensions.get::()) + .expect("request should carry meta"); + + assert_eq!( + meta.0.get(SESSION_ID_HEADER), + Some(&json!("test-session-id")) + ); + if expect_trace_context { + assert_eq!( + meta.0.get("traceparent"), + Some(&json!( + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" + )) + ); + assert_eq!(meta.0.get("tracestate"), Some(&json!("vendor=value"))); + } else { + assert!(!meta.0.contains_key("traceparent")); + assert!(!meta.0.contains_key("tracestate")); + } + } + #[test] fn test_client_info_advertises_mcp_apps_ui_extension() { let client = new_client(GoosePlatform::GooseDesktop); diff --git a/crates/goose/src/otel/mod.rs b/crates/goose/src/otel/mod.rs index 95eb89aa4d..9ecd468a4f 100644 --- a/crates/goose/src/otel/mod.rs +++ b/crates/goose/src/otel/mod.rs @@ -1 +1,2 @@ pub mod otlp; +pub mod trace_context; diff --git a/crates/goose/src/otel/trace_context.rs b/crates/goose/src/otel/trace_context.rs new file mode 100644 index 0000000000..c9759f3534 --- /dev/null +++ b/crates/goose/src/otel/trace_context.rs @@ -0,0 +1,126 @@ +//! W3C trace context carried in ACP and MCP `_meta` (MCP SEP-414 keys). + +use opentelemetry::propagation::TextMapPropagator; +use opentelemetry::trace::TraceContextExt; +use opentelemetry::Context; +use opentelemetry_sdk::propagation::TraceContextPropagator; +use serde_json::{Map, Value}; +use std::collections::HashMap; + +const TRACEPARENT: &str = "traceparent"; +const TRACESTATE: &str = "tracestate"; + +/// Returns the remote context from `_meta.traceparent`/`tracestate`, or `None` if absent or invalid. +/// Only the lowercase SEP-414 keys are read, so mixed-case duplicates cannot race. +pub fn extract_from_meta(meta: Option<&Map>) -> Option { + let meta = meta?; + let carrier: HashMap = [TRACEPARENT, TRACESTATE] + .into_iter() + .filter_map(|key| Some((key.to_string(), meta.get(key)?.as_str()?.to_string()))) + .collect(); + let cx = TraceContextPropagator::new().extract_with_context(&Context::new(), &carrier); + cx.span().span_context().is_valid().then_some(cx) +} + +/// Writes the current OpenTelemetry context into `meta`; no-op when there is no valid span context. +pub fn inject_current(meta: &mut Map) { + let mut carrier = HashMap::new(); + TraceContextPropagator::new().inject_context(&Context::current(), &mut carrier); + meta.extend( + carrier + .into_iter() + .filter(|(_, value)| !value.is_empty()) + .map(|(key, value)| (key, Value::String(value))), + ); +} + +#[cfg(test)] +mod tests { + use super::*; + use opentelemetry::trace::{ + SpanContext, SpanId, TraceFlags, TraceId, TraceState, TracerProvider, + }; + use opentelemetry_sdk::trace::SdkTracerProvider; + use serde_json::json; + use test_case::test_case; + use tracing_opentelemetry::OpenTelemetrySpanExt; + use tracing_subscriber::layer::SubscriberExt; + + const TRACE_ID: &str = "4bf92f3577b34da6a3ce929d0e0e4736"; + const SPAN_ID: &str = "00f067aa0ba902b7"; + const PARENT: &str = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"; + + fn meta(value: Value) -> Map { + value.as_object().unwrap().clone() + } + + fn remote_context() -> Context { + Context::new().with_remote_span_context(SpanContext::new( + TraceId::from_hex(TRACE_ID).unwrap(), + SpanId::from_hex(SPAN_ID).unwrap(), + TraceFlags::SAMPLED, + true, + "vendor=value".parse::().unwrap(), + )) + } + + #[test] + fn extracts_remote_context() { + let value = json!({ "traceparent": PARENT, "tracestate": "vendor=value" }); + let cx = extract_from_meta(Some(&meta(value))).unwrap(); + let span_context = cx.span().span_context().clone(); + + assert_eq!(span_context.trace_id().to_string(), TRACE_ID); + assert_eq!(span_context.span_id().to_string(), SPAN_ID); + assert!(span_context.is_remote()); + assert!(span_context.is_sampled()); + assert_eq!(span_context.trace_state().header(), "vendor=value"); + } + + #[test_case(None; "no meta")] + #[test_case(Some(json!({ "tracestate": "vendor=value" })); "no traceparent")] + #[test_case(Some(json!({ "traceparent": "not-a-traceparent" })); "malformed")] + #[test_case(Some(json!({ "traceparent": "00-00000000000000000000000000000000-00f067aa0ba902b7-01" })); "zero trace id")] + #[test_case(Some(json!({ "traceparent": 42 })); "not a string")] + #[test_case(Some(json!({ "TraceParent": PARENT })); "non lowercase key")] + fn ignores_missing_or_invalid_context(value: Option) { + let meta = value.map(meta); + assert!(extract_from_meta(meta.as_ref()).is_none()); + } + + #[test] + fn injects_nothing_without_current_context() { + let mut meta = Map::new(); + inject_current(&mut meta); + assert!(meta.is_empty()); + } + + #[test] + fn relays_attached_remote_context_unchanged() { + let _guard = remote_context().attach(); + let mut meta = Map::new(); + inject_current(&mut meta); + + assert_eq!(meta[TRACEPARENT], PARENT); + assert_eq!(meta[TRACESTATE], "vendor=value"); + } + + #[test] + fn injects_entered_tracing_span_as_parent() { + let provider = SdkTracerProvider::builder().build(); + let subscriber = tracing_subscriber::registry() + .with(tracing_opentelemetry::layer().with_tracer(provider.tracer("test"))); + + tracing::subscriber::with_default(subscriber, || { + let _remote = remote_context().attach(); + let span = tracing::info_span!("tool"); + let _entered = span.enter(); + let span_id = span.context().span().span_context().span_id(); + let mut meta = Map::new(); + inject_current(&mut meta); + + assert_eq!(meta[TRACEPARENT], format!("00-{TRACE_ID}-{span_id}-01"),); + assert_eq!(meta[TRACESTATE], "vendor=value"); + }); + } +} diff --git a/crates/goose/tests/acp_test_data/openai_inspect_context_tool_call.txt b/crates/goose/tests/acp_test_data/openai_inspect_context_tool_call.txt new file mode 100644 index 0000000000..ea674af0d5 --- /dev/null +++ b/crates/goose/tests/acp_test_data/openai_inspect_context_tool_call.txt @@ -0,0 +1,10 @@ +data: {"id":"chatcmpl-trace","object":"chat.completion.chunk","created":1766709751,"model":"gpt-5-nano","choices":[{"index":0,"delta":{"role":"assistant","content":null,"tool_calls":[{"index":0,"id":"call_inspect_context","type":"function","function":{"name":"mcp-fixture__inspect_context","arguments":""}}]},"finish_reason":null}]} + +data: {"id":"chatcmpl-trace","object":"chat.completion.chunk","created":1766709751,"model":"gpt-5-nano","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{}"}}]},"finish_reason":null}]} + +data: {"id":"chatcmpl-trace","object":"chat.completion.chunk","created":1766709751,"model":"gpt-5-nano","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]} + +data: {"id":"chatcmpl-trace","object":"chat.completion.chunk","created":1766709751,"model":"gpt-5-nano","choices":[],"usage":{"prompt_tokens":100,"completion_tokens":10,"total_tokens":110}} + +data: [DONE] + diff --git a/crates/goose/tests/acp_trace_context_test.rs b/crates/goose/tests/acp_trace_context_test.rs new file mode 100644 index 0000000000..fbd1fb7a62 --- /dev/null +++ b/crates/goose/tests/acp_trace_context_test.rs @@ -0,0 +1,253 @@ +#![cfg(feature = "otel")] +#![recursion_limit = "256"] +#[allow(dead_code)] +#[path = "acp_common_tests/mod.rs"] +mod common_tests; + +use agent_client_protocol::schema::v1::{ + ContentBlock, McpServer, McpServerHttp, PromptRequest, SessionUpdate, StopReason, TextContent, + ToolCallContent, +}; +use common_tests::fixtures::server::AcpServerConnection; +use common_tests::fixtures::{ + run_test, Connection, OpenAiFixture, Session, SessionData, TestConnectionConfig, +}; +use common_tests::TURN_CONTEXT_OPEN; +use goose_test_support::mcp::ContextReport; +use goose_test_support::McpFixture; +use opentelemetry::trace::{SpanId, TraceId, TracerProvider as _}; +use opentelemetry_sdk::trace::{InMemorySpanExporter, SdkTracerProvider, SpanData}; +use std::sync::OnceLock; +use std::time::Duration; +use tracing_subscriber::filter::LevelFilter; +use tracing_subscriber::layer::SubscriberExt; +use tracing_subscriber::Layer as _; + +const TRACESTATE: &str = "vendor=value"; + +/// The ACP server runs on its own runtime threads, so spans are only visible through a global subscriber. +fn exporter() -> &'static InMemorySpanExporter { + static EXPORTER: OnceLock = OnceLock::new(); + EXPORTER.get_or_init(|| { + let exporter = InMemorySpanExporter::default(); + let provider = SdkTracerProvider::builder() + .with_simple_exporter(exporter.clone()) + .build(); + let layer = tracing_opentelemetry::layer() + .with_tracer(provider.tracer("acp-trace-context-test")) + .with_filter(LevelFilter::INFO); + tracing::subscriber::set_global_default(tracing_subscriber::registry().with(layer)) + .unwrap(); + exporter + }) +} + +fn exported_spans() -> Vec { + exporter().get_finished_spans().unwrap() +} + +fn attribute<'a>(span: &'a SpanData, key: &str) -> Option> { + span.attributes + .iter() + .find(|kv| kv.key.as_str() == key) + .map(|kv| kv.value.as_str()) +} + +fn is_operation(span: &SpanData, operation: &str) -> bool { + attribute(span, "gen_ai.operation.name").as_deref() == Some(operation) +} + +fn is_turn(span: &SpanData) -> bool { + is_operation(span, "invoke_agent") +} + +async fn wait_for_span(description: &str, matches: impl Fn(&SpanData) -> bool) -> SpanData { + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + loop { + if let Some(span) = exported_spans().into_iter().find(|span| matches(span)) { + return span; + } + assert!( + tokio::time::Instant::now() < deadline, + "timed out waiting for {description} span" + ); + tokio::time::sleep(Duration::from_millis(20)).await; + } +} + +fn random_remote_parent() -> (TraceId, SpanId) { + let trace_id = TraceId::from_bytes(*uuid::Uuid::new_v4().as_bytes()); + let span_bytes = uuid::Uuid::new_v4().as_bytes()[..8].try_into().unwrap(); + (trace_id, SpanId::from_bytes(span_bytes)) +} + +fn prompt_meta(traceparent: &str) -> serde_json::Map { + serde_json::json!({ + "traceparent": traceparent, + "tracestate": TRACESTATE, + }) + .as_object() + .unwrap() + .clone() +} + +async fn send_prompt( + conn: &AcpServerConnection, + session: &impl Session, + prompt: &str, + meta: serde_json::Map, +) { + let response = conn + .cx() + .send_request( + PromptRequest::new( + session.session_id().clone(), + vec![ContentBlock::Text(TextContent::new(prompt))], + ) + .meta(meta), + ) + .block_task() + .await + .unwrap(); + assert_eq!(response.stop_reason, StopReason::EndTurn); +} + +async fn wait_for_context_report(session: &impl Session) -> ContextReport { + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + loop { + let report = session.session_updates().into_iter().find_map(|update| { + let SessionUpdate::ToolCallUpdate(update) = update else { + return None; + }; + update + .fields + .content? + .into_iter() + .find_map(|content| match content { + ToolCallContent::Content(content) => match content.content { + ContentBlock::Text(text) => serde_json::from_str(&text.text).ok(), + _ => None, + }, + _ => None, + }) + }); + if let Some(report) = report { + return report; + } + assert!( + tokio::time::Instant::now() < deadline, + "timed out waiting for the inspect_context tool result" + ); + tokio::time::sleep(Duration::from_millis(20)).await; + } +} + +#[test] +fn test_prompt_trace_context_reaches_mcp_tool_call() { + exporter(); + run_test(async move { + let (trace_id, remote_span_id) = random_remote_parent(); + let traceparent = format!("00-{trace_id}-{remote_span_id}-01"); + let prompt = "Use the inspect_context tool."; + let mcp = McpFixture::new().await; + let openai = OpenAiFixture::new( + vec![ + ( + format!("{prompt}{TURN_CONTEXT_OPEN}"), + include_str!("acp_test_data/openai_inspect_context_tool_call.txt"), + ), + ( + "requestTraceparent".to_string(), + include_str!("acp_test_data/openai_basic.txt"), + ), + ], + AcpServerConnection::expected_session_id(), + ) + .await; + let mut conn = AcpServerConnection::new( + TestConnectionConfig { + mcp_servers: vec![McpServer::Http(McpServerHttp::new("mcp-fixture", &mcp.url))], + ..Default::default() + }, + openai, + ) + .await; + let SessionData { session, .. } = conn.new_session().await.unwrap(); + + send_prompt(&conn, &session, prompt, prompt_meta(&traceparent)).await; + let report = wait_for_context_report(&session).await; + + let in_trace = |span: &SpanData| span.span_context.trace_id() == trace_id; + let turn = wait_for_span("invoke_agent", |span| in_trace(span) && is_turn(span)).await; + let tool = wait_for_span("execute_tool", |span| { + in_trace(span) && is_operation(span, "execute_tool") + }) + .await; + + assert_eq!(turn.parent_span_id, remote_span_id); + assert!(turn.parent_span_is_remote); + assert_eq!( + report.request_traceparent, + Some(format!("00-{trace_id}-{}-01", tool.span_context.span_id())), + "MCP tools/call must carry the tool span as its parent" + ); + assert_eq!(report.request_tracestate.as_deref(), Some(TRACESTATE)); + }); +} + +#[test] +fn test_prompt_with_invalid_trace_context_is_ignored() { + exporter(); + run_test(async { + // Uppercase hex is invalid W3C, but leaves ids we can check were not adopted. + let invalid_trace_id = "4BF92F3577B34DA6A3CE929D0E0E4736"; + let invalid_span_id = "00F067AA0BA902B7"; + let prompt = "what is 1+1"; + let openai = OpenAiFixture::new( + vec![( + format!("{prompt}{TURN_CONTEXT_OPEN}"), + include_str!("acp_test_data/openai_basic.txt"), + )], + AcpServerConnection::expected_session_id(), + ) + .await; + let mut conn = AcpServerConnection::new(TestConnectionConfig::default(), openai).await; + let SessionData { session, .. } = conn.new_session().await.unwrap(); + // Session ids repeat across tests, so only consider spans exported by this test. + let earlier_spans = exported_spans().len(); + + send_prompt( + &conn, + &session, + prompt, + prompt_meta(&format!("00-{invalid_trace_id}-{invalid_span_id}-01")), + ) + .await; + + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + let turn = loop { + if let Some(turn) = exported_spans() + .into_iter() + .skip(earlier_spans) + .find(is_turn) + { + break turn; + } + assert!( + tokio::time::Instant::now() < deadline, + "timed out waiting for invoke_agent span" + ); + tokio::time::sleep(Duration::from_millis(20)).await; + }; + + assert!(!turn.parent_span_is_remote); + assert_ne!( + turn.span_context.trace_id(), + TraceId::from_hex(invalid_trace_id).unwrap() + ); + assert_ne!( + turn.parent_span_id, + SpanId::from_hex(invalid_span_id).unwrap() + ); + }); +}