diff --git a/src/uipath/runtime/governance/runtime.py b/src/uipath/runtime/governance/runtime.py index c9cf6996..78973b68 100644 --- a/src/uipath/runtime/governance/runtime.py +++ b/src/uipath/runtime/governance/runtime.py @@ -59,6 +59,10 @@ UiPathStreamOptions, ) from uipath.runtime.events import UiPathRuntimeEvent +from uipath.runtime.governance._audit.traces import ( + SPAN_TYPE_GOVERNANCE, + UIPATH_SOURCE, +) from uipath.runtime.governance.native.evaluator import GovernanceEvaluator from uipath.runtime.governance.native.models import PolicyIndex from uipath.runtime.result import UiPathRuntimeResult @@ -106,6 +110,11 @@ def _governance_root_span(agent_name: str, runtime_id: str) -> Iterator[None]: # child (same trace_id). Otherwise we open a root span (new # trace_id). with tracer.start_as_current_span("uipath.governance.run") as span: + span.set_attribute("uipath.custom_instrumentation", True) + span.set_attribute("type", SPAN_TYPE_GOVERNANCE) + span.set_attribute("span_type", SPAN_TYPE_GOVERNANCE) + span.set_attribute("uipath.source", UIPATH_SOURCE) + # Span attributes for downstream consumers. ``agent_name`` # and ``runtime_id`` are the primary keys an operator # filters on; ``source`` identifies which producer emitted diff --git a/tests/test_governance_runtime.py b/tests/test_governance_runtime.py index b2ca3e3d..4dce0f99 100644 --- a/tests/test_governance_runtime.py +++ b/tests/test_governance_runtime.py @@ -10,7 +10,8 @@ from __future__ import annotations -from typing import Any +from contextlib import contextmanager +from typing import Any, Iterator import pytest from uipath.core.governance import EnforcementMode @@ -628,3 +629,86 @@ def _blocked_import(name: str, *args: Any, **kwargs: Any) -> Any: assert result == "result" assert delegate.execute_calls == [({"x": 1}, None)] + + +# --------------------------------------------------------------------------- +# _governance_root_span — export markers +# --------------------------------------------------------------------------- + + +@contextmanager +def _recording_tracer_provider() -> Iterator[Any]: + """Install an in-memory tracer provider globally and yield its exporter.""" + from opentelemetry import trace + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) + + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + + original = trace.get_tracer_provider() + trace._TRACER_PROVIDER = provider # type: ignore[attr-defined] + try: + yield exporter + finally: + trace._TRACER_PROVIDER = original # type: ignore[attr-defined] + + +def _governance_run_span(exporter: Any) -> Any: + """Return the single ``uipath.governance.run`` span the exporter saw.""" + spans = [ + s for s in exporter.get_finished_spans() if s.name == "uipath.governance.run" + ] + assert len(spans) == 1 + return spans[0] + + +async def test_root_span_is_marked_for_export() -> None: + """The span must survive host-side export filters.""" + from opentelemetry import trace + + runtime = UiPathGovernedRuntime( + _StubDelegate(), + PolicyIndex(), + EnforcementMode.AUDIT, + agent_name="HR Assistant", + runtime_id="rt-1", + ) + + with _recording_tracer_provider() as exporter: + host_tracer = trace.get_tracer("test.host") + with host_tracer.start_as_current_span("host exchange") as host_span: + host_span_id = host_span.get_span_context().span_id + assert await runtime.execute({"x": 1}) == "result" + + span = _governance_run_span(exporter) + + attributes = span.attributes or {} + + assert attributes["uipath.custom_instrumentation"] is True + assert attributes["type"] == "governance" + assert attributes["span_type"] == "governance" + assert attributes["uipath.source"] == "Governance" + assert attributes["uipath_governance.agent_name"] == "HR Assistant" + assert attributes["uipath_governance.runtime_id"] == "rt-1" + + assert span.parent is not None + assert span.parent.span_id == host_span_id + + +async def test_root_span_is_marked_for_export_when_streaming() -> None: + """``stream`` marks the span the same way ``execute`` does.""" + runtime = UiPathGovernedRuntime( + _StubDelegate(), PolicyIndex(), EnforcementMode.AUDIT + ) + + with _recording_tracer_provider() as exporter: + events = [event async for event in runtime.stream({"x": 1})] + span = _governance_run_span(exporter) + + assert events == ["a", "b"] + assert (span.attributes or {})["uipath.custom_instrumentation"] is True