fix: avoid OpenTelemetry context detach errors in mock runtimes

demo/mock_invoice_runtime.py

@@ -174,7 +174,7 @@ async def stream(
         invoice_id = str(payload.get("invoice_id", "INV-2024-08731"))
         currency = str(payload.get("currency", "USD"))
 
-        with self.tracer.start_as_current_span(
+        root_span = self.tracer.start_span(
             "invoice_processing.execute",
             attributes={
                 "uipath.runtime.name": "InvoiceProcessingRuntime",
@@ -183,7 +183,8 @@ async def stream(
                 "uipath.input.invoice_id": invoice_id,
                 "uipath.input.currency": currency,
             },
-        ):
+        )
+        try:
             # 1. Ingest document — parse the invoice PDF
             yield _state("ingest_document", S)
             with self.tracer.start_as_current_span(
@@ -391,6 +392,8 @@ async def stream(
                 C,
                 {"payment_reference": payment_ref},
             )
+        finally:
+            root_span.end()
 
         yield UiPathRuntimeResult(
             output={

demo/mock_movies_runtime.py

@@ -419,7 +419,7 @@ async def stream(
 
         message_id = str(uuid4())
 
-        with self.tracer.start_as_current_span(
+        root_span = self.tracer.start_span(
             "movies.execute",
             attributes={
                 "uipath.runtime.name": "MoviesRuntime",
@@ -428,18 +428,20 @@ async def stream(
                 "uipath.input.message.length": len(message),
                 "uipath.input.turn_keywords": ",".join(turn.get("keywords", [])),
             },
-        ):
+        )
+        try:
             # --- Model (first call — decides to search) ---
             yield self._node_state("model", S)
-            with self.tracer.start_as_current_span(
+            model_span = self.tracer.start_span(
                 "model",
                 attributes={
                     "uipath.step.kind": "model",
                     "uipath.input.model": "claude-3-7-sonnet-latest",
                     "uipath.output.tool_calls_count": len(tool_defs),
                     "uipath.output.tool_names": ",".join(t["name"] for t in tool_defs),
                 },
-            ):
+            )
+            try:
                 # message_start for tool-call message
                 yield UiPathRuntimeMessageEvent(
                     payload=UiPathConversationMessageEvent(
@@ -453,6 +455,8 @@ async def stream(
                 # Emit tool call events
                 async for evt in self._emit_tool_calls(message_id, tool_defs):
                     yield evt
+            finally:
+                model_span.end()
             yield self._node_state("model", C)
 
             # --- Tools (execute search) ---
@@ -473,18 +477,23 @@ async def stream(
 
             # --- Model (second call — streaming response with search results) ---
             yield self._node_state("model", S)
-            with self.tracer.start_as_current_span(
+            model_final_span = self.tracer.start_span(
                 "model.final",
                 attributes={
                     "uipath.step.kind": "model",
                     "uipath.input.has_tool_results": True,
                     "uipath.output.reply.length": len(reply),
                     "uipath.output.finish_reason": "end_turn",
                 },
-            ):
+            )
+            try:
                 async for evt in self._emit_streaming_response(reply):
                     yield evt
+            finally:
+                model_final_span.end()
             yield self._node_state("model", C)
+        finally:
+            root_span.end()
 
         yield UiPathRuntimeResult(
             output={"reply": reply},

demo/mock_pharma_runtime.py

@@ -186,7 +186,7 @@ async def stream(
         doc_id = str(payload.get("document_id", "SOP-2024-0142"))
         review_type = str(payload.get("review_type", "initial"))
 
-        with self.tracer.start_as_current_span(
+        root_span = self.tracer.start_span(
             "pharma_compliance.execute",
             attributes={
                 "uipath.runtime.name": "PharmaComplianceReview",
@@ -195,7 +195,8 @@ async def stream(
                 "uipath.input.document_id": doc_id,
                 "uipath.input.review_type": review_type,
             },
-        ):
+        )
+        try:
             # 1. Extract metadata
             yield _state("extract_metadata", S)
             with self.tracer.start_as_current_span(
@@ -351,6 +352,8 @@ async def stream(
                 )
                 span.set_attribute("uipath.output.recommendation", recommendation)
             yield _state("generate_report", C)
+        finally:
+            root_span.end()
 
         yield UiPathRuntimeResult(
             output={

demo/mock_support_runtime.py

@@ -458,15 +458,18 @@ async def _stream_phase2(
 
         # --- Model (streaming final response) ---
         yield self._node_state("model", S)
-        with self.tracer.start_as_current_span(
+        model_span = self.tracer.start_span(
             "model.2",
             attributes={
                 "uipath.step.kind": "model",
                 "uipath.output.reply.length": len(reply),
             },
-        ):
+        )
+        try:
             async for evt in self._emit_streaming_response(reply):
                 yield evt
+        finally:
+            model_span.end()
         yield self._node_state("model", C)
 
         # --- Middleware: TodoList (routes to __end__) ---
@@ -500,15 +503,18 @@ async def stream(
             self._suspended_turn = None
             self._suspended_message_id = None
 
-            with self.tracer.start_as_current_span(
+            resume_span = self.tracer.start_span(
                 "support_chat.resume",
                 attributes={
                     "uipath.runtime.name": "SupportChatRuntime",
                     "uipath.step.kind": "resume",
                 },
-            ):
+            )
+            try:
                 async for evt in self._stream_phase2(turn):
                     yield evt
+            finally:
+                resume_span.end()
             return
 
         # --- Normal flow ---
@@ -531,7 +537,7 @@ async def stream(
         # Shared message_id for the assistant turn (tool calls + response)
         message_id = str(uuid4())
 
-        with self.tracer.start_as_current_span(
+        root_span = self.tracer.start_span(
             "support_chat.execute",
             attributes={
                 "uipath.runtime.name": "SupportChatRuntime",
@@ -540,7 +546,8 @@ async def stream(
                 "uipath.input.message.length": len(message),
                 "uipath.input.turn_keywords": ",".join(turn.get("keywords", [])),
             },
-        ):
+        )
+        try:
             # --- Middleware: PatchToolCalls ---
             yield self._node_state("PatchToolCallsMiddleware.before_agent", S)
             with self.tracer.start_as_current_span(
@@ -567,14 +574,15 @@ async def stream(
 
             # --- Model (first call — decides to use tools) ---
             yield self._node_state("model", S)
-            with self.tracer.start_as_current_span(
+            model_span = self.tracer.start_span(
                 "model",
                 attributes={
                     "uipath.step.kind": "model",
                     "uipath.output.tool_calls_count": len(tool_defs),
                     "uipath.output.tool_names": ",".join(t["name"] for t in tool_defs),
                 },
-            ):
+            )
+            try:
                 # message_start for tool-call message
                 yield UiPathRuntimeMessageEvent(
                     payload=UiPathConversationMessageEvent(
@@ -588,6 +596,8 @@ async def stream(
                 # Emit tool call events from the model
                 async for evt in self._emit_tool_calls(message_id, tool_defs):
                     yield evt
+            finally:
+                model_span.end()
             yield self._node_state("model", C)
 
             # --- Middleware: TodoList (routes to tools) ---
@@ -658,15 +668,18 @@ async def stream(
 
             # --- Model (second call — streaming final response with tool results) ---
             yield self._node_state("model", S)
-            with self.tracer.start_as_current_span(
+            model2_span = self.tracer.start_span(
                 "model.2",
                 attributes={
                     "uipath.step.kind": "model",
                     "uipath.output.reply.length": len(reply),
                 },
-            ):
+            )
+            try:
                 async for evt in self._emit_streaming_response(reply):
                     yield evt
+            finally:
+                model2_span.end()
             yield self._node_state("model", C)
 
             # --- Middleware: TodoList (routes to __end__) ---
@@ -680,6 +693,8 @@ async def stream(
             ):
                 await asyncio.sleep(0.1)
             yield self._node_state("TodoListMiddleware.after_model", C)
+        finally:
+            root_span.end()
 
         yield UiPathRuntimeResult(
             output={"reply": reply},

demo/mock_telemetry_runtime.py

@@ -279,15 +279,16 @@ async def stream(
         if is_resuming:
             self.current_step_index += 1
 
-        with self.tracer.start_as_current_span(
+        root_span = self.tracer.start_span(
             "mock-runtime.execute",
             attributes={
                 "uipath.runtime.name": "MockRuntime",
                 "uipath.runtime.type": "agent",
                 "uipath.runtime.entrypoint": self.entrypoint,
                 "uipath.input.message.length": len(message),
             },
-        ):
+        )
+        try:
             while self.current_step_index < len(steps):
                 step = steps[self.current_step_index]
 
@@ -327,34 +328,42 @@ async def stream(
 
                 elif step == "researcher":
                     yield _state("researcher", S)
-                    with self.tracer.start_as_current_span(
+                    researcher_span = self.tracer.start_span(
                         "researcher",
                         attributes={
                             "uipath.step.kind": "subgraph",
                             "uipath.input.task": "web_research",
                             "uipath.input.query": message[:100],
                         },
-                    ) as span:
+                    )
+                    try:
                         async for evt in self._run_subgraph("researcher", "researcher"):
                             yield evt
-                        span.set_attribute("uipath.output.sources_found", 3)
-                        span.set_attribute("uipath.output.summary.length", 450)
+                        researcher_span.set_attribute("uipath.output.sources_found", 3)
+                        researcher_span.set_attribute(
+                            "uipath.output.summary.length", 450
+                        )
+                    finally:
+                        researcher_span.end()
                     yield _state("researcher", C)
 
                 elif step == "coder":
                     yield _state("coder", S)
-                    with self.tracer.start_as_current_span(
+                    coder_span = self.tracer.start_span(
                         "coder",
                         attributes={
                             "uipath.step.kind": "subgraph",
                             "uipath.input.task": "code_generation",
                             "uipath.input.language": "python",
                         },
-                    ) as span:
+                    )
+                    try:
                         async for evt in self._run_subgraph("coder", "coder"):
                             yield evt
-                        span.set_attribute("uipath.output.lines_generated", 47)
-                        span.set_attribute("uipath.output.tests_passed", True)
+                        coder_span.set_attribute("uipath.output.lines_generated", 47)
+                        coder_span.set_attribute("uipath.output.tests_passed", True)
+                    finally:
+                        coder_span.end()
                     yield _state("coder", C)
 
                 elif step == "output":
@@ -399,6 +408,8 @@ async def stream(
 
             # All steps completed — reset
             self.current_step_index = 0
+        finally:
+            root_span.end()
 
         yield UiPathRuntimeResult(
             output={

demo/mock_template_runtime.py

@@ -114,13 +114,14 @@ async def stream(
         """Stream events from the JSON file."""
         logger.info(f"MockTemplateRuntime: streaming {len(self._events)} events")
 
-        with self.tracer.start_as_current_span(
+        root_span = self.tracer.start_span(
             "template.stream",
             attributes={
                 "uipath.runtime.name": "MockTemplateRuntime",
                 "uipath.event.count": len(self._events),
             },
-        ):
+        )
+        try:
             for i, event_data in enumerate(self._events):
                 event_type = event_data.get("event_type")
 
@@ -166,6 +167,8 @@ async def stream(
                 except Exception as e:
                     logger.error(f"Error processing event {i}: {e}", exc_info=True)
                     continue
+        finally:
+            root_span.end()
 
         logger.info("MockTemplateRuntime: streaming completed")