feat(agents): promote tool-call JSON repair to a ToolCallJsonRepair capability

code_puppy/agents/_builder.py

@@ -23,6 +23,7 @@
     build_speculative_code_mode,
 )
 from code_puppy.agents._compaction import HistoryCompaction
+from code_puppy.agents._json_repair import build_tool_call_json_repair
 from code_puppy.agents._model_message_transform import build_model_message_transform
 from code_puppy.agents._subagent_recursion import build_subagent_recursion_guard
 from code_puppy.agents._output_limits import (
@@ -685,8 +686,12 @@ def _new_pydantic_agent(toolsets: List[Any]) -> PydanticAgent:
             # hook (after_tool_execute), so its position is inert; the
             # response clamp runs before_model_request after compaction and
             # steering. The plugin transform wraps the final model request.
+            # ToolCallJsonRepair rides the before_tool_validate seam, which
+            # never interacts with the history hooks, so its position is
+            # inert too.
             capabilities=[
                 *build_tool_output_limits(),
+                *build_tool_call_json_repair(),
                 history_compaction,
                 ProcessHistory(steer_processor),
                 build_response_clamp(),

code_puppy/agents/_json_repair.py

@@ -0,0 +1,110 @@
+"""Tool-call JSON repair as a first-class pydantic-ai capability.
+
+Promotes the ``patch_tool_call_json_repair`` monkeypatch (see
+``code_puppy.pydantic_patches``) to the ``before_tool_validate`` capability
+seam. LLMs sometimes emit slightly broken JSON in tool-call arguments
+(trailing commas, missing quotes, unclosed braces); repairing the raw args
+before validation prevents an unnecessary retry round-trip.
+
+**Delivery model -- explicit-when-ours, fallback-for-guests** (the same
+split ``Instrumentation`` uses for tracing): agents built by code_puppy's
+own construction sites carry an explicit :class:`ToolCallJsonRepair`
+capability, and the monkeypatch detects it (via the run's
+``ToolManager.root_capability`` tree) and steps aside. Raw pydantic-ai
+agents built by plugins (e.g. wiggum's judge, which registers real
+read-only tools) carry no capability, so the patch keeps repairing for
+them exactly as before.
+
+Parity notes versus the eager patch:
+
+* **Same repair, same custody.** The hook receives the *live*
+  ``ToolCallPart`` -- the object recorded in the run's message state -- so
+  assigning ``call.args`` here lands the repaired JSON in message history
+  exactly as the patch's in-place mutation did, while the returned args
+  feed validation.
+* **Unknown/unavailable tools are a bounded divergence.**
+  ``ToolManager._resolve_tool`` raises before ``before_tool_validate``
+  fires, so a call to a nonexistent tool no longer gets its recorded args
+  repaired (the patch repaired first, resolution failed after). The call
+  fails with the identical ``ModelRetry`` either way. (In-app, the
+  ``patch_tool_call_callbacks`` pre_tool_call writeback still records a
+  repaired dict view for such calls, so history bytes match the old
+  behavior anyway.)
+* **Output tools were never covered.** On pydantic-ai 2.31.0, tool-based
+  structured output validates through ``validate_output_tool_call`` -- a
+  method the patch never wrapped -- so skipping the ``kind == 'output'``
+  carve-out here (pydantic-ai excludes output tools from tool hooks
+  anyway) changes nothing.
+
+``json_repair`` is an optional dependency: :func:`build_tool_call_json_repair`
+returns ``[]`` when it is missing (mirroring the patch's quiet skip and the
+``build_tool_output_limits`` conditional-splice pattern), and the hook
+no-ops defensively if a capability instance is constructed anyway.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, List
+
+from pydantic_ai import RunContext
+from pydantic_ai.capabilities import AbstractCapability
+from pydantic_ai.messages import ToolCallPart
+from pydantic_ai.tools import ToolDefinition
+
+from code_puppy.pydantic_patches import _repair_tool_call_json
+
+try:  # pragma: no cover - exercised via the None-fallback tests
+    import json_repair
+except ImportError:  # pragma: no cover - optional dependency absent
+    json_repair = None  # type: ignore[assignment]
+
+__all__ = ["ToolCallJsonRepair", "build_tool_call_json_repair"]
+
+
+@dataclass
+class ToolCallJsonRepair(AbstractCapability[Any]):
+    """Repair malformed JSON tool-call arguments before validation.
+
+    Stateless: safe to share across construction passes and spec-construct
+    (the default ``get_serialization_name``/``from_spec`` apply).
+    """
+
+    async def before_tool_validate(
+        self,
+        ctx: RunContext[Any],
+        *,
+        call: ToolCallPart,
+        tool_def: ToolDefinition,
+        args: str | dict[str, Any],
+    ) -> str | dict[str, Any]:
+        """Return repaired raw args; mirror the repair onto ``call.args``.
+
+        Only string args are candidates (dict args already parsed upstream).
+        The repair itself is the patch's shared helper, so both tiers apply
+        the same policy: strict parse first, repair only on failure, and
+        reject any repair that changes the top-level object shape. Repair
+        failures fall back to the original args, which then earn the retry
+        the model would have gotten anyway.
+        """
+        if json_repair is None or not isinstance(args, str) or not args:
+            return args
+        repaired = _repair_tool_call_json(args)
+        if repaired != args:
+            # Same in-place custody the patch performed: this is the
+            # ToolCallPart recorded in the run's message state, so history
+            # shows the repaired JSON the tool actually ran with.
+            call.args = repaired
+            return repaired
+        return args
+
+
+def build_tool_call_json_repair() -> List[ToolCallJsonRepair]:
+    """Build the repair capability, or ``[]`` when ``json_repair`` is absent.
+
+    Returned as a list so callers can splice it into ``capabilities=[...]``
+    unconditionally, exactly like ``build_tool_output_limits``.
+    """
+    if json_repair is None:
+        return []
+    return [ToolCallJsonRepair()]

code_puppy/pydantic_patches.py

@@ -247,6 +247,35 @@ def patch_message_history_cleaning() -> bool:
         )
 
 
+def _run_owns_json_repair(tool_manager: Any) -> bool:
+    """True when the run's capability tree contains ``ToolCallJsonRepair``.
+
+    Walks the run's ``ToolManager.root_capability`` with the public
+    ``AbstractCapability.apply`` visitor (the same traversal the run layer
+    uses for its own explicit-capability checks). Any failure reads as
+    ``False``; the patch then repairs eagerly, which is the old behavior
+    and idempotent alongside the capability (repairing already-repaired
+    JSON is a no-op).
+    """
+    try:
+        root = getattr(tool_manager, "root_capability", None)
+        if root is None:
+            return False
+        from code_puppy.agents._json_repair import ToolCallJsonRepair
+
+        found = False
+
+        def _visit(leaf: Any) -> None:
+            nonlocal found
+            if isinstance(leaf, ToolCallJsonRepair):
+                found = True
+
+        root.apply(_visit)
+        return found
+    except Exception:
+        return False
+
+
 def patch_openai_chat_video_url() -> bool:
     """Let OpenAI-compatible Chat Completions endpoints receive video parts.
 
@@ -330,6 +359,16 @@ def patch_tool_call_json_repair() -> bool:
     json_repair on the raw arguments before validation, preventing
     unnecessary retries. Repairs are attempted only after strict parsing fails,
     and are accepted only when they preserve the required object shape.
+
+    **Fallback tier only.** Agents built by code_puppy's own construction
+    sites carry an explicit ``ToolCallJsonRepair`` capability (see
+    ``code_puppy.agents._json_repair``), which repairs on the
+    ``before_tool_validate`` seam; when it is present in the run's
+    capability tree this patch steps aside. The eager repair below remains
+    for *guest* agents: raw pydantic-ai agents built by plugins (e.g.
+    wiggum's tool-wielding judge) that never pass through our builders.
+    Same explicit-when-ours, fallback-for-guests split as Logfire
+    instrumentation.
     """
     try:
         import json_repair  # noqa: F401  (optional-dependency gate)
@@ -345,8 +384,13 @@ def patch_tool_call_json_repair() -> bool:
 
         async def _patched_validate_tool_call(self, call, **kwargs):
             """Repair malformed JSON args before pydantic-ai validates them."""
-            # Only attempt repair if args is a string (JSON)
-            if isinstance(call.args, str) and call.args:
+            # Only attempt repair if args is a string (JSON), and only when
+            # no ToolCallJsonRepair capability owns the job for this run.
+            if (
+                isinstance(call.args, str)
+                and call.args
+                and not _run_owns_json_repair(self)
+            ):
                 call.args = _repair_tool_call_json(call.args)
 
             # Drop a stray {"arguments": ...} envelope (valid JSON, so the

code_puppy/tools/subagent_invocation.py

@@ -460,6 +460,7 @@ async def _invoke_agent_impl(
                 build_speculative_code_mode,
             )
             from code_puppy.agents._compaction import HistoryCompaction
+            from code_puppy.agents._json_repair import build_tool_call_json_repair
             from code_puppy.agents._subagent_recursion import (
                 build_subagent_recursion_guard,
             )
@@ -485,8 +486,11 @@ async def _invoke_agent_impl(
                 toolsets=mcp_servers,
                 # HistoryCompaction hits before_model_request (the seam the
                 # deprecated `history_processors=` kwarg fed, removed in
-                # pydantic-ai v2).
+                # pydantic-ai v2). ToolCallJsonRepair rides the
+                # before_tool_validate seam (position inert relative to the
+                # history hooks).
                 capabilities=[
+                    *build_tool_call_json_repair(),
                     HistoryCompaction(agent_config),
                     build_model_message_transform(agent_name),
                     # Recursion guards ride the wrap_tool_execute seam so a