From cd11a25b77163078ac8e3db6f37d31a715d739f9 Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Wed, 19 Aug 2026 22:10:10 +0530 Subject: [PATCH] feat(governance): stamp trace_id on governance customEvents Resolve the run's trace id (UIPATH_TRACE_ID, else the live OTel span) as a dashed GUID on the hook thread and stamp it on governance.rule.denied and governance.hook.summary, so each event carries the id its LLMOps spans are stored under. Resolved synchronously before the background telemetry dispatch so it survives the trace-context header being dropped on that thread. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../runtime/governance/_audit/track_events.py | 23 ++++++++ tests/test_track_events_sink.py | 57 +++++++++++++++++++ 2 files changed, 80 insertions(+) diff --git a/src/uipath/runtime/governance/_audit/track_events.py b/src/uipath/runtime/governance/_audit/track_events.py index 8a2b745a..06a6db5d 100644 --- a/src/uipath/runtime/governance/_audit/track_events.py +++ b/src/uipath/runtime/governance/_audit/track_events.py @@ -22,6 +22,7 @@ from __future__ import annotations import logging +import os from typing import Any, Callable from uipath.core.governance import EnforcementMode @@ -77,6 +78,25 @@ def _resolve_operation_id() -> str | None: return f"{span_ctx.trace_id:032x}" +_UIPATH_TRACE_ID_ENV = "UIPATH_TRACE_ID" + + +def _resolve_trace_id() -> str | None: + """The run's trace id as a dashed GUID, matching the span store's TraceId. + + Prefers ``UIPATH_TRACE_ID`` (what the exporter binds spans to), else the + live OTel span. Resolved on the hook thread so it survives the background + telemetry dispatch. ``None`` when neither yields a valid 32-hex id. + """ + for raw in (os.environ.get(_UIPATH_TRACE_ID_ENV), _resolve_operation_id()): + if not raw: + continue + hex_id = raw.replace("-", "").lower() + if len(hex_id) == 32 and all(c in "0123456789abcdef" for c in hex_id): + return f"{hex_id[0:8]}-{hex_id[8:12]}-{hex_id[12:16]}-{hex_id[16:20]}-{hex_id[20:32]}" + return None + + def _mode_str(mode: Any) -> str: """Coerce an enforcement-mode field to its canonical uppercase string.""" if isinstance(mode, EnforcementMode): @@ -272,6 +292,9 @@ def _common_payload(self, event: AuditEvent) -> dict[str, Any]: payload["agent_name"] = event.agent_name payload["hook"] = event.hook payload["timestamp"] = event.timestamp.isoformat() + trace_id = _resolve_trace_id() + if trace_id is not None: + payload["trace_id"] = trace_id return payload def _emit_rule_denied(self, event: AuditEvent) -> None: diff --git a/tests/test_track_events_sink.py b/tests/test_track_events_sink.py index ee38f311..b739180b 100644 --- a/tests/test_track_events_sink.py +++ b/tests/test_track_events_sink.py @@ -414,3 +414,60 @@ def test_operation_id_none_when_no_active_span( """No live OTel span → ``operation_id=None`` (consumer fills in its own).""" sink.emit(_rule_event(matched=True)) assert capture.calls[0]["operation_id"] is None + + +# --------------------------------------------------------------------------- +# trace_id payload field — the id the LLMOps span store persists/fetches by +# (UIPATH_TRACE_ID when set, else the live OTel span), as a dashed GUID +# --------------------------------------------------------------------------- + + +def test_rule_denied_trace_id_prefers_uipath_trace_id( + sink: TrackEventAuditSink, capture: _Capture, monkeypatch: pytest.MonkeyPatch +) -> None: + """``data.trace_id`` = UIPATH_TRACE_ID (dashed) — the exact id the span + exporter binds spans to, so the event and its spans share one id. + """ + monkeypatch.setenv("UIPATH_TRACE_ID", "406d33561ff4450382a877a4fba3f838") + sink.emit(_rule_event(matched=True)) + assert ( + capture.calls[0]["data"]["trace_id"] == "406d3356-1ff4-4503-82a8-77a4fba3f838" + ) + + +def test_rule_denied_trace_id_falls_back_to_live_span( + sink: TrackEventAuditSink, capture: _Capture, monkeypatch: pytest.MonkeyPatch +) -> None: + """No UIPATH_TRACE_ID → the live OTel span's trace id, dashed to a GUID.""" + from opentelemetry.sdk.trace import TracerProvider + + monkeypatch.delenv("UIPATH_TRACE_ID", raising=False) + tracer = TracerProvider().get_tracer("test") + with tracer.start_as_current_span("agent-run") as span: + hex_id = f"{span.get_span_context().trace_id:032x}" + sink.emit(_rule_event(matched=True)) + + expected = ( + f"{hex_id[0:8]}-{hex_id[8:12]}-{hex_id[12:16]}-{hex_id[16:20]}-{hex_id[20:32]}" + ) + assert capture.calls[0]["data"]["trace_id"] == expected + + +def test_hook_summary_payload_carries_trace_id( + sink: TrackEventAuditSink, capture: _Capture, monkeypatch: pytest.MonkeyPatch +) -> None: + """The trace id is stamped on hook summaries too (shared ``_common_payload``).""" + monkeypatch.setenv("UIPATH_TRACE_ID", "406d33561ff4450382a877a4fba3f838") + sink.emit(_hook_event()) + assert ( + capture.calls[0]["data"]["trace_id"] == "406d3356-1ff4-4503-82a8-77a4fba3f838" + ) + + +def test_trace_id_absent_when_no_env_and_no_active_span( + sink: TrackEventAuditSink, capture: _Capture, monkeypatch: pytest.MonkeyPatch +) -> None: + """No UIPATH_TRACE_ID and no live span → ``trace_id`` omitted (never null).""" + monkeypatch.delenv("UIPATH_TRACE_ID", raising=False) + sink.emit(_rule_event(matched=True)) + assert "trace_id" not in capture.calls[0]["data"]