feat(guardrails): add PRE+POST tool scope support for PII and HarmfulContent middleware [AL-445]

[apetraru-uipath]

What changed?

Added PRE+POST tool scope support for UiPathPIIDetectionMiddleware and UiPathHarmfulContentMiddleware, matching the behavior of UiPath low-code agents where guardrails are evaluated against both tool input (PRE) and tool output (POST).

Key changes:

• pii_detection.py / harmful_content.py: Added stage: GuardrailExecutionStage = PRE_AND_POST parameter. TOOL scope now evaluates the guardrail against tool args (PRE) and against the tool's return value (POST).
• _base.py: Extracted _extract_tool_output_data, _extract_tool_input_data, and _create_tool_wrap_hook into BuiltInGuardrailMiddlewareMixin to eliminate copy-paste duplication between the two middlewares. The inner closure was further extracted into _run_tool_guardrail(request, handler) to make the PRE+POST logic directly testable.
• samples/joke-agent/graph.py: Added UiPathHarmfulContentMiddleware on TOOL scope and explicit stage=PRE_AND_POST to both TOOL-scope middlewares.
• Version bump: uipath-langchain 0.11.3 → 0.11.4.

How has this been tested?

• 27 new unit tests (tests/guardrails/middlewares/test_tool_wrap.py): Cover all branches of _extract_tool_input_data, _extract_tool_output_data, and _run_tool_guardrail — tool-name skipping, PRE-only, POST-only, PRE+POST stage routing, GuardrailBlockException → AgentRuntimeError conversion, AgentRuntimeError re-raise, generic exception logging, request/result modification, and empty-output skipping.
• 4 new e2e parity test runs (test_guardrails_in_langgraph.py): test_tool_pii_post_block and test_tool_harmful_content_post_block, each running against both middleware and decorator flavors, verifying that violations in tool output trigger BlockAction correctly.
• Hook wiring tests extended with TOOL-scope assertions for both affected middlewares.
• Validated end-to-end by running samples/joke-agent — HTTP logs show two guardrail API calls per tool invocation (PRE + POST).
• ruff check and mypy pass on all changed files.

Are there any breaking changes?

• Under Feature Flag
• None
• DB migrations required
• API removals or changes

[valentinabojan]

Took a closer look at the new BuiltInGuardrailMiddlewareMixin._run_tool_guardrail. The refactor and test coverage are solid. Posting 5 inline comments on _base.py:

• Sync HTTP call from async context (now fires up to 2× per tool invocation under the new default)
• _extract_tool_output_data only inspects messages[0] of a Command
• _tool_stage / _tool_names are annotation-only on the mixin — AttributeError trap for future TOOL-scope subclasses
• input_data extracted unconditionally even on POST-only path (wasted work)
• PRE/POST asymmetry on the isinstance(..., dict) guard around modified_*

No approval blocker — flagging as COMMENT.

[valentinabojan]

_evaluate_guardrail (line 57) is a synchronous HTTP call to the UiPath guardrails API, but _run_tool_guardrail is async. With the new default of PRE_AND_POST, every tool invocation now blocks the event loop twice — once here and again on line 157. Under any concurrent agent load this is a real throughput regression vs. the old PRE-only path.

Cheapest fix is to push the sync call into a worker thread (needs import asyncio at the top):

Suggested change

	result = self._evaluate_guardrail(input_data)
	result = await asyncio.to_thread(self._evaluate_guardrail, input_data)

And the same swap on line 157. Long-term, an async evaluate_guardrail on the SDK side would be cleaner.

[valentinabojan]

A Command(update={"messages": [...]}) can carry multiple messages, and the ToolMessage isn't guaranteed to be first. If a flow emits e.g. [AIMessage(...), ToolMessage(<real content>)], this branch reads messages[0], sees it isn't a ToolMessage, sets content = None, and POST silently gets no content to validate.

Iterate to find the first ToolMessage instead:

Suggested change

	messages[0].content
	tool_msg = next(
	(m for m in messages if isinstance(m, ToolMessage)), None
	)
	content = tool_msg.content if tool_msg is not None else None

[valentinabojan]

_tool_stage and _tool_names are annotation-only here — no class-level binding. Today both subclasses (pii_detection.py, harmful_content.py) set them in __init__, so it works. But the moment a future TOOL-scope subclass forgets, _run_tool_guardrail will AttributeError at runtime — not at construction — and only when a specific tool is called. The mixin docstring also only mentions _guardrail / _name / action as required attrs, which makes this easy to miss.

Suggested change

	_tool_stage: GuardrailExecutionStage
	_tool_names: list[str] | None = None
	_tool_stage: GuardrailExecutionStage = GuardrailExecutionStage.PRE_AND_POST

And extend the docstring (lines 38–41) to list these alongside the others.

[valentinabojan]

input_data is extracted on every invocation, even when stage=POST and the value is never read. Cheap call, but pointless work — move it inside the PRE block (right after try: on line 132). Then a POST-only middleware skips the dict copy / str() conversion entirely.

[valentinabojan]

Asymmetry with line 135: PRE only re-wraps the request when isinstance(modified_input, dict), but POST forwards any non-None to create_modified_tool_result. The helper actually handles both str and dict (_utils.py:71-82), so the difference is intentional — but it's load-bearing in a way that's not obvious from the call site, and the next refactor pass is likely to "tidy it up" and break things.

At minimum, add a one-line comment:

Suggested change

	if modified_output is not None:
	if modified_output is not None:
	# POST: ToolMessage.content accepts both str and dict
	tool_result = create_modified_tool_result(tool_result, modified_output)

Or unify the contract by letting create_modified_tool_request accept str too, and gating both sites the same way.

[apetraru-uipath]

The original UiPathChat was a no-op stub (pass). The change was needed to fix a real bug in the joke agent output.

Root cause: When LangChain tracing callbacks are registered, ainvoke uses the streaming API path. Each SSE chunk from the UiPath
LLM gateway includes model_name, id, and created in its generation_info. LangChain accumulates chunks via AIMessageChunk.add →
merge_dicts, which string-concatenates values for these keys. With 14 SSE chunks, the final AI message had:

response_metadata.model_name = "gpt-4o-2024-08-06gpt-4o-2024-08-06gpt-4o-2024-08-06..." # × 14
response_metadata.finish_reason = "stopstop" # sent in 2 final chunks

The fix strips model_name, id, created from intermediate chunks (they only belong on the final chunk), and deduplicates when the
API sends finish_reason twice. The parent class _UpstreamUiPathChat includes these fields in every chunk by design (for standalone
streaming use), so the override is the correct layer to fix this without patching the upstream client.

The change is minimal: a module-level _strip_chunk_metadata helper + overrides of the two streaming methods that call super() and
filter each chunk through it. No behavior changes for non-streaming invocations.

[sonarqubecloud]

Quality Gate passed

Issues
1 New issue
0 Accepted issues

Measures
0 Security Hotspots
97.8% Coverage on New Code
0.2% Duplication on New Code

See analysis details on SonarQube Cloud