Skip to content

Commit 84fbb52

Browse files
valentinabojanValentina Bojanclaude
authored
feat(guardrails): send execution source and job key headers (#1739)
Co-authored-by: Valentina Bojan <valentina.bojan@uipath.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
1 parent 8cfb822 commit 84fbb52

15 files changed

Lines changed: 306 additions & 41 deletions

File tree

‎packages/uipath-platform/pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[project]
22
name = "uipath-platform"
3-
version = "0.1.77"
3+
version = "0.1.78"
44
description = "HTTP client library for programmatic access to UiPath Platform"
55
readme = { file = "README.md", content-type = "text/markdown" }
66
requires-python = ">=3.11"

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
)
1717
from ._config import UiPathApiConfig, UiPathConfig
1818
from ._endpoints_manager import EndpointManager
19-
from ._execution_context import UiPathExecutionContext
19+
from ._execution_context import ExecutionSourceContext, UiPathExecutionContext
2020
from ._external_application_service import ExternalApplicationService
2121
from ._folder_context import FolderContext, header_folder
2222
from ._http_config import get_ca_bundle_path, get_httpx_client_kwargs
@@ -61,6 +61,7 @@
6161
"BaseService",
6262
"UiPathApiConfig",
6363
"UiPathExecutionContext",
64+
"ExecutionSourceContext",
6465
"ExternalApplicationService",
6566
"FolderContext",
6667
"TokenData",

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

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,34 @@
1+
from contextvars import ContextVar, Token
12
from os import environ as env
23

34
from uipath.platform.common.constants import ENV_JOB_ID, ENV_JOB_KEY, ENV_ROBOT_KEY
45

6+
_execution_source: ContextVar[str | None] = ContextVar("execution_source", default=None)
7+
8+
9+
class ExecutionSourceContext:
10+
"""Scope the execution source for the duration of a run.
11+
12+
Carries the source (e.g. ``runtime``/``playground``/``eval``) via a context
13+
variable and releases it on exit so it stays correctly scoped in concurrent
14+
runs. The CLI enters this with ``UiPathRuntimeContext.execution_source`` so
15+
platform clients can read it via
16+
:attr:`UiPathExecutionContext.execution_source`.
17+
"""
18+
19+
def __init__(self, execution_source: str | None) -> None:
20+
self._execution_source = execution_source
21+
self._token: Token[str | None] | None = None
22+
23+
def __enter__(self) -> "ExecutionSourceContext":
24+
self._token = _execution_source.set(self._execution_source)
25+
return self
26+
27+
def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
28+
if self._token is not None:
29+
_execution_source.reset(self._token)
30+
self._token = None
31+
532

