feat(guardrails): send execution source and job key headers

packages/uipath-platform/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-platform"
-version = "0.1.77"
+version = "0.1.78"
 description = "HTTP client library for programmatic access to UiPath Platform"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

packages/uipath-platform/src/uipath/platform/common/__init__.py

@@ -16,7 +16,7 @@
 )
 from ._config import UiPathApiConfig, UiPathConfig
 from ._endpoints_manager import EndpointManager
-from ._execution_context import UiPathExecutionContext
+from ._execution_context import ExecutionSourceContext, UiPathExecutionContext
 from ._external_application_service import ExternalApplicationService
 from ._folder_context import FolderContext, header_folder
 from ._http_config import get_ca_bundle_path, get_httpx_client_kwargs
@@ -61,6 +61,7 @@
     "BaseService",
     "UiPathApiConfig",
     "UiPathExecutionContext",
+    "ExecutionSourceContext",
     "ExternalApplicationService",
     "FolderContext",
     "TokenData",

packages/uipath-platform/src/uipath/platform/common/_execution_context.py

@@ -1,7 +1,34 @@
+from contextvars import ContextVar, Token
 from os import environ as env
 
 from uipath.platform.common.constants import ENV_JOB_ID, ENV_JOB_KEY, ENV_ROBOT_KEY
 
