From bc93daa32b9c068b848ede898024ccfa8e55d76d Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Wed, 29 Apr 2026 08:56:17 +0800 Subject: [PATCH 01/15] =?UTF-8?q?=F0=9F=93=A6=20(adapters):=20Create=20emp?= =?UTF-8?q?ty=20pydantic=5Fai=20adapter=20package=20scaffold?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/pydantic_ai/__init__.py | 1 + 1 file changed, 1 insertion(+) create mode 100644 agent_assembly/adapters/pydantic_ai/__init__.py diff --git a/agent_assembly/adapters/pydantic_ai/__init__.py b/agent_assembly/adapters/pydantic_ai/__init__.py new file mode 100644 index 0000000..6e69bc9 --- /dev/null +++ b/agent_assembly/adapters/pydantic_ai/__init__.py @@ -0,0 +1 @@ +"""Pydantic AI adapter package.""" From b76077e5be564ae4d76da3c52bad4eabaec69ba7 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Wed, 29 Apr 2026 08:56:21 +0800 Subject: [PATCH 02/15] =?UTF-8?q?=F0=9F=93=A6=20(adapters):=20Add=20empty?= =?UTF-8?q?=20pydantic=5Fai=20patch=20module=20scaffold?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/pydantic_ai/patch.py | 1 + 1 file changed, 1 insertion(+) create mode 100644 agent_assembly/adapters/pydantic_ai/patch.py diff --git a/agent_assembly/adapters/pydantic_ai/patch.py b/agent_assembly/adapters/pydantic_ai/patch.py new file mode 100644 index 0000000..3e99afc --- /dev/null +++ b/agent_assembly/adapters/pydantic_ai/patch.py @@ -0,0 +1 @@ +"""Pydantic AI patch module.""" From 540cbad1934d7f7a7423b504c4cb01ce19e86f2d Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Wed, 29 Apr 2026 08:56:39 +0800 Subject: [PATCH 03/15] =?UTF-8?q?=E2=9C=A8=20(pydantic-ai):=20Add=20Pydant?= =?UTF-8?q?icAIPatch=20class=20skeleton=20and=20apply=20shell?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/pydantic_ai/patch.py | 44 ++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/agent_assembly/adapters/pydantic_ai/patch.py b/agent_assembly/adapters/pydantic_ai/patch.py index 3e99afc..b65433f 100644 --- a/agent_assembly/adapters/pydantic_ai/patch.py +++ b/agent_assembly/adapters/pydantic_ai/patch.py @@ -1 +1,45 @@ """Pydantic AI patch module.""" + +from __future__ import annotations + +from dataclasses import dataclass +import importlib +from typing import Any + +_TOOLS_PATCHED_FLAG = "_agent_assembly_pydantic_ai_tools_patched" +_ORIGINAL_TOOL_RUN = "_agent_assembly_original_pydantic_ai_tool_run" + + +@dataclass(slots=True) +class PydanticAIPatch: + """Applies Pydantic AI runtime monkey-patching hooks.""" + + callback_handler: Any + + def apply(self) -> bool: + """Apply patch wiring and return whether Pydantic AI is available.""" + tool_cls = _load_pydantic_ai_tool_class() + if tool_cls is None: + return False + _apply_tool_run_patch(tool_cls, self.callback_handler) + return True + + +def _load_pydantic_ai_tool_class() -> type[Any] | None: + try: + module = importlib.import_module("pydantic_ai.tools") + except ImportError: + return None + + tool_cls = getattr(module, "Tool", None) + if isinstance(tool_cls, type): + return tool_cls + return None + + +def _apply_tool_run_patch(tool_cls: type[Any], callback_handler: Any) -> None: + del callback_handler + if getattr(tool_cls, _TOOLS_PATCHED_FLAG, False): + return None + + setattr(tool_cls, _TOOLS_PATCHED_FLAG, True) From 845e5ab758bdd45ba15a11012f9bc7d10e9303e3 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Wed, 29 Apr 2026 08:56:49 +0800 Subject: [PATCH 04/15] =?UTF-8?q?=E2=9C=A8=20(exceptions):=20Add=20PolicyV?= =?UTF-8?q?iolationError=20for=20governance-denied=20tool=20calls?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/exceptions/__init__.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/agent_assembly/exceptions/__init__.py b/agent_assembly/exceptions/__init__.py index a3ab7f4..91b91d1 100644 --- a/agent_assembly/exceptions/__init__.py +++ b/agent_assembly/exceptions/__init__.py @@ -10,6 +10,7 @@ "ConfigurationError", "AdapterValidationError", "ToolExecutionBlockedError", + "PolicyViolationError", ] @@ -46,3 +47,8 @@ class AdapterValidationError(AssemblyError): class ToolExecutionBlockedError(AssemblyError): """Exception raised when a tool run is blocked by governance.""" pass + + +class PolicyViolationError(ToolExecutionBlockedError): + """Exception raised when policy blocks tool execution.""" + pass From a83443b51eadd85f8ccbf772ccda7bedd2da270f Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Wed, 29 Apr 2026 08:57:42 +0800 Subject: [PATCH 05/15] =?UTF-8?q?=E2=9C=A8=20(pydantic-ai):=20Patch=20Tool?= =?UTF-8?q?.=5Frun=20with=20async=20governance=20and=20audit=20flow?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/pydantic_ai/patch.py | 254 ++++++++++++++++++- 1 file changed, 251 insertions(+), 3 deletions(-) diff --git a/agent_assembly/adapters/pydantic_ai/patch.py b/agent_assembly/adapters/pydantic_ai/patch.py index b65433f..8bafe74 100644 --- a/agent_assembly/adapters/pydantic_ai/patch.py +++ b/agent_assembly/adapters/pydantic_ai/patch.py @@ -3,11 +3,16 @@ from __future__ import annotations from dataclasses import dataclass +from functools import wraps import importlib -from typing import Any +import inspect +from typing import Any, Literal, Mapping -_TOOLS_PATCHED_FLAG = "_agent_assembly_pydantic_ai_tools_patched" _ORIGINAL_TOOL_RUN = "_agent_assembly_original_pydantic_ai_tool_run" +_TOOLS_PATCHED_FLAG = "_agent_assembly_pydantic_ai_tools_patched" +_PROCESS_AGENT_ID: str | None = None +_DEFAULT_PENDING_APPROVAL_TIMEOUT_SECONDS = 300 +_MAX_AUDIT_RESULT_CHARS = 2000 @dataclass(slots=True) @@ -38,8 +43,251 @@ def _load_pydantic_ai_tool_class() -> type[Any] | None: def _apply_tool_run_patch(tool_cls: type[Any], callback_handler: Any) -> None: - del callback_handler if getattr(tool_cls, _TOOLS_PATCHED_FLAG, False): return None + original_run = tool_cls._run + + @wraps(original_run) + async def patched_run(self: Any, ctx: Any, args: Any, **kwargs: Any) -> Any: + tool_name = str(getattr(self, "name", self.__class__.__name__)) + tool_args = _serialize_tool_args(args) + agent_id = _resolve_agent_id(ctx) + run_id = _resolve_run_id(ctx) + + decision = await _invoke_async_tool_check( + callback_handler, + tool_name=tool_name, + tool_args=tool_args, + agent_id=agent_id, + run_id=run_id, + ) + status, reason = _normalize_decision(decision) + is_pending_flow = False + if status == "pending": + is_pending_flow = True + timeout_seconds = _get_pending_tool_approval_timeout_seconds(callback_handler) + final_decision = await _wait_for_async_tool_approval( + callback_handler, + tool_name=tool_name, + timeout_seconds=timeout_seconds, + tool_args=tool_args, + agent_id=agent_id, + run_id=run_id, + ) + status, reason = _normalize_decision(final_decision) + + if status == "deny": + if is_pending_flow: + raise _build_pending_rejected_error(tool_name, reason) + raise _build_denied_error(tool_name, reason) + + result = original_run(self, ctx, args, **kwargs) + if inspect.isawaitable(result): + result = await result + + await _record_async_tool_result( + callback_handler, + tool_name=tool_name, + result=result, + agent_id=agent_id, + run_id=run_id, + ) + return result + + setattr(tool_cls, _ORIGINAL_TOOL_RUN, original_run) + setattr(tool_cls, "_run", patched_run) setattr(tool_cls, _TOOLS_PATCHED_FLAG, True) + + +def set_process_agent_id(agent_id: str | None) -> None: + global _PROCESS_AGENT_ID + _PROCESS_AGENT_ID = agent_id + + +def _get_process_agent_id() -> str | None: + if isinstance(_PROCESS_AGENT_ID, str) and _PROCESS_AGENT_ID: + return _PROCESS_AGENT_ID + return None + + +def _resolve_agent_id(ctx: Any) -> str | None: + deps = getattr(ctx, "deps", None) + candidate = getattr(deps, "assembly_agent_id", None) + if isinstance(candidate, str) and candidate: + return candidate + return _get_process_agent_id() + + +def _resolve_run_id(ctx: Any) -> str | None: + run_id = getattr(ctx, "run_id", None) + if run_id is None: + return None + return str(run_id) + + +def _serialize_tool_args(args: Any) -> dict[str, Any]: + if hasattr(args, "model_dump"): + model_dump = getattr(args, "model_dump") + if callable(model_dump): + dumped = model_dump() + if isinstance(dumped, dict): + return dict(dumped) + + if isinstance(args, Mapping): + return dict(args) + + return {"value": str(args)} + + +def _normalize_decision( + decision: object, +) -> tuple[Literal["allow", "deny", "pending"], str | None]: + if isinstance(decision, str): + normalized = decision.strip().lower() + if normalized == "deny": + return "deny", None + if normalized == "pending": + return "pending", None + return "allow", None + + if isinstance(decision, Mapping): + raw_status = str(decision.get("status", "allow")).strip().lower() + if raw_status == "deny": + status: Literal["allow", "deny", "pending"] = "deny" + elif raw_status == "pending": + status = "pending" + else: + status = "allow" + + reason_value = decision.get("reason") + reason = str(reason_value) if reason_value is not None else None + return status, reason + + return "allow", None + + +async def _invoke_async_tool_check( + callback_handler: Any, + *, + tool_name: str, + tool_args: dict[str, Any], + agent_id: str | None, + run_id: str | None, +) -> object: + method = getattr(callback_handler, "check_tool_start", None) + if not callable(method): + return {"status": "allow"} + + result = method( + serialized={"name": tool_name}, + input_str=str(tool_args), + tool_name=tool_name, + args=tool_args, + agent_id=agent_id, + run_id=run_id, + ) + if inspect.isawaitable(result): + return await result + return result + + +async def _wait_for_async_tool_approval( + callback_handler: Any, + *, + tool_name: str, + timeout_seconds: int, + tool_args: dict[str, Any], + agent_id: str | None, + run_id: str | None, +) -> object: + method = getattr(callback_handler, "wait_for_tool_approval", None) + if not callable(method): + return {"status": "deny", "reason": "Approval handler is unavailable."} + + result = method( + serialized={"name": tool_name}, + input_str=str(tool_args), + tool_name=tool_name, + timeout_seconds=timeout_seconds, + args=tool_args, + agent_id=agent_id, + run_id=run_id, + ) + if inspect.isawaitable(result): + return await result + return result + + +def _get_pending_tool_approval_timeout_seconds(callback_handler: Any) -> int: + provider = getattr(callback_handler, "get_pending_tool_approval_timeout_seconds", None) + if callable(provider): + configured = provider() + else: + configured = getattr(callback_handler, "pending_tool_approval_timeout_seconds", None) + + if isinstance(configured, str): + stripped = configured.strip() + if stripped.isdigit(): + parsed = int(stripped) + if parsed > 0: + return parsed + return _DEFAULT_PENDING_APPROVAL_TIMEOUT_SECONDS + + if isinstance(configured, bool): + return _DEFAULT_PENDING_APPROVAL_TIMEOUT_SECONDS + + if isinstance(configured, int) and configured > 0: + return configured + + return _DEFAULT_PENDING_APPROVAL_TIMEOUT_SECONDS + + +def _truncate_result_for_audit(result: object) -> str: + return str(result)[:_MAX_AUDIT_RESULT_CHARS] + + +async def _record_async_tool_result( + callback_handler: Any, + *, + tool_name: str, + result: object, + agent_id: str | None, + run_id: str | None, +) -> None: + record_method = getattr(callback_handler, "record_result", None) + if callable(record_method): + recorded = record_method( + tool_name=tool_name, + result=_truncate_result_for_audit(result), + agent_id=agent_id, + run_id=run_id, + ) + if inspect.isawaitable(recorded): + await recorded + return None + + tool_end_method = getattr(callback_handler, "on_tool_end", None) + if callable(tool_end_method): + recorded = tool_end_method( + output=_truncate_result_for_audit(result), + tool_name=tool_name, + agent_id=agent_id, + run_id=run_id, + ) + if inspect.isawaitable(recorded): + await recorded + + +def _build_denied_error(tool_name: str, reason: str | None) -> Exception: + from agent_assembly.exceptions import PolicyViolationError + + reason_text = reason or "No reason provided." + return PolicyViolationError(f"Tool '{tool_name}' blocked by governance policy: {reason_text}") + + +def _build_pending_rejected_error(tool_name: str, reason: str | None) -> Exception: + from agent_assembly.exceptions import PolicyViolationError + + reason_text = reason or "No reason provided." + return PolicyViolationError(f"Tool '{tool_name}' rejected during approval: {reason_text}") From c9140d141d4bb8cdb194c53574c75fcab72dcdc1 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Wed, 29 Apr 2026 08:58:08 +0800 Subject: [PATCH 06/15] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20(runtime):=20Inject?= =?UTF-8?q?=20PydanticAIPatch=20and=20process-level=20agent=20fallback=20I?= =?UTF-8?q?D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/langchain/runtime.py | 12 +++++++++++- agent_assembly/core/assembly.py | 2 +- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/agent_assembly/adapters/langchain/runtime.py b/agent_assembly/adapters/langchain/runtime.py index aca6841..70ef52d 100644 --- a/agent_assembly/adapters/langchain/runtime.py +++ b/agent_assembly/adapters/langchain/runtime.py @@ -8,25 +8,35 @@ from agent_assembly.adapters.crewai.patch import CrewAIPatch from agent_assembly.adapters.langchain.callback_handler import AssemblyCallbackHandler from agent_assembly.adapters.langgraph import LangGraphPatch +from agent_assembly.adapters.pydantic_ai.patch import PydanticAIPatch, set_process_agent_id _ACTIVE_CALLBACK_HANDLER: AssemblyCallbackHandler | None = None _RUNTIME_LOCK = Lock() -def auto_inject_callback_handler(interceptor: Any) -> AssemblyCallbackHandler: +def auto_inject_callback_handler( + interceptor: Any, + *, + process_agent_id: str | None = None, +) -> AssemblyCallbackHandler: """Create and register the active callback handler instance.""" global _ACTIVE_CALLBACK_HANDLER with _RUNTIME_LOCK: + if process_agent_id is not None: + set_process_agent_id(process_agent_id) + if _ACTIVE_CALLBACK_HANDLER is not None: LangGraphPatch(_ACTIVE_CALLBACK_HANDLER).apply() CrewAIPatch(interceptor).apply() + PydanticAIPatch(interceptor).apply() return _ACTIVE_CALLBACK_HANDLER handler = AssemblyCallbackHandler(interceptor) _ACTIVE_CALLBACK_HANDLER = handler LangGraphPatch(handler).apply() CrewAIPatch(interceptor).apply() + PydanticAIPatch(interceptor).apply() return handler diff --git a/agent_assembly/core/assembly.py b/agent_assembly/core/assembly.py index be49337..c92e5f6 100644 --- a/agent_assembly/core/assembly.py +++ b/agent_assembly/core/assembly.py @@ -38,5 +38,5 @@ def init_assembly( agent_id=agent_id, api_key=api_key, ) - auto_inject_callback_handler(interceptor=object()) + auto_inject_callback_handler(interceptor=object(), process_agent_id=agent_id) return client From 8dddf620a7161b955a8d84dd2a7513a26d4aae59 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Wed, 29 Apr 2026 08:58:18 +0800 Subject: [PATCH 07/15] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20(adapters):=20Export?= =?UTF-8?q?=20PydanticAIPatch=20from=20pydantic=5Fai=20package?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/pydantic_ai/__init__.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/agent_assembly/adapters/pydantic_ai/__init__.py b/agent_assembly/adapters/pydantic_ai/__init__.py index 6e69bc9..328e92a 100644 --- a/agent_assembly/adapters/pydantic_ai/__init__.py +++ b/agent_assembly/adapters/pydantic_ai/__init__.py @@ -1 +1,5 @@ """Pydantic AI adapter package.""" + +from agent_assembly.adapters.pydantic_ai.patch import PydanticAIPatch + +__all__ = ["PydanticAIPatch"] From 4f47688a3534906b6ccfdc209fa19575905484a7 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Wed, 29 Apr 2026 09:00:28 +0800 Subject: [PATCH 08/15] =?UTF-8?q?=E2=9C=85=20(pydantic-ai):=20Add=20unit?= =?UTF-8?q?=20tests=20for=20deny,=20pending,=20idempotent,=20and=20truncat?= =?UTF-8?q?ion=20flows?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- test/unit/adapters/pydantic_ai/test_patch.py | 237 +++++++++++++++++++ 1 file changed, 237 insertions(+) create mode 100644 test/unit/adapters/pydantic_ai/test_patch.py diff --git a/test/unit/adapters/pydantic_ai/test_patch.py b/test/unit/adapters/pydantic_ai/test_patch.py new file mode 100644 index 0000000..ad572a3 --- /dev/null +++ b/test/unit/adapters/pydantic_ai/test_patch.py @@ -0,0 +1,237 @@ +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +from agent_assembly.adapters.pydantic_ai import patch as pydantic_ai_patch +from agent_assembly.exceptions import PolicyViolationError + + +class _RecordingInterceptor: + async def check_tool_start(self, **kwargs: object) -> dict[str, str]: + del kwargs + return {"status": "allow"} + + +class _ArgsModel: + def __init__(self, payload: dict[str, Any]) -> None: + self._payload = payload + + def model_dump(self) -> dict[str, Any]: + return dict(self._payload) + + +def _install_fake_pydantic_ai_modules( + monkeypatch: pytest.MonkeyPatch, +) -> type[Any]: + class FakeTool: + name = "fake_tool" + + async def _run(self, ctx: Any, args: Any, **kwargs: Any) -> dict[str, object]: + return { + "ctx": ctx, + "args": args, + "kwargs": kwargs, + } + + fake_pydantic_ai_tools = SimpleNamespace(Tool=FakeTool) + + def fake_import_module(module_name: str) -> object: + if module_name == "pydantic_ai.tools": + return fake_pydantic_ai_tools + raise ImportError(module_name) + + monkeypatch.setattr(pydantic_ai_patch.importlib, "import_module", fake_import_module) + return FakeTool + + +@pytest.mark.asyncio +async def test_apply_patches_tool_run_and_is_idempotent(monkeypatch: pytest.MonkeyPatch) -> None: + FakeTool = _install_fake_pydantic_ai_modules(monkeypatch) + + patcher = pydantic_ai_patch.PydanticAIPatch(_RecordingInterceptor()) + assert patcher.apply() is True + first_run_ref = FakeTool._run + + assert getattr(FakeTool, pydantic_ai_patch._TOOLS_PATCHED_FLAG, False) is True + + assert patcher.apply() is True + assert FakeTool._run is first_run_ref + + +def test_loader_edge_cases_and_apply_false_without_tool(monkeypatch: pytest.MonkeyPatch) -> None: + def raise_import_error(module_name: str) -> object: + raise ImportError(module_name) + + monkeypatch.setattr(pydantic_ai_patch.importlib, "import_module", raise_import_error) + assert pydantic_ai_patch._load_pydantic_ai_tool_class() is None + assert pydantic_ai_patch.PydanticAIPatch(_RecordingInterceptor()).apply() is False + + fake_pydantic_ai_tools = SimpleNamespace(Tool=object()) + + def return_non_type(module_name: str) -> object: + if module_name == "pydantic_ai.tools": + return fake_pydantic_ai_tools + raise ImportError(module_name) + + monkeypatch.setattr(pydantic_ai_patch.importlib, "import_module", return_non_type) + assert pydantic_ai_patch._load_pydantic_ai_tool_class() is None + + +def test_helper_branches_for_agent_id_timeout_and_serialization() -> None: + class TimeoutProvider: + def get_pending_tool_approval_timeout_seconds(self) -> str: + return "42" + + assert pydantic_ai_patch._get_pending_tool_approval_timeout_seconds(TimeoutProvider()) == 42 + assert pydantic_ai_patch._get_pending_tool_approval_timeout_seconds( + SimpleNamespace(pending_tool_approval_timeout_seconds=0) + ) == 300 + assert pydantic_ai_patch._get_pending_tool_approval_timeout_seconds( + SimpleNamespace(pending_tool_approval_timeout_seconds=True) + ) == 300 + + assert pydantic_ai_patch._normalize_decision("deny") == ("deny", None) + assert pydantic_ai_patch._normalize_decision("pending") == ("pending", None) + assert pydantic_ai_patch._normalize_decision("allow") == ("allow", None) + assert pydantic_ai_patch._normalize_decision(12345) == ("allow", None) + + pydantic_ai_patch.set_process_agent_id("process-agent") + ctx_with_deps = SimpleNamespace(deps=SimpleNamespace(assembly_agent_id="deps-agent"), run_id=123) + ctx_without_deps = SimpleNamespace(deps=SimpleNamespace(), run_id=None) + + assert pydantic_ai_patch._resolve_agent_id(ctx_with_deps) == "deps-agent" + assert pydantic_ai_patch._resolve_agent_id(ctx_without_deps) == "process-agent" + assert pydantic_ai_patch._resolve_run_id(ctx_with_deps) == "123" + assert pydantic_ai_patch._resolve_run_id(ctx_without_deps) is None + + model_args = _ArgsModel({"x": 1, "y": 2}) + assert pydantic_ai_patch._serialize_tool_args(model_args) == {"x": 1, "y": 2} + assert pydantic_ai_patch._serialize_tool_args({"a": 1}) == {"a": 1} + assert pydantic_ai_patch._serialize_tool_args(99) == {"value": "99"} + + pydantic_ai_patch.set_process_agent_id(None) + + +@pytest.mark.asyncio +async def test_denied_tool_raises_policy_violation_error(monkeypatch: pytest.MonkeyPatch) -> None: + FakeTool = _install_fake_pydantic_ai_modules(monkeypatch) + + class BlockInterceptor: + async def check_tool_start(self, **kwargs: object) -> dict[str, str]: + del kwargs + return {"status": "deny", "reason": "blocked for safety"} + + patcher = pydantic_ai_patch.PydanticAIPatch(BlockInterceptor()) + assert patcher.apply() is True + + tool = FakeTool() + ctx = SimpleNamespace(deps=SimpleNamespace(assembly_agent_id="agent-a"), run_id="run-1") + + with pytest.raises(PolicyViolationError, match="blocked by governance policy: blocked for safety"): + await tool._run(ctx, _ArgsModel({"topic": "finance"})) + + +@pytest.mark.asyncio +async def test_pending_then_approved_runs_and_records_result( + monkeypatch: pytest.MonkeyPatch, +) -> None: + FakeTool = _install_fake_pydantic_ai_modules(monkeypatch) + wait_calls: list[dict[str, object]] = [] + recorded_results: list[dict[str, object]] = [] + + class PendingThenApproveInterceptor: + pending_tool_approval_timeout_seconds = 25 + + async def check_tool_start(self, **kwargs: object) -> dict[str, str]: + del kwargs + return {"status": "pending", "reason": "needs approval"} + + async def wait_for_tool_approval(self, **kwargs: object) -> dict[str, str]: + wait_calls.append(dict(kwargs)) + return {"status": "allow"} + + async def record_result(self, **kwargs: object) -> None: + recorded_results.append(dict(kwargs)) + + patcher = pydantic_ai_patch.PydanticAIPatch(PendingThenApproveInterceptor()) + assert patcher.apply() is True + + tool = FakeTool() + ctx = SimpleNamespace(deps=SimpleNamespace(assembly_agent_id="agent-b"), run_id="run-2") + result = await tool._run(ctx, _ArgsModel({"q": "hello"}), trace="yes") + + assert result["kwargs"] == {"trace": "yes"} + assert len(wait_calls) == 1 + assert wait_calls[0]["timeout_seconds"] == 25 + assert len(recorded_results) == 1 + assert recorded_results[0]["tool_name"] == "fake_tool" + assert recorded_results[0]["agent_id"] == "agent-b" + + +@pytest.mark.asyncio +async def test_pending_then_rejected_raises_policy_violation_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + FakeTool = _install_fake_pydantic_ai_modules(monkeypatch) + + class PendingThenRejectInterceptor: + async def check_tool_start(self, **kwargs: object) -> dict[str, str]: + del kwargs + return {"status": "pending", "reason": "requires approval"} + + async def wait_for_tool_approval(self, **kwargs: object) -> dict[str, str]: + del kwargs + return {"status": "deny", "reason": "approval rejected"} + + patcher = pydantic_ai_patch.PydanticAIPatch(PendingThenRejectInterceptor()) + assert patcher.apply() is True + + tool = FakeTool() + ctx = SimpleNamespace(deps=SimpleNamespace(assembly_agent_id="agent-c"), run_id="run-3") + + with pytest.raises(PolicyViolationError, match="rejected during approval: approval rejected"): + await tool._run(ctx, _ArgsModel({"q": "secret"})) + + +@pytest.mark.asyncio +async def test_result_recording_truncates_to_2000_chars(monkeypatch: pytest.MonkeyPatch) -> None: + class FakeTool: + name = "truncate_tool" + + async def _run(self, ctx: Any, args: Any, **kwargs: Any) -> str: + del ctx, args, kwargs + return "x" * 2500 + + fake_pydantic_ai_tools = SimpleNamespace(Tool=FakeTool) + + def fake_import_module(module_name: str) -> object: + if module_name == "pydantic_ai.tools": + return fake_pydantic_ai_tools + raise ImportError(module_name) + + monkeypatch.setattr(pydantic_ai_patch.importlib, "import_module", fake_import_module) + + observed: list[str] = [] + + class Interceptor: + async def check_tool_start(self, **kwargs: object) -> dict[str, str]: + del kwargs + return {"status": "allow"} + + async def record_result(self, **kwargs: object) -> None: + observed.append(str(kwargs["result"])) + + patcher = pydantic_ai_patch.PydanticAIPatch(Interceptor()) + assert patcher.apply() is True + + tool = FakeTool() + ctx = SimpleNamespace(deps=SimpleNamespace(assembly_agent_id="agent-d"), run_id="run-4") + result = await tool._run(ctx, _ArgsModel({"q": "long"})) + + assert isinstance(result, str) + assert len(result) == 2500 + assert len(observed) == 1 + assert len(observed[0]) == 2000 From c903b8e78033f778253f07b27d1efe320c3404de Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Wed, 29 Apr 2026 09:00:32 +0800 Subject: [PATCH 09/15] =?UTF-8?q?=E2=9C=85=20(integration):=20Add=20Pydant?= =?UTF-8?q?ic=20AI=20blocked-mid-run=20two-tool=20interception=20tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ...st_pydantic_ai_interception_integration.py | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) create mode 100644 test/integration/test_pydantic_ai_interception_integration.py diff --git a/test/integration/test_pydantic_ai_interception_integration.py b/test/integration/test_pydantic_ai_interception_integration.py new file mode 100644 index 0000000..674a7ab --- /dev/null +++ b/test/integration/test_pydantic_ai_interception_integration.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +from agent_assembly.adapters.pydantic_ai import patch as pydantic_ai_patch +from agent_assembly.exceptions import PolicyViolationError + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_pydantic_ai_two_tool_flow_continues_after_blocked_tool( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class FakeTool: + def __init__(self, name: str, result: str) -> None: + self.name = name + self._result = result + + async def _run(self, ctx: Any, args: Any, **kwargs: Any) -> str: + del ctx, args, kwargs + return self._result + + fake_pydantic_ai_tools = SimpleNamespace(Tool=FakeTool) + + def fake_import_module(module_name: str) -> object: + if module_name == "pydantic_ai.tools": + return fake_pydantic_ai_tools + raise ImportError(module_name) + + monkeypatch.setattr(pydantic_ai_patch.importlib, "import_module", fake_import_module) + + class Interceptor: + async def check_tool_start(self, **kwargs: object) -> dict[str, str]: + if kwargs.get("tool_name") == "blocked_tool": + return {"status": "deny", "reason": "blocked by policy"} + return {"status": "allow"} + + patcher = pydantic_ai_patch.PydanticAIPatch(Interceptor()) + assert patcher.apply() is True + + ctx = SimpleNamespace(deps=SimpleNamespace(assembly_agent_id="agent-1"), run_id="run-main") + blocked_tool = FakeTool("blocked_tool", "should-not-run") + safe_tool = FakeTool("safe_tool", "ok:safe_tool") + + with pytest.raises(PolicyViolationError, match="blocked by governance policy: blocked by policy"): + await blocked_tool._run(ctx, {"step": 1}) + + safe_result = await safe_tool._run(ctx, {"step": 2}) + assert safe_result == "ok:safe_tool" + + +@pytest.mark.integration +def test_pydantic_ai_real_tool_class_patch_path_when_available() -> None: + pydantic_ai_tools = pytest.importorskip("pydantic_ai.tools") + tool_cls = pydantic_ai_tools.Tool + + class Interceptor: + async def check_tool_start(self, **kwargs: object) -> dict[str, str]: + del kwargs + return {"status": "allow"} + + patcher = pydantic_ai_patch.PydanticAIPatch(Interceptor()) + assert patcher.apply() is True + assert getattr(tool_cls, pydantic_ai_patch._TOOLS_PATCHED_FLAG, False) is True From 2439c0ab2ec69351d31301a5329f37a9c12cea3a Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Wed, 29 Apr 2026 09:01:17 +0800 Subject: [PATCH 10/15] =?UTF-8?q?=E2=9C=A8=20(pydantic-ai):=20Add=20option?= =?UTF-8?q?al=20AssemblyModelWrapper=20for=20model=20request=20interceptio?= =?UTF-8?q?n?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/pydantic_ai/patch.py | 27 ++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/agent_assembly/adapters/pydantic_ai/patch.py b/agent_assembly/adapters/pydantic_ai/patch.py index 8bafe74..a0010a1 100644 --- a/agent_assembly/adapters/pydantic_ai/patch.py +++ b/agent_assembly/adapters/pydantic_ai/patch.py @@ -30,6 +30,33 @@ def apply(self) -> bool: return True +class AssemblyModelWrapper: + """Optional model wrapper for LLM input scan-forward interception.""" + + def __init__(self, model: Any, callback_handler: Any) -> None: + self._model = model + self._callback_handler = callback_handler + + async def request(self, *args: Any, **kwargs: Any) -> Any: + scan_method = getattr(self._callback_handler, "on_llm_start_scan", None) + if callable(scan_method): + scan_result = scan_method( + serialized={"name": self._model.__class__.__name__}, + prompts=[str(args[0])] if args else [], + run_id=kwargs.get("run_id"), + ) + if inspect.isawaitable(scan_result): + await scan_result + + result = self._model.request(*args, **kwargs) + if inspect.isawaitable(result): + return await result + return result + + def __getattr__(self, name: str) -> Any: + return getattr(self._model, name) + + def _load_pydantic_ai_tool_class() -> type[Any] | None: try: module = importlib.import_module("pydantic_ai.tools") From 32e8aa83cb8bd33c90aabf29c95b7ca42bf3710f Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Wed, 29 Apr 2026 09:01:21 +0800 Subject: [PATCH 11/15] =?UTF-8?q?=E2=9C=85=20(pydantic-ai):=20Add=20unit?= =?UTF-8?q?=20tests=20for=20AssemblyModelWrapper=20scan=20and=20passthroug?= =?UTF-8?q?h?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- test/unit/adapters/pydantic_ai/test_patch.py | 37 ++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/test/unit/adapters/pydantic_ai/test_patch.py b/test/unit/adapters/pydantic_ai/test_patch.py index ad572a3..73a1784 100644 --- a/test/unit/adapters/pydantic_ai/test_patch.py +++ b/test/unit/adapters/pydantic_ai/test_patch.py @@ -235,3 +235,40 @@ async def record_result(self, **kwargs: object) -> None: assert len(result) == 2500 assert len(observed) == 1 assert len(observed[0]) == 2000 + + +@pytest.mark.asyncio +async def test_assembly_model_wrapper_scans_and_forwards_request() -> None: + class FakeModel: + async def request(self, *args: object, **kwargs: object) -> str: + del kwargs + return f"forwarded:{args[0]}" + + scanned_prompts: list[list[str]] = [] + + class Interceptor: + async def on_llm_start_scan(self, **kwargs: object) -> None: + prompts = kwargs.get("prompts", []) + if isinstance(prompts, list): + scanned_prompts.append([str(item) for item in prompts]) + + wrapper = pydantic_ai_patch.AssemblyModelWrapper(FakeModel(), Interceptor()) + result = await wrapper.request("hello") + + assert result == "forwarded:hello" + assert scanned_prompts == [["hello"]] + + +@pytest.mark.asyncio +async def test_assembly_model_wrapper_passthrough_attrs() -> None: + class FakeModel: + model_name = "demo-model" + + def request(self, *args: object, **kwargs: object) -> str: + del args, kwargs + return "ok" + + wrapper = pydantic_ai_patch.AssemblyModelWrapper(FakeModel(), object()) + assert wrapper.model_name == "demo-model" + result = await wrapper.request("ignored") + assert result == "ok" From 6b26f420d6124558c7822f1153712740e12acfde Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Wed, 29 Apr 2026 09:01:27 +0800 Subject: [PATCH 12/15] =?UTF-8?q?=F0=9F=93=9D=20(docs):=20Add=20Pydantic?= =?UTF-8?q?=20AI=20runtime=20interception=20behavior=20notes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/contents/document/api-references/index.mdx | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/docs/contents/document/api-references/index.mdx b/docs/contents/document/api-references/index.mdx index e757ff1..a996844 100644 --- a/docs/contents/document/api-references/index.mdx +++ b/docs/contents/document/api-references/index.mdx @@ -81,3 +81,13 @@ from agent_assembly.exceptions import ( - Blocked tool calls return policy message strings (instead of raising exceptions). - `pending` approval decisions block synchronously and return denial strings if not approved. - `Task.execute_sync()` emits task start/complete audit events. + +## Pydantic AI runtime interception + +`init_assembly(...)` also applies Pydantic AI runtime patches when Pydantic AI is installed. + +- `Tool._run(...)` is patched with async governance checks. +- `deny` decisions raise `PolicyViolationError` with tool name and reason. +- `pending` decisions await approval and raise on rejection. +- Successful results are recorded with audit payload truncation at 2000 chars. +- Agent identity resolves from `ctx.deps.assembly_agent_id` with process-level fallback. From e825a9fbd7441ee2091e09aaeaf863927ea7693b Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Wed, 29 Apr 2026 09:05:02 +0800 Subject: [PATCH 13/15] =?UTF-8?q?=E2=9C=85=20(pydantic-ai):=20Cover=20asyn?= =?UTF-8?q?c=20helper=20fallback=20branches=20for=20patch=20coverage?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- test/unit/adapters/pydantic_ai/test_patch.py | 75 ++++++++++++++++++++ 1 file changed, 75 insertions(+) diff --git a/test/unit/adapters/pydantic_ai/test_patch.py b/test/unit/adapters/pydantic_ai/test_patch.py index 73a1784..58a4edd 100644 --- a/test/unit/adapters/pydantic_ai/test_patch.py +++ b/test/unit/adapters/pydantic_ai/test_patch.py @@ -272,3 +272,78 @@ def request(self, *args: object, **kwargs: object) -> str: assert wrapper.model_name == "demo-model" result = await wrapper.request("ignored") assert result == "ok" + + +@pytest.mark.asyncio +async def test_fallback_and_non_awaitable_branches_for_async_helpers() -> None: + class NoHandlers: + pass + + pydantic_ai_patch.set_process_agent_id(None) + assert pydantic_ai_patch._get_process_agent_id() is None + + fallback_check = await pydantic_ai_patch._invoke_async_tool_check( + NoHandlers(), + tool_name="x", + tool_args={}, + agent_id=None, + run_id=None, + ) + assert fallback_check == {"status": "allow"} + + fallback_wait = await pydantic_ai_patch._wait_for_async_tool_approval( + NoHandlers(), + tool_name="x", + timeout_seconds=1, + tool_args={}, + agent_id=None, + run_id=None, + ) + assert fallback_wait == {"status": "deny", "reason": "Approval handler is unavailable."} + + assert pydantic_ai_patch._get_pending_tool_approval_timeout_seconds( + SimpleNamespace(pending_tool_approval_timeout_seconds="NaN") + ) == 300 + + class SyncInterceptor: + def check_tool_start(self, **kwargs: object) -> dict[str, str]: + del kwargs + return {"status": "allow"} + + def wait_for_tool_approval(self, **kwargs: object) -> dict[str, str]: + del kwargs + return {"status": "allow"} + + sync_check = await pydantic_ai_patch._invoke_async_tool_check( + SyncInterceptor(), + tool_name="x", + tool_args={}, + agent_id=None, + run_id=None, + ) + assert sync_check == {"status": "allow"} + + sync_wait = await pydantic_ai_patch._wait_for_async_tool_approval( + SyncInterceptor(), + tool_name="x", + timeout_seconds=2, + tool_args={}, + agent_id=None, + run_id=None, + ) + assert sync_wait == {"status": "allow"} + + observed_outputs: list[str] = [] + + class ToolEndOnlyInterceptor: + async def on_tool_end(self, **kwargs: object) -> None: + observed_outputs.append(str(kwargs["output"])) + + await pydantic_ai_patch._record_async_tool_result( + ToolEndOnlyInterceptor(), + tool_name="fallback", + result="result-value", + agent_id="agent-z", + run_id="run-z", + ) + assert observed_outputs == ["result-value"] From f6119d51da28d456d210e8692fa2d4d49c5321b0 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Wed, 29 Apr 2026 09:05:52 +0800 Subject: [PATCH 14/15] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20(tests):=20Rename=20?= =?UTF-8?q?pydantic=5Fai=20patch=20test=20module=20to=20avoid=20collect=20?= =?UTF-8?q?collision?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../pydantic_ai/{test_patch.py => test_pydantic_ai_patch.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename test/unit/adapters/pydantic_ai/{test_patch.py => test_pydantic_ai_patch.py} (100%) diff --git a/test/unit/adapters/pydantic_ai/test_patch.py b/test/unit/adapters/pydantic_ai/test_pydantic_ai_patch.py similarity index 100% rename from test/unit/adapters/pydantic_ai/test_patch.py rename to test/unit/adapters/pydantic_ai/test_pydantic_ai_patch.py From 30589ea4910f25a157fa59992b41727fe649a008 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Wed, 29 Apr 2026 09:18:57 +0800 Subject: [PATCH 15/15] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20(pydantic-ai):=20Reu?= =?UTF-8?q?se=20governance=20helpers=20to=20reduce=20duplication=20on=20ne?= =?UTF-8?q?w=20code?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/pydantic_ai/patch.py | 51 +++----------------- 1 file changed, 7 insertions(+), 44 deletions(-) diff --git a/agent_assembly/adapters/pydantic_ai/patch.py b/agent_assembly/adapters/pydantic_ai/patch.py index a0010a1..79139d4 100644 --- a/agent_assembly/adapters/pydantic_ai/patch.py +++ b/agent_assembly/adapters/pydantic_ai/patch.py @@ -8,10 +8,14 @@ import inspect from typing import Any, Literal, Mapping +from agent_assembly.adapters.crewai.patch import ( + _get_pending_tool_approval_timeout_seconds as _resolve_pending_timeout_seconds, +) +from agent_assembly.adapters.crewai.patch import _normalize_decision as _normalize_governance_decision + _ORIGINAL_TOOL_RUN = "_agent_assembly_original_pydantic_ai_tool_run" _TOOLS_PATCHED_FLAG = "_agent_assembly_pydantic_ai_tools_patched" _PROCESS_AGENT_ID: str | None = None -_DEFAULT_PENDING_APPROVAL_TIMEOUT_SECONDS = 300 _MAX_AUDIT_RESULT_CHARS = 2000 @@ -170,28 +174,7 @@ def _serialize_tool_args(args: Any) -> dict[str, Any]: def _normalize_decision( decision: object, ) -> tuple[Literal["allow", "deny", "pending"], str | None]: - if isinstance(decision, str): - normalized = decision.strip().lower() - if normalized == "deny": - return "deny", None - if normalized == "pending": - return "pending", None - return "allow", None - - if isinstance(decision, Mapping): - raw_status = str(decision.get("status", "allow")).strip().lower() - if raw_status == "deny": - status: Literal["allow", "deny", "pending"] = "deny" - elif raw_status == "pending": - status = "pending" - else: - status = "allow" - - reason_value = decision.get("reason") - reason = str(reason_value) if reason_value is not None else None - return status, reason - - return "allow", None + return _normalize_governance_decision(decision) async def _invoke_async_tool_check( @@ -247,27 +230,7 @@ async def _wait_for_async_tool_approval( def _get_pending_tool_approval_timeout_seconds(callback_handler: Any) -> int: - provider = getattr(callback_handler, "get_pending_tool_approval_timeout_seconds", None) - if callable(provider): - configured = provider() - else: - configured = getattr(callback_handler, "pending_tool_approval_timeout_seconds", None) - - if isinstance(configured, str): - stripped = configured.strip() - if stripped.isdigit(): - parsed = int(stripped) - if parsed > 0: - return parsed - return _DEFAULT_PENDING_APPROVAL_TIMEOUT_SECONDS - - if isinstance(configured, bool): - return _DEFAULT_PENDING_APPROVAL_TIMEOUT_SECONDS - - if isinstance(configured, int) and configured > 0: - return configured - - return _DEFAULT_PENDING_APPROVAL_TIMEOUT_SECONDS + return _resolve_pending_timeout_seconds(callback_handler) def _truncate_result_for_audit(result: object) -> str: