Apps/ImageApi
1
0 Fork 0
Code Issues 7 Pull Requests Releases Wiki Activity

Tracing: LLM clients + streaming spans on tracing, drop cx threading (#140)

Browse Source 
LlmClient methods no longer take an OpenTelemetry Context. llm.request is
now a tracing span (otel.kind=client, llm.provider/operation/model) that
parents on the caller's ambient span; traceparent injection reads the
span's OTel context via OpenTelemetrySpanExt.

The src/ai feature spans (handlers, insight generation, chat, persona
chat, daily summaries, TTS) move to #[instrument] / info_span!, and every
LLM-bearing tokio::spawn carries its parent span with .instrument(): the
streaming chat turn (chat_turn_stream_inner), async chat / persona turns,
background insight jobs and TTS jobs. Sub-spans in the agentic loops
(ai.chat.loop, ai.agentic.loop) instrument each LLM / tool await.

DAO calls keep their explicit Context (derived from the tracing span)
until the DAO migration (#141).

tests/otel_llm_spans.rs pins nesting + traceparent for the blocking and
streaming paths across a spawn on a multi-thread runtime.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Cameron CordesandClaude Opus 5.5 committed 2026-09-29 17:02:40 -04:00
1 parent 7a59afc62b
commit 1c508cdabd
20 files changed
+973 -1030
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/ai/insight_chat.rs
+95 -87
View File
@@ -1,7 +1,6 @@
use anyhow::{Result, anyhow, bail};
use chrono::{TimeZone, Utc};
use opentelemetry::KeyValue;
use opentelemetry::trace::{Span, Status, TraceContextExt, Tracer};
use opentelemetry::trace::Status;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use tokio::sync::Mutex as TokioMutex;
@@ -15,9 +14,10 @@ use crate::ai::turn_registry::TurnEntry;
use crate::ai::turn_registry::TurnRegistry;
use crate::database::InsightDao;
use crate::database::models::InsertPhotoInsight;
use crate::otel::global_tracer;
use crate::utils::{normalize_path, retry_with_backoff};
use futures::stream::{BoxStream, StreamExt};
use tracing::Instrument;
use tracing_opentelemetry::OpenTelemetrySpanExt;
use uuid::Uuid;
pub const DEFAULT_MAX_ITERATIONS: usize = 6;
@@ -200,13 +200,12 @@ impl InsightChatService {
})
}
#[tracing::instrument(name = "ai.insight.chat_turn", skip_all)]
pub async fn chat_turn(&self, req: ChatTurnRequest) -> Result<ChatTurnResult> {
let tracer = global_tracer();
let parent_cx = opentelemetry::Context::new();
let mut span = tracer.start_with_context("ai.insight.chat_turn", &parent_cx);
span.set_attribute(KeyValue::new("file_path", req.file_path.clone()));
span.set_attribute(KeyValue::new("library_id", req.library_id as i64));
span.set_attribute(KeyValue::new("amend", req.amend));
let span = tracing::Span::current();
span.set_attribute("file_path", req.file_path.clone());
span.set_attribute("library_id", req.library_id as i64);
span.set_attribute("amend", req.amend);
if req.user_message.trim().is_empty() {
bail!("user_message must not be empty");
@@ -220,7 +219,7 @@ impl InsightChatService {
.clone()
.filter(|s| !s.trim().is_empty())
.unwrap_or_else(|| "default".to_string());
span.set_attribute(KeyValue::new("persona_id", active_persona.clone()));
span.set_attribute("persona_id", active_persona.clone());
let normalized = normalize_path(&req.file_path);
@@ -278,13 +277,13 @@ impl InsightChatService {
.unwrap_or_else(|| stored_backend.clone());
validate_cross_replay(&stored_backend, &effective_backend)?;
let kind = BackendKind::parse(&effective_backend)?;
span.set_attribute(KeyValue::new("backend", kind.as_str()));
span.set_attribute("backend", kind.as_str());
let max_iterations = req
.max_iterations
.unwrap_or(DEFAULT_MAX_ITERATIONS)
.clamp(1, env_max_iterations());
span.set_attribute(KeyValue::new("max_iterations", max_iterations as i64));
span.set_attribute("max_iterations", max_iterations as i64);
let stored_model = insight.model_version.clone();
let overrides = SamplingOverrides {
@@ -302,7 +301,7 @@ impl InsightChatService {
};
let backend = self.generator.resolve_backend(kind, &overrides).await?;
let model_used = backend.model().to_string();
span.set_attribute(KeyValue::new("model", model_used.clone()));
span.set_attribute("model", model_used.clone());
// 5. Decide vision + tool set. In hybrid (describe-then-inline) mode
// we omit `describe_photo`. Otherwise trust the stored history:
@@ -333,7 +332,7 @@ impl InsightChatService {
let budget_bytes = budget_tokens.saturating_mul(BYTES_PER_TOKEN);
let truncated = apply_context_budget(&mut messages, budget_bytes);
if truncated {
span.set_attribute(KeyValue::new("history_truncated", true));
span.set_attribute("history_truncated", true);
}
// Everything in `messages` at this point is replayed history that is
// ALREADY in the tree; only what the turn appends past this mark
@@ -359,12 +358,12 @@ impl InsightChatService {
// accumulate across turns.
let original_system_content = annotate_system_with_budget(&mut messages, max_iterations);
let insight_cx = parent_cx.with_span(span);
// 8. Agentic loop — same shape as insight_generator's, but capped
// tighter and dispatching tools through the shared executor.
let loop_span = tracer.start_with_context("ai.chat.loop", &insight_cx);
let loop_cx = insight_cx.with_span(loop_span);
// Each LLM / tool await runs inside `loop_span` so their spans
// nest under it.
let loop_span = tracing::info_span!("ai.chat.loop");
let loop_cx = loop_span.context();
let mut tool_calls_made = 0usize;
let mut iterations_used = 0usize;
let mut last_prompt_eval_count: Option<i32> = None;
@@ -377,7 +376,8 @@ impl InsightChatService {
let (response, prompt_tokens, eval_tokens) = backend
.chat()
.chat_with_tools(&loop_cx, messages.clone(), tools.clone())
.chat_with_tools(messages.clone(), tools.clone())
.instrument(loop_span.clone())
.await?;
last_prompt_eval_count = prompt_tokens;
last_eval_count = eval_tokens;
@@ -416,6 +416,7 @@ impl InsightChatService {
&active_persona,
&loop_cx,
)
.instrument(loop_span.clone())
.await;
messages.push(ChatMessage::tool_result(result));
}
@@ -434,7 +435,8 @@ impl InsightChatService {
let synthetic_idx = push_synthetic_final_prompt(&mut messages);
let (final_response, prompt_tokens, eval_tokens) = backend
.chat()
.chat_with_tools(&loop_cx, messages.clone(), vec![])
.chat_with_tools(messages.clone(), vec![])
.instrument(loop_span.clone())
.await?;
last_prompt_eval_count = prompt_tokens;
last_eval_count = eval_tokens;
@@ -449,7 +451,8 @@ impl InsightChatService {
// return / persist as the reply (the raw transcript keeps it).
let final_content = crate::ai::llm_client::strip_think_blocks(&final_content);
loop_cx.span().set_status(Status::Ok);
loop_span.set_status(Status::Ok);
drop(loop_span);
// Drop the per-turn iteration-budget note from the system message
// before we persist so it doesn't snowball on each subsequent turn.
@@ -492,7 +495,6 @@ impl InsightChatService {
let title_raw = backend
.chat()
.generate(
&loop_cx,
&title_prompt,
Some(
"You are my long term memory assistant. Use only the information provided. Do not invent details.",
@@ -779,14 +781,22 @@ impl InsightChatService {
///
/// The stream takes ownership of the service via `Arc<Self>` (passed by
/// the caller) so it can live past the handler's await boundary.
///
/// Must be called from inside the request's span: the turn span is
/// created here, eagerly, so it parents on the caller. The returned
/// stream is polled by the HTTP body writer, outside any span.
pub fn chat_turn_stream(
self: Arc<Self>,
req: ChatTurnRequest,
cx: opentelemetry::Context,
) -> BoxStream<'static, ChatStreamEvent> {
let span = tracing::info_span!(
"ai.chat.turn.stream",
file_path = %req.file_path,
library_id = req.library_id,
);
let svc = self;
let s = async_stream::stream! {
match svc.chat_turn_stream_inner(req, cx, Ok).await {
match svc.chat_turn_stream_inner(req, span, Ok).await {
Ok(mut rx) => {
while let Some(ev) = rx.recv().await {
yield ev;
@@ -808,7 +818,7 @@ impl InsightChatService {
async fn chat_turn_stream_inner<F>(
self: Arc<Self>,
req: ChatTurnRequest,
cx: opentelemetry::Context,
span: tracing::Span,
_ev_mapper: F,
) -> Result<tokio::sync::mpsc::Receiver<ChatStreamEvent>>
where
@@ -816,12 +826,22 @@ impl InsightChatService {
{
let (tx, rx) = tokio::sync::mpsc::channel::<ChatStreamEvent>(64);
let svc = self.clone();
tokio::spawn(async move {
let result = svc.run_streaming_turn(req, cx, tx.clone()).await;
if let Err(e) = result {
let _ = tx.send(ChatStreamEvent::Error(format!("{}", e))).await;
// Carry the turn span into the task explicitly — tracing context is
// thread-local, so without `.instrument` every LLM / DB span inside
// the turn would orphan onto a fresh trace.
tokio::spawn(
async move {
let result = svc.run_streaming_turn(req, tx.clone()).await;
let span = tracing::Span::current();
if let Err(e) = result {
span.set_status(Status::error(format!("{e}")));
let _ = tx.send(ChatStreamEvent::Error(format!("{}", e))).await;
} else {
span.set_status(Status::Ok);
}
}
});
.instrument(span),
);
Ok(rx)
}
@@ -843,35 +863,38 @@ impl InsightChatService {
let svc = self.clone();
let entry_clone = entry.clone();
let turn_id_for_span = turn_id.clone();
let library_id = req.library_id;
let handle = tokio::spawn(async move {
// Span covering the whole spawned turn execution. Created here (not
// in the HTTP handler) because the dispatch span ends at the 202
// response, long before this work runs.
let tracer = global_tracer();
let mut span = tracer.start("ai.chat.turn.execute");
span.set_attribute(KeyValue::new("turn_id", turn_id_for_span));
span.set_attribute(KeyValue::new("library_id", library_id as i64));
let result = svc
.run_streaming_turn_with_entry(req, entry_clone.clone(), &mut span)
.await;
if let Err(ref e) = result {
span.set_attribute(KeyValue::new("status", "error"));
span.set_status(Status::error(format!("{e}")));
// Push the terminal event BEFORE flipping status: a replay
// reader treats a terminal status with no buffered tail as
// "closed", so the Error must be in the buffer first.
let _ = entry_clone
.push_event(ChatStreamEvent::Error(format!("{}", e)))
// Span covering the whole spawned turn execution. Created before the
// spawn so it parents on the dispatching request (same trace as the
// 202), and carried into the task with `.instrument`. It outlives the
// dispatch span, which ends at the 202 response.
let turn_span = tracing::info_span!(
"ai.chat.turn.execute",
turn_id = %turn_id,
library_id = req.library_id,
);
let handle = tokio::spawn(
async move {
let span = tracing::Span::current();
let result = svc
.run_streaming_turn_with_entry(req, entry_clone.clone(), &span)
.await;
entry_clone.set_terminal_status(crate::ai::turn_registry::TurnStatus::Error);
} else {
span.set_attribute(KeyValue::new("status", "done"));
span.set_status(Status::Ok);
if let Err(ref e) = result {
span.set_attribute("status", "error");
span.set_status(Status::error(format!("{e}")));
// Push the terminal event BEFORE flipping status: a replay
// reader treats a terminal status with no buffered tail as
// "closed", so the Error must be in the buffer first.
let _ = entry_clone
.push_event(ChatStreamEvent::Error(format!("{}", e)))
.await;
entry_clone.set_terminal_status(crate::ai::turn_registry::TurnStatus::Error);
} else {
span.set_attribute("status", "done");
span.set_status(Status::Ok);
}
}
});
.instrument(turn_span),
);
// Install the abort handle so DELETE can actually stop the task.
entry.set_abort_handle(handle.abort_handle());
@@ -885,7 +908,7 @@ impl InsightChatService {
self: Arc<Self>,
req: ChatTurnRequest,
entry: Arc<TurnEntry>,
turn_span: &mut opentelemetry::global::BoxedSpan,
turn_span: &tracing::Span,
) -> Result<()> {
if req.user_message.trim().is_empty() {
bail!("user_message must not be empty");
@@ -930,7 +953,7 @@ impl InsightChatService {
normalized: String,
insight: crate::database::models::PhotoInsight,
entry: Arc<TurnEntry>,
turn_span: &mut opentelemetry::global::BoxedSpan,
turn_span: &tracing::Span,
) -> Result<()> {
let active_persona = req
.persona_id
@@ -988,10 +1011,8 @@ impl InsightChatService {
let model_used = backend.model().to_string();
// llm.* attributes on the turn-level span (issue #130): clients can't
// know local-vs-hybrid, but the turn knows its resolved backend.
turn_span.set_attribute(KeyValue::new("llm.backend", kind.as_str()));
turn_span.set_attribute(KeyValue::new("llm.model", model_used.clone()));
let turn_cx = opentelemetry::Context::new()
.with_remote_span_context(turn_span.span_context().clone());
turn_span.set_attribute("llm.backend", kind.as_str());
turn_span.set_attribute("llm.model", model_used.clone());
let local_first_user_has_image = messages
.iter()
@@ -1041,7 +1062,6 @@ impl InsightChatService {
&active_persona,
max_iterations,
&entry,
&turn_cx,
)
.await?;
let AgenticLoopOutcome {
@@ -1158,7 +1178,7 @@ impl InsightChatService {
req: ChatTurnRequest,
normalized: String,
entry: Arc<TurnEntry>,
turn_span: &mut opentelemetry::global::BoxedSpan,
turn_span: &tracing::Span,
) -> Result<()> {
let active_persona = req
.persona_id
@@ -1185,10 +1205,8 @@ impl InsightChatService {
let backend = self.generator.resolve_backend(kind, &overrides).await?;
let model_used = backend.model().to_string();
// llm.* attributes on the turn-level span (issue #130).
turn_span.set_attribute(KeyValue::new("llm.backend", kind.as_str()));
turn_span.set_attribute(KeyValue::new("llm.model", model_used.clone()));
let turn_cx = opentelemetry::Context::new()
.with_remote_span_context(turn_span.span_context().clone());
turn_span.set_attribute("llm.backend", kind.as_str());
turn_span.set_attribute("llm.model", model_used.clone());
let image_base64: Option<String> = self.generator.load_image_as_base64(&normalized).ok();
@@ -1208,7 +1226,7 @@ impl InsightChatService {
let visual_block = if !backend.images_inline {
match image_base64.as_deref() {
Some(b64) => match backend.local().describe_image(&turn_cx, b64).await {
Some(b64) => match backend.local().describe_image(b64).await {
Ok(desc) => {
format!("Visual description (from local vision model):\n{}\n", desc)
}
@@ -1263,7 +1281,6 @@ impl InsightChatService {
&active_persona,
max_iterations,
&entry,
&turn_cx,
)
.await?;
let AgenticLoopOutcome {
@@ -1352,7 +1369,6 @@ impl InsightChatService {
active_persona: &str,
max_iterations: usize,
entry: &Arc<TurnEntry>,
cx: &opentelemetry::Context,
) -> Result<AgenticLoopOutcome> {
crate::ai::insight_chat::run_streaming_agentic_loop_with_entry(
&self.generator,
@@ -1365,7 +1381,6 @@ impl InsightChatService {
active_persona,
max_iterations,
entry,
cx,
)
.await
}
@@ -1373,7 +1388,6 @@ impl InsightChatService {
async fn run_streaming_turn(
self: Arc<Self>,
req: ChatTurnRequest,
cx: opentelemetry::Context,
tx: tokio::sync::mpsc::Sender<ChatStreamEvent>,
) -> Result<()> {
if req.user_message.trim().is_empty() {
@@ -1410,10 +1424,10 @@ impl InsightChatService {
};
if req.regenerate || existing_insight.is_none() {
return self.run_bootstrap_streaming(req, normalized, cx, tx).await;
return self.run_bootstrap_streaming(req, normalized, tx).await;
}
let insight = existing_insight.expect("just checked Some above");
self.run_continuation_streaming(req, normalized, cx, insight, tx)
self.run_continuation_streaming(req, normalized, insight, tx)
.await
}
@@ -1425,7 +1439,6 @@ impl InsightChatService {
&self,
req: ChatTurnRequest,
normalized: String,
cx: opentelemetry::Context,
insight: crate::database::models::PhotoInsight,
tx: tokio::sync::mpsc::Sender<ChatStreamEvent>,
) -> Result<()> {
@@ -1523,7 +1536,6 @@ impl InsightChatService {
&active_persona,
max_iterations,
&tx,
&cx,
)
.await?;
let AgenticLoopOutcome {
@@ -1635,7 +1647,6 @@ impl InsightChatService {
&self,
req: ChatTurnRequest,
normalized: String,
cx: opentelemetry::Context,
tx: tokio::sync::mpsc::Sender<ChatStreamEvent>,
) -> Result<()> {
let active_persona = req
@@ -1692,7 +1703,7 @@ impl InsightChatService {
// images_inline backends send images directly to the chat model.
let visual_block = if !backend.images_inline {
match image_base64.as_deref() {
Some(b64) => match backend.local().describe_image(&cx, b64).await {
Some(b64) => match backend.local().describe_image(b64).await {
Ok(desc) => {
format!("Visual description (from local vision model):\n{}\n", desc)
}
@@ -1759,7 +1770,6 @@ impl InsightChatService {
&active_persona,
max_iterations,
&tx,
&cx,
)
.await?;
let AgenticLoopOutcome {
@@ -1842,7 +1852,6 @@ impl InsightChatService {
active_persona: &str,
max_iterations: usize,
tx: &tokio::sync::mpsc::Sender<ChatStreamEvent>,
cx: &opentelemetry::Context,
) -> Result<AgenticLoopOutcome> {
let mut tool_calls_made = 0usize;
let mut iterations_used = 0usize;
@@ -1861,7 +1870,7 @@ impl InsightChatService {
let mut stream = backend
.chat()
.chat_with_tools_stream(cx, messages.clone(), tools.clone())
.chat_with_tools_stream(messages.clone(), tools.clone())
.await?;
let mut final_message: Option<ChatMessage> = None;
@@ -1955,7 +1964,7 @@ impl InsightChatService {
let synthetic_idx = push_synthetic_final_prompt(messages);
let mut stream = backend
.chat()
.chat_with_tools_stream(cx, messages.clone(), vec![])
.chat_with_tools_stream(messages.clone(), vec![])
.await?;
let mut final_message: Option<ChatMessage> = None;
while let Some(ev) = stream.next().await {
@@ -2335,7 +2344,6 @@ pub async fn run_streaming_agentic_loop_with_entry(
active_persona: &str,
max_iterations: usize,
entry: &Arc<TurnEntry>,
cx: &opentelemetry::Context,
) -> Result<AgenticLoopOutcome> {
let mut tool_calls_made = 0usize;
let mut iterations_used = 0usize;
@@ -2365,7 +2373,7 @@ pub async fn run_streaming_agentic_loop_with_entry(
let mut stream = backend
.chat()
.chat_with_tools_stream(cx, messages.clone(), tools.clone())
.chat_with_tools_stream(messages.clone(), tools.clone())
.await?;
let mut final_message: Option<ChatMessage> = None;
@@ -2448,7 +2456,7 @@ pub async fn run_streaming_agentic_loop_with_entry(
let synthetic_idx = push_synthetic_final_prompt(messages);
let mut stream = backend
.chat()
.chat_with_tools_stream(cx, messages.clone(), vec![])
.chat_with_tools_stream(messages.clone(), vec![])
.await?;
let mut final_message: Option<ChatMessage> = None;
while let Some(ev) = stream.next().await {