feat(tracer): enhance ASGI compatibility by capturing context during tracing

Author: viniciusdsmello

src/openlayer/lib/tracing/tracer.py

@@ -496,6 +496,7 @@ def __init__(self):
                         self._token = None
                         self._output_chunks = []
                         self._trace_initialized = False
+                        self._captured_context = None  # Capture context for ASGI compatibility
 
                     def __iter__(self):
                         return self
@@ -517,6 +518,10 @@ def __next__(self):
                                 func_kwargs=func_kwargs,
                                 context_kwarg=context_kwarg,
                             )
+                            # Capture context after trace is initialized for ASGI compatibility
+                            # This ensures we can finalize in the correct context even if
+                            # later __next__ calls happen in different worker threads
+                            self._captured_context = contextvars.copy_context()
                             self._trace_initialized = True
 
                         try:
@@ -525,23 +530,21 @@ def __next__(self):
                             return chunk
                         except StopIteration:
                             # Finalize trace when generator is exhausted
+                            # Use captured context to ensure we have access to the trace
                             output = _join_output_chunks(self._output_chunks)
-                            _finalize_sync_generator_step(
-                                step=self._step,
-                                token=self._token,
-                                is_root_step=self._is_root_step,
-                                step_name=step_name,
-                                inputs=self._inputs,
-                                output=output,
-                                inference_pipeline_id=inference_pipeline_id,
-                                on_flush_failure=on_flush_failure,
-                            )
-                            raise
-                        except Exception as exc:
-                            # Handle exceptions
-                            if self._step:
-                                _log_step_exception(self._step, exc)
-                                output = _join_output_chunks(self._output_chunks)
+                            if self._captured_context:
+                                self._captured_context.run(
+                                    _finalize_sync_generator_step,
+                                    step=self._step,
+                                    token=self._token,
+                                    is_root_step=self._is_root_step,
+                                    step_name=step_name,
+                                    inputs=self._inputs,
+                                    output=output,
+                                    inference_pipeline_id=inference_pipeline_id,
+                                    on_flush_failure=on_flush_failure,
+                                )
+                            else:
                                 _finalize_sync_generator_step(
                                     step=self._step,
                                     token=self._token,
@@ -553,6 +556,35 @@ def __next__(self):
                                     on_flush_failure=on_flush_failure,
                                 )
                             raise
+                        except Exception as exc:
+                            # Handle exceptions
+                            if self._step:
+                                _log_step_exception(self._step, exc)
+                                output = _join_output_chunks(self._output_chunks)
+                                if self._captured_context:
+                                    self._captured_context.run(
+                                        _finalize_sync_generator_step,
+                                        step=self._step,
+                                        token=self._token,
+                                        is_root_step=self._is_root_step,
+                                        step_name=step_name,
+                                        inputs=self._inputs,
+                                        output=output,
+                                        inference_pipeline_id=inference_pipeline_id,
+                                        on_flush_failure=on_flush_failure,
+                                    )
+                                else:
+                                    _finalize_sync_generator_step(
+                                        step=self._step,
+                                        token=self._token,
+                                        is_root_step=self._is_root_step,
+                                        step_name=step_name,
+                                        inputs=self._inputs,
+                                        output=output,
+                                        inference_pipeline_id=inference_pipeline_id,
+                                        on_flush_failure=on_flush_failure,
+                                    )
+                            raise
 
                 return TracedSyncGenerator()