633
class UiPathExecutionContext:
734
"""Manages the execution context for UiPath automation processes.
@@ -76,3 +103,13 @@ def robot_key(self) -> str | None:
76103
raise ValueError(f"Robot key is not set ({ENV_ROBOT_KEY})")
77104

78105
return self._robot_key
106+
107+
@property
108+
def execution_source(self) -> str | None:
109+
"""Get the execution source for the current run.
110+
111+
Identifies the run context (e.g. ``runtime``/``playground``/``eval``),
112+
derived from the CLI command and carried via
113+
:class:`ExecutionSourceContext`. Returns ``None`` when not set.
114+
"""
115+
return _execution_source.get()

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
HEADER_PROCESS_KEY = "x-uipath-processkey"
4040
HEADER_TRACE_ID = "x-uipath-traceid"
4141
HEADER_AGENTHUB_CONFIG = "x-uipath-agenthub-config"
42+
HEADER_GUARDRAILS_SOURCE = "x-uipath-guardrails-source"
4243
HEADER_LLMGATEWAY_BYO_CONNECTION_ID = "x-uipath-llmgateway-byoisconnectionid"
4344
HEADER_SW_LOCK_KEY = "x-uipath-sw-lockkey"
4445
HEADER_LICENSING_CONTEXT = "x-uipath-licensing-context"

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

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,9 @@
1212
from ..common._base_service import BaseService
1313
from ..common._config import UiPathApiConfig
1414
from ..common._execution_context import UiPathExecutionContext
15+
from ..common._job_context import header_job_key
1516
from ..common._models import Endpoint, RequestSpec
17+
from ..common.constants import HEADER_GUARDRAILS_SOURCE
1618
from ..errors import EnrichedException
1719
from .guardrails import BuiltInValidatorGuardrail
1820

@@ -123,9 +125,21 @@ def evaluate_guardrail(
123125
endpoint=Endpoint("/agentsruntime_/api/execution/guardrails/validate"),
124126
json=payload,
125127
)
126-
# Include trace context headers for server-side span correlation
128+
# Include trace context headers for server-side span correlation, plus
129+
# the execution source (x-uipath-guardrails-source) and job key headers
130+
# for licensing/metering correlation. The execution source is read from
131+
# the execution context, propagated from the runtime context.
127132
trace_headers = build_trace_context_headers()
128-
request_headers = {**(spec.headers or {}), **trace_headers}
133+
source_headers: dict[str, str] = {}
134+
execution_source = self._execution_context.execution_source
135+
if execution_source:
136+
source_headers[HEADER_GUARDRAILS_SOURCE] = execution_source
137+
request_headers = {
138+
**(spec.headers or {}),
139+
**trace_headers,
140+
**source_headers,
141+
**header_job_key(),
142+
}
129143
span_id = None
130144
try:
131145
response = self.request(

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
from uipath.platform.common import ExecutionSourceContext, UiPathExecutionContext
2+
3+
4+
def test_execution_source_none_by_default() -> None:
5+
assert UiPathExecutionContext().execution_source is None
6+
7+
8+
def test_execution_source_set_within_context() -> None:
9+
ctx = UiPathExecutionContext()
10+
11+
with ExecutionSourceContext("runtime"):
12+
assert ctx.execution_source == "runtime"
13+
14+
assert ctx.execution_source is None
15+
16+
17+
def test_execution_source_context_restores_previous_value() -> None:
18+
ctx = UiPathExecutionContext()
19+
20+
with ExecutionSourceContext("eval"):
21+
assert ctx.execution_source == "eval"
22+
with ExecutionSourceContext("playground"):
23+
assert ctx.execution_source == "playground"
24+
assert ctx.execution_source == "eval"
25+
26+
assert ctx.execution_source is None

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

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
)
1111

1212
from uipath.platform import UiPathApiConfig, UiPathExecutionContext
13+
from uipath.platform.common import ExecutionSourceContext
1314
from uipath.platform.guardrails import (
1415
BuiltInValidatorGuardrail,
1516
EnumListParameterValue,
@@ -356,6 +357,103 @@ def capture_request(request):
356357
# header merging works even when no active span exists)
357358
assert "content-type" in headers
358359

360+
def test_evaluate_guardrail_sends_source_and_job_key_headers(
361+
self,
362+
httpx_mock: HTTPXMock,
363+
service: GuardrailsService,
364+
base_url: str,
365+
org: str,
366+
tenant: str,
367+
monkeypatch: pytest.MonkeyPatch,
368+
) -> None:
369+
"""Outgoing request includes execution source and job key headers."""
370+
monkeypatch.setenv("UIPATH_JOB_KEY", "job-123")
371+
372+
captured_request = None
373+
374+
def capture_request(request):
375+
nonlocal captured_request
376+
captured_request = request
377+
return httpx.Response(
378+
status_code=200,
379+
json={"result": "PASSED", "details": "OK"},
380+
)
381+
382+
httpx_mock.add_callback(
383+
method="POST",
384+
url=f"{base_url}{org}{tenant}/agentsruntime_/api/execution/guardrails/validate",
385+
callback=capture_request,
386+
)
387+
388+
pii_guardrail = BuiltInValidatorGuardrail(
389+
id="test-id",
390+
name="PII guardrail",
391+
description="Test",
392+
enabled_for_evals=True,
393+
selector=GuardrailSelector(
394+
scopes=[GuardrailScope.TOOL], match_names=["tool1"]
395+
),
396+
guardrail_type="builtInValidator",
397+
validator_type="pii_detection",
398+
validator_parameters=[],
399+
)
400+
401+
with ExecutionSourceContext("runtime"):
402+
service.evaluate_guardrail("test input", pii_guardrail)
403+
404+
assert captured_request is not None
405+
headers = dict(captured_request.headers)
406+
assert headers.get("x-uipath-guardrails-source") == "runtime"
407+
assert headers.get("x-uipath-jobkey") == "job-123"
408+
409+
def test_evaluate_guardrail_omits_source_and_job_key_when_unset(
410+
self,
411+
httpx_mock: HTTPXMock,
412+
service: GuardrailsService,
413+
base_url: str,
414+
org: str,
415+
tenant: str,
416+
monkeypatch: pytest.MonkeyPatch,
417+
) -> None:
418+
"""Source/job key headers are absent when unset."""
419+
monkeypatch.delenv("UIPATH_JOB_KEY", raising=False)
420+
421+
captured_request = None
422+
423+
def capture_request(request):
424+
nonlocal captured_request
425+
captured_request = request
426+
return httpx.Response(
427+
status_code=200,
428+
json={"result": "PASSED", "details": "OK"},
429+
)
430+
431+
httpx_mock.add_callback(
432+
method="POST",
433+
url=f"{base_url}{org}{tenant}/agentsruntime_/api/execution/guardrails/validate",
434+
callback=capture_request,
435+
)
436+
437+
pii_guardrail = BuiltInValidatorGuardrail(
438+
id="test-id",
439+
name="PII guardrail",
440+
description="Test",
441+
enabled_for_evals=True,
442+
selector=GuardrailSelector(
443+
scopes=[GuardrailScope.TOOL], match_names=["tool1"]
444+
),
445+
guardrail_type="builtInValidator",
446+
validator_type="pii_detection",
447+
validator_parameters=[],
448+
)
449+
450+
service.evaluate_guardrail("test input", pii_guardrail)
451+
452+
assert captured_request is not None
453+
headers = dict(captured_request.headers)
454+
assert "x-uipath-guardrails-source" not in headers
455+
assert "x-uipath-jobkey" not in headers
456+
359457
def test_evaluate_guardrail_extracts_span_id_from_traceparent(
360458
self,
361459
httpx_mock: HTTPXMock,

‎packages/uipath-platform/uv.lock‎

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/uipath/pyproject.toml‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
11
[project]
22
name = "uipath"
3-
version = "2.11.12"
3+
version = "2.11.13"
44
description = "Python SDK and CLI for UiPath Platform, enabling programmatic interaction with automation services, process management, and deployment tools."
55
readme = { file = "README.md", content-type = "text/markdown" }
66
requires-python = ">=3.11"
77
dependencies = [
88
"uipath-core>=0.5.21, <0.6.0",
9-
"uipath-runtime>=0.11.0, <0.12.0",
10-
"uipath-platform>=0.1.76, <0.2.0",
9+
"uipath-runtime>=0.11.4, <0.12.0",
10+
"uipath-platform>=0.1.78, <0.2.0",
1111
"click>=8.3.1",
1212
"httpx>=0.28.1",
1313
"pyjwt>=2.10.1",

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

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,11 @@
1111
from uipath.core.tracing import UiPathTraceManager
1212
from uipath.eval.mocks import UiPathMockRuntime
1313
from uipath.eval.mocks._mock_runtime import load_simulation_config
14-
from uipath.platform.common import ResourceOverwritesContext, UiPathConfig
14+
from uipath.platform.common import (
15+
ExecutionSourceContext,
16+
ResourceOverwritesContext,
17+
UiPathConfig,
18+
)
1519
from uipath.runtime import (
1620
UiPathExecuteOptions,
1721
UiPathRuntimeContext,
@@ -122,14 +126,15 @@ def debug(
122126
async def execute_debug_runtime():
123127
trace_manager = UiPathTraceManager()
124128

125-
with UiPathRuntimeContext.with_defaults(
129+
ctx = UiPathRuntimeContext.with_defaults(
126130
input=input,
127131
input_file=input_file,
128132
output_file=output_file,
129133
resume=resume,
130134
trace_manager=trace_manager,
131135
command="debug",
132-
) as ctx:
136+
)
137+
with ExecutionSourceContext(ctx.execution_source), ctx:
133138
factory: UiPathRuntimeFactoryProtocol | None = None
134139

135140
try:

0 commit comments

Comments
 (0)