Replace private pydantic-ai access with public Agent.toolsets surface

The _get_function_tools helper reached into agent._function_tools (0.2.x)
and agent._function_toolset (1.x) — both private dataclass fields. The
tool-dispatch and requires-approval lookup paths additionally probed
speculative private attributes (_wrapped_function, callable_function,
_core_tool, wrapped_tool) on Tool objects that do not exist on the
public Tool dataclass.

Refactor to walk the public Agent.toolsets property (documented to
include the auto-built function toolset for tools registered directly
on the agent), pick FunctionToolset instances, and read the public
FunctionToolset.tools dict. Use the public Tool.function and
Tool.requires_approval fields for callable extraction and fallback
approval checks. Verified against pydantic-ai 1.87.0 source (our
minimum pin).

Update test mocks in test_pydantic_ai_agents.py and
test_nested_approval_gates.py to set inst.toolsets via a
MagicMock(spec=FunctionToolset) helper so isinstance() recognises the
fake. The integration test
test_check_tool_requires_approval_with_real_pydantic_ai_agent
exercises the helper end-to-end against a real Agent instance.

Author: JSv4

opencontractserver/llms/agents/pydantic_ai_agents.py

@@ -29,6 +29,7 @@
     ToolReturnPart,
     UserPromptPart,
 )
+from pydantic_ai.toolsets import FunctionToolset
 from pydantic_graph import End
 
 from opencontractserver.constants.context_guardrails import COMPACTION_SUMMARY_PREFIX
