From a9a8afc9f20a37b863f551a3be7e926127f22b88 Mon Sep 17 00:00:00 2001 From: Bryant Date: Fri, 17 Jul 2026 13:54:55 +0800 Subject: [PATCH 1/3] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20(adapters):=20Add=20?= =?UTF-8?q?=5Fshared=20async=20tool-governance=20helper=20module?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extract the async tool-governance flow (arg serialization, verdict normalization, pending-approval round-trip, deny-by-raise, spawn-scoped original call, result recording) that Google ADK and Pydantic AI adapters duplicated verbatim into one framework-agnostic helper. Preserves the AAASM-4734 fail-closed-under-enforce contract by threading enforce and reusing the CrewAI _missing_interceptor_decision fallback. Refs AAASM-4746 --- agent_assembly/adapters/_shared/__init__.py | 6 + .../adapters/_shared/tool_governance.py | 244 ++++++++++++++++++ 2 files changed, 250 insertions(+) create mode 100644 agent_assembly/adapters/_shared/__init__.py create mode 100644 agent_assembly/adapters/_shared/tool_governance.py diff --git a/agent_assembly/adapters/_shared/__init__.py b/agent_assembly/adapters/_shared/__init__.py new file mode 100644 index 00000000..7d63bc80 --- /dev/null +++ b/agent_assembly/adapters/_shared/__init__.py @@ -0,0 +1,6 @@ +"""Internal helpers shared across framework adapters. + +This package holds framework-agnostic governance plumbing that would otherwise be +copy-pasted verbatim into each adapter's ``patch`` module. It is private (leading +underscore): adapters import from it, but it is not part of the public SDK API. +""" diff --git a/agent_assembly/adapters/_shared/tool_governance.py b/agent_assembly/adapters/_shared/tool_governance.py new file mode 100644 index 00000000..54e9dae0 --- /dev/null +++ b/agent_assembly/adapters/_shared/tool_governance.py @@ -0,0 +1,244 @@ +"""Shared async tool-governance flow for framework adapters. + +The Google ADK and Pydantic AI adapters intercept tool execution through +framework-specific hook points, but the governance logic they run once a tool +call is intercepted is identical: serialize the args, ask the interceptor for a +verdict, honour a ``pending`` approval round-trip, deny by raising when the +verdict is ``deny``, otherwise run the original inside a spawn-context scope and +record the result. That shared body — previously duplicated verbatim in both +adapters (the cross-file duplication SonarCloud flagged on PR #269, AAASM-4746) — +lives here so each adapter keeps only its framework-specific glue. + +The fail-closed-under-enforce contract from AAASM-4734 is preserved exactly: the +``enforce`` flag is threaded through ``_normalize_decision`` and a +``check_tool_start``-less interceptor falls back to +``_missing_interceptor_decision`` (both re-used from the CrewAI leaf helpers), so +an unrecognized verdict or a missing ``check_tool_start`` still denies under +enforce. +""" + +from __future__ import annotations + +import inspect +from collections.abc import Callable, Mapping +from typing import TYPE_CHECKING, Any, Literal + +if TYPE_CHECKING: + from agent_assembly.exceptions import PolicyViolationError + +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 ( + _missing_interceptor_decision, +) +from agent_assembly.adapters.crewai.patch import ( + _normalize_decision as _normalize_governance_decision, +) +from agent_assembly.core.spawn import _SPAWN_CTX, SpawnContext, spawn_context_scope + +_MAX_AUDIT_RESULT_CHARS = 2000 + + +def _current_spawn_depth() -> int: + current = _SPAWN_CTX.get() + return (current.depth + 1) if current is not None else 1 + + +def _serialize_tool_args(args: Any) -> dict[str, Any]: + if hasattr(args, "model_dump"): + model_dump = 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, + *, + enforce: bool = False, +) -> tuple[Literal["allow", "deny", "pending"], str | None]: + return _normalize_governance_decision(decision, enforce=enforce) + + +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 _missing_interceptor_decision(callback_handler) + + 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: + return _resolve_pending_timeout_seconds(callback_handler) + + +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) -> PolicyViolationError: + 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) -> PolicyViolationError: + from agent_assembly.exceptions import PolicyViolationError + + reason_text = reason or "No reason provided." + return PolicyViolationError(f"Tool '{tool_name}' rejected during approval: {reason_text}") + + +async def run_governed_async_tool( + callback_handler: Any, + *, + enforce: bool, + tool_name: str, + tool_args: dict[str, Any], + agent_id: str | None, + run_id: str | None, + invoke_original: Callable[[], Any], +) -> Any: + """Run one intercepted async tool call through the governance flow. + + Applies the pre-execution check (with a ``pending`` approval round-trip), + raises on ``deny``, then runs ``invoke_original`` inside a spawn-context scope + and records the result. ``invoke_original`` is a zero-argument callable that + the adapter supplies to call the framework's original method with its own + signature; it may return an awaitable, which is awaited here. + + Raises: + PolicyViolationError: When the (final) verdict is ``deny`` — a + pending-rejection message when the deny followed an approval + round-trip, otherwise the blocked-by-policy message. + """ + 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, enforce=enforce) + 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, enforce=enforce) + + if status == "deny": + if is_pending_flow: + raise _build_pending_rejected_error(tool_name, reason) + raise _build_denied_error(tool_name, reason) + + spawn_ctx = SpawnContext( + parent_agent_id=agent_id or "", + depth=_current_spawn_depth(), + spawned_by_tool=tool_name, + delegation_reason=f"tool:{tool_name}", + ) + with spawn_context_scope(spawn_ctx): + result = invoke_original() + 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 From 017cb7d276ffa5d1033f15c79bfa0849a174513b Mon Sep 17 00:00:00 2001 From: Bryant Date: Fri, 17 Jul 2026 13:56:34 +0800 Subject: [PATCH 2/3] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20(google=5Fadk):=20Rout?= =?UTF-8?q?e=20tool=20governance=20through=20=5Fshared=20helper?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the adapter's local copies of the async tool-governance flow and its leaf helpers with the shared _shared.tool_governance module; the patched run_async now supplies only the framework-specific invoke_original callable. Shared helpers are re-imported (and listed in __all__) so the adapter's unit tests still reach them through this module. Behavior unchanged. Refs AAASM-4746 --- agent_assembly/adapters/google_adk/patch.py | 225 +++----------------- 1 file changed, 34 insertions(+), 191 deletions(-) diff --git a/agent_assembly/adapters/google_adk/patch.py b/agent_assembly/adapters/google_adk/patch.py index e24b0700..724ed6fa 100644 --- a/agent_assembly/adapters/google_adk/patch.py +++ b/agent_assembly/adapters/google_adk/patch.py @@ -3,33 +3,43 @@ from __future__ import annotations import importlib as importlib -import inspect -from collections.abc import Mapping from dataclasses import dataclass from functools import wraps -from typing import TYPE_CHECKING, Any, Literal - -if TYPE_CHECKING: - from agent_assembly.exceptions import PolicyViolationError - -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 ( - _interceptor_enforces, - _missing_interceptor_decision, +from typing import Any + +from agent_assembly.adapters._shared.tool_governance import ( + _current_spawn_depth, + _get_pending_tool_approval_timeout_seconds, + _invoke_async_tool_check, + _normalize_decision, + _record_async_tool_result, + _serialize_tool_args, + _wait_for_async_tool_approval, + run_governed_async_tool, ) -from agent_assembly.adapters.crewai.patch import ( - _normalize_decision as _normalize_governance_decision, -) -from agent_assembly.core.spawn import _SPAWN_CTX, SpawnContext, spawn_context_scope +from agent_assembly.adapters.crewai.patch import _interceptor_enforces +from agent_assembly.core.spawn import SpawnContext, spawn_context_scope + +# The shared governance helpers are imported here (not just used internally) so +# the adapter's unit tests can reach them through this module — listing them in +# __all__ marks the re-export as intentional. See AAASM-4746 (dedup). +__all__ = [ + "GoogleADKPatch", + "_current_spawn_depth", + "_get_pending_tool_approval_timeout_seconds", + "_invoke_async_tool_check", + "_normalize_decision", + "_record_async_tool_result", + "_serialize_tool_args", + "_wait_for_async_tool_approval", + "set_process_agent_id", +] _ORIGINAL_TOOL_RUN_ASYNC = "_agent_assembly_original_google_adk_tool_run_async" _TOOLS_PATCHED_FLAG = "_agent_assembly_google_adk_tools_patched" _ORIGINAL_AGENT_RUN_ASYNC = "_agent_assembly_original_google_adk_agent_run_async" _AGENT_PATCHED_FLAG = "_agent_assembly_google_adk_agent_patched" _PROCESS_AGENT_ID: str | None = None -_MAX_AUDIT_RESULT_CHARS = 2000 @dataclass(slots=True) @@ -122,11 +132,6 @@ def _load_google_adk_base_agent_class() -> type[Any] | None: return None -def _current_spawn_depth() -> int: - current = _SPAWN_CTX.get() - return (current.depth + 1) if current is not None else 1 - - def _apply_agent_run_async_patch(agent_cls: type[Any], process_agent_id: str | None) -> None: if getattr(agent_cls, _AGENT_PATCHED_FLAG, False): return None @@ -181,52 +186,18 @@ async def patched_run_async(self: Any, *, args: Any, tool_context: Any, **kwargs agent_id = _resolve_agent_id(tool_context) run_id = _resolve_run_id(tool_context) - 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, enforce=enforce) - 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, enforce=enforce) - - if status == "deny": - if is_pending_flow: - raise _build_pending_rejected_error(tool_name, reason) - raise _build_denied_error(tool_name, reason) - - spawn_ctx = SpawnContext( - parent_agent_id=agent_id or "", - depth=_current_spawn_depth(), - spawned_by_tool=tool_name, - delegation_reason=f"tool:{tool_name}", - ) - with spawn_context_scope(spawn_ctx): - result = original_run_async(self, args=args, tool_context=tool_context, **kwargs) - if inspect.isawaitable(result): - result = await result + def _invoke_original() -> Any: + return original_run_async(self, args=args, tool_context=tool_context, **kwargs) - await _record_async_tool_result( + return await run_governed_async_tool( callback_handler, + enforce=enforce, tool_name=tool_name, - result=result, + tool_args=tool_args, agent_id=agent_id, run_id=run_id, + invoke_original=_invoke_original, ) - return result setattr(tool_cls, _ORIGINAL_TOOL_RUN_ASYNC, original_run_async) tool_cls.run_async = patched_run_async @@ -275,131 +246,3 @@ def _resolve_run_id(tool_context: Any) -> str | None: if invocation_id is None: return None return str(invocation_id) - - -def _serialize_tool_args(args: Any) -> dict[str, Any]: - if hasattr(args, "model_dump"): - model_dump = 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, - *, - enforce: bool = False, -) -> tuple[Literal["allow", "deny", "pending"], str | None]: - return _normalize_governance_decision(decision, enforce=enforce) - - -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 _missing_interceptor_decision(callback_handler) - - 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: - return _resolve_pending_timeout_seconds(callback_handler) - - -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) -> PolicyViolationError: - 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) -> PolicyViolationError: - 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 579fafe6ce7302620ad105daef8e7b52edcae72b Mon Sep 17 00:00:00 2001 From: Bryant Date: Fri, 17 Jul 2026 13:58:22 +0800 Subject: [PATCH 3/3] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20(pydantic=5Fai):=20Rou?= =?UTF-8?q?te=20tool=20governance=20through=20=5Fshared=20helper?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the adapter's local copies of the async tool-governance flow and its leaf helpers with the shared _shared.tool_governance module; both the Tool._run and toolset call_tool patches now supply only a framework-specific invoke_original callable. Shared helpers are re-imported (and listed in __all__) so the adapter's unit tests still reach them through this module. Behavior unchanged. Refs AAASM-4746 --- agent_assembly/adapters/pydantic_ai/patch.py | 271 +++---------------- 1 file changed, 41 insertions(+), 230 deletions(-) diff --git a/agent_assembly/adapters/pydantic_ai/patch.py b/agent_assembly/adapters/pydantic_ai/patch.py index 5a9e12a2..cbf59058 100644 --- a/agent_assembly/adapters/pydantic_ai/patch.py +++ b/agent_assembly/adapters/pydantic_ai/patch.py @@ -4,25 +4,38 @@ import importlib as importlib import inspect -from collections.abc import Mapping from dataclasses import dataclass from functools import wraps -from typing import TYPE_CHECKING, Any, Literal - -if TYPE_CHECKING: - from agent_assembly.exceptions import PolicyViolationError - -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 ( - _interceptor_enforces, - _missing_interceptor_decision, -) -from agent_assembly.adapters.crewai.patch import ( - _normalize_decision as _normalize_governance_decision, +from typing import Any + +from agent_assembly.adapters._shared.tool_governance import ( + _current_spawn_depth, + _get_pending_tool_approval_timeout_seconds, + _invoke_async_tool_check, + _normalize_decision, + _record_async_tool_result, + _serialize_tool_args, + _wait_for_async_tool_approval, + run_governed_async_tool, ) -from agent_assembly.core.spawn import _SPAWN_CTX, SpawnContext, spawn_context_scope +from agent_assembly.adapters.crewai.patch import _interceptor_enforces +from agent_assembly.core.spawn import SpawnContext, spawn_context_scope + +# The shared governance helpers are imported here (not just used internally) so +# the adapter's unit tests can reach them through this module — listing them in +# __all__ marks the re-export as intentional. See AAASM-4746 (dedup). +__all__ = [ + "AssemblyModelWrapper", + "PydanticAIPatch", + "_current_spawn_depth", + "_get_pending_tool_approval_timeout_seconds", + "_invoke_async_tool_check", + "_normalize_decision", + "_record_async_tool_result", + "_serialize_tool_args", + "_wait_for_async_tool_approval", + "set_process_agent_id", +] _ORIGINAL_TOOL_RUN = "_agent_assembly_original_pydantic_ai_tool_run" _ORIGINAL_TOOLSET_CALL_TOOL = "_agent_assembly_original_pydantic_ai_toolset_call_tool" @@ -31,7 +44,6 @@ _ORIGINAL_AGENT_RUN_SYNC = "_agent_assembly_original_pydantic_ai_agent_run_sync" _AGENT_PATCHED_FLAG = "_agent_assembly_pydantic_ai_agent_patched" _PROCESS_AGENT_ID: str | None = None -_MAX_AUDIT_RESULT_CHARS = 2000 @dataclass(slots=True) @@ -224,11 +236,6 @@ def _load_pydantic_ai_agent_class() -> type[Any] | None: return None -def _current_spawn_depth() -> int: - current = _SPAWN_CTX.get() - return (current.depth + 1) if current is not None else 1 - - def _apply_agent_run_patch(agent_cls: type[Any], process_agent_id: str | None) -> None: if getattr(agent_cls, _AGENT_PATCHED_FLAG, False): return None @@ -304,52 +311,18 @@ async def patched_run(self: Any, ctx: Any, args: Any, **kwargs: Any) -> Any: 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, enforce=enforce) - 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, enforce=enforce) - - if status == "deny": - if is_pending_flow: - raise _build_pending_rejected_error(tool_name, reason) - raise _build_denied_error(tool_name, reason) + def _invoke_original() -> Any: + return original_run(self, ctx, args, **kwargs) - spawn_ctx = SpawnContext( - parent_agent_id=agent_id or "", - depth=_current_spawn_depth(), - spawned_by_tool=tool_name, - delegation_reason=f"tool:{tool_name}", - ) - with spawn_context_scope(spawn_ctx): - result = original_run(self, ctx, args, **kwargs) - if inspect.isawaitable(result): - result = await result - - await _record_async_tool_result( + return await run_governed_async_tool( callback_handler, + enforce=enforce, tool_name=tool_name, - result=result, + tool_args=tool_args, agent_id=agent_id, run_id=run_id, + invoke_original=_invoke_original, ) - return result setattr(tool_cls, _ORIGINAL_TOOL_RUN, original_run) tool_cls._run = patched_run @@ -396,52 +369,18 @@ async def patched_call_tool(self: Any, name: Any, tool_args: Any, ctx: Any, tool 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=serialized_args, - agent_id=agent_id, - run_id=run_id, - ) - status, reason = _normalize_decision(decision, enforce=enforce) - 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=serialized_args, - agent_id=agent_id, - run_id=run_id, - ) - status, reason = _normalize_decision(final_decision, enforce=enforce) - - if status == "deny": - if is_pending_flow: - raise _build_pending_rejected_error(tool_name, reason) - raise _build_denied_error(tool_name, reason) - - spawn_ctx = SpawnContext( - parent_agent_id=agent_id or "", - depth=_current_spawn_depth(), - spawned_by_tool=tool_name, - delegation_reason=f"tool:{tool_name}", - ) - with spawn_context_scope(spawn_ctx): - result = original_call_tool(self, name, tool_args, ctx, tool, **kwargs) - if inspect.isawaitable(result): - result = await result + def _invoke_original() -> Any: + return original_call_tool(self, name, tool_args, ctx, tool, **kwargs) - await _record_async_tool_result( + return await run_governed_async_tool( callback_handler, + enforce=enforce, tool_name=tool_name, - result=result, + tool_args=serialized_args, agent_id=agent_id, run_id=run_id, + invoke_original=_invoke_original, ) - return result setattr(toolset_cls, _ORIGINAL_TOOLSET_CALL_TOOL, original_call_tool) toolset_cls.call_tool = patched_call_tool @@ -490,131 +429,3 @@ def _resolve_run_id(ctx: Any) -> str | 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 = 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, - *, - enforce: bool = False, -) -> tuple[Literal["allow", "deny", "pending"], str | None]: - return _normalize_governance_decision(decision, enforce=enforce) - - -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 _missing_interceptor_decision(callback_handler) - - 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: - return _resolve_pending_timeout_seconds(callback_handler) - - -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) -> PolicyViolationError: - 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) -> PolicyViolationError: - from agent_assembly.exceptions import PolicyViolationError - - reason_text = reason or "No reason provided." - return PolicyViolationError(f"Tool '{tool_name}' rejected during approval: {reason_text}")