fix(bindings): restore trace context across the Python chat-processor boundary (3 orphan traces → 1 waterfall)

lib/bindings/python/rust/llm/routed_engine.rs

@@ -1,9 +1,12 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+use opentelemetry::trace::{SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId, TraceState};
 use pyo3::prelude::*;
 use pythonize::{depythonize, pythonize};
 use tokio_stream::StreamExt;
+use tracing::Instrument;
+use tracing_opentelemetry::OpenTelemetrySpanExt;
 
 use dynamo_llm::entrypoint::PrefillRoutedEngine;
 use dynamo_llm::protocols::common::preprocessor::PreprocessedRequest;
@@ -55,41 +58,88 @@ impl RoutedEngine {
         };
         let inner = self.inner.clone();
 
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            let mut stream = inner.generate(request_context).await.map_err(to_pyerr)?;
-            let task_context = stream.context();
-            let (tx, rx) = tokio::sync::mpsc::channel::<RsAnnotated<PyObject>>(32);
+        // Restore the caller's trace context across the Python boundary.
+        //
+        // When the frontend hands a request to a Python chat processor, the
+        // Rust side captures its `DistributedTraceContext` (the identity of
+        // the `http-request` span) into the PyContext (engine.rs:
+        // `Context::new(ctx, get_distributed_tracing_context(), ..)`). The
+        // Python processor then calls back into this method — but the future
+        // below runs on a fresh tokio task where `Span::current()` is empty.
+        // Without re-parenting, every downstream dispatch span
+        // (`router.route_request` → worker `handle_payload` → engine spans)
+        // becomes an orphan trace root: one request fans out into three
+        // unrelated traces (http-request / prefill / decode) instead of one.
+        //
+        // Re-parent a `routed_engine.generate` span to the captured context
+        // so the whole PD pipeline lands back on the request's trace.
+        let dispatch_span = match context.as_ref().and_then(|c| c.trace_context()) {
+            Some(tc) => {
+                let span = tracing::info_span!(
+                    "routed_engine.generate",
+                    request_id = request_context.id()
+                );
+                if let (Ok(trace_id), Ok(span_id)) = (
+                    TraceId::from_hex(&tc.trace_id),
+                    SpanId::from_hex(&tc.span_id),
+                ) {
+                    let span_context = SpanContext::new(
+                        trace_id,
+                        span_id,
+                        TraceFlags::SAMPLED,
+                        true, // is_remote
+                        TraceState::default(),
+                    );
+                    let _ = span.set_parent(
+                        opentelemetry::Context::new().with_remote_span_context(span_context),
+                    );
+                }
+                span
+            }
+            // No trace context on the PyContext (e.g. tests constructing a
+            // bare Context, or tracing disabled): keep current behavior.
+            None => tracing::Span::current(),
+        };
 
-            tokio::spawn(async move {
-                loop {
-                    let response = tokio::select! {
-                        _ = tx.closed() => {
-                            task_context.stop_generating();
-                            break;
-                        }
-                        response = stream.next() => response,
-                    };
+        pyo3_async_runtimes::tokio::future_into_py(
+            py,
+            async move {
+                let mut stream = inner.generate(request_context).await.map_err(to_pyerr)?;
+                let task_context = stream.context();
+                let (tx, rx) = tokio::sync::mpsc::channel::<RsAnnotated<PyObject>>(32);
 
-                    let Some(response) = response else {
-                        break;
-                    };
+                tokio::spawn(async move {
+                    loop {
+                        let response = tokio::select! {
+                            _ = tx.closed() => {
+                                task_context.stop_generating();
+                                break;
+                            }
+                            response = stream.next() => response,
+                        };
 
-                    let py_response = Python::with_gil(|py| {
-                        response.map_data(|data| {
-                            pythonize(py, &data)
-                                .map(|obj| obj.unbind())
-                                .map_err(|e| format!("pythonize failed: {e}"))
-                        })
-                    });
+                        let Some(response) = response else {
+                            break;
+                        };
 
-                    if tx.send(py_response).await.is_err() {
-                        task_context.stop_generating();
-                        break;
+                        let py_response = Python::with_gil(|py| {
+                            response.map_data(|data| {
+                                pythonize(py, &data)
+                                    .map(|obj| obj.unbind())
+                                    .map_err(|e| format!("pythonize failed: {e}"))
+                            })
+                        });
+
+                        if tx.send(py_response).await.is_err() {
+                            task_context.stop_generating();
+                            break;
+                        }
                     }
-                }
-            });
+                });
 
-            Ok(crate::AsyncResponseStream::new(rx, true))
-        })
+                Ok(crate::AsyncResponseStream::new(rx, true))
+            }
+            .instrument(dispatch_span),
+        )
     }
 }