@@ -148,18 +149,16 @@ async def similarity_search(
 def _get_function_tools(agent: PydanticAIAgent) -> dict:
     """Return the function-tools dict from a pydantic-ai Agent.
 
-    Handles both pydantic-ai 0.2.x (``agent._function_tools``) and
-    1.x (``agent._function_toolset.tools``).
+    Uses only the public surface: ``Agent.toolsets`` (documented property
+    that includes the auto-built function toolset for tools registered
+    directly on the agent) and ``FunctionToolset.tools`` (public dict of
+    tool name -> ``Tool``).
     """
-    # pydantic-ai 0.2.x
-    ft = getattr(agent, "_function_tools", None)
-    if ft is not None:
-        return ft
-    # pydantic-ai 1.x
-    toolset = getattr(agent, "_function_toolset", None)
-    if toolset is not None:
-        return getattr(toolset, "tools", {})
-    return {}
+    merged: dict = {}
+    for toolset in agent.toolsets:
+        if isinstance(toolset, FunctionToolset):
+            merged.update(toolset.tools)
+    return merged
 
 
 @dataclasses.dataclass
@@ -1605,27 +1604,17 @@ async def _maybe_await(call_result):  # noqa: D401 – small helper
                         # Don't retry here, fall through to registry lookup
 
                 if not tool_executed:
-                    # Resort to pydantic-ai registry – may return Tool object.
+                    # Resort to pydantic-ai registry – returns a public ``Tool``.
                     tool_obj = _get_function_tools(self.pydantic_ai_agent).get(
                         tool_name
                     )
                     if tool_obj is None:
                         raise ValueError(f"Tool '{tool_name}' not found for execution")
 
-                    # Try common attributes to reach the underlying callable.
-                    candidate = None
-                    for attr in (
-                        "function",
-                        "_wrapped_function",
-                        "callable_function",
-                    ):
-                        candidate = getattr(tool_obj, attr, None)
-                        if callable(candidate):
-                            break
-
-                    if candidate is None or not callable(candidate):
+                    candidate = tool_obj.function
+                    if not callable(candidate):
                         raise TypeError(
-                            "Tool object is not callable and no inner function found"
+                            f"Tool '{tool_name}' has a non-callable function"
                         )
 
                     logger.info(
@@ -1933,36 +1922,24 @@ def _check_tool_requires_approval(self, tool_name: str) -> bool:
                     if hasattr(tool, "requires_approval"):
                         return tool.requires_approval
 
-        # Check tools registered with pydantic-ai agent
+        # Check tools registered with pydantic-ai agent. Tools registered as
+        # plain async callables (our common case) carry their CoreTool on the
+        # underlying function, not on the Tool object — pydantic-ai 1.x's
+        # Tool.requires_approval defaults to False unless the caller passes it
+        # in, so we must consult the function attribute first.
         function_tools = _get_function_tools(self.pydantic_ai_agent)
         if function_tools:
             tool_obj = function_tools.get(tool_name)
-            if tool_obj:
-                # Check various possible attributes where the CoreTool might be stored
-                for attr in ("core_tool", "_core_tool", "wrapped_tool"):
-                    core_tool = getattr(tool_obj, attr, None)
-                    if core_tool and hasattr(core_tool, "requires_approval"):
-                        return core_tool.requires_approval
-
-                # Check the wrapped function (must come before the native
-                # tool_obj.requires_approval check because pydantic-ai 1.x
-                # Tool has a native requires_approval field that defaults to
-                # False, shadowing the custom attribute on the function).
-                for attr in ("function", "_wrapped_function", "callable_function"):
-                    func = getattr(tool_obj, attr, None)
-                    if func:
-                        # Check if the function has a core_tool attribute
-                        if hasattr(func, "core_tool") and hasattr(
-                            func.core_tool, "requires_approval"
-                        ):
-                            return func.core_tool.requires_approval
-                        # Check if the function itself has requires_approval
-                        if hasattr(func, "requires_approval"):
-                            return func.requires_approval
-
-                # Fall back to the tool object's own requires_approval
-                if hasattr(tool_obj, "requires_approval"):
-                    return tool_obj.requires_approval
+            if tool_obj is not None:
+                func = tool_obj.function
+                core_tool = getattr(func, "core_tool", None)
+                if core_tool is not None and getattr(
+                    core_tool, "requires_approval", False
+                ):
+                    return True
+                if getattr(func, "requires_approval", False):
+                    return True
+                return tool_obj.requires_approval
 
         # Default to not requiring approval
         return False

opencontractserver/tests/test_nested_approval_gates.py

@@ -87,6 +87,18 @@ async def __aexit__(self, exc_type, exc, tb):
         return False
 
 
+def _fake_function_toolset(entries: dict[str, Any]) -> Any:
+    """Return a MagicMock that satisfies isinstance(_, FunctionToolset) and
+    exposes a .tools dict, mimicking the public toolset surface that
+    pydantic-ai exposes via Agent.toolsets.
+    """
+    from pydantic_ai.toolsets import FunctionToolset
+
+    ts = MagicMock(spec=FunctionToolset)
+    ts.tools = entries
+    return ts
+
+
 # ---------------------------------------------------------------------------
 # Mock sub-agent that yields configurable stream events
 # ---------------------------------------------------------------------------
@@ -180,7 +192,7 @@ async def _create_corpus_agent(self, sub_agent_events=None):
         ) as mock_agent_cls:
             inst = MagicMock()
             inst.iter = MagicMock(return_value=_IterCtx())
-            inst._function_tools = {}
+            inst.toolsets = []
             inst.run = AsyncMock(
                 return_value=types.SimpleNamespace(
                     data="ok", sources=[], usage=lambda: None
@@ -413,11 +425,14 @@ async def _spy_tool(ctx, **kwargs):
                 )
             )
             inst.iter = MagicMock(return_value=_IterCtx())
-            # resume_with_approval uses pydantic_ai_agent._function_tools
-            # as fallback to find the tool callable.
-            inst._function_tools = {
-                "ask_document": types.SimpleNamespace(function=_spy_tool),
-            }
+            # resume_with_approval falls back to pydantic-ai's public
+            # function-toolset registry (Agent.toolsets) to find the tool
+            # callable when config.tools doesn't carry it.
+            inst.toolsets = [
+                _fake_function_toolset(
+                    {"ask_document": types.SimpleNamespace(function=_spy_tool)}
+                )
+            ]
             mock_agent_cls.return_value = inst
 
             agent = await UnifiedAgentFactory.create_corpus_agent(
@@ -499,9 +514,11 @@ async def _capture_bypass_tool(ctx, **kwargs):
                 )
                 return {"result": "done"}
 
-            inst._function_tools = {
-                "test_tool": types.SimpleNamespace(function=_capture_bypass_tool),
-            }
+            inst.toolsets = [
+                _fake_function_toolset(
+                    {"test_tool": types.SimpleNamespace(function=_capture_bypass_tool)}
+                )
+            ]
             mock_agent_cls.return_value = inst
 
             agent = await UnifiedAgentFactory.create_corpus_agent(
@@ -569,9 +586,11 @@ async def test_bypass_flag_reset_on_tool_error(self):
             async def _failing_tool(ctx, **kwargs):
                 raise RuntimeError("Tool exploded")
 
-            inst._function_tools = {
-                "exploding_tool": types.SimpleNamespace(function=_failing_tool),
-            }
+            inst.toolsets = [
+                _fake_function_toolset(
+                    {"exploding_tool": types.SimpleNamespace(function=_failing_tool)}
+                )
+            ]
             mock_agent_cls.return_value = inst
 
             agent = await UnifiedAgentFactory.create_corpus_agent(
@@ -615,7 +634,7 @@ async def _failing_tool(ctx, **kwargs):
     async def test_resume_executes_tool_from_config_tools(self):
         """resume_with_approval should find and execute a tool present in
         config.tools (the wrapper_fn path) rather than falling back to
-        pydantic_ai_agent._function_tools."""
+        pydantic-ai's function-toolset registry."""
         from opencontractserver.conversations.models import ChatMessage, Conversation
         from opencontractserver.llms.agents.agent_factory import (
             UnifiedAgentFactory,
@@ -640,9 +659,9 @@ async def _config_spy(ctx, **kwargs):
                 )
             )
             inst.iter = MagicMock(return_value=_IterCtx())
-            # Leave _function_tools empty so the only way to find the tool
-            # is via config.tools.
-            inst._function_tools = {}
+            # Leave the function toolset empty so the only way to find the
+            # tool is via config.tools.
+            inst.toolsets = []
             mock_agent_cls.return_value = inst
 
             agent = await UnifiedAgentFactory.create_corpus_agent(
@@ -683,7 +702,7 @@ async def _config_spy(ctx, **kwargs):
 
     async def test_resume_config_tools_typeerror_falls_through(self):
         """When a config.tools wrapper raises TypeError, resume_with_approval
-        should fall through to the _function_tools registry lookup."""
+        should fall through to the function-toolset registry lookup."""
         from opencontractserver.conversations.models import ChatMessage, Conversation
         from opencontractserver.llms.agents.agent_factory import (
             UnifiedAgentFactory,
@@ -712,9 +731,11 @@ async def _fallback_tool(ctx, **kwargs):
                 )
             )
             inst.iter = MagicMock(return_value=_IterCtx())
-            inst._function_tools = {
-                "dual_tool": types.SimpleNamespace(function=_fallback_tool),
-            }
+            inst.toolsets = [
+                _fake_function_toolset(
+                    {"dual_tool": types.SimpleNamespace(function=_fallback_tool)}
+                )
+            ]
             mock_agent_cls.return_value = inst
 
             agent = await UnifiedAgentFactory.create_corpus_agent(

opencontractserver/tests/test_pydantic_ai_agents.py

@@ -859,7 +859,7 @@ def test_tool():
         config = AgentConfig(user_id=self.user.id, tools=[mock_tool])
 
         mock_pydantic_agent = MagicMock()
-        mock_pydantic_agent._function_tools = {}
+        mock_pydantic_agent.toolsets = []
         mock_agent_cls.return_value = mock_pydantic_agent
 
         agent = PydanticAICoreAgent(
@@ -885,7 +885,7 @@ def test_check_tool_requires_approval_default_false(
         config = AgentConfig(user_id=self.user.id)
 
         mock_pydantic_agent = MagicMock()
-        mock_pydantic_agent._function_tools = {}
+        mock_pydantic_agent.toolsets = []
         mock_agent_cls.return_value = mock_pydantic_agent
 
         agent = PydanticAICoreAgent(
@@ -1048,7 +1048,7 @@ async def test_resume_with_approval_approved_success(
         )
 
         mock_pydantic_agent = MagicMock()
-        mock_pydantic_agent._function_tools = {}
+        mock_pydantic_agent.toolsets = []
 
         agent = PydanticAICoreAgent(
             config=config,
@@ -1203,7 +1203,7 @@ async def test_resume_with_approval_parses_json_string_args(
         )
 
         mock_pydantic_agent = MagicMock()
-        mock_pydantic_agent._function_tools = {}
+        mock_pydantic_agent.toolsets = []
 
         agent = PydanticAICoreAgent(
             config=config,