From 84e2be7a7f133686d89b2bf28804812789603e8b Mon Sep 17 00:00:00 2001 From: Andrei Petraru Date: Sun, 20 Sep 2026 22:22:00 +0300 Subject: [PATCH 1/3] feat(guardrails): let tool-scope judge guardrails read the files a tool call mentions Tool-scope guardrails no longer skip attachments. Before the tool runs, the judge receives references for the attachments named in the call's arguments; after it runs, for those named in the result. Only ids the run's registry holds are forwarded, with the registry's name and type, so a call that mentions no file forwards none, and a mention of a file the run never held is skipped rather than sent to the backend for lookup. Agent and LLM scope keep reading the run's registry; both paths build references through one collector. AL-593 Co-Authored-By: Claude Fable 5.1 --- .../agent/guardrails/attachment_refs.py | 149 +++++++++++-- .../agent/guardrails/guardrail_nodes.py | 59 ++++- .../agent/guardrails/test_attachment_refs.py | 178 ++++++++++++++- .../agent/guardrails/test_guardrail_nodes.py | 204 ++++++++++++++++-- .../test_tool_guardrails_subgraph_run.py | 117 ++++++++++ 5 files changed, 662 insertions(+), 45 deletions(-) create mode 100644 tests/agent/guardrails/test_tool_guardrails_subgraph_run.py diff --git a/src/uipath_langchain/agent/guardrails/attachment_refs.py b/src/uipath_langchain/agent/guardrails/attachment_refs.py index eace068f1..a0353cab0 100644 --- a/src/uipath_langchain/agent/guardrails/attachment_refs.py +++ b/src/uipath_langchain/agent/guardrails/attachment_refs.py @@ -1,18 +1,30 @@ -"""Project a run's job-attachment registry into guardrail attachment references. +"""Project a run's attachments into guardrail attachment references. -Any built-in guardrail forwards the run's attachments unless it is scoped to prompts. The -runtime forwards id, file name and mime type only; the backend's feature flag decides -whether they are used at all, and the backend decides which validators and file types it -can inspect and resolves the id through Orchestrator. +Two sources feed the same reference shape: -Nothing in this module raises: the guardrail node re-raises any exception, which would end -the run over a single malformed attachment. +* Agent- and LLM-scope guardrails judge the conversation so far, so they read the whole + job-attachment registry (:func:`resolve_guardrail_attachments`). +* Tool-scope guardrails judge one tool call, so they read only the attachments that call + mentions: the ``{"ID": ...}`` objects in its arguments before the tool runs, or in its + result afterwards (:func:`resolve_referenced_attachments`). A tool call that names no + file forwards nothing, even when the run holds files elsewhere; otherwise every tool + call would ship every file to the backend. + +Any built-in guardrail forwards references unless it is scoped to prompts. The runtime +forwards id, file name and mime type only; the backend's feature flag decides whether +they are used at all, the backend decides which validators and file types it can +inspect, and it resolves the id through Orchestrator. + +Nothing in this module raises: the guardrail node re-raises any exception, which would +end the run over a single malformed attachment. """ import logging import uuid +from collections.abc import Iterable, Iterator, Mapping from typing import Any +from pydantic import BaseModel from uipath.platform.attachments import Attachment from uipath.platform.guardrails import BuiltInValidatorGuardrail, GuardrailAttachment @@ -24,6 +36,11 @@ #: ``appliesTo`` guardrail parameter; only ``Prompts`` excludes files (default is ``Both``). _APPLIES_TO_PARAMETER = "appliesto" _PROMPTS_ONLY = "prompts" +#: Wire key of a job attachment reference as the model and the tools exchange it. +_ID_KEY = "ID" +#: Bounds for scanning a tool payload, which can be arbitrarily large or deep. +_MAX_SCAN_DEPTH = 32 +_MAX_SCAN_NODES = 10_000 def _scope_includes_files(guardrail: BuiltInValidatorGuardrail) -> bool: @@ -44,8 +61,9 @@ async def resolve_guardrail_attachments( job_attachments: dict[str, Attachment], guardrail: BuiltInValidatorGuardrail, ) -> list[GuardrailAttachment]: - """Return up to five attachment references for the guardrail, or an empty list. + """Return up to five references for every attachment the run knows about. + For Agent- and LLM-scope guardrails, which evaluate the conversation as a whole. Empty when the guardrail is scoped to prompts or the run has no attachments. Never raises. """ @@ -57,14 +75,119 @@ async def resolve_guardrail_attachments( guardrail.name, ) return [] + return _collect(job_attachments.values()) + + +def resolve_referenced_attachments( + data: Any, + job_attachments: dict[str, Attachment] | None, + guardrail: BuiltInValidatorGuardrail, +) -> list[GuardrailAttachment]: + """Return up to five references for the attachments ``data`` mentions. + + For Tool-scope guardrails: ``data`` is the tool call's arguments (before the tool + runs) or its parsed result (after). A mention is a mapping with an ``ID`` that parses + as a UUID, the shape the model emits and the tool wrapper expands. Only ids the run's + registry holds are forwarded, with the registry's name and type: the registry is the + set of files this run legitimately has (agent input plus files its tools returned), + so a mention the run never held, or a non-attachment resource id, is skipped rather + than sent to the backend for lookup. Empty when nothing is mentioned or the guardrail + is scoped to prompts. Never raises. + """ + if data is None: + return [] + if not _scope_includes_files(guardrail): + logger.debug( + "Guardrail '%s' is scoped to prompts; skipping attachment resolution.", + guardrail.name, + ) + return [] + try: + mentions = list(_iter_attachment_mentions(data)) + except Exception: + logger.warning( + "Could not scan the tool payload for attachments; guardrail '%s' evaluates " + "without files.", + guardrail.name, + exc_info=True, + ) + return [] + registry = job_attachments or {} + resolved = (_resolve_mention(mention, registry) for mention in mentions) + return _collect(attachment for attachment in resolved if attachment is not None) + + +def _iter_attachment_mentions(data: Any) -> Iterator[Any]: + """Yield every attachment-shaped value in ``data``, depth first, within bounds.""" + stack: list[tuple[Any, int]] = [(data, 0)] + visited = 0 + while stack: + value, depth = stack.pop() + visited += 1 + if visited > _MAX_SCAN_NODES: + logger.debug( + "Stopped scanning the tool payload for attachments after %d values.", + _MAX_SCAN_NODES, + ) + return + if depth > _MAX_SCAN_DEPTH: + continue + if isinstance(value, Attachment): + yield value + elif isinstance(value, BaseModel): + stack.append((value.model_dump(by_alias=True), depth + 1)) + elif isinstance(value, Mapping): + if _is_mention(value): + yield value + else: + stack.extend((item, depth + 1) for item in value.values()) + elif isinstance(value, (list, tuple, set, frozenset)): + stack.extend((item, depth + 1) for item in value) + + +def _is_mention(value: Mapping[Any, Any]) -> bool: + raw_id = value.get(_ID_KEY) + if raw_id is None or isinstance(raw_id, bool): + return False + try: + uuid.UUID(str(raw_id)) + except ValueError: + return False + return True + + +def _resolve_mention(mention: Any, registry: dict[str, Attachment]) -> Any | None: + """Look one mention up in the run's registry; None when the run never held it.""" + try: + raw_id = mention.id if isinstance(mention, Attachment) else mention[_ID_KEY] + attachment_id = str(uuid.UUID(str(raw_id))) + known = registry.get(attachment_id) + if known is not None: + return known + logger.debug( + "Skipping attachment reference '%s': this run does not hold it.", + attachment_id, + ) + except Exception: + logger.debug( + "Skipping a malformed attachment reference in the tool payload.", + exc_info=True, + ) + return None + +def _collect(attachments: Iterable[Any]) -> list[GuardrailAttachment]: + """Build references, dropping unusable and duplicate ones, capped to the API limit.""" references: list[GuardrailAttachment] = [] - for attachment in job_attachments.values(): + seen: set[str] = set() + for attachment in attachments: reference = _to_reference(attachment) - if reference is not None: - references.append(reference) - if len(references) == _MAX_ATTACHMENTS: - break + if reference is None or reference.id in seen: + continue + seen.add(reference.id) + references.append(reference) + if len(references) == _MAX_ATTACHMENTS: + break return references diff --git a/src/uipath_langchain/agent/guardrails/guardrail_nodes.py b/src/uipath_langchain/agent/guardrails/guardrail_nodes.py index be96eb695..dd9e19698 100644 --- a/src/uipath_langchain/agent/guardrails/guardrail_nodes.py +++ b/src/uipath_langchain/agent/guardrails/guardrail_nodes.py @@ -23,6 +23,7 @@ from uipath_langchain.agent.guardrails.attachment_refs import ( resolve_guardrail_attachments, + resolve_referenced_attachments, ) from uipath_langchain.agent.guardrails.types import ExecutionStage from uipath_langchain.agent.guardrails.utils import ( @@ -37,9 +38,6 @@ logger = logging.getLogger(__name__) -#: Scopes whose guardrails may inspect attached files (tool scope excluded on purpose). -_ATTACHMENT_SCOPES = frozenset({GuardrailScope.AGENT, GuardrailScope.LLM}) - def _evaluate_deterministic_guardrail( state: AgentGuardrailsGraphState, @@ -168,6 +166,48 @@ def _create_validation_command( ) +async def _resolve_attachments( + state: AgentGuardrailsGraphState, + guardrail: BuiltInValidatorGuardrail, + scope: GuardrailScope, + execution_stage: ExecutionStage, + input_data_extractor: Callable[[AgentGuardrailsGraphState], dict[str, Any]] | None, + output_data_extractor: Callable[[AgentGuardrailsGraphState], dict[str, Any]] | None, +) -> list[GuardrailAttachment]: + """Attachment references for one built-in guardrail evaluation. + + Agent and LLM scope judge the conversation, so they read the run's whole attachment + registry. Tool scope judges one call, so it reads only the files that call's + arguments (pre-execution) or result (post-execution) mention: a tool call without a + file forwards nothing even when the run has files elsewhere. Never raises. + """ + registry = state.inner_state.job_attachments + if scope != GuardrailScope.TOOL: + return await resolve_guardrail_attachments(registry, guardrail) + + if execution_stage == ExecutionStage.PRE_EXECUTION: + extractor, source_name = input_data_extractor, "arguments" + else: + extractor, source_name = output_data_extractor, "result" + # A mention always carries the literal ``ID`` key; skip parsing plain-text + # results, which the output extractor would otherwise warn about. + if not state.messages or "ID" not in get_message_content(state.messages[-1]): + return [] + if extractor is None: + return [] + try: + source = extractor(state) + except Exception: + logger.warning( + "Could not read the tool %s for guardrail '%s'; evaluating without files.", + source_name, + guardrail.name, + exc_info=True, + ) + return [] + return resolve_referenced_attachments(source, registry, guardrail) + + def _create_guardrail_node( guardrail: BaseGuardrail, scope: GuardrailScope, @@ -235,12 +275,13 @@ async def node( else: metadata["payload"]["output"] = payload - attachments = ( - await resolve_guardrail_attachments( - state.inner_state.job_attachments, guardrail - ) - if scope in _ATTACHMENT_SCOPES - else [] + attachments = await _resolve_attachments( + state, + guardrail, + scope, + execution_stage, + input_data_extractor, + output_data_extractor, ) result = await _evaluate_builtin_guardrail( diff --git a/tests/agent/guardrails/test_attachment_refs.py b/tests/agent/guardrails/test_attachment_refs.py index 738c209d5..b8e67b976 100644 --- a/tests/agent/guardrails/test_attachment_refs.py +++ b/tests/agent/guardrails/test_attachment_refs.py @@ -1,10 +1,11 @@ -"""Tests for projecting the job-attachment registry into guardrail attachment refs.""" +"""Tests for projecting a run's attachments into guardrail attachment references.""" import uuid from typing import Any from unittest.mock import MagicMock import pytest +from pydantic import BaseModel from uipath.platform.attachments import Attachment from uipath.platform.guardrails import BuiltInValidatorGuardrail from uipath.platform.guardrails.guardrails import EnumParameterValue @@ -12,6 +13,7 @@ from uipath_langchain.agent.guardrails.attachment_refs import ( _MAX_ATTACHMENTS, resolve_guardrail_attachments, + resolve_referenced_attachments, ) _UUID = "7f2c1e44-0b3a-4a1e-9d55-2f9a1c3b8e10" @@ -178,3 +180,177 @@ async def test_resolves_when_the_scope_parameter_is_malformed(self, monkeypatch) result = await resolve_guardrail_attachments(_registry(), guardrail) assert [r.file_name for r in result] == ["a.csv"] + + +def _other_attachment(name: str = "b.pdf", mime: str = "application/pdf"): + attachment_id = str(uuid.uuid4()) + return attachment_id, Attachment( + id=uuid.UUID(attachment_id), full_name=name, mime_type=mime + ) + + +class TestResolveReferencedAttachments: + """Tool scope: only the attachments a tool call mentions are forwarded.""" + + async def test_resolves_the_attachment_the_tool_call_names(self): + """The model passes ``{"ID": ...}`` only; name and type come from the registry.""" + result = resolve_referenced_attachments( + {"attachment": {"ID": _UUID}, "question": "summarize"}, + _registry(), + _judge(), + ) + + assert [r.model_dump(by_alias=True) for r in result] == [ + {"id": _UUID, "fileName": "a.csv", "mimeType": "text/csv"} + ] + + async def test_forwards_nothing_when_the_call_names_no_file(self): + """The run holds a file, but this call does not touch it: forwarding it would + ship every file on every tool call.""" + result = resolve_referenced_attachments( + {"query": "cats", "limit": 5}, _registry(), _judge() + ) + + assert result == [] + + async def test_finds_mentions_nested_in_lists_and_dicts(self): + other_id, other = _other_attachment() + registry = _registry() | {other_id: other} + data = { + "files": [{"ID": _UUID}], + "options": {"inner": {"ref": {"ID": other_id}, "flag": True}}, + } + + result = resolve_referenced_attachments(data, registry, _judge()) + + assert sorted(r.id for r in result) == sorted([_UUID, other_id]) + + async def test_skips_an_id_the_run_never_held_even_with_inline_name_and_type( + self, + ): + """A tool result (or a prompt-injected argument) can name any id; only files the + run legitimately holds are sent to the backend, which would otherwise look the + id up in Orchestrator and judge a file the run never had.""" + other_id = str(uuid.uuid4()) + data = {"file": {"ID": other_id, "FullName": "out.csv", "MimeType": "text/csv"}} + + assert resolve_referenced_attachments(data, _registry(), _judge()) == [] + + async def test_uses_the_registry_name_and_type_not_the_inline_ones(self): + data = {"file": {"ID": _UUID, "FullName": "renamed.csv", "MimeType": "x/y"}} + + result = resolve_referenced_attachments(data, _registry(), _judge()) + + assert [(r.file_name, r.mime_type) for r in result] == [("a.csv", "text/csv")] + + async def test_skips_a_uuid_that_is_not_one_of_the_run_attachments(self): + """A UUID under ``ID`` that is not an attachment (a queue item, a job) must not + reach the backend, which would look it up in Orchestrator.""" + data = {"item": {"ID": str(uuid.uuid4())}} + + assert resolve_referenced_attachments(data, _registry(), _judge()) == [] + + @pytest.mark.parametrize("raw_id", ["queue-item-42", 12, True, None, ""]) + async def test_skips_a_non_uuid_id(self, raw_id): + data = {"item": {"ID": raw_id, "FullName": "a.csv", "MimeType": "text/csv"}} + + assert resolve_referenced_attachments(data, _registry(), _judge()) == [] + + async def test_deduplicates_the_same_attachment(self): + data = {"a": {"ID": _UUID}, "b": [{"ID": _UUID.upper()}]} + + result = resolve_referenced_attachments(data, _registry(), _judge()) + + assert [r.id for r in result] == [_UUID] + + async def test_caps_at_the_api_limit(self): + registry = dict(_other_attachment() for _ in range(_MAX_ATTACHMENTS + 2)) + data = {"files": [{"ID": attachment_id} for attachment_id in registry]} + + result = resolve_referenced_attachments(data, registry, _judge()) + + assert len(result) == _MAX_ATTACHMENTS + + @pytest.mark.parametrize("applies_to", ["Prompts", "prompts"]) + async def test_returns_empty_when_scoped_to_prompts(self, applies_to): + result = resolve_referenced_attachments( + {"attachment": {"ID": _UUID}}, _registry(), _scoped_judge(applies_to) + ) + + assert result == [] + + async def test_accepts_attachment_instances_and_models(self): + """Arguments may already carry expanded objects, not only wire dicts; they are + still looked up in the registry by id.""" + + class ToolArgs(BaseModel): + attachment: Attachment + + registry = _registry() + data = { + "direct": registry[_UUID], + "wrapped": ToolArgs(attachment=registry[_UUID]), + } + + result = resolve_referenced_attachments(data, registry, _judge()) + + assert [r.id for r in result] == [_UUID] + + async def test_bounds_the_scan_on_deep_and_large_payloads(self): + """A tool result can be arbitrarily large; the scan stops instead of stalling the + guardrail node, and never raises.""" + deep: dict[str, Any] = {"ID": _UUID} + for _ in range(100): + deep = {"child": deep} + wide = {"rows": [{"n": i} for i in range(20_000)], "file": {"ID": _UUID}} + + assert resolve_referenced_attachments(deep, _registry(), _judge()) == [] + wide_result = resolve_referenced_attachments(wide, _registry(), _judge()) + assert isinstance(wide_result, list) + + @pytest.mark.parametrize("data", [None, 42, "just text", object(), [1, "two"]]) + async def test_never_raises_on_non_structured_payloads(self, data): + assert resolve_referenced_attachments(data, _registry(), _judge()) == [] + + async def test_matches_the_registry_path_for_the_same_attachment(self): + """Agent/LLM scope and tool scope must send the backend identical references.""" + registry = _registry() + + via_registry = await resolve_guardrail_attachments(registry, _judge()) + via_mention = resolve_referenced_attachments( + {"attachment": {"ID": _UUID}}, registry, _judge() + ) + + assert via_registry == via_mention + + +class _UnscannablePayload(dict[str, Any]): + """A mapping whose lookups blow up, as a broken tool result might.""" + + def get(self, key, default=None): # noqa: D401 + raise RuntimeError("boom") + + +class _MentionWithUnreadableId(dict[str, Any]): + """Looks like a mention to the scanner but fails when the id is read.""" + + def get(self, key, default=None): + return _UUID if key == "ID" else default + + def __getitem__(self, key): + raise RuntimeError("boom") + + +class TestResolveReferencedAttachmentsErrorPaths: + async def test_returns_empty_when_the_payload_cannot_be_scanned(self): + """The guardrail node re-raises, so a broken payload must degrade to no files.""" + data = {"result": _UnscannablePayload(ID=_UUID)} + + assert resolve_referenced_attachments(data, _registry(), _judge()) == [] + + async def test_skips_a_mention_whose_id_cannot_be_read(self): + data = {"file": _MentionWithUnreadableId(), "other": {"ID": _UUID}} + + result = resolve_referenced_attachments(data, _registry(), _judge()) + + assert [r.id for r in result] == [_UUID] diff --git a/tests/agent/guardrails/test_guardrail_nodes.py b/tests/agent/guardrails/test_guardrail_nodes.py index fc030bb85..cf8eff31c 100644 --- a/tests/agent/guardrails/test_guardrail_nodes.py +++ b/tests/agent/guardrails/test_guardrail_nodes.py @@ -14,6 +14,7 @@ from uipath.platform.guardrails import BuiltInValidatorGuardrail from uipath_langchain.agent.guardrails.guardrail_nodes import ( + _create_guardrail_node, create_agent_init_guardrail_node, create_agent_terminate_guardrail_node, create_llm_guardrail_node, @@ -1003,7 +1004,8 @@ async def test_tool_guardrail_payload_populated_post_execution(self, monkeypatch class TestGuardrailNodeAttachments: - """Agent- and LLM-scope nodes forward the run's job attachments to the judge.""" + """Built-in guardrail nodes forward attachment references to the judge: the run's + registry at Agent and LLM scope, the files one tool call mentions at Tool scope.""" _UUID = "7f2c1e44-0b3a-4a1e-9d55-2f9a1c3b8e10" @@ -1087,18 +1089,64 @@ async def test_llm_node_forwards_resolved_attachments(self, monkeypatch): assert fake.guardrails.last_attachments == [attachment] - @pytest.mark.asyncio - async def test_tool_scope_node_never_resolves_attachments(self, monkeypatch): - """Tool scope is excluded by product decision: a tool-scope judge would ship file - contents on every tool call.""" - from unittest.mock import AsyncMock + def _tool_pre_state(self, args): + return AgentGuardrailsGraphState( + messages=[ + AIMessage( + content="", + tool_calls=[{"name": "my_tool", "args": args, "id": "c1"}], + ) + ], + inner_state=InnerAgentGuardrailsGraphState( + job_attachments=self._state_with_attachment().inner_state.job_attachments + ), + ) + + def _tool_post_state(self, content): + return AgentGuardrailsGraphState( + messages=[ + AIMessage( + content="", + tool_calls=[{"name": "my_tool", "args": {}, "id": "c1"}], + ), + ToolMessage(content=content, tool_call_id="c1"), + ], + inner_state=InnerAgentGuardrailsGraphState( + job_attachments=self._state_with_attachment().inner_state.job_attachments + ), + ) + @pytest.mark.asyncio + async def test_tool_pre_node_forwards_attachments_referenced_in_tool_args( + self, monkeypatch + ): + """Before the tool runs, the judge reads the file the call names. The model + passes ``{"ID": ...}`` only, so name and type come from the run's registry.""" fake = _patch_uipath(monkeypatch, reason="ok") - resolver = AsyncMock(return_value=[]) - monkeypatch.setattr( - "uipath_langchain.agent.guardrails.guardrail_nodes.resolve_guardrail_attachments", - resolver, + args = {"attachment": {"ID": self._UUID}, "question": "summarize"} + + _, node = create_tool_guardrail_node( + guardrail=self._judge_guardrail(), + execution_stage=ExecutionStage.PRE_EXECUTION, + success_node="ok", + failure_node="nope", + tool_name="my_tool", ) + cmd = await node(self._tool_pre_state(args)) + + assert cmd.goto == "ok" + assert json.loads(fake.guardrails.last_text) == args + assert [ + a.model_dump(by_alias=True) for a in fake.guardrails.last_attachments + ] == [{"id": self._UUID, "fileName": "Tickets.csv", "mimeType": "text/csv"}] + + @pytest.mark.asyncio + async def test_tool_pre_node_ignores_registry_when_args_reference_nothing( + self, monkeypatch + ): + """A tool call without a file forwards nothing even though the run holds one; + otherwise every tool call would ship every file to the backend.""" + fake = _patch_uipath(monkeypatch, reason="ok") _, node = create_tool_guardrail_node( guardrail=self._judge_guardrail(), @@ -1107,22 +1155,134 @@ async def test_tool_scope_node_never_resolves_attachments(self, monkeypatch): failure_node="nope", tool_name="my_tool", ) - state = AgentGuardrailsGraphState( - messages=[ - AIMessage( - content="", - tool_calls=[{"name": "my_tool", "args": {"q": 1}, "id": "c1"}], - ) - ], - inner_state=InnerAgentGuardrailsGraphState( - job_attachments=self._state_with_attachment().inner_state.job_attachments - ), + await node(self._tool_pre_state({"q": 1})) + + assert fake.guardrails.last_attachments == [] + + @pytest.mark.asyncio + async def test_tool_post_node_forwards_attachment_returned_by_the_tool( + self, monkeypatch + ): + """After the tool runs, the judge reads the file the result names.""" + fake = _patch_uipath(monkeypatch, reason="ok") + content = json.dumps( + { + "file": { + "ID": self._UUID, + "FullName": "Tickets.csv", + "MimeType": "text/csv", + } + } ) - await node(state) - resolver.assert_not_awaited() + _, node = create_tool_guardrail_node( + guardrail=self._judge_guardrail(), + execution_stage=ExecutionStage.POST_EXECUTION, + success_node="ok", + failure_node="nope", + tool_name="my_tool", + ) + await node(self._tool_post_state(content)) + + assert fake.guardrails.last_text == content + assert [a.id for a in fake.guardrails.last_attachments] == [self._UUID] + + @pytest.mark.asyncio + async def test_tool_post_node_with_plain_text_result_forwards_nothing( + self, monkeypatch + ): + fake = _patch_uipath(monkeypatch, reason="ok") + + _, node = create_tool_guardrail_node( + guardrail=self._judge_guardrail(), + execution_stage=ExecutionStage.POST_EXECUTION, + success_node="ok", + failure_node="nope", + tool_name="my_tool", + ) + await node(self._tool_post_state("tool output")) + + assert fake.guardrails.last_attachments == [] + + @pytest.mark.asyncio + async def test_tool_node_without_extractors_forwards_nothing(self, monkeypatch): + """A tool-scope built-in node built without the argument/result extractors has + no source to scan and must not fall back to the whole registry.""" + from uipath.platform.guardrails import GuardrailScope + + fake = _patch_uipath(monkeypatch, reason="ok") + + _, node = _create_guardrail_node( + self._judge_guardrail(), + GuardrailScope.TOOL, + ExecutionStage.PRE_EXECUTION, + lambda state: "payload", + "ok", + "nope", + ) + cmd = await node(self._tool_pre_state({"attachment": {"ID": self._UUID}})) + + assert cmd.goto == "ok" + assert fake.guardrails.last_attachments == [] + + @pytest.mark.asyncio + async def test_tool_post_node_evaluates_without_files_when_the_result_cannot_be_read( + self, monkeypatch + ): + """A file must never fail the run: if the result extractor blows up, the judge + still sees the text payload, just without attachments.""" + fake = _patch_uipath(monkeypatch, reason="ok") + + def broken(_state): + raise RuntimeError("boom") + + monkeypatch.setattr( + "uipath_langchain.agent.guardrails.guardrail_nodes._extract_tool_output_data", + broken, + ) + + _, node = create_tool_guardrail_node( + guardrail=self._judge_guardrail(), + execution_stage=ExecutionStage.POST_EXECUTION, + success_node="ok", + failure_node="nope", + tool_name="my_tool", + ) + cmd = await node(self._tool_post_state('{"ID": "not-json-but-mentions-ID"')) + + assert cmd.goto == "ok" + assert fake.guardrails.last_text == '{"ID": "not-json-but-mentions-ID"' assert fake.guardrails.last_attachments == [] + @pytest.mark.asyncio + @pytest.mark.parametrize( + "factory", [create_agent_init_guardrail_node, create_llm_guardrail_node] + ) + async def test_agent_and_llm_pre_nodes_forward_files_scoped_attachments( + self, monkeypatch, factory + ): + """Pre-execution at Agent and LLM scope with ``appliesTo = Files`` still hands + the run's files to the judge (real resolver, nothing patched).""" + from uipath.platform.guardrails.guardrails import EnumParameterValue + + fake = _patch_uipath(monkeypatch, reason="ok") + guardrail = self._judge_guardrail() + guardrail.validator_parameters = [ + EnumParameterValue.model_validate( + {"$parameterType": "enum", "id": "appliesTo", "value": "Files"} + ) + ] + + _, node = factory( + guardrail=guardrail, + execution_stage=ExecutionStage.PRE_EXECUTION, + success_node="ok", + failure_node="nope", + ) + await node(self._state_with_attachment()) + + assert [a.id for a in fake.guardrails.last_attachments] == [self._UUID] + @pytest.mark.asyncio async def test_attachment_rejection_falls_back_to_text_only(self, monkeypatch): """A 400 on a request that carried attachments must not kill the run: the backend diff --git a/tests/agent/guardrails/test_tool_guardrails_subgraph_run.py b/tests/agent/guardrails/test_tool_guardrails_subgraph_run.py new file mode 100644 index 000000000..06336d455 --- /dev/null +++ b/tests/agent/guardrails/test_tool_guardrails_subgraph_run.py @@ -0,0 +1,117 @@ +"""A compiled tool-guardrail subgraph, run end to end with a judge that reads files.""" + +import json +import uuid +from typing import Any + +import pytest +from langchain_core.messages import AIMessage, ToolMessage +from langchain_core.tools import StructuredTool +from pydantic import BaseModel +from uipath.core.guardrails import GuardrailValidationResultType +from uipath.platform.attachments import Attachment +from uipath.platform.guardrails import BuiltInValidatorGuardrail + +from tests.agent.guardrails.test_guardrail_nodes import _patch_uipath +from uipath_langchain.agent.exceptions import AgentRuntimeError, AgentRuntimeErrorCode +from uipath_langchain.agent.guardrails.actions import BlockAction +from uipath_langchain.agent.react.guardrails.guardrails_subgraph import ( + create_tool_guardrails_subgraph, +) +from uipath_langchain.agent.tools.tool_node import UiPathToolNode + +_UUID = "7f2c1e44-0b3a-4a1e-9d55-2f9a1c3b8e10" +_TOOL = "analyze_files" + + +class _Args(BaseModel): + attachment: dict[str, Any] + question: str + + +def _judge() -> BuiltInValidatorGuardrail: + return BuiltInValidatorGuardrail.model_validate( + { + "$guardrailType": "builtInValidator", + "id": "judge-files", + "name": "Judge files", + "description": "Blocks poisoned files before a tool reads them.", + "enabledForEvals": True, + "selector": {"scopes": ["Tool"], "matchNames": [_TOOL]}, + "validatorType": "llm_as_judge", + "validatorParameters": [], + } + ) + + +def _build(monkeypatch, *, result: GuardrailValidationResultType): + calls: list[dict[str, Any]] = [] + + async def analyze(attachment: dict[str, Any], question: str) -> str: + calls.append({"attachment": attachment, "question": question}) + return "The file lists 12 tickets." + + tool = StructuredTool.from_function( + coroutine=analyze, name=_TOOL, description="Reads a file.", args_schema=_Args + ) + graph = create_tool_guardrails_subgraph( + tool_node=(_TOOL, UiPathToolNode(tool)), + guardrails=[(_judge(), BlockAction("unsafe file"))], + ) + fake = _patch_uipath(monkeypatch, result=result, reason="judged") + return graph, fake, calls + + +def _input() -> dict[str, Any]: + args = {"attachment": {"ID": _UUID}, "question": "summarize"} + return { + "messages": [ + AIMessage( + content="", tool_calls=[{"name": _TOOL, "args": args, "id": "c1"}] + ) + ], + "inner_state": { + "job_attachments": { + _UUID: Attachment( + id=uuid.UUID(_UUID), full_name="Tickets.csv", mime_type="text/csv" + ) + } + }, + } + + +@pytest.mark.asyncio +async def test_pre_execution_judge_blocks_before_the_tool_runs(monkeypatch): + """The judge saw the referenced file and the tool implementation never ran.""" + graph, fake, calls = _build( + monkeypatch, result=GuardrailValidationResultType.VALIDATION_FAILED + ) + + payload = _input() + + with pytest.raises(AgentRuntimeError) as raised: + await graph.ainvoke(payload) + + assert raised.value.error_info.code == AgentRuntimeError.full_code( + AgentRuntimeErrorCode.TERMINATION_GUARDRAIL_VIOLATION + ) + assert calls == [] + assert fake.guardrails.call_count == 1 + assert [a.id for a in fake.guardrails.last_attachments] == [_UUID] + assert json.loads(fake.guardrails.last_text)["attachment"] == {"ID": _UUID} + + +@pytest.mark.asyncio +async def test_passing_judge_lets_the_tool_run_and_judges_its_result(monkeypatch): + graph, fake, calls = _build( + monkeypatch, result=GuardrailValidationResultType.PASSED + ) + + final = await graph.ainvoke(_input()) + + assert [call["question"] for call in calls] == ["summarize"] + assert fake.guardrails.call_count == 2 + assert isinstance(final["messages"][-1], ToolMessage) + # The post-execution judge read the plain-text result and no file with it. + assert fake.guardrails.last_text == "The file lists 12 tickets." + assert fake.guardrails.last_attachments == [] From 0c6aeb1916d73fccf63b88d76ce0bebcfdee73fa Mon Sep 17 00:00:00 2001 From: Andrei Petraru Date: Tue, 22 Sep 2026 00:26:36 +0300 Subject: [PATCH 2/3] chore: bump to 0.18.13 0.18.12 shipped from main on 2026-09-21 without this change. Co-Authored-By: Claude Fable 5.1 --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d2ac71462..e74b80dfd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath-langchain" -version = "0.18.12" +version = "0.18.13" description = "Python SDK that enables developers to build and deploy LangGraph agents to the UiPath Cloud Platform" readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.11" diff --git a/uv.lock b/uv.lock index b6cf4ec82..a1a2cb0fa 100644 --- a/uv.lock +++ b/uv.lock @@ -4828,7 +4828,7 @@ wheels = [ [[package]] name = "uipath-langchain" -version = "0.18.12" +version = "0.18.13" source = { editable = "." } dependencies = [ { name = "a2a-sdk" }, From 093f3e483f05549a2c402874197037f10c73ce23 Mon Sep 17 00:00:00 2001 From: Andrei Petraru Date: Tue, 22 Sep 2026 11:05:10 +0300 Subject: [PATCH 3/3] fix(guardrails): judge the current tool call, not the first call with that name Review finding on #1103: with two calls to the same tool in one AI message, the tool-scope guardrail read "the first call named tool_name", so the second pre-execution evaluation judged (and resolved attachments for) the first call, and once the first ToolMessage was appended it read nothing at all. The node now selects the call the way the tool node does (extract_current_tool_call_index) before execution, and the call the last ToolMessage answers after execution; payload, deterministic input extraction and attachment resolution all use that selection. Co-Authored-By: Claude Fable 5.1 --- .../agent/guardrails/guardrail_nodes.py | 72 +++++++++++++--- .../agent/guardrails/test_guardrail_nodes.py | 85 +++++++++++++++++++ 2 files changed, 144 insertions(+), 13 deletions(-) diff --git a/src/uipath_langchain/agent/guardrails/guardrail_nodes.py b/src/uipath_langchain/agent/guardrails/guardrail_nodes.py index dd9e19698..5a918934d 100644 --- a/src/uipath_langchain/agent/guardrails/guardrail_nodes.py +++ b/src/uipath_langchain/agent/guardrails/guardrail_nodes.py @@ -4,6 +4,7 @@ import re from typing import Any, Callable +from langchain_core.messages import AIMessage, ToolMessage from langgraph.types import Command from uipath.core.guardrails import ( DeterministicGuardrail, @@ -33,6 +34,10 @@ get_message_content, ) from uipath_langchain.agent.react.types import AgentGuardrailsGraphState +from uipath_langchain.agent.react.utils import ( + extract_current_tool_call_index, + find_latest_ai_message, +) from ..exceptions import AgentRuntimeError, AgentRuntimeErrorCode @@ -387,6 +392,13 @@ def _payload_generator(state: AgentGuardrailsGraphState) -> str: ) +def _tool_call_field(tool_call: Any, field: str) -> Any: + """Read ``field`` from a tool call given as a dict or an object.""" + if isinstance(tool_call, dict): + return tool_call.get(field) + return getattr(tool_call, field, None) + + def create_tool_guardrail_node( guardrail: BaseGuardrail, execution_stage: ExecutionStage, @@ -409,6 +421,51 @@ def create_tool_guardrail_node( A tuple of (node_name, node_function) for the guardrail evaluation node. """ + def _current_call_args(state: AgentGuardrailsGraphState) -> dict[str, Any]: + """Arguments of the tool call this evaluation is about. + + One AI message can carry several calls to the same tool, executed one after + another with a ToolMessage appended after each. The call under evaluation is + therefore not "the first call named ``tool_name``" but, before execution, the + first one without a ToolMessage yet (the same selection the tool node makes), + and after execution, the one the last ToolMessage answers. Falls back to the + first matching call when the history has no ToolMessage bookkeeping. + """ + messages = state.messages + if not messages: + return {} + + ai_message = find_latest_ai_message(messages) + if ai_message is None or not ai_message.tool_calls: + return {} + + selected: Any = None + if execution_stage == ExecutionStage.PRE_EXECUTION: + try: + index = extract_current_tool_call_index(messages, tool_name) + except AgentRuntimeError: + index = None + if index is not None and index < len(ai_message.tool_calls): + selected = ai_message.tool_calls[index] + else: # POST_EXECUTION + last_message = messages[-1] + if isinstance(last_message, ToolMessage): + selected = next( + ( + call + for call in ai_message.tool_calls + if _tool_call_field(call, "id") == last_message.tool_call_id + ), + None, + ) + + if selected is None: + return _extract_tool_args_from_message(ai_message, tool_name) + + return _extract_tool_args_from_message( + AIMessage(content="", tool_calls=[selected]), tool_name + ) + def _payload_generator(state: AgentGuardrailsGraphState) -> str: """Extract tool call arguments for the specified tool name. @@ -422,24 +479,13 @@ def _payload_generator(state: AgentGuardrailsGraphState) -> str: return "" if execution_stage == ExecutionStage.PRE_EXECUTION: - last_message = state.messages[-1] - args_dict = _extract_tool_args_from_message(last_message, tool_name) - return json.dumps(args_dict) + return json.dumps(_current_call_args(state)) return get_message_content(state.messages[-1]) # Create closures for input/output data extraction (for deterministic guardrails) def _input_data_extractor(state: AgentGuardrailsGraphState) -> dict[str, Any]: - if execution_stage == ExecutionStage.PRE_EXECUTION: - if len(state.messages) < 1: - return {} - message = state.messages[-1] - else: # POST_EXECUTION - if len(state.messages) < 2: - return {} - message = state.messages[-2] - - return _extract_tool_args_from_message(message, tool_name) + return _current_call_args(state) def _output_data_extractor(state: AgentGuardrailsGraphState) -> dict[str, Any]: return _extract_tool_output_data(state) diff --git a/tests/agent/guardrails/test_guardrail_nodes.py b/tests/agent/guardrails/test_guardrail_nodes.py index cf8eff31c..5ae0fafd3 100644 --- a/tests/agent/guardrails/test_guardrail_nodes.py +++ b/tests/agent/guardrails/test_guardrail_nodes.py @@ -1116,6 +1116,91 @@ def _tool_post_state(self, content): ), ) + @pytest.mark.asyncio + async def test_tool_pre_node_judges_the_current_call_when_a_tool_is_called_twice( + self, monkeypatch + ): + """One AI message, two calls to the same tool: after the first call's ToolMessage + lands, the second evaluation must judge the second call's arguments and file, + not the first call's (the tool node selects the call the same way).""" + fake = _patch_uipath(monkeypatch, reason="ok") + other = "00000000-0000-4000-8000-000000000001" + first_args = {"attachment": {"ID": other}, "question": "first"} + second_args = {"attachment": {"ID": self._UUID}, "question": "second"} + state = AgentGuardrailsGraphState( + messages=[ + AIMessage( + content="", + tool_calls=[ + {"name": "my_tool", "args": first_args, "id": "c1"}, + {"name": "my_tool", "args": second_args, "id": "c2"}, + ], + ), + ToolMessage(content="first done", tool_call_id="c1"), + ], + inner_state=InnerAgentGuardrailsGraphState( + job_attachments=self._state_with_attachment().inner_state.job_attachments + ), + ) + + _, node = create_tool_guardrail_node( + guardrail=self._judge_guardrail(), + execution_stage=ExecutionStage.PRE_EXECUTION, + success_node="ok", + failure_node="nope", + tool_name="my_tool", + ) + cmd = await node(state) + + assert cmd.goto == "ok" + assert json.loads(fake.guardrails.last_text) == second_args + assert [a.id for a in fake.guardrails.last_attachments] == [self._UUID] + + @pytest.mark.asyncio + async def test_tool_post_node_judges_the_answered_call_when_a_tool_is_called_twice( + self, monkeypatch + ): + """After the second of two calls returns, the post evaluation reads that + call's result (the last ToolMessage), which here names the file.""" + fake = _patch_uipath(monkeypatch, reason="ok") + second_result = json.dumps( + { + "file": { + "ID": self._UUID, + "FullName": "Tickets.csv", + "MimeType": "text/csv", + } + } + ) + state = AgentGuardrailsGraphState( + messages=[ + AIMessage( + content="", + tool_calls=[ + {"name": "my_tool", "args": {"n": 1}, "id": "c1"}, + {"name": "my_tool", "args": {"n": 2}, "id": "c2"}, + ], + ), + ToolMessage(content="first done", tool_call_id="c1"), + ToolMessage(content=second_result, tool_call_id="c2"), + ], + inner_state=InnerAgentGuardrailsGraphState( + job_attachments=self._state_with_attachment().inner_state.job_attachments + ), + ) + + _, node = create_tool_guardrail_node( + guardrail=self._judge_guardrail(), + execution_stage=ExecutionStage.POST_EXECUTION, + success_node="ok", + failure_node="nope", + tool_name="my_tool", + ) + await node(state) + + assert fake.guardrails.last_text == second_result + assert [a.id for a in fake.guardrails.last_attachments] == [self._UUID] + @pytest.mark.asyncio async def test_tool_pre_node_forwards_attachments_referenced_in_tool_args( self, monkeypatch