tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 8 Packages Projects Releases Wiki Activity

feat(otel): propagate session.id to spans and log records (#7490)

Browse Source 
Signed-off-by: Adrian Cole <adrian@tetrate.io>
This commit is contained in:
Adrian Cole
2026-03-16 13:06:31 +08:00
committed by GitHub
parent 10d597de3b
commit 2631095f20
20 changed files with 193 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-test-support/Cargo.toml
+2
View File
@@ -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
crates/goose-test-support/src/lib.rs
+1
View File
@@ -1,4 +1,5 @@
pub mod mcp;
pub mod otel;
pub mod session;
pub use mcp::{McpFixture, FAKE_CODE, TEST_IMAGE_B64};
crates/goose-test-support/src/otel.rs
+62
View File
@@ -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<dyn MeterProvider + Send + Sync>);
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<dyn MeterProvider + Send + Sync>,
}
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,
}
}
crates/goose/Cargo.toml
+3 -1
View File
@@ -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]]
crates/goose/src/agents/agent.rs
+8 -7
View File
@@ -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) {
crates/goose/src/agents/reply_parts.rs
+6 -2
View File
@@ -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<dyn Provider>,
session_id: &str,
crates/goose/src/otel/otlp.rs
+6 -64
View File
@@ -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<OtlpLogsLayer> {
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<dyn MeterProvider>
struct SavedMeterProvider(Arc<dyn MeterProvider + Send + Sync>);
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<dyn MeterProvider + Send + Sync>,
}
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);
crates/goose/src/providers/base.rs
+8
View File
@@ -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<MessageStream, ProviderError>;
/// 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,
crates/goose/src/providers/bedrock.rs
-4
View File
@@ -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,
crates/goose/src/providers/codex.rs
-4
View File
@@ -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,
crates/goose/src/providers/cursor_agent.rs
-4
View File
@@ -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,
crates/goose/src/providers/gemini_cli.rs
-4
View File
@@ -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,
crates/goose/src/providers/litellm.rs
-1
View File
@@ -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,
crates/goose/src/providers/sagemaker_tgi.rs
-4
View File
@@ -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,
crates/goose/src/providers/snowflake.rs
-4
View File
@@ -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,
crates/goose/src/providers/venice.rs
-4
View File
@@ -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,
crates/goose/src/tracing/rate_limiter.rs
+4 -4
View File
@@ -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,
crates/goose/tests/session_id_propagation_test.rs
+36
View File
@@ -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;