From 9432cc4b5b6085ca2abbe03a8e0cf23e923fa86c Mon Sep 17 00:00:00 2001 From: MacOS Date: Tue, 18 Aug 2026 16:35:30 +0800 Subject: [PATCH 1/2] feat(agents): add optional AgentFuse tool middleware --- .../praisonaiagents/agent/chat_mixin.py | 7 +- .../praisonaiagents/agent/execution_mixin.py | 2 +- .../praisonaiagents/agent/tool_execution.py | 1 + .../praisonaiagents/hooks/agentfuse.py | 88 +++++++++++++++ src/praisonai-agents/pyproject.toml | 3 + .../unit/hooks/test_agentfuse_middleware.py | 106 ++++++++++++++++++ 6 files changed, 205 insertions(+), 2 deletions(-) create mode 100644 src/praisonai-agents/praisonaiagents/hooks/agentfuse.py create mode 100644 src/praisonai-agents/tests/unit/hooks/test_agentfuse_middleware.py diff --git a/src/praisonai-agents/praisonaiagents/agent/chat_mixin.py b/src/praisonai-agents/praisonaiagents/agent/chat_mixin.py index a5643af7b1..544c908f37 100644 --- a/src/praisonai-agents/praisonaiagents/agent/chat_mixin.py +++ b/src/praisonai-agents/praisonaiagents/agent/chat_mixin.py @@ -4439,7 +4439,12 @@ async def _achat_completion(self, response, tools, reasoning_steps=False): # after-context aggregation) are fired once, guarded, inside # execute_tool_async — no inline duplicate dispatch here. # Pass the tools list to honor task-scoped tools - result = await self.execute_tool_async(function_name, arguments, tools_override=tools) + result = await self.execute_tool_async( + function_name, + arguments, + tool_call_id=getattr(tool_call, "id", None), + tools_override=tools, + ) results.append(result) except Exception as e: diff --git a/src/praisonai-agents/praisonaiagents/agent/execution_mixin.py b/src/praisonai-agents/praisonaiagents/agent/execution_mixin.py index e7b76e5bfd..74b8d93c83 100644 --- a/src/praisonai-agents/praisonaiagents/agent/execution_mixin.py +++ b/src/praisonai-agents/praisonaiagents/agent/execution_mixin.py @@ -1519,6 +1519,7 @@ async def _execute_tool_async_via_middleware( run_id=getattr(self, '_current_run_id', 'unknown'), session_id=getattr(self, '_session_id', None) or 'default', tool_name=function_name, + metadata={"tool_call_id": tool_call_id}, ), ) @@ -2122,4 +2123,3 @@ async def _emit_retry_hook_async(self, tool_name, attempt, delay_ms, error, max_ except Exception as e: # Don't let hook failures break retry logic logging.debug(f"Failed to emit async retry hook: {e}") - diff --git a/src/praisonai-agents/praisonaiagents/agent/tool_execution.py b/src/praisonai-agents/praisonaiagents/agent/tool_execution.py index 4a9bf3b0af..8db7652735 100644 --- a/src/praisonai-agents/praisonaiagents/agent/tool_execution.py +++ b/src/praisonai-agents/praisonaiagents/agent/tool_execution.py @@ -489,6 +489,7 @@ def execute_tool(self, function_name: str, arguments: Dict[str, Any], tool_call_ run_id=getattr(self, '_current_run_id', 'unknown'), session_id=getattr(self, '_session_id', None) or 'default', tool_name=function_name, + metadata={"tool_call_id": tool_call_id}, ), ) diff --git a/src/praisonai-agents/praisonaiagents/hooks/agentfuse.py b/src/praisonai-agents/praisonaiagents/hooks/agentfuse.py new file mode 100644 index 0000000000..4fa208396a --- /dev/null +++ b/src/praisonai-agents/praisonaiagents/hooks/agentfuse.py @@ -0,0 +1,88 @@ +"""Optional AgentFuse middleware for pre-dispatch tool decisions.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +from .middleware import ToolRequest, ToolResponse + + +class AgentFuseToolMiddleware: + """Evaluate AgentFuse immediately before PraisonAI's tool handler chain.""" + + _hook_type = "wrap_tool_call" + + def __init__(self, guard: Any) -> None: + try: + from dhms_agentfuse import RuntimeGuard, ToolCallRequest + except ImportError as exc: + raise ImportError( + "AgentFuse middleware requires the 'agentfuse' optional dependency: " + "pip install 'praisonaiagents[agentfuse]'" + ) from exc + + if not isinstance(guard, RuntimeGuard): + raise TypeError("guard must be a dhms_agentfuse.RuntimeGuard") + + self._guard = guard + self._request_type = ToolCallRequest + self._decisions: dict[str, Any] = {} + + def decision_for(self, tool_call_id: str) -> Any | None: + """Return the completed policy decision for a tool call, if available.""" + return self._decisions.get(tool_call_id) + + def __call__( + self, + request: ToolRequest, + call_next: Callable[[ToolRequest], ToolResponse], + ) -> ToolResponse: + context = request.context + tool_call_id = ( + context.metadata.get("tool_call_id") if context is not None else None + ) + if not tool_call_id: + return ToolResponse( + tool_name=request.tool_name, + result={ + "status": "blocked", + "policy_denied": True, + "tool_failure": False, + "reason_code": "missing_tool_call_id", + "host_execution": { + "outcome": "not_executed", + "handler_started": False, + }, + }, + context=context, + ) + + decision = self._guard.evaluate( + self._request_type( + tool_call_id=tool_call_id, + tool_name=request.tool_name, + arguments=request.arguments, + safe_metadata={"integration": "praisonaiagents"}, + ) + ) + self._decisions[tool_call_id] = decision + + if decision.action == "block": + return ToolResponse( + tool_name=request.tool_name, + result={ + "status": "blocked", + "policy_denied": True, + "tool_failure": False, + "reason_code": decision.reason_code, + "agentfuse_decision": decision.to_safe_dict(), + "host_execution": { + "outcome": "not_executed", + "handler_started": False, + }, + }, + context=context, + ) + + return call_next(request) diff --git a/src/praisonai-agents/pyproject.toml b/src/praisonai-agents/pyproject.toml index 967ee770f6..3e892fd637 100644 --- a/src/praisonai-agents/pyproject.toml +++ b/src/praisonai-agents/pyproject.toml @@ -21,6 +21,9 @@ dependencies = [ ] [project.optional-dependencies] +agentfuse = [ + "dhms-agentfuse==3.7.3", +] mcp = [ "mcp>=1.20.0", "fastapi>=0.115.0", diff --git a/src/praisonai-agents/tests/unit/hooks/test_agentfuse_middleware.py b/src/praisonai-agents/tests/unit/hooks/test_agentfuse_middleware.py new file mode 100644 index 0000000000..602d90daa1 --- /dev/null +++ b/src/praisonai-agents/tests/unit/hooks/test_agentfuse_middleware.py @@ -0,0 +1,106 @@ +"""Real Agent dispatch tests for the optional AgentFuse tool middleware.""" + +from __future__ import annotations + +import pytest +from dhms_agentfuse import RuntimeGuard +from praisonaiagents import Agent +from praisonaiagents.hooks.agentfuse import AgentFuseToolMiddleware + + +def _sync_agent(guard: RuntimeGuard): + handler_calls: list[str] = [] + + def protected_write(value: str) -> str: + handler_calls.append(value) + return "write completed" + + middleware = AgentFuseToolMiddleware(guard) + agent = Agent( + name="agentfuse-sync-test", + instructions="Exercise one inert test tool.", + tools=[protected_write], + hooks=[middleware], + approval=True, + ) + return agent, middleware, handler_calls + + +def test_sync_allow_dispatches_once_and_preserves_identity(): + agent, middleware, calls = _sync_agent( + RuntimeGuard(allow_tools={"protected_write"}) + ) + + result = agent.execute_tool( + "protected_write", {"value": "synthetic-value"}, "praison-allow-001" + ) + + assert result == "write completed" + assert calls == ["synthetic-value"] + assert middleware.decision_for("praison-allow-001").tool_call_id == ( + "praison-allow-001" + ) + + +def test_sync_block_returns_non_execution_without_dispatch(): + agent, middleware, calls = _sync_agent(RuntimeGuard(deny_tools={"protected_write"})) + + result = agent.execute_tool( + "protected_write", {"value": "synthetic-value"}, "praison-block-001" + ) + + assert calls == [] + assert result["status"] == "blocked" + assert result["tool_failure"] is False + assert result["host_execution"] == { + "outcome": "not_executed", + "handler_started": False, + } + assert result["agentfuse_decision"]["tool_call_id"] == "praison-block-001" + assert middleware.decision_for("praison-block-001").action == "block" + + +def test_sync_policy_failure_fails_closed_without_dispatch(): + def failing_policy(tool_call): + del tool_call + raise RuntimeError("synthetic policy failure") + + agent, middleware, calls = _sync_agent(RuntimeGuard(policy=failing_policy)) + + result = agent.execute_tool( + "protected_write", {"value": "synthetic-value"}, "praison-failure-001" + ) + + assert calls == [] + assert result["status"] == "blocked" + assert result["agentfuse_decision"]["reason_code"] == "policy_exception" + assert result["host_execution"]["handler_started"] is False + assert middleware.decision_for("praison-failure-001").reason_code == ( + "policy_exception" + ) + + +@pytest.mark.asyncio +async def test_async_dispatch_uses_same_guard_and_preserves_identity(): + handler_calls: list[str] = [] + + async def protected_write(value: str) -> str: + handler_calls.append(value) + return "write completed" + + middleware = AgentFuseToolMiddleware(RuntimeGuard(deny_tools={"protected_write"})) + agent = Agent( + name="agentfuse-async-test", + instructions="Exercise one inert test tool.", + tools=[protected_write], + hooks=[middleware], + approval=True, + ) + + result = await agent.execute_tool_async( + "protected_write", {"value": "synthetic-value"}, "praison-async-001" + ) + + assert handler_calls == [] + assert result["host_execution"]["outcome"] == "not_executed" + assert result["agentfuse_decision"]["tool_call_id"] == "praison-async-001" From 19f5f4a15be4362038adcbe6adb8d12e771c3212 Mon Sep 17 00:00:00 2001 From: MacOS Date: Tue, 18 Aug 2026 18:07:43 +0800 Subject: [PATCH 2/2] refactor(agents): keep middleware identity plumbing generic --- .../praisonaiagents/agent/execution_mixin.py | 1 + .../praisonaiagents/hooks/agentfuse.py | 88 --------------- src/praisonai-agents/pyproject.toml | 3 - .../unit/agent/test_achat_unified_dispatch.py | 2 + .../unit/hooks/test_agentfuse_middleware.py | 106 ------------------ .../hooks/test_tool_call_identity_context.py | 68 +++++++++++ 6 files changed, 71 insertions(+), 197 deletions(-) delete mode 100644 src/praisonai-agents/praisonaiagents/hooks/agentfuse.py delete mode 100644 src/praisonai-agents/tests/unit/hooks/test_agentfuse_middleware.py create mode 100644 src/praisonai-agents/tests/unit/hooks/test_tool_call_identity_context.py diff --git a/src/praisonai-agents/praisonaiagents/agent/execution_mixin.py b/src/praisonai-agents/praisonaiagents/agent/execution_mixin.py index 74b8d93c83..f3d03085a5 100644 --- a/src/praisonai-agents/praisonaiagents/agent/execution_mixin.py +++ b/src/praisonai-agents/praisonaiagents/agent/execution_mixin.py @@ -2123,3 +2123,4 @@ async def _emit_retry_hook_async(self, tool_name, attempt, delay_ms, error, max_ except Exception as e: # Don't let hook failures break retry logic logging.debug(f"Failed to emit async retry hook: {e}") + diff --git a/src/praisonai-agents/praisonaiagents/hooks/agentfuse.py b/src/praisonai-agents/praisonaiagents/hooks/agentfuse.py deleted file mode 100644 index 4fa208396a..0000000000 --- a/src/praisonai-agents/praisonaiagents/hooks/agentfuse.py +++ /dev/null @@ -1,88 +0,0 @@ -"""Optional AgentFuse middleware for pre-dispatch tool decisions.""" - -from __future__ import annotations - -from collections.abc import Callable -from typing import Any - -from .middleware import ToolRequest, ToolResponse - - -class AgentFuseToolMiddleware: - """Evaluate AgentFuse immediately before PraisonAI's tool handler chain.""" - - _hook_type = "wrap_tool_call" - - def __init__(self, guard: Any) -> None: - try: - from dhms_agentfuse import RuntimeGuard, ToolCallRequest - except ImportError as exc: - raise ImportError( - "AgentFuse middleware requires the 'agentfuse' optional dependency: " - "pip install 'praisonaiagents[agentfuse]'" - ) from exc - - if not isinstance(guard, RuntimeGuard): - raise TypeError("guard must be a dhms_agentfuse.RuntimeGuard") - - self._guard = guard - self._request_type = ToolCallRequest - self._decisions: dict[str, Any] = {} - - def decision_for(self, tool_call_id: str) -> Any | None: - """Return the completed policy decision for a tool call, if available.""" - return self._decisions.get(tool_call_id) - - def __call__( - self, - request: ToolRequest, - call_next: Callable[[ToolRequest], ToolResponse], - ) -> ToolResponse: - context = request.context - tool_call_id = ( - context.metadata.get("tool_call_id") if context is not None else None - ) - if not tool_call_id: - return ToolResponse( - tool_name=request.tool_name, - result={ - "status": "blocked", - "policy_denied": True, - "tool_failure": False, - "reason_code": "missing_tool_call_id", - "host_execution": { - "outcome": "not_executed", - "handler_started": False, - }, - }, - context=context, - ) - - decision = self._guard.evaluate( - self._request_type( - tool_call_id=tool_call_id, - tool_name=request.tool_name, - arguments=request.arguments, - safe_metadata={"integration": "praisonaiagents"}, - ) - ) - self._decisions[tool_call_id] = decision - - if decision.action == "block": - return ToolResponse( - tool_name=request.tool_name, - result={ - "status": "blocked", - "policy_denied": True, - "tool_failure": False, - "reason_code": decision.reason_code, - "agentfuse_decision": decision.to_safe_dict(), - "host_execution": { - "outcome": "not_executed", - "handler_started": False, - }, - }, - context=context, - ) - - return call_next(request) diff --git a/src/praisonai-agents/pyproject.toml b/src/praisonai-agents/pyproject.toml index 3e892fd637..967ee770f6 100644 --- a/src/praisonai-agents/pyproject.toml +++ b/src/praisonai-agents/pyproject.toml @@ -21,9 +21,6 @@ dependencies = [ ] [project.optional-dependencies] -agentfuse = [ - "dhms-agentfuse==3.7.3", -] mcp = [ "mcp>=1.20.0", "fastapi>=0.115.0", diff --git a/src/praisonai-agents/tests/unit/agent/test_achat_unified_dispatch.py b/src/praisonai-agents/tests/unit/agent/test_achat_unified_dispatch.py index 438fc94207..9295cb062a 100644 --- a/src/praisonai-agents/tests/unit/agent/test_achat_unified_dispatch.py +++ b/src/praisonai-agents/tests/unit/agent/test_achat_unified_dispatch.py @@ -104,6 +104,7 @@ def run(self, **kwargs): assert not hasattr(browser_tool, "__name__") tool_call = SimpleNamespace( + id="legacy-tool-call-001", function=SimpleNamespace(name="browserbase", arguments="{}") ) response = SimpleNamespace( @@ -117,6 +118,7 @@ def run(self, **kwargs): mock_exec.assert_awaited_once() assert mock_exec.await_args[0][0] == "browserbase" + assert mock_exec.await_args.kwargs["tool_call_id"] == "legacy-tool-call-001" assert results is not None diff --git a/src/praisonai-agents/tests/unit/hooks/test_agentfuse_middleware.py b/src/praisonai-agents/tests/unit/hooks/test_agentfuse_middleware.py deleted file mode 100644 index 602d90daa1..0000000000 --- a/src/praisonai-agents/tests/unit/hooks/test_agentfuse_middleware.py +++ /dev/null @@ -1,106 +0,0 @@ -"""Real Agent dispatch tests for the optional AgentFuse tool middleware.""" - -from __future__ import annotations - -import pytest -from dhms_agentfuse import RuntimeGuard -from praisonaiagents import Agent -from praisonaiagents.hooks.agentfuse import AgentFuseToolMiddleware - - -def _sync_agent(guard: RuntimeGuard): - handler_calls: list[str] = [] - - def protected_write(value: str) -> str: - handler_calls.append(value) - return "write completed" - - middleware = AgentFuseToolMiddleware(guard) - agent = Agent( - name="agentfuse-sync-test", - instructions="Exercise one inert test tool.", - tools=[protected_write], - hooks=[middleware], - approval=True, - ) - return agent, middleware, handler_calls - - -def test_sync_allow_dispatches_once_and_preserves_identity(): - agent, middleware, calls = _sync_agent( - RuntimeGuard(allow_tools={"protected_write"}) - ) - - result = agent.execute_tool( - "protected_write", {"value": "synthetic-value"}, "praison-allow-001" - ) - - assert result == "write completed" - assert calls == ["synthetic-value"] - assert middleware.decision_for("praison-allow-001").tool_call_id == ( - "praison-allow-001" - ) - - -def test_sync_block_returns_non_execution_without_dispatch(): - agent, middleware, calls = _sync_agent(RuntimeGuard(deny_tools={"protected_write"})) - - result = agent.execute_tool( - "protected_write", {"value": "synthetic-value"}, "praison-block-001" - ) - - assert calls == [] - assert result["status"] == "blocked" - assert result["tool_failure"] is False - assert result["host_execution"] == { - "outcome": "not_executed", - "handler_started": False, - } - assert result["agentfuse_decision"]["tool_call_id"] == "praison-block-001" - assert middleware.decision_for("praison-block-001").action == "block" - - -def test_sync_policy_failure_fails_closed_without_dispatch(): - def failing_policy(tool_call): - del tool_call - raise RuntimeError("synthetic policy failure") - - agent, middleware, calls = _sync_agent(RuntimeGuard(policy=failing_policy)) - - result = agent.execute_tool( - "protected_write", {"value": "synthetic-value"}, "praison-failure-001" - ) - - assert calls == [] - assert result["status"] == "blocked" - assert result["agentfuse_decision"]["reason_code"] == "policy_exception" - assert result["host_execution"]["handler_started"] is False - assert middleware.decision_for("praison-failure-001").reason_code == ( - "policy_exception" - ) - - -@pytest.mark.asyncio -async def test_async_dispatch_uses_same_guard_and_preserves_identity(): - handler_calls: list[str] = [] - - async def protected_write(value: str) -> str: - handler_calls.append(value) - return "write completed" - - middleware = AgentFuseToolMiddleware(RuntimeGuard(deny_tools={"protected_write"})) - agent = Agent( - name="agentfuse-async-test", - instructions="Exercise one inert test tool.", - tools=[protected_write], - hooks=[middleware], - approval=True, - ) - - result = await agent.execute_tool_async( - "protected_write", {"value": "synthetic-value"}, "praison-async-001" - ) - - assert handler_calls == [] - assert result["host_execution"]["outcome"] == "not_executed" - assert result["agentfuse_decision"]["tool_call_id"] == "praison-async-001" diff --git a/src/praisonai-agents/tests/unit/hooks/test_tool_call_identity_context.py b/src/praisonai-agents/tests/unit/hooks/test_tool_call_identity_context.py new file mode 100644 index 0000000000..26cfa608f7 --- /dev/null +++ b/src/praisonai-agents/tests/unit/hooks/test_tool_call_identity_context.py @@ -0,0 +1,68 @@ +"""Tests for tool-call identity exposed through middleware context.""" + +from __future__ import annotations + +import pytest +from praisonaiagents import Agent +from praisonaiagents.hooks import wrap_tool_call + + +def test_sync_tool_middleware_receives_original_tool_call_id(): + observed_ids: list[str | None] = [] + handler_calls: list[str] = [] + + @wrap_tool_call + def capture_identity(request, call_next): + observed_ids.append(request.context.metadata.get("tool_call_id")) + return call_next(request) + + def inert_tool(value: str) -> str: + handler_calls.append(value) + return "completed" + + agent = Agent( + name="sync-middleware-identity", + instructions="Exercise one inert test tool.", + tools=[inert_tool], + hooks=[capture_identity], + approval=True, + ) + + result = agent.execute_tool( + "inert_tool", {"value": "synthetic-value"}, "sync-tool-call-001" + ) + + assert result == "completed" + assert handler_calls == ["synthetic-value"] + assert observed_ids == ["sync-tool-call-001"] + + +@pytest.mark.asyncio +async def test_async_tool_middleware_receives_original_tool_call_id(): + observed_ids: list[str | None] = [] + handler_calls: list[str] = [] + + @wrap_tool_call + def capture_identity(request, call_next): + observed_ids.append(request.context.metadata.get("tool_call_id")) + return call_next(request) + + async def inert_tool(value: str) -> str: + handler_calls.append(value) + return "completed" + + agent = Agent( + name="async-middleware-identity", + instructions="Exercise one inert test tool.", + tools=[inert_tool], + hooks=[capture_identity], + approval=True, + ) + + result = await agent.execute_tool_async( + "inert_tool", {"value": "synthetic-value"}, "async-tool-call-001" + ) + + assert result == "completed" + assert handler_calls == ["synthetic-value"] + assert observed_ids == ["async-tool-call-001"]