From f850c535f1f2b69023604e57b9a83472522f4f9f Mon Sep 17 00:00:00 2001 From: Cameron Cordes Date: Mon, 28 Sep 2026 21:41:38 -0400 Subject: [PATCH] Otel: thread request context through streaming chat path (#129 review) The legacy SSE streaming path (chat_turn_stream -> run_streaming_turn -> run_bootstrap/continuation_streaming) built its LLM spans from Context::current(), which is always the empty root context here: the codebase threads OTel contexts explicitly and never attaches one, there is no context-aware tokio layer, and the turn runs on a tokio::spawn. So every llm.request span on that path was an orphaned root span, uncorrelated with the HTTP request. Thread the handler's extracted request context down through chat_turn_stream -> chat_turn_stream_inner -> run_streaming_turn -> run_bootstrap_streaming / run_continuation_streaming, and use it for the describe_image and agentic-loop calls. Mirrors the non-streaming path, which already threads parent_cx -> loop_cx. Also rename the block-scoped db 'cx' in run_streaming_turn to db_cx so it no longer shadows the new parameter. --- src/ai/handlers.rs | 4 +++- src/ai/insight_chat.rs | 27 ++++++++++++++------------- 2 files changed, 17 insertions(+), 14 deletions(-) diff --git a/src/ai/handlers.rs b/src/ai/handlers.rs index b00c7b2..82a5406 100644 --- a/src/ai/handlers.rs +++ b/src/ai/handlers.rs @@ -1667,10 +1667,12 @@ pub struct ChatSwitchBranchRequest { /// Returns `text/event-stream` with one event per chat stream event. #[post("/insights/chat/stream")] pub async fn chat_stream_handler( + http_request: HttpRequest, claims: Claims, request: web::Json, app_state: web::Data, ) -> HttpResponse { + let parent_context = extract_context_from_request(&http_request); let mut request = request.0; // Stored owner defaults fill what the client omitted (request → pref → // hard default; the hard default still applies downstream). @@ -1716,7 +1718,7 @@ pub async fn chat_stream_handler( }; let service = app_state.insight_chat.clone(); - let events = service.chat_turn_stream(chat_req); + let events = service.chat_turn_stream(chat_req, parent_context); // Map ChatStreamEvent → SSE frame bytes. let sse_stream = futures::stream::StreamExt::map(events, |ev| { diff --git a/src/ai/insight_chat.rs b/src/ai/insight_chat.rs index f4dd3da..b11a1e8 100644 --- a/src/ai/insight_chat.rs +++ b/src/ai/insight_chat.rs @@ -782,10 +782,11 @@ impl InsightChatService { pub fn chat_turn_stream( self: Arc, req: ChatTurnRequest, + cx: opentelemetry::Context, ) -> BoxStream<'static, ChatStreamEvent> { let svc = self; let s = async_stream::stream! { - match svc.chat_turn_stream_inner(req, Ok).await { + match svc.chat_turn_stream_inner(req, cx, Ok).await { Ok(mut rx) => { while let Some(ev) = rx.recv().await { yield ev; @@ -807,6 +808,7 @@ impl InsightChatService { async fn chat_turn_stream_inner( self: Arc, req: ChatTurnRequest, + cx: opentelemetry::Context, _ev_mapper: F, ) -> Result> where @@ -815,7 +817,7 @@ impl InsightChatService { let (tx, rx) = tokio::sync::mpsc::channel::(64); let svc = self.clone(); tokio::spawn(async move { - let result = svc.run_streaming_turn(req, tx.clone()).await; + let result = svc.run_streaming_turn(req, cx, tx.clone()).await; if let Err(e) = result { let _ = tx.send(ChatStreamEvent::Error(format!("{}", e))).await; } @@ -1371,6 +1373,7 @@ impl InsightChatService { async fn run_streaming_turn( self: Arc, req: ChatTurnRequest, + cx: opentelemetry::Context, tx: tokio::sync::mpsc::Sender, ) -> Result<()> { if req.user_message.trim().is_empty() { @@ -1400,17 +1403,17 @@ impl InsightChatService { // library bootstraps cleanly, even when another library has // an insight for the same rel_path. let existing_insight = { - let cx = opentelemetry::Context::new(); + let db_cx = opentelemetry::Context::new(); let mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao"); - dao.get_current_insight_for_library(&cx, req.library_id, &normalized) + dao.get_current_insight_for_library(&db_cx, req.library_id, &normalized) .map_err(|e| anyhow!("failed to load insight: {:?}", e))? }; if req.regenerate || existing_insight.is_none() { - return self.run_bootstrap_streaming(req, normalized, tx).await; + return self.run_bootstrap_streaming(req, normalized, cx, tx).await; } let insight = existing_insight.expect("just checked Some above"); - self.run_continuation_streaming(req, normalized, insight, tx) + self.run_continuation_streaming(req, normalized, cx, insight, tx) .await } @@ -1422,6 +1425,7 @@ impl InsightChatService { &self, req: ChatTurnRequest, normalized: String, + cx: opentelemetry::Context, insight: crate::database::models::PhotoInsight, tx: tokio::sync::mpsc::Sender, ) -> Result<()> { @@ -1519,7 +1523,7 @@ impl InsightChatService { &active_persona, max_iterations, &tx, - &opentelemetry::Context::current(), + &cx, ) .await?; let AgenticLoopOutcome { @@ -1631,6 +1635,7 @@ impl InsightChatService { &self, req: ChatTurnRequest, normalized: String, + cx: opentelemetry::Context, tx: tokio::sync::mpsc::Sender, ) -> Result<()> { let active_persona = req @@ -1687,11 +1692,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(&opentelemetry::Context::current(), b64) - .await - { + Some(b64) => match backend.local().describe_image(&cx, b64).await { Ok(desc) => { format!("Visual description (from local vision model):\n{}\n", desc) } @@ -1758,7 +1759,7 @@ impl InsightChatService { &active_persona, max_iterations, &tx, - &opentelemetry::Context::current(), + &cx, ) .await?; let AgenticLoopOutcome {