Skip to content

Commit 47f08b8

Browse files
cristipufucristipufu
authored
Merge pull request UiPath#74 from UiPath/fix/otel-context-detach-in-mock-runtimes
fix: avoid OpenTelemetry context detach errors in mock runtimes
2 parents 244c2f9 + a308d01 commit 47f08b8

6 files changed

Lines changed: 76 additions & 32 deletions

‎demo/mock_invoice_runtime.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,7 @@ async def stream(
174174
invoice_id = str(payload.get("invoice_id", "INV-2024-08731"))
175175
currency = str(payload.get("currency", "USD"))
176176

177-
with self.tracer.start_as_current_span(
177+
root_span = self.tracer.start_span(
178178
"invoice_processing.execute",
179179
attributes={
180180
"uipath.runtime.name": "InvoiceProcessingRuntime",
@@ -183,7 +183,8 @@ async def stream(
183183
"uipath.input.invoice_id": invoice_id,
184184
"uipath.input.currency": currency,
185185
},
186-
):
186+
)
187+
try:
187188
# 1. Ingest document — parse the invoice PDF
188189
yield _state("ingest_document", S)
189190
with self.tracer.start_as_current_span(
@@ -391,6 +392,8 @@ async def stream(
391392
C,
392393
{"payment_reference": payment_ref},
393394
)
395+
finally:
396+
root_span.end()
394397

395398
yield UiPathRuntimeResult(
396399
output={

‎demo/mock_movies_runtime.py‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -419,7 +419,7 @@ async def stream(
419419

420420
message_id = str(uuid4())
421421

422-
with self.tracer.start_as_current_span(
422+
root_span = self.tracer.start_span(
423423
"movies.execute",
424424
attributes={
425425
"uipath.runtime.name": "MoviesRuntime",
@@ -428,18 +428,20 @@ async def stream(
428428
"uipath.input.message.length": len(message),
429429
"uipath.input.turn_keywords": ",".join(turn.get("keywords", [])),
430430
},
431-
):
431+
)
432+
try:
432433
# --- Model (first call — decides to search) ---
433434
yield self._node_state("model", S)
434-
with self.tracer.start_as_current_span(
435+
model_span = self.tracer.start_span(
435436
"model",
436437
attributes={
437438
"uipath.step.kind": "model",
438439
"uipath.input.model": "claude-3-7-sonnet-latest",
439440
"uipath.output.tool_calls_count": len(tool_defs),
440441
"uipath.output.tool_names": ",".join(t["name"] for t in tool_defs),
441442
},
442-
):
443+
)
444+
try:
443445
# message_start for tool-call message
444446
yield UiPathRuntimeMessageEvent(
445447
payload=UiPathConversationMessageEvent(
@@ -453,6 +455,8 @@ async def stream(
453455
# Emit tool call events
454456
async for evt in self._emit_tool_calls(message_id, tool_defs):
455457
yield evt
458+
finally:
459+
model_span.end()
456460
yield self._node_state("model", C)
457461

458462
# --- Tools (execute search) ---
@@ -473,18 +477,23 @@ async def stream(
473477

474478
# --- Model (second call — streaming response with search results) ---
475479
yield self._node_state("model", S)
476-
with self.tracer.start_as_current_span(
480+
model_final_span = self.tracer.start_span(
477481
"model.final",
478482
attributes={
479483
"uipath.step.kind": "model",
480484
"uipath.input.has_tool_results": True,
481485
"uipath.output.reply.length": len(reply),
482486
"uipath.output.finish_reason": "end_turn",
483487
},
484-
):
488+
)
489+
try:
485490
async for evt in self._emit_streaming_response(reply):
486491
yield evt
492+
finally:
493+
model_final_span.end()
487494
yield self._node_state("model", C)
495+
finally:
496+
root_span.end()
488497

489498
yield UiPathRuntimeResult(
490499
output={"reply": reply},

‎demo/mock_pharma_runtime.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -186,7 +186,7 @@ async def stream(
186186
doc_id = str(payload.get("document_id", "SOP-2024-0142"))
187187
review_type = str(payload.get("review_type", "initial"))
188188

189-
with self.tracer.start_as_current_span(
189+
root_span = self.tracer.start_span(
190190
"pharma_compliance.execute",
191191
attributes={
192192
"uipath.runtime.name": "PharmaComplianceReview",
@@ -195,7 +195,8 @@ async def stream(
195195
"uipath.input.document_id": doc_id,
196196
"uipath.input.review_type": review_type,
197197
},
198-
):
198+
)
199+
try:
199200
# 1. Extract metadata
200201
yield _state("extract_metadata", S)
201202
with self.tracer.start_as_current_span(
@@ -351,6 +352,8 @@ async def stream(
351352
)
352353
span.set_attribute("uipath.output.recommendation", recommendation)
353354
yield _state("generate_report", C)
355+
finally:
356+
root_span.end()
354357

355358
yield UiPathRuntimeResult(
356359
output={

‎demo/mock_support_runtime.py‎

Lines changed: 25 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -458,15 +458,18 @@ async def _stream_phase2(
458458

459459
# --- Model (streaming final response) ---
460460
yield self._node_state("model", S)
461-
with self.tracer.start_as_current_span(
461+
model_span = self.tracer.start_span(
462462
"model.2",
463463
attributes={
464464
"uipath.step.kind": "model",
465465
"uipath.output.reply.length": len(reply),
466466
},
467-
):
467+
)
468+
try:
468469
async for evt in self._emit_streaming_response(reply):
469470
yield evt
471+
finally:
472+
model_span.end()
470473
yield self._node_state("model", C)
471474

472475
# --- Middleware: TodoList (routes to __end__) ---
@@ -500,15 +503,18 @@ async def stream(
500503
self._suspended_turn = None
501504
self._suspended_message_id = None
502505

503-
with self.tracer.start_as_current_span(
506+
resume_span = self.tracer.start_span(
504507
"support_chat.resume",
505508
attributes={
506509
"uipath.runtime.name": "SupportChatRuntime",
507510
"uipath.step.kind": "resume",
508511
},
509-
):
512+
)
513+
try:
510514
async for evt in self._stream_phase2(turn):
511515
yield evt
516+
finally:
517+
resume_span.end()
512518
return
513519

514520
# --- Normal flow ---
@@ -531,7 +537,7 @@ async def stream(
531537
# Shared message_id for the assistant turn (tool calls + response)
532538
message_id = str(uuid4())
533539

534-
with self.tracer.start_as_current_span(
540+
root_span = self.tracer.start_span(
535541
"support_chat.execute",
536542
attributes={
537543
"uipath.runtime.name": "SupportChatRuntime",
@@ -540,7 +546,8 @@ async def stream(
540546
"uipath.input.message.length": len(message),
541547
"uipath.input.turn_keywords": ",".join(turn.get("keywords", [])),
542548
},
543-
):
549+
)
550+
try:
544551
# --- Middleware: PatchToolCalls ---
545552
yield self._node_state("PatchToolCallsMiddleware.before_agent", S)
546553
with self.tracer.start_as_current_span(
@@ -567,14 +574,15 @@ async def stream(
567574

568575
# --- Model (first call — decides to use tools) ---
569576
yield self._node_state("model", S)
570-
with self.tracer.start_as_current_span(
577+
model_span = self.tracer.start_span(
571578
"model",
572579
attributes={
573580
"uipath.step.kind": "model",
574581
"uipath.output.tool_calls_count": len(tool_defs),
575582
"uipath.output.tool_names": ",".join(t["name"] for t in tool_defs),
576583
},
577-
):
584+
)
585+
try:
578586
# message_start for tool-call message
579587
yield UiPathRuntimeMessageEvent(
580588
payload=UiPathConversationMessageEvent(
@@ -588,6 +596,8 @@ async def stream(
588596
# Emit tool call events from the model
589597
async for evt in self._emit_tool_calls(message_id, tool_defs):
590598
yield evt
599+
finally:
600+
model_span.end()
591601
yield self._node_state("model", C)
592602

593603
# --- Middleware: TodoList (routes to tools) ---
@@ -658,15 +668,18 @@ async def stream(
658668

659669
# --- Model (second call — streaming final response with tool results) ---
660670
yield self._node_state("model", S)
661-
with self.tracer.start_as_current_span(
671+
model2_span = self.tracer.start_span(
662672
"model.2",
663673
attributes={
664674
"uipath.step.kind": "model",
665675
"uipath.output.reply.length": len(reply),
666676
},
667-
):
677+
)
678+
try:
668679
async for evt in self._emit_streaming_response(reply):
669680
yield evt
681+
finally:
682+
model2_span.end()
670683
yield self._node_state("model", C)
671684

672685
# --- Middleware: TodoList (routes to __end__) ---
@@ -680,6 +693,8 @@ async def stream(
680693
):
681694
await asyncio.sleep(0.1)
682695
yield self._node_state("TodoListMiddleware.after_model", C)
696+
finally:
697+
root_span.end()
683698

684699
yield UiPathRuntimeResult(
685700
output={"reply": reply},

‎demo/mock_telemetry_runtime.py‎

Lines changed: 21 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -279,15 +279,16 @@ async def stream(
279279
if is_resuming:
280280
self.current_step_index += 1
281281

282-
with self.tracer.start_as_current_span(
282+
root_span = self.tracer.start_span(
283283
"mock-runtime.execute",
284284
attributes={
285285
"uipath.runtime.name": "MockRuntime",
286286
"uipath.runtime.type": "agent",
287287
"uipath.runtime.entrypoint": self.entrypoint,
288288
"uipath.input.message.length": len(message),
289289
},
290-
):
290+
)
291+
try:
291292
while self.current_step_index < len(steps):
292293
step = steps[self.current_step_index]
293294

@@ -327,34 +328,42 @@ async def stream(
327328

328329
elif step == "researcher":
329330
yield _state("researcher", S)
330-
with self.tracer.start_as_current_span(
331+
researcher_span = self.tracer.start_span(
331332
"researcher",
332333
attributes={
333334
"uipath.step.kind": "subgraph",
334335
"uipath.input.task": "web_research",
335336
"uipath.input.query": message[:100],
336337
},
337-
) as span:
338+
)
339+
try:
338340
async for evt in self._run_subgraph("researcher", "researcher"):
339341
yield evt
340-
span.set_attribute("uipath.output.sources_found", 3)
341-
span.set_attribute("uipath.output.summary.length", 450)
342+
researcher_span.set_attribute("uipath.output.sources_found", 3)
343+
researcher_span.set_attribute(
344+
"uipath.output.summary.length", 450
345+
)
346+
finally:
347+
researcher_span.end()
342348
yield _state("researcher", C)
343349

344350
elif step == "coder":
345351
yield _state("coder", S)
346-
with self.tracer.start_as_current_span(
352+
coder_span = self.tracer.start_span(
347353
"coder",
348354
attributes={
349355
"uipath.step.kind": "subgraph",
350356
"uipath.input.task": "code_generation",
351357
"uipath.input.language": "python",
352358
},
353-
) as span:
359+
)
360+
try:
354361
async for evt in self._run_subgraph("coder", "coder"):
355362
yield evt
356-
span.set_attribute("uipath.output.lines_generated", 47)
357-
span.set_attribute("uipath.output.tests_passed", True)
363+
coder_span.set_attribute("uipath.output.lines_generated", 47)
364+
coder_span.set_attribute("uipath.output.tests_passed", True)
365+
finally:
366+
coder_span.end()
358367
yield _state("coder", C)
359368

360369
elif step == "output":
@@ -399,6 +408,8 @@ async def stream(
399408

400409
# All steps completed — reset
401410
self.current_step_index = 0
411+
finally:
412+
root_span.end()
402413

403414
yield UiPathRuntimeResult(
404415
output={

‎demo/mock_template_runtime.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -114,13 +114,14 @@ async def stream(
114114
"""Stream events from the JSON file."""
115115
logger.info(f"MockTemplateRuntime: streaming {len(self._events)} events")
116116

117-
with self.tracer.start_as_current_span(
117+
root_span = self.tracer.start_span(
118118
"template.stream",
119119
attributes={
120120
"uipath.runtime.name": "MockTemplateRuntime",
121121
"uipath.event.count": len(self._events),
122122
},
123-
):
123+
)
124+
try:
124125
for i, event_data in enumerate(self._events):
125126
event_type = event_data.get("event_type")
126127

@@ -166,6 +167,8 @@ async def stream(
166167
except Exception as e:
167168
logger.error(f"Error processing event {i}: {e}", exc_info=True)
168169
continue
170+
finally:
171+
root_span.end()
169172

170173
logger.info("MockTemplateRuntime: streaming completed")
171174

0 commit comments

Comments
 (0)