fix(pydantic-ai): address remaining review minors (page 4)

Follow-up to the #359 review pass — the earlier commit covered the blockers
+ page-3 majors; these are the page-4 minors:

- model.py:239 — BuiltinToolReturnPart now fires AFTER_TOOL (symmetric with
  the built-in TOOL_CALL). A provider-executed built-in tool carries both its
  call and result inline in the response, so AFTER_TOOL for built-ins fires
  from on_response, not the next request.
- model.py:371 — _coerce_args preserves malformed JSON as {"_raw": args}
  instead of dropping it to {}, so an arg-based policy can still scan it (a
  malformed payload must not slip past governance).
- model.py:229 — tool_name fallback to "unknown" now logs a warning via a
  shared _tool_name() helper (used by all three tool-part call sites).

Not changed: runtime/factory.py:271 kwargs.get("evaluator") — this is the
same magic-string kwarg LangChain #899's factory uses (factory.py:281); it's
the host-dispatch contract shared by all adapters, so kept for parity.

Tests: added built-in-tool-result AFTER_TOOL coverage; updated _coerce_args
variant assertion. 82 pass.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: aditik0303

packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/model.py

@@ -45,6 +45,7 @@
 from pydantic_ai import Agent
 from pydantic_ai.messages import (
     BuiltinToolCallPart,
+    BuiltinToolReturnPart,
     ModelRequest,
     TextPart,
     ToolCallPart,
@@ -200,24 +201,37 @@ def on_request(self, messages: Any) -> None:
         for part in parts:
             if isinstance(part, ToolReturnPart):
                 self._after_tool(
-                    part.tool_name or "unknown",
+                    _tool_name(part),
                     part.content,
                     tool_call_id=getattr(part, "tool_call_id", None),
                 )
 
     # ----- after the model call ---------------------------------------
 
     def on_response(self, response: Any) -> None:
-        """Fire AFTER_MODEL (response text) + TOOL_CALL (each tool-call part)."""
+        """Fire AFTER_MODEL (response text) + TOOL_CALL / AFTER_TOOL parts.
+
+        A provider-executed **built-in** tool carries both its call
+        (``BuiltinToolCallPart``) and its result (``BuiltinToolReturnPart``)
+        inline in the model response, so AFTER_TOOL for built-in tools is fired
+        here — symmetric with the built-in TOOL_CALL — rather than on the next
+        request (where only user-tool ``ToolReturnPart``s arrive).
+        """
         parts = getattr(response, "parts", None) or []
         self._after_model(self._response_text(parts))
         for part in parts:
             if isinstance(part, (ToolCallPart, BuiltinToolCallPart)):
                 self._tool_call(
-                    part.tool_name or "unknown",
+                    _tool_name(part),
                     part.args,
                     tool_call_id=getattr(part, "tool_call_id", None),
                 )
+            elif isinstance(part, BuiltinToolReturnPart):
+                self._after_tool(
+                    _tool_name(part),
+                    part.content,
+                    tool_call_id=getattr(part, "tool_call_id", None),
+                )
 
     # ----- individual evaluate_* wrappers (block-propagate, else swallow) --
 
@@ -380,6 +394,22 @@ def _pieces() -> Iterable[str]:
     return _stringify(content)
 
 
+def _tool_name(part: Any) -> str:
+    """Return ``part.tool_name`` or ``"unknown"``, logging the fallback.
+
+    A missing tool name means TOOL_CALL / AFTER_TOOL can't be attributed to a
+    real tool, so surface it rather than silently reporting ``"unknown"``.
+    """
+    name = getattr(part, "tool_name", None)
+    if name:
+        return name
+    logger.warning(
+        "governance: %s carries no tool_name; reporting 'unknown'",
+        type(part).__name__,
+    )
+    return "unknown"
+
+
 def _coerce_args(args: Any) -> Dict[str, Any]:
     """Normalise ``ToolCallPart.args`` (dict / JSON string / None) to a dict."""
     if args is None:
@@ -391,7 +421,9 @@ def _coerce_args(args: Any) -> Dict[str, Any]:
             parsed = json.loads(args)
             return parsed if isinstance(parsed, dict) else {"_": parsed}
         except (TypeError, ValueError):
-            return {}
+            # Preserve the raw string so an arg-based policy can still scan it;
+            # a malformed payload must not be a way to slip past governance.
+            return {"_raw": args}
     return {}
 
 

packages/uipath-pydantic-ai/tests/governance/test_model.py

@@ -18,6 +18,8 @@
 import pytest
 from pydantic_ai import Agent
 from pydantic_ai.messages import (
+    BuiltinToolCallPart,
+    BuiltinToolReturnPart,
     ModelRequest,
     ModelResponse,
     TextPart,
@@ -225,6 +227,31 @@ def test_on_response_fires_after_model_and_tool_call():
     assert tool_call["session_state"]["tool_calls"] == 1
 
 
+def test_on_response_fires_after_tool_for_builtin_tool_result():
+    """A provider-executed built-in tool carries its result inline in the
+    response (BuiltinToolReturnPart); AFTER_TOOL must fire for it — symmetric
+    with the built-in TOOL_CALL."""
+    ev = FakeEvaluator()
+    cb = _make_callbacks(ev)
+    response = ModelResponse(
+        parts=[
+            BuiltinToolCallPart(
+                tool_name="web_search", args={"q": "x"}, tool_call_id="b1"
+            ),
+            BuiltinToolReturnPart(
+                tool_name="web_search", content={"results": "found"}, tool_call_id="b1"
+            ),
+        ]
+    )
+    cb.on_response(response)
+    hooks = _hooks(ev)
+    assert "tool_call" in hooks and "after_tool" in hooks
+    after_tool = [kw for h, kw in ev.calls if h == "after_tool"][0]
+    assert after_tool["tool_name"] == "web_search"
+    assert after_tool["tool_call_id"] == "b1"
+    assert "found" in after_tool["tool_result"]
+
+
 def test_on_response_coerces_json_string_args():
     ev = FakeEvaluator()
     cb = _make_callbacks(ev)
@@ -393,7 +420,8 @@ def test_coerce_args_variants():
     assert _coerce_args({"a": 1}) == {"a": 1}
     assert _coerce_args('{"a": 1}') == {"a": 1}
     assert _coerce_args(None) == {}
-    assert _coerce_args("not json") == {}
+    # malformed JSON is preserved (not dropped) so arg-based policies can scan it
+    assert _coerce_args("not json") == {"_raw": "not json"}
 
 
 def test_block_in_before_model_propagates():