Skip to content

fix(traced): fix parent for traced #596

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ionmincu merged 1 commit into from
Sep 18, 2025
Merged

fix(traced): fix parent for traced #596

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[project]
name = "uipath"
version = "2.1.54"
version = "2.1.55"
description = "Python SDK and CLI for UiPath Platform, enabling programmatic interaction with automation services, process management, and deployment tools."
readme = { file = "README.md", content-type = "text/markdown" }
requires-python = ">=3.10"
Expand Down
264 changes: 130 additions & 134 deletions src/uipath/tracing/_traced.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,25 +2,28 @@
import inspect
import json
import logging
from contextvars import ContextVar
from functools import wraps
from typing import Any, Callable, List, Optional, Tuple

from opentelemetry import trace
from opentelemetry import context, trace
from opentelemetry.trace import set_span_in_context

from ._utils import _SpanUtils

logger = logging.getLogger(__name__)

_tracer_instance: Optional[trace.Tracer] = None
# ContextVar to track the currently active span for nesting
_active_traced_span: ContextVar[Optional[trace.Span]] = ContextVar(
"_active_traced_span", default=None
)


def get_tracer() -> trace.Tracer:
"""Lazily initializes and returns the tracer instance."""
global _tracer_instance
if _tracer_instance is None:
logger.warning(
"Initializing tracer instance. This should only be done once per process."
)
_tracer_instance = trace.get_tracer(__name__)
return _tracer_instance