+_execution_source: ContextVar[str | None] = ContextVar("execution_source", default=None)
+
+
+class ExecutionSourceContext:
+    """Scope the execution source for the duration of a run.
+
+    Carries the source (e.g. ``runtime``/``playground``/``eval``) via a context
+    variable and releases it on exit so it stays correctly scoped in concurrent
+    runs. The CLI enters this with ``UiPathRuntimeContext.execution_source`` so
+    platform clients can read it via
+    :attr:`UiPathExecutionContext.execution_source`.
+    """
+
+    def __init__(self, execution_source: str | None) -> None:
+        self._execution_source = execution_source
+        self._token: Token[str | None] | None = None
+
+    def __enter__(self) -> "ExecutionSourceContext":
+        self._token = _execution_source.set(self._execution_source)
+        return self
+
+    def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
+        if self._token is not None:
+            _execution_source.reset(self._token)
+            self._token = None
+
 
 class UiPathExecutionContext:
     """Manages the execution context for UiPath automation processes.
@@ -76,3 +103,13 @@ def robot_key(self) -> str | None:
             raise ValueError(f"Robot key is not set ({ENV_ROBOT_KEY})")
 
         return self._robot_key
+
+    @property
+    def execution_source(self) -> str | None:
+        """Get the execution source for the current run.
+
+        Identifies the run context (e.g. ``runtime``/``playground``/``eval``),
+        derived from the CLI command and carried via
+        :class:`ExecutionSourceContext`. Returns ``None`` when not set.
+        """
+        return _execution_source.get()

packages/uipath-platform/src/uipath/platform/common/constants.py

@@ -39,6 +39,7 @@
 HEADER_PROCESS_KEY = "x-uipath-processkey"
 HEADER_TRACE_ID = "x-uipath-traceid"
 HEADER_AGENTHUB_CONFIG = "x-uipath-agenthub-config"
+HEADER_GUARDRAILS_SOURCE = "x-uipath-guardrails-source"
 HEADER_LLMGATEWAY_BYO_CONNECTION_ID = "x-uipath-llmgateway-byoisconnectionid"
 HEADER_SW_LOCK_KEY = "x-uipath-sw-lockkey"
 HEADER_LICENSING_CONTEXT = "x-uipath-licensing-context"

packages/uipath-platform/src/uipath/platform/guardrails/_guardrails_service.py

@@ -12,7 +12,9 @@
 from ..common._base_service import BaseService
 from ..common._config import UiPathApiConfig
 from ..common._execution_context import UiPathExecutionContext
+from ..common._job_context import header_job_key
 from ..common._models import Endpoint, RequestSpec
+from ..common.constants import HEADER_GUARDRAILS_SOURCE
 from ..errors import EnrichedException
 from .guardrails import BuiltInValidatorGuardrail
 
@@ -123,9 +125,21 @@ def evaluate_guardrail(
             endpoint=Endpoint("/agentsruntime_/api/execution/guardrails/validate"),
             json=payload,
         )
-        # Include trace context headers for server-side span correlation
+        # Include trace context headers for server-side span correlation, plus
+        # the execution source (x-uipath-guardrails-source) and job key headers
+        # for licensing/metering correlation. The execution source is read from
+        # the execution context, propagated from the runtime context.
         trace_headers = build_trace_context_headers()
-        request_headers = {**(spec.headers or {}), **trace_headers}
+        source_headers: dict[str, str] = {}
+        execution_source = self._execution_context.execution_source
+        if execution_source:
+            source_headers[HEADER_GUARDRAILS_SOURCE] = execution_source
+        request_headers = {
+            **(spec.headers or {}),
+            **trace_headers,
+            **source_headers,
+            **header_job_key(),
+        }
         span_id = None
         try:
             response = self.request(

packages/uipath-platform/tests/common/test_execution_context.py

@@ -0,0 +1,26 @@
+from uipath.platform.common import ExecutionSourceContext, UiPathExecutionContext
+
+
+def test_execution_source_none_by_default() -> None:
+    assert UiPathExecutionContext().execution_source is None
+
+
+def test_execution_source_set_within_context() -> None:
+    ctx = UiPathExecutionContext()
+
+    with ExecutionSourceContext("runtime"):
+        assert ctx.execution_source == "runtime"
+
+    assert ctx.execution_source is None
+
+
+def test_execution_source_context_restores_previous_value() -> None:
+    ctx = UiPathExecutionContext()
+
+    with ExecutionSourceContext("eval"):
+        assert ctx.execution_source == "eval"
+        with ExecutionSourceContext("playground"):
+            assert ctx.execution_source == "playground"
+        assert ctx.execution_source == "eval"
+
+    assert ctx.execution_source is None

packages/uipath-platform/tests/services/test_guardrails_service.py

@@ -10,6 +10,7 @@
 )
 
 from uipath.platform import UiPathApiConfig, UiPathExecutionContext
+from uipath.platform.common import ExecutionSourceContext
 from uipath.platform.guardrails import (
     BuiltInValidatorGuardrail,
     EnumListParameterValue,
@@ -356,6 +357,103 @@ def capture_request(request):
             # header merging works even when no active span exists)
             assert "content-type" in headers
 
+        def test_evaluate_guardrail_sends_source_and_job_key_headers(
+            self,
+            httpx_mock: HTTPXMock,
+            service: GuardrailsService,
+            base_url: str,
+            org: str,
+            tenant: str,
+            monkeypatch: pytest.MonkeyPatch,
+        ) -> None:
+            """Outgoing request includes execution source and job key headers."""
+            monkeypatch.setenv("UIPATH_JOB_KEY", "job-123")
+
+            captured_request = None
+
+            def capture_request(request):
+                nonlocal captured_request
+                captured_request = request
+                return httpx.Response(
+                    status_code=200,
+                    json={"result": "PASSED", "details": "OK"},
+                )
+
+            httpx_mock.add_callback(
+                method="POST",
+                url=f"{base_url}{org}{tenant}/agentsruntime_/api/execution/guardrails/validate",
+                callback=capture_request,
+            )
+
+            pii_guardrail = BuiltInValidatorGuardrail(
+                id="test-id",
+                name="PII guardrail",
+                description="Test",
+                enabled_for_evals=True,
+                selector=GuardrailSelector(
+                    scopes=[GuardrailScope.TOOL], match_names=["tool1"]
+                ),
+                guardrail_type="builtInValidator",
+                validator_type="pii_detection",
+                validator_parameters=[],
+            )
+
+            with ExecutionSourceContext("runtime"):
+                service.evaluate_guardrail("test input", pii_guardrail)
+
+            assert captured_request is not None
+            headers = dict(captured_request.headers)
+            assert headers.get("x-uipath-guardrails-source") == "runtime"
+            assert headers.get("x-uipath-jobkey") == "job-123"
+
+        def test_evaluate_guardrail_omits_source_and_job_key_when_unset(
+            self,
+            httpx_mock: HTTPXMock,
+            service: GuardrailsService,
+            base_url: str,
+            org: str,
+            tenant: str,
+            monkeypatch: pytest.MonkeyPatch,
+        ) -> None:
+            """Source/job key headers are absent when unset."""
+            monkeypatch.delenv("UIPATH_JOB_KEY", raising=False)
+
+            captured_request = None
+
+            def capture_request(request):
+                nonlocal captured_request
+                captured_request = request
+                return httpx.Response(
+                    status_code=200,
+                    json={"result": "PASSED", "details": "OK"},
+                )
+
+            httpx_mock.add_callback(
+                method="POST",
+                url=f"{base_url}{org}{tenant}/agentsruntime_/api/execution/guardrails/validate",
+                callback=capture_request,
+            )
+
+            pii_guardrail = BuiltInValidatorGuardrail(
+                id="test-id",
+                name="PII guardrail",
+                description="Test",
+                enabled_for_evals=True,
+                selector=GuardrailSelector(
+                    scopes=[GuardrailScope.TOOL], match_names=["tool1"]
+                ),
+                guardrail_type="builtInValidator",
+                validator_type="pii_detection",
+                validator_parameters=[],
+            )
+
+            service.evaluate_guardrail("test input", pii_guardrail)
+
+            assert captured_request is not None
+            headers = dict(captured_request.headers)
+            assert "x-uipath-guardrails-source" not in headers
+            assert "x-uipath-jobkey" not in headers
+
         def test_evaluate_guardrail_extracts_span_id_from_traceparent(
             self,
             httpx_mock: HTTPXMock,

packages/uipath/pyproject.toml

@@ -1,13 +1,13 @@
 [project]
 name = "uipath"
-version = "2.11.12"
+version = "2.11.13"
 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.11"
 dependencies = [
   "uipath-core>=0.5.21, <0.6.0",
-  "uipath-runtime>=0.11.0, <0.12.0",
-  "uipath-platform>=0.1.76, <0.2.0",
+  "uipath-runtime>=0.11.4, <0.12.0",
+  "uipath-platform>=0.1.78, <0.2.0",
   "click>=8.3.1",
   "httpx>=0.28.1",
   "pyjwt>=2.10.1",

packages/uipath/src/uipath/_cli/cli_debug.py

@@ -11,7 +11,11 @@
 from uipath.core.tracing import UiPathTraceManager
 from uipath.eval.mocks import UiPathMockRuntime
 from uipath.eval.mocks._mock_runtime import load_simulation_config
-from uipath.platform.common import ResourceOverwritesContext, UiPathConfig
+from uipath.platform.common import (
+    ExecutionSourceContext,
+    ResourceOverwritesContext,
+    UiPathConfig,
+)
 from uipath.runtime import (
     UiPathExecuteOptions,
     UiPathRuntimeContext,
@@ -122,14 +126,15 @@ def debug(
             async def execute_debug_runtime():
                 trace_manager = UiPathTraceManager()
 
-                with UiPathRuntimeContext.with_defaults(
+                ctx = UiPathRuntimeContext.with_defaults(
                     input=input,
                     input_file=input_file,
                     output_file=output_file,
                     resume=resume,
                     trace_manager=trace_manager,
                     command="debug",
-                ) as ctx:
+                )
+                with ExecutionSourceContext(ctx.execution_source), ctx:
                     factory: UiPathRuntimeFactoryProtocol | None = None
 
                     try:

packages/uipath/src/uipath/_cli/cli_dev.py

@@ -7,13 +7,20 @@
 from uipath._cli._utils._debug import setup_debugging
 from uipath._cli.middlewares import Middlewares
 from uipath.core.tracing import UiPathTraceManager
+from uipath.platform.common import ExecutionSourceContext
 from uipath.runtime import UiPathRuntimeContext, UiPathRuntimeFactoryRegistry
 
 from ._telemetry import track_command
 
 console = ConsoleLogger()
 
 
+def _create_dev_context_and_factory(trace_manager: UiPathTraceManager):
+    """Build the dev runtime context and its factory."""
+    context = UiPathRuntimeContext(trace_manager=trace_manager, command="dev")
+    return context, UiPathRuntimeFactoryRegistry.get(context=context)
+
+
 def _check_dev_dependency(interface: str) -> None:
     """Check if uipath-dev is installed and raise helpful error if not."""
     import importlib.util
@@ -80,17 +87,14 @@ async def run_terminal() -> None:
             factory = None
             try:
                 trace_manager = UiPathTraceManager()
-                factory = UiPathRuntimeFactoryRegistry.get(
-                    context=UiPathRuntimeContext(
-                        trace_manager=trace_manager, command="dev"
-                    )
-                )
+                context, factory = _create_dev_context_and_factory(trace_manager)
 
                 app = UiPathDeveloperConsole(
                     runtime_factory=factory, trace_manager=trace_manager
                 )
 
-                await app.run_async()
+                with ExecutionSourceContext(context.execution_source):
+                    await app.run_async()
 
             except KeyboardInterrupt:
                 console.info("Debug session interrupted by user")
@@ -124,11 +128,7 @@ def signal_handler(sig, frame):
 
             try:
                 trace_manager = UiPathTraceManager()
-                factory = UiPathRuntimeFactoryRegistry.get(
-                    context=UiPathRuntimeContext(
-                        trace_manager=trace_manager, command="dev"
-                    )
-                )
+                context, factory = _create_dev_context_and_factory(trace_manager)
 
                 app = UiPathDeveloperServer(
                     runtime_factory=factory,
@@ -140,13 +140,17 @@ def signal_handler(sig, frame):
                     ),
                 )
 
-                server_task = asyncio.create_task(app.run_async())
-                shutdown_task = asyncio.create_task(shutdown_event.wait())
+                # Enter the execution source context before creating the server
+                # task so request tasks spawned during the run inherit it.
+                with ExecutionSourceContext(context.execution_source):
+                    server_task = asyncio.create_task(app.run_async())
+                    shutdown_task = asyncio.create_task(shutdown_event.wait())
 
-                # Wait for either server to complete or shutdown signal
-                done, pending = await asyncio.wait(
-                    {server_task, shutdown_task}, return_when=asyncio.FIRST_COMPLETED
-                )
+                    # Wait for either server to complete or shutdown signal
+                    done, pending = await asyncio.wait(
+                        {server_task, shutdown_task},
+                        return_when=asyncio.FIRST_COMPLETED,
+                    )
 
                 for task in pending:
                     task.cancel()