vrr/goose
2
0
Fork 0
mirror of https://github.com/block/goose.git synced 2026-10-18 09:30:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 13

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

Browse source
This commit is contained in:
Sergo Beruashvili 2026-10-12 03:24:58 +00:00 • committed by GitHub
parent 3295f5e0f6
commit 1c3ac1ba19
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 498 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

6
crates/goose-test-support/src/mcp.rs
View file

@ -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(

33
crates/goose/src/acp/server.rs
View file

@ -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?;

95
crates/goose/src/agents/mcp_client.rs
View file

@ -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::<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,
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::<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);

1
crates/goose/src/otel/mod.rs
View file

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

126
crates/goose/src/otel/trace_context.rs Normal file
View file

@ -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");
});
}
}

10
crates/goose/tests/acp_test_data/openai_inspect_context_tool_call.txt Normal file
View file

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

253
crates/goose/tests/acp_trace_context_test.rs Normal file
View file

@ -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()
);
});
}