feat: propagate W3C trace context from ACP prompts to MCP tool calls (#12697)

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<String>,
     pub request_session_id: Option<String>,
     pub request_working_dir: Option<String>,
+    pub request_traceparent: Option<String>,
+    pub request_tracestate: Option<String>,
     pub roots: Vec<String>,
 }
 
@@ -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(

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
+        let reply = agent.reply(user_message, session_config, Some(cancel_token.clone()));
-            .reply(user_message, session_config, Some(cancel_token.clone()))
+        #[cfg(feature = "otel")]
-            .await
+        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
+        let forward = self.forward_agent_stream(
-            .forward_agent_stream(
+            cx,
-                cx,
+            &args.session_id,
-                &args.session_id,
+            &session_id,
-                &session_id,
+            &agent,
-                &agent,
+            &cancel_token,
-                &cancel_token,
+            stream,
-                stream,
+        );
-            )
+        #[cfg(feature = "otel")]
-            .await;
+        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?;

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(
+            let extensions =
-                Extensions::new(),
+                inject_trace_context_into_extensions(inject_session_context_into_extensions(
-                Some(&ctx.session_id),
+                    Extensions::new(),
-                ctx.working_dir_str(),
+                    Some(&ctx.session_id),
-                ctx.tool_call_request_id.as_deref(),
+                    ctx.working_dir_str(),
-            );
+                    ctx.tool_call_request_id.as_deref(),
+                ));
             if let Some(meta) = extensions.get::<MetaObject>() {
                 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::<MetaObject>()
+        .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 =
-                req.extensions,
+                inject_trace_context_into_extensions(inject_session_context_into_extensions(
-                session_id,
+                    req.extensions,
-                working_dir,
+                    session_id,
-                tool_call_request_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::<TraceState>().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::<MetaObject>())
+            .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);

crates/goose/src/otel/mod.rs

@@ -1 +1,2 @@
 pub mod otlp;
+pub mod trace_context;

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<String, Value>>) -> Option<Context> {
+    let meta = meta?;
+    let carrier: HashMap<String, String> = [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<String, Value>) {
+    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<String, Value> {
+        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::<TraceState>().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<Value>) {
+        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");
+        });
+    }
+}

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]
+

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<InMemorySpanExporter> = 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<SpanData> {
+    exporter().get_finished_spans().unwrap()
+}
+
+fn attribute<'a>(span: &'a SpanData, key: &str) -> Option<std::borrow::Cow<'a, str>> {
+    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<String, serde_json::Value> {
+    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<String, serde_json::Value>,
+) {
+    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()
+        );
+    });
+}