From 2631095f20d99c357bb04504b7e5c2afcfe5ebe3 Mon Sep 17 00:00:00 2001 From: Adrian Cole <64215+codefromthecrypt@users.noreply.github.com> Date: Mon, 16 Mar 2026 13:06:31 +0800 Subject: [PATCH] feat(otel): propagate session.id to spans and log records (#7490) Signed-off-by: Adrian Cole --- Cargo.lock | 64 ++++++++++++----- Cargo.toml | 12 +++- crates/goose-test-support/Cargo.toml | 2 + crates/goose-test-support/src/lib.rs | 1 + crates/goose-test-support/src/otel.rs | 62 ++++++++++++++++ crates/goose/Cargo.toml | 4 +- crates/goose/src/agents/agent.rs | 15 ++-- crates/goose/src/agents/reply_parts.rs | 8 ++- crates/goose/src/otel/otlp.rs | 70 ++----------------- crates/goose/src/providers/base.rs | 8 +++ crates/goose/src/providers/bedrock.rs | 4 -- crates/goose/src/providers/codex.rs | 4 -- crates/goose/src/providers/cursor_agent.rs | 4 -- crates/goose/src/providers/gemini_cli.rs | 4 -- crates/goose/src/providers/litellm.rs | 1 - crates/goose/src/providers/sagemaker_tgi.rs | 4 -- crates/goose/src/providers/snowflake.rs | 4 -- crates/goose/src/providers/venice.rs | 4 -- crates/goose/src/tracing/rate_limiter.rs | 8 +-- .../tests/session_id_propagation_test.rs | 36 ++++++++++ 20 files changed, 193 insertions(+), 126 deletions(-) create mode 100644 crates/goose-test-support/src/otel.rs diff --git a/Cargo.lock b/Cargo.lock index e56eab1e52..ff00adc2bf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4299,6 +4299,7 @@ dependencies = [ "tokio-stream", "tokio-util", "tracing", + "tracing-futures", "tracing-opentelemetry", "tracing-subscriber", "tree-sitter", @@ -4506,6 +4507,8 @@ name = "goose-test-support" version = "1.27.0" dependencies = [ "axum 0.7.9", + "env-lock", + "opentelemetry", "rmcp 1.2.0", "serde_json", "tokio", @@ -6483,8 +6486,7 @@ dependencies = [ [[package]] name = "opentelemetry" version = "0.31.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b84bcd6ae87133e903af7ef497404dda70c60d0ea14895fc8a5e6722754fc2a0" +source = "git+https://github.com/open-telemetry/opentelemetry-rust?rev=345cd74a#345cd74a9c88ad1a47435d3d063c12d47235e803" dependencies = [ "futures-core", "futures-sink", @@ -6497,8 +6499,7 @@ dependencies = [ [[package]] name = "opentelemetry-appender-tracing" version = "0.31.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef6a1ac5ca3accf562b8c306fa8483c85f4390f768185ab775f242f7fe8fdcc2" +source = "git+https://github.com/open-telemetry/opentelemetry-rust?rev=345cd74a#345cd74a9c88ad1a47435d3d063c12d47235e803" dependencies = [ "opentelemetry", "tracing", @@ -6509,21 +6510,19 @@ dependencies = [ [[package]] name = "opentelemetry-http" version = "0.31.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7a6d09a73194e6b66df7c8f1b680f156d916a1a942abf2de06823dd02b7855d" +source = "git+https://github.com/open-telemetry/opentelemetry-rust?rev=345cd74a#345cd74a9c88ad1a47435d3d063c12d47235e803" dependencies = [ "async-trait", "bytes", "http 1.4.0", "opentelemetry", - "reqwest 0.12.28", + "reqwest 0.13.2", ] [[package]] name = "opentelemetry-otlp" version = "0.31.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a2366db2dca4d2ad033cad11e6ee42844fd727007af5ad04a1730f4cb8163bf" +source = "git+https://github.com/open-telemetry/opentelemetry-rust?rev=345cd74a#345cd74a9c88ad1a47435d3d063c12d47235e803" dependencies = [ "http 1.4.0", "opentelemetry", @@ -6531,18 +6530,17 @@ dependencies = [ "opentelemetry-proto", "opentelemetry_sdk", "prost", - "reqwest 0.12.28", + "reqwest 0.13.2", "thiserror 2.0.18", "tokio", "tonic", - "tracing", + "tonic-types", ] [[package]] name = "opentelemetry-proto" version = "0.31.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7175df06de5eaee9909d4805a3d07e28bb752c34cab57fa9cff549da596b30f" +source = "git+https://github.com/open-telemetry/opentelemetry-rust?rev=345cd74a#345cd74a9c88ad1a47435d3d063c12d47235e803" dependencies = [ "opentelemetry", "opentelemetry_sdk", @@ -6554,8 +6552,7 @@ dependencies = [ [[package]] name = "opentelemetry-stdout" version = "0.31.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bc8887887e169414f637b18751487cce4e095be787d23fad13c454e2fb1b3811" +source = "git+https://github.com/open-telemetry/opentelemetry-rust?rev=345cd74a#345cd74a9c88ad1a47435d3d063c12d47235e803" dependencies = [ "chrono", "opentelemetry", @@ -6565,16 +6562,17 @@ dependencies = [ [[package]] name = "opentelemetry_sdk" version = "0.31.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e14ae4f5991976fd48df6d843de219ca6d31b01daaab2dad5af2badeded372bd" +source = "git+https://github.com/open-telemetry/opentelemetry-rust?rev=345cd74a#345cd74a9c88ad1a47435d3d063c12d47235e803" dependencies = [ "futures-channel", "futures-executor", "futures-util", "opentelemetry", "percent-encoding", + "portable-atomic", "rand 0.9.2", "thiserror 2.0.18", + "tokio", ] [[package]] @@ -7208,6 +7206,15 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "prost-types" +version = "0.14.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8991c4cbdb8bc5b11f0b074ffe286c30e523de90fee5ba8132f1399f23cb3dd7" +dependencies = [ + "prost", +] + [[package]] name = "psl-types" version = "2.0.11" @@ -10443,6 +10450,17 @@ dependencies = [ "tonic", ] +[[package]] +name = "tonic-types" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a875a902255423d34c1f20838ab374126db8eb41625b7947a1d54113b0b7399" +dependencies = [ + "prost", + "prost-types", + "tonic", +] + [[package]] name = "tower" version = "0.5.3" @@ -10542,6 +10560,18 @@ dependencies = [ "valuable", ] +[[package]] +name = "tracing-futures" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97d095ae15e245a057c8e8451bab9b3ee1e1f68e9ba2b4fbc18d0ac5237835f2" +dependencies = [ + "futures", + "futures-task", + "pin-project", + "tracing", +] + [[package]] name = "tracing-log" version = "0.2.0" diff --git a/Cargo.toml b/Cargo.toml index 524a0a508a..1b8e533adb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -71,8 +71,9 @@ url = "2.5.8" opentelemetry = "0.31" opentelemetry_sdk = { version = "0.31", features = ["metrics"] } opentelemetry-otlp = "0.31" -opentelemetry-appender-tracing = "0.31" +opentelemetry-appender-tracing = { version = "0.31", features = ["experimental_span_attributes"] } opentelemetry-stdout = { version = "0.31", features = ["trace", "metrics", "logs"] } +tracing-futures = { version = "0.2", features = ["futures-03"] } tracing-opentelemetry = "0.32" rayon = "1.10" @@ -88,4 +89,11 @@ tree-sitter-swift = "0.7" tree-sitter-typescript = "0.23" [patch.crates-io] -v8 = { path = "vendor/v8" } \ No newline at end of file +v8 = { path = "vendor/v8" } +# TODO: switch to released version in opentelemetry 0.32.0 +# https://github.com/open-telemetry/opentelemetry-rust/issues/3408 +opentelemetry = { git = "https://github.com/open-telemetry/opentelemetry-rust", rev = "345cd74a" } +opentelemetry_sdk = { git = "https://github.com/open-telemetry/opentelemetry-rust", rev = "345cd74a" } +opentelemetry-appender-tracing = { git = "https://github.com/open-telemetry/opentelemetry-rust", rev = "345cd74a" } +opentelemetry-otlp = { git = "https://github.com/open-telemetry/opentelemetry-rust", rev = "345cd74a" } +opentelemetry-stdout = { git = "https://github.com/open-telemetry/opentelemetry-rust", rev = "345cd74a" } \ No newline at end of file diff --git a/crates/goose-test-support/Cargo.toml b/crates/goose-test-support/Cargo.toml index 3b11067c32..8e7d4f2477 100644 --- a/crates/goose-test-support/Cargo.toml +++ b/crates/goose-test-support/Cargo.toml @@ -9,6 +9,8 @@ description.workspace = true [dependencies] axum = "0.7" +env-lock = { workspace = true } +opentelemetry = { workspace = true } rmcp = { workspace = true, features = ["server", "macros", "transport-streamable-http-server"] } serde_json.workspace = true tokio.workspace = true diff --git a/crates/goose-test-support/src/lib.rs b/crates/goose-test-support/src/lib.rs index be80c4a878..6cc0b629fa 100644 --- a/crates/goose-test-support/src/lib.rs +++ b/crates/goose-test-support/src/lib.rs @@ -1,4 +1,5 @@ pub mod mcp; +pub mod otel; pub mod session; pub use mcp::{McpFixture, FAKE_CODE, TEST_IMAGE_B64}; diff --git a/crates/goose-test-support/src/otel.rs b/crates/goose-test-support/src/otel.rs new file mode 100644 index 0000000000..4eed3370f0 --- /dev/null +++ b/crates/goose-test-support/src/otel.rs @@ -0,0 +1,62 @@ +use opentelemetry::global; +use opentelemetry::metrics::{Meter, MeterProvider}; +use opentelemetry::InstrumentationScope; +use std::env; +use std::sync::Arc; + +struct SavedMeterProvider(Arc); + +impl MeterProvider for SavedMeterProvider { + fn meter_with_scope(&self, scope: InstrumentationScope) -> Meter { + self.0.meter_with_scope(scope) + } +} + +pub struct OtelTestGuard { + pub _env: env_lock::EnvGuard<'static>, + prev_tracer: global::GlobalTracerProvider, + prev_meter: Arc, +} + +impl Drop for OtelTestGuard { + fn drop(&mut self) { + global::set_tracer_provider(self.prev_tracer.clone()); + global::set_meter_provider(SavedMeterProvider(self.prev_meter.clone())); + } +} + +pub fn clear_otel_env(overrides: &[(&'static str, &'static str)]) -> OtelTestGuard { + let prev_tracer = global::tracer_provider(); + let prev_meter = global::meter_provider(); + + let mut keys: Vec<&'static str> = vec![ + "OTEL_EXPORTER_OTLP_ENDPOINT", + "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT", + "OTEL_EXPORTER_OTLP_METRICS_ENDPOINT", + "OTEL_EXPORTER_OTLP_METRICS_TEMPORALITY_PREFERENCE", + "OTEL_EXPORTER_OTLP_TIMEOUT", + "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", + "OTEL_LOG_LEVEL", + "OTEL_LOGS_EXPORTER", + "OTEL_METRICS_EXPORTER", + "OTEL_RESOURCE_ATTRIBUTES", + "OTEL_SDK_DISABLED", + "OTEL_SERVICE_NAME", + "OTEL_TRACES_EXPORTER", + ]; + for &(k, _) in overrides { + if !keys.contains(&k) { + keys.push(k); + } + } + + let guard = env_lock::lock_env(keys.into_iter().map(|k| (k, None::<&str>))); + for &(k, v) in overrides { + env::set_var(k, v); + } + OtelTestGuard { + _env: guard, + prev_tracer, + prev_meter, + } +} diff --git a/crates/goose/Cargo.toml b/crates/goose/Cargo.toml index 68b49a2ad0..c6c6d17198 100644 --- a/crates/goose/Cargo.toml +++ b/crates/goose/Cargo.toml @@ -54,10 +54,11 @@ webbrowser = { workspace = true } lazy_static = "1.5.0" tracing = { workspace = true } tracing-subscriber = { workspace = true } +tracing-futures = { workspace = true } tracing-opentelemetry = { workspace = true } opentelemetry = { workspace = true } -opentelemetry-appender-tracing = { workspace = true } opentelemetry_sdk = { workspace = true } +opentelemetry-appender-tracing = { workspace = true } opentelemetry-otlp = { workspace = true } opentelemetry-stdout = { workspace = true } keyring = { version = "3.6.2", features = [ @@ -163,6 +164,7 @@ ctor = "0.2.9" test-case = { workspace = true } env-lock = { workspace = true } rmcp = { workspace = true, features = ["transport-streamable-http-server"] } +opentelemetry_sdk = { workspace = true, features = ["testing"] } goose-test-support = { path = "../goose-test-support" } [[example]] diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 2c62b0b773..e35ace9969 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use anyhow::{anyhow, Context, Result}; use futures::stream::BoxStream; use futures::{stream, FutureExt, Stream, StreamExt, TryStreamExt}; +use tracing_futures::Instrument; use uuid::Uuid; use super::container::Container; @@ -492,7 +493,7 @@ impl Agent { } /// Dispatch a single tool call to the appropriate client - #[instrument(skip(self, tool_call, request_id, session), fields(input, output, session_id = %session.id))] + #[instrument(skip(self, tool_call, request_id, cancellation_token, session), fields(input, output, session.id = %session.id))] pub async fn dispatch_tool_call( &self, tool_call: CallToolRequestParams, @@ -875,8 +876,8 @@ impl Agent { } #[instrument( - skip(self, user_message, session_config), - fields(user_message, trace_input) + skip(self, user_message, session_config, cancel_token), + fields(user_message, trace_input, session.id = %session_config.id) )] pub async fn reply( &self, @@ -1120,9 +1121,8 @@ impl Agent { } let working_dir = session.working_dir.clone(); - Ok(Box::pin(async_stream::try_stream! { - let reply_stream_span = tracing::info_span!(target: "goose::agents::agent", "reply_stream"); - let _stream_guard = reply_stream_span.enter(); + let reply_stream_span = tracing::info_span!(target: "goose::agents::agent", "reply_stream", session.id = %session_config.id); + let inner = Box::pin(async_stream::try_stream! { let mut turns_taken = 0u32; let max_turns = session_config.max_turns.unwrap_or_else(|| { Config::global() @@ -1690,7 +1690,8 @@ impl Agent { if !last_assistant_text.is_empty() { tracing::info!(target: "goose::agents::agent", trace_output = last_assistant_text.as_str()); } - })) + }.instrument(reply_stream_span)); + Ok(inner) } pub async fn extend_system_prompt(&self, key: String, instruction: String) { diff --git a/crates/goose/src/agents/reply_parts.rs b/crates/goose/src/agents/reply_parts.rs index 26b795698b..b31ac44aba 100644 --- a/crates/goose/src/agents/reply_parts.rs +++ b/crates/goose/src/agents/reply_parts.rs @@ -205,8 +205,12 @@ impl Agent { Ok((tools, toolshim_tools, system_prompt)) } - /// Stream a response from the LLM provider. - /// Handles toolshim transformations if needed + // Don't add gen_ai.request.model here — provider.get_model_config() + // returns the wrong model for LeadWorkerProvider. + #[tracing::instrument( + skip(provider, session_id, system_prompt, messages, tools, toolshim_tools), + fields(session.id = %session_id) + )] pub(crate) async fn stream_response_from_provider( provider: Arc, session_id: &str, diff --git a/crates/goose/src/otel/otlp.rs b/crates/goose/src/otel/otlp.rs index 4efae0f1eb..19c69e7e45 100644 --- a/crates/goose/src/otel/otlp.rs +++ b/crates/goose/src/otel/otlp.rs @@ -131,8 +131,8 @@ pub fn init_otlp_layers( if let Ok(layer) = create_otlp_metrics_layer() { layers.push(layer.with_filter(create_otlp_metrics_filter()).boxed()); } - if let Ok(layer) = create_otlp_logs_layer() { - layers.push(layer.with_filter(create_otlp_logs_filter()).boxed()); + if let Ok(bridge) = create_otlp_logs_layer() { + layers.push(bridge.with_filter(create_otlp_logs_filter()).boxed()); } if !layers.is_empty() { @@ -241,7 +241,9 @@ fn create_otlp_logs_layer() -> OtlpResult { ExporterType::None => return Err("Logs exporter set to none".into()), }; - let bridge = OpenTelemetryTracingBridge::new(&logger_provider); + let bridge = OpenTelemetryTracingBridge::builder(&logger_provider) + .with_span_attribute_allowlist(["session.id"]) + .build(); *LOGGER_PROVIDER.lock().unwrap_or_else(|e| e.into_inner()) = Some(logger_provider); Ok(bridge) @@ -368,70 +370,10 @@ pub fn shutdown_otlp() { #[cfg(test)] mod tests { use super::*; - use opentelemetry::metrics::{Meter, MeterProvider}; - use opentelemetry::InstrumentationScope; + use goose_test_support::otel::clear_otel_env; use opentelemetry_sdk::metrics::Temporality; - use std::sync::Arc; use test_case::test_case; - // set_meter_provider requires P: MeterProvider, not Arc - struct SavedMeterProvider(Arc); - - impl MeterProvider for SavedMeterProvider { - fn meter_with_scope(&self, scope: InstrumentationScope) -> Meter { - self.0.meter_with_scope(scope) - } - } - - struct OtelTestGuard { - _env: env_lock::EnvGuard<'static>, - prev_tracer: global::GlobalTracerProvider, - prev_meter: Arc, - } - - impl Drop for OtelTestGuard { - fn drop(&mut self) { - global::set_tracer_provider(self.prev_tracer.clone()); - global::set_meter_provider(SavedMeterProvider(self.prev_meter.clone())); - } - } - - fn clear_otel_env(overrides: &[(&'static str, &'static str)]) -> OtelTestGuard { - let prev_tracer = global::tracer_provider(); - let prev_meter = global::meter_provider(); - - let mut keys: Vec<&'static str> = vec![ - "OTEL_EXPORTER_OTLP_ENDPOINT", - "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT", - "OTEL_EXPORTER_OTLP_METRICS_ENDPOINT", - "OTEL_EXPORTER_OTLP_METRICS_TEMPORALITY_PREFERENCE", - "OTEL_EXPORTER_OTLP_TIMEOUT", - "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", - "OTEL_LOG_LEVEL", - "OTEL_LOGS_EXPORTER", - "OTEL_METRICS_EXPORTER", - "OTEL_RESOURCE_ATTRIBUTES", - "OTEL_SDK_DISABLED", - "OTEL_SERVICE_NAME", - "OTEL_TRACES_EXPORTER", - ]; - for &(k, _) in overrides { - if !keys.contains(&k) { - keys.push(k); - } - } - - let guard = env_lock::lock_env(keys.into_iter().map(|k| (k, None::<&str>))); - for &(k, v) in overrides { - env::set_var(k, v); - } - OtelTestGuard { - _env: guard, - prev_tracer, - prev_meter, - } - } - #[test] fn exporter_type_from_env_value() { assert_eq!(ExporterType::from_env_value("otlp"), ExporterType::Otlp); diff --git a/crates/goose/src/providers/base.rs b/crates/goose/src/providers/base.rs index 7c863b89cd..97ef2bdd38 100644 --- a/crates/goose/src/providers/base.rs +++ b/crates/goose/src/providers/base.rs @@ -459,6 +459,10 @@ pub trait Provider: Send + Sync { fn get_name(&self) -> &str; /// Primary streaming method that all providers must implement. + /// + /// Note: Do not add `#[instrument]` here — the call sites (`complete` and + /// `stream_response_from_provider`) create the telemetry span so that + /// `session.id` is set once rather than in every provider. async fn stream( &self, model_config: &ModelConfig, @@ -469,6 +473,10 @@ pub trait Provider: Send + Sync { ) -> Result; /// Complete with a specific model config. + #[tracing::instrument( + skip(self, model_config, session_id, system, messages, tools), + fields(session.id = %session_id, gen_ai.request.model = %model_config.model_name) + )] async fn complete( &self, model_config: &ModelConfig, diff --git a/crates/goose/src/providers/bedrock.rs b/crates/goose/src/providers/bedrock.rs index 0e9f8c09e7..72f07850ce 100644 --- a/crates/goose/src/providers/bedrock.rs +++ b/crates/goose/src/providers/bedrock.rs @@ -356,10 +356,6 @@ impl Provider for BedrockProvider { Ok(BEDROCK_KNOWN_MODELS.iter().map(|s| s.to_string()).collect()) } - #[tracing::instrument( - skip(self, model_config, system, messages, tools), - fields(model_config, input, output, input_tokens, output_tokens, total_tokens) - )] async fn stream( &self, model_config: &ModelConfig, diff --git a/crates/goose/src/providers/codex.rs b/crates/goose/src/providers/codex.rs index fbc591e5d5..96705daa08 100644 --- a/crates/goose/src/providers/codex.rs +++ b/crates/goose/src/providers/codex.rs @@ -668,10 +668,6 @@ impl Provider for CodexProvider { self.model.clone() } - #[tracing::instrument( - skip(self, model_config, system, messages, tools), - fields(model_config, input, output, input_tokens, output_tokens, total_tokens) - )] async fn stream( &self, model_config: &ModelConfig, diff --git a/crates/goose/src/providers/cursor_agent.rs b/crates/goose/src/providers/cursor_agent.rs index e9790a3648..ceb2577694 100644 --- a/crates/goose/src/providers/cursor_agent.rs +++ b/crates/goose/src/providers/cursor_agent.rs @@ -322,10 +322,6 @@ impl Provider for CursorAgentProvider { .collect()) } - #[tracing::instrument( - skip(self, model_config, system, messages, tools), - fields(model_config, input, output, input_tokens, output_tokens, total_tokens) - )] async fn stream( &self, model_config: &ModelConfig, diff --git a/crates/goose/src/providers/gemini_cli.rs b/crates/goose/src/providers/gemini_cli.rs index 5f1e549b39..27cdf965ef 100644 --- a/crates/goose/src/providers/gemini_cli.rs +++ b/crates/goose/src/providers/gemini_cli.rs @@ -196,10 +196,6 @@ impl Provider for GeminiCliProvider { .collect()) } - #[tracing::instrument( - skip(self, model_config, system, messages, _tools), - fields(model_config, input, output, input_tokens, output_tokens, total_tokens) - )] async fn stream( &self, model_config: &ModelConfig, diff --git a/crates/goose/src/providers/litellm.rs b/crates/goose/src/providers/litellm.rs index 78149bce57..8008e0659e 100644 --- a/crates/goose/src/providers/litellm.rs +++ b/crates/goose/src/providers/litellm.rs @@ -184,7 +184,6 @@ impl Provider for LiteLLMProvider { self.model.clone() } - #[tracing::instrument(skip_all, name = "provider_complete")] async fn stream( &self, model_config: &ModelConfig, diff --git a/crates/goose/src/providers/sagemaker_tgi.rs b/crates/goose/src/providers/sagemaker_tgi.rs index 1777bfb52f..89fe5424cf 100644 --- a/crates/goose/src/providers/sagemaker_tgi.rs +++ b/crates/goose/src/providers/sagemaker_tgi.rs @@ -309,10 +309,6 @@ impl Provider for SageMakerTgiProvider { self.model.clone() } - #[tracing::instrument( - skip(self, model_config, system, messages, tools), - fields(model_config, input, output, input_tokens, output_tokens, total_tokens) - )] async fn stream( &self, model_config: &ModelConfig, diff --git a/crates/goose/src/providers/snowflake.rs b/crates/goose/src/providers/snowflake.rs index 7cb6cd574e..6dc3742894 100644 --- a/crates/goose/src/providers/snowflake.rs +++ b/crates/goose/src/providers/snowflake.rs @@ -340,10 +340,6 @@ impl Provider for SnowflakeProvider { .collect()) } - #[tracing::instrument( - skip(self, model_config, system, messages, tools), - fields(model_config, input, output, input_tokens, output_tokens, total_tokens) - )] async fn stream( &self, model_config: &ModelConfig, diff --git a/crates/goose/src/providers/venice.rs b/crates/goose/src/providers/venice.rs index 3312bcfea7..5440632eaf 100644 --- a/crates/goose/src/providers/venice.rs +++ b/crates/goose/src/providers/venice.rs @@ -271,10 +271,6 @@ impl Provider for VeniceProvider { Ok(models) } - #[tracing::instrument( - skip(self, model_config, system, messages, tools), - fields(model_config, input, output, input_tokens, output_tokens, total_tokens) - )] async fn stream( &self, model_config: &ModelConfig, diff --git a/crates/goose/src/tracing/rate_limiter.rs b/crates/goose/src/tracing/rate_limiter.rs index f1a9c9b7e0..d72e4edd85 100644 --- a/crates/goose/src/tracing/rate_limiter.rs +++ b/crates/goose/src/tracing/rate_limiter.rs @@ -49,10 +49,10 @@ impl RateLimitedTelemetrySender { match event { TelemetryEvent::Span(span_data) => { - Self::process_span(span_data).await; + Self::process_span(span_data); } TelemetryEvent::Metric(metric_data) => { - Self::process_metric(metric_data).await; + Self::process_metric(metric_data); } } @@ -79,7 +79,7 @@ impl RateLimitedTelemetrySender { self.sender.send(TelemetryEvent::Metric(metric_data)) } - async fn process_span(span_data: SpanData) { + fn process_span(span_data: SpanData) { let span = tracing::info_span!("telemetry_span", name = %span_data.name); let _enter = span.enter(); @@ -92,7 +92,7 @@ impl RateLimitedTelemetrySender { } } - async fn process_metric(metric_data: MetricData) { + fn process_metric(metric_data: MetricData) { info!( metric_name = %metric_data.name, metric_value = metric_data.value, diff --git a/crates/goose/tests/session_id_propagation_test.rs b/crates/goose/tests/session_id_propagation_test.rs index 1ace2e6ae7..0f68e9851a 100644 --- a/crates/goose/tests/session_id_propagation_test.rs +++ b/crates/goose/tests/session_id_propagation_test.rs @@ -4,9 +4,14 @@ use goose::providers::api_client::{ApiClient, AuthMethod}; use goose::providers::base::Provider; use goose::providers::openai::OpenAiProvider; use goose::session_context::SESSION_ID_HEADER; +use opentelemetry::logs::AnyValue; +use opentelemetry::Key; +use opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge; +use opentelemetry_sdk::logs::{InMemoryLogExporterBuilder, SdkLoggerProvider}; use serde_json::json; use std::sync::Arc; use std::sync::Mutex; +use tracing_subscriber::prelude::*; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, Request, ResponseTemplate}; @@ -104,6 +109,37 @@ async fn make_request(provider: &dyn Provider, session_id: &str) { .unwrap(); } +#[tokio::test] +async fn test_session_id_propagates_to_log_records() { + let exporter = InMemoryLogExporterBuilder::default().build(); + let provider = SdkLoggerProvider::builder() + .with_simple_exporter(exporter.clone()) + .build(); + + let layer = OpenTelemetryTracingBridge::builder(&provider) + .with_span_attribute_allowlist(["session.id"]) + .build(); + let subscriber = tracing_subscriber::registry().with(layer); + let _guard = tracing::subscriber::set_default(subscriber); + + let span = tracing::info_span!("test", session.id = "test-session-42"); + let _enter = span.enter(); + tracing::info!("hello from test"); + drop(_enter); + drop(_guard); + + provider.force_flush().unwrap(); + let logs = exporter.get_emitted_logs().unwrap(); + assert_eq!(logs.len(), 1); + let log = &logs[0]; + + let has_session_id = log.record.attributes_iter().any(|(k, v)| { + k == &Key::new("session.id") + && matches!(v, AnyValue::String(s) if s.as_str() == "test-session-42") + }); + assert!(has_session_id); +} + #[tokio::test] async fn test_session_id_propagation_to_llm() { let (_, capture, provider) = setup_mock_server().await;