Expand Down Expand Up @@ -53,24 +56,24 @@ def register_current_span_provider(
"""
cls._current_span_provider = current_span_provider

@classmethod
def get_parent_context(cls):
"""Get the parent context using the registered current span provider.
@staticmethod
def get_parent_context():
# Always use the currently active OTel span if valid (recursion / children)
current_span = trace.get_current_span()
if current_span is not None and current_span.get_span_context().is_valid:
return set_span_in_context(current_span)

Returns:
Context object with the current span set, or None if no provider is registered.
"""
if cls._current_span_provider is not None:
# Only for the very top-level call, fallback to LangGraph span
if TracingManager._current_span_provider is not None:
try:
current_span = cls._current_span_provider()
if current_span is not None:
from opentelemetry.trace import set_span_in_context

return set_span_in_context(current_span)
external_span = TracingManager._current_span_provider()
if external_span is not None:
return set_span_in_context(external_span)
except Exception as e:
logger.warning(f"Error getting current span from provider: {e}")
return None
return None

# Last fallback
return context.get_current()

@classmethod
def register_traced_function(cls, original_func, decorated_func, params):
Expand Down Expand Up @@ -176,176 +179,169 @@ def _opentelemetry_traced(
"""Default tracer implementation using OpenTelemetry."""

def decorator(func):
trace_name = name if name is not None else func.__name__
trace_name = name or func.__name__

def get_parent_context():
"""Return a context object for starting the new span."""
current_span = _active_traced_span.get()
if current_span is not None and current_span.get_span_context().is_valid:
return set_span_in_context(current_span)

if TracingManager._current_span_provider is not None:
try:
external_span = TracingManager._current_span_provider()
if external_span is not None:
return set_span_in_context(external_span)
except Exception as e:
logger.warning(f"Error getting current span from provider: {e}")

return context.get_current()

# --------- Sync wrapper ---------
@wraps(func)
def sync_wrapper(*args, **kwargs):
context = TracingManager.get_parent_context()

with get_tracer().start_as_current_span(
trace_name, context=context
) as span:
default_span_type = "function_call_sync"
span.set_attribute(
"span_type",
span_type if span_type is not None else default_span_type,
)
ctx = get_parent_context()
span_cm = get_tracer().start_as_current_span(trace_name, context=ctx)
span = span_cm.__enter__()
token = _active_traced_span.set(span)
try:
span.set_attribute("span_type", span_type or "function_call_sync")
if run_type is not None:
span.set_attribute("run_type", run_type)

# Format arguments for tracing
inputs = _SpanUtils.format_args_for_trace_json(
inspect.signature(func), *args, **kwargs
)
# Apply input processor if provided
if input_processor is not None:
if input_processor:
processed_inputs = input_processor(json.loads(inputs))
inputs = json.dumps(processed_inputs, default=str)
span.set_attribute("inputs", inputs)
try:
result = func(*args, **kwargs)
# Process output if processor is provided
output = result
if output_processor is not None:
output = output_processor(result)
span.set_attribute("output", json.dumps(output, default=str))
return result
except Exception as e:
span.record_exception(e)
span.set_status(
trace.status.Status(trace.status.StatusCode.ERROR, str(e))
)
raise

result = func(*args, **kwargs)
output = output_processor(result) if output_processor else result
span.set_attribute("output", json.dumps(output, default=str))
return result
except Exception as e:
span.record_exception(e)
span.set_status(
trace.status.Status(trace.status.StatusCode.ERROR, str(e))
)
raise
finally:
_active_traced_span.reset(token)
span_cm.__exit__(None, None, None)

# --------- Async wrapper ---------
@wraps(func)
async def async_wrapper(*args, **kwargs):
context = TracingManager.get_parent_context()

with get_tracer().start_as_current_span(
trace_name, context=context
) as span:
default_span_type = "function_call_async"
span.set_attribute(
"span_type",
span_type if span_type is not None else default_span_type,
)
ctx = get_parent_context()
span_cm = get_tracer().start_as_current_span(trace_name, context=ctx)
span = span_cm.__enter__()
token = _active_traced_span.set(span)
try:
span.set_attribute("span_type", span_type or "function_call_async")
if run_type is not None:
span.set_attribute("run_type", run_type)

# Format arguments for tracing
inputs = _SpanUtils.format_args_for_trace_json(
inspect.signature(func), *args, **kwargs
)
# Apply input processor if provided
if input_processor is not None:
if input_processor:
processed_inputs = input_processor(json.loads(inputs))
inputs = json.dumps(processed_inputs, default=str)
span.set_attribute("inputs", inputs)
try:
result = await func(*args, **kwargs)
# Process output if processor is provided
output = result
if output_processor is not None:
output = output_processor(result)
span.set_attribute("output", json.dumps(output, default=str))
return result
except Exception as e:
span.record_exception(e)
span.set_status(
trace.status.Status(trace.status.StatusCode.ERROR, str(e))
)
raise

result = await func(*args, **kwargs)
output = output_processor(result) if output_processor else result
span.set_attribute("output", json.dumps(output, default=str))
return result
except Exception as e:
span.record_exception(e)
span.set_status(
trace.status.Status(trace.status.StatusCode.ERROR, str(e))
)
raise
finally:
_active_traced_span.reset(token)
span_cm.__exit__(None, None, None)

# --------- Generator wrapper ---------
@wraps(func)
def generator_wrapper(*args, **kwargs):
context = TracingManager.get_parent_context()

with get_tracer().start_as_current_span(
trace_name, context=context
) as span:
span.get_span_context()
default_span_type = "function_call_generator_sync"
ctx = get_parent_context()
span_cm = get_tracer().start_as_current_span(trace_name, context=ctx)
span = span_cm.__enter__()
token = _active_traced_span.set(span)
try:
span.set_attribute(
"span_type",
span_type if span_type is not None else default_span_type,
"span_type", span_type or "function_call_generator_sync"
)
if run_type is not None:
span.set_attribute("run_type", run_type)

# Format arguments for tracing
inputs = _SpanUtils.format_args_for_trace_json(
inspect.signature(func), *args, **kwargs
)
# Apply input processor if provided
if input_processor is not None:
if input_processor:
processed_inputs = input_processor(json.loads(inputs))
inputs = json.dumps(processed_inputs, default=str)
span.set_attribute("inputs", inputs)

outputs = []
try:
for item in func(*args, **kwargs):
outputs.append(item)
span.add_event(f"Yielded: {item}") # Add event for each yield
yield item

# Process output if processor is provided
output_to_record = outputs
if output_processor is not None:
output_to_record = output_processor(outputs)
span.set_attribute(
"output", json.dumps(output_to_record, default=str)
)
except Exception as e:
span.record_exception(e)
span.set_status(
trace.status.Status(trace.status.StatusCode.ERROR, str(e))
)
raise
for item in func(*args, **kwargs):
outputs.append(item)
span.add_event(f"Yielded: {item}")
yield item
output = output_processor(outputs) if output_processor else outputs
span.set_attribute("output", json.dumps(output, default=str))
except Exception as e:
span.record_exception(e)
span.set_status(
trace.status.Status(trace.status.StatusCode.ERROR, str(e))
)
raise
finally:
_active_traced_span.reset(token)
span_cm.__exit__(None, None, None)

# --------- Async generator wrapper ---------
@wraps(func)
async def async_generator_wrapper(*args, **kwargs):
context = TracingManager.get_parent_context()

with get_tracer().start_as_current_span(
trace_name, context=context
) as span:
default_span_type = "function_call_generator_async"
ctx = get_parent_context()
span_cm = get_tracer().start_as_current_span(trace_name, context=ctx)
span = span_cm.__enter__()
token = _active_traced_span.set(span)
try:
span.set_attribute(
"span_type",
span_type if span_type is not None else default_span_type,
"span_type", span_type or "function_call_generator_async"
)
if run_type is not None:
span.set_attribute("run_type", run_type)

# Format arguments for tracing
inputs = _SpanUtils.format_args_for_trace_json(
inspect.signature(func), *args, **kwargs
)
# Apply input processor if provided
if input_processor is not None:
if input_processor:
processed_inputs = input_processor(json.loads(inputs))
inputs = json.dumps(processed_inputs, default=str)
span.set_attribute("inputs", inputs)

outputs = []
try:
async for item in func(*args, **kwargs):
outputs.append(item)
span.add_event(f"Yielded: {item}") # Add event for each yield
yield item

# Process output if processor is provided
output_to_record = outputs
if output_processor is not None:
output_to_record = output_processor(outputs)
span.set_attribute(
"output", json.dumps(output_to_record, default=str)
)
except Exception as e:
span.record_exception(e)
span.set_status(
trace.status.Status(trace.status.StatusCode.ERROR, str(e))
)
raise
async for item in func(*args, **kwargs):
outputs.append(item)
span.add_event(f"Yielded: {item}")
yield item
output = output_processor(outputs) if output_processor else outputs
span.set_attribute("output", json.dumps(output, default=str))
except Exception as e:
span.record_exception(e)
span.set_status(
trace.status.Status(trace.status.StatusCode.ERROR, str(e))
)
raise
finally:
_active_traced_span.reset(token)
span_cm.__exit__(None, None, None)

if inspect.iscoroutinefunction(func):
return async_wrapper
Expand Down
Loading