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"]