improve span coverage

crates/goose/src/agents/reply_parts.rs

@@ -8,7 +8,9 @@ use tracing::debug;
 use super::super::agents::Agent;
 use crate::conversation::message::{Message, MessageContent, ToolRequest};
 use crate::conversation::Conversation;
-use crate::providers::base::{stream_from_single_message, MessageStream, Provider, ProviderUsage};
+use crate::providers::base::{
+    instrumented_stream, stream_from_single_message, MessageStream, Provider, ProviderUsage,
+};
 use crate::providers::errors::ProviderError;
 use crate::providers::toolshim::{
     augment_message_with_tool_calls, convert_tool_messages_to_text,
@@ -152,15 +154,22 @@ impl Agent {
             }
         };
 
-        // If there was an error creating the stream, return a stream that yields that error
         let mut stream = match stream_result {
-            Ok(s) => s,
+            Ok(s) => {
+                if crate::tracing::is_langfuse_enabled() {
+                    instrumented_stream(
+                        s,
+                        provider.get_active_model_name(),
+                        system_prompt.clone(),
+                        messages_for_provider.messages().to_vec(),
+                        tools.iter().map(|t| t.name.to_string()).collect(),
+                    )
+                } else {
+                    s
+                }
+            }
             Err(e) => {
-                // Return a stream that immediately yields the error
-                // This allows the error to be caught by existing error handling in agent.rs
-                return Ok(Box::pin(try_stream! {
-                    yield Err(e)?;
-                }));
+                return Ok(Box::pin(futures::stream::once(async move { Err(e) })));
             }
         };
 

crates/goose/src/providers/anthropic.rs

@@ -187,10 +187,6 @@ impl Provider for AnthropicProvider {
         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 complete_with_model(
         &self,
         model_config: &ModelConfig,

crates/goose/src/providers/base.rs

@@ -1,5 +1,5 @@
 use anyhow::Result;
-use futures::Stream;
+use futures::{Stream, StreamExt};
 use serde::{Deserialize, Serialize};
 
 use super::errors::ProviderError;
@@ -533,6 +533,70 @@ pub fn stream_from_single_message(message: Message, usage: ProviderUsage) -> Mes
     Box::pin(stream)
 }
 
+pub fn instrumented_stream(
+    stream: MessageStream,
+    model_name: String,
+    system: String,
+    messages: Vec<Message>,
+    tools: Vec<String>,
+) -> MessageStream {
+    Box::pin(async_stream::stream! {
+        let span = tracing::info_span!(
+            target: "goose::providers",
+            "complete_with_model",
+            model_config = tracing::field::Empty,
+            input = tracing::field::Empty,
+            output = tracing::field::Empty,
+        );
+
+        {
+            let _enter = span.enter();
+            span.record("model_config", model_name.as_str());
+
+            let input_data = serde_json::json!({
+                "system": system,
+                "messages": messages,
+                "tools": tools,
+            });
+
+            match serde_json::to_string(&input_data) {
+                Ok(input_json) => {
+                    let truncated = safe_truncate(&input_json, 128_000);
+                    span.record("input", truncated.as_str());
+                }
+                Err(e) => {
+                    tracing::warn!("Failed to serialize input for tracing: {}", e);
+                }
+            }
+        }
+
+        const OUTPUT_PREVIEW_LIMIT: usize = 64_000;
+        let mut output_preview = String::new();
+
+        tokio::pin!(stream);
+        while let Some(result) = stream.next().await {
+            if let Ok((Some(msg), _usage_opt)) = &result {
+                if output_preview.len() < OUTPUT_PREVIEW_LIMIT {
+                    let text = msg.as_concat_text();
+                    if !text.is_empty() {
+                        let remaining = OUTPUT_PREVIEW_LIMIT.saturating_sub(output_preview.len());
+                        let chunk = safe_truncate(&text, remaining);
+                        output_preview.push_str(&chunk);
+                    }
+                }
+            }
+            yield result;
+        }
+
+        {
+            let _enter = span.enter();
+            if !output_preview.is_empty() {
+                span.record("output", output_preview.as_str());
+            }
+        }
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/goose/src/tracing/langfuse_layer.rs

@@ -152,20 +152,31 @@ impl BatchManager for LangfuseBatchManager {
     }
 }
 
-pub fn create_langfuse_observer() -> Option<ObservationLayer> {
+pub fn is_langfuse_enabled() -> bool {
     let public_key = env::var("LANGFUSE_PUBLIC_KEY")
         .or_else(|_| env::var("LANGFUSE_INIT_PROJECT_PUBLIC_KEY"))
-        .unwrap_or_default(); // Use empty string if not found
+        .unwrap_or_default();
 
     let secret_key = env::var("LANGFUSE_SECRET_KEY")
         .or_else(|_| env::var("LANGFUSE_INIT_PROJECT_SECRET_KEY"))
-        .unwrap_or_default(); // Use empty string if not found
+        .unwrap_or_default();
 
-    // Return None if either key is empty
-    if public_key.is_empty() || secret_key.is_empty() {
+    !public_key.is_empty() && !secret_key.is_empty()
+}
+
+pub fn create_langfuse_observer() -> Option<ObservationLayer> {
+    if !is_langfuse_enabled() {
         return None;
     }
 
+    let public_key = env::var("LANGFUSE_PUBLIC_KEY")
+        .or_else(|_| env::var("LANGFUSE_INIT_PROJECT_PUBLIC_KEY"))
+        .unwrap_or_default();
+
+    let secret_key = env::var("LANGFUSE_SECRET_KEY")
+        .or_else(|_| env::var("LANGFUSE_INIT_PROJECT_SECRET_KEY"))
+        .unwrap_or_default();
+
     let base_url = env::var("LANGFUSE_URL").unwrap_or_else(|_| DEFAULT_LANGFUSE_URL.to_string());
 
     let batch_manager = Arc::new(Mutex::new(LangfuseBatchManager::new(

crates/goose/src/tracing/mod.rs

@@ -3,7 +3,7 @@ mod observation_layer;
 pub mod otlp_layer;
 pub mod rate_limiter;
 
-pub use langfuse_layer::{create_langfuse_observer, LangfuseBatchManager};
+pub use langfuse_layer::{create_langfuse_observer, is_langfuse_enabled, LangfuseBatchManager};
 pub use observation_layer::{
     flatten_metadata, map_level, BatchManager, ObservationLayer, SpanData, SpanTracker,
 };