Merge pull request #6108 from MervinPraison/claude/issue-5061-20261028

fix: let after_agent plugins block or rewrite the final response

src/praisonai-agents/praisonaiagents/agent/tool_execution.py

@@ -2068,6 +2068,41 @@ class ToolExecutionMixin:
             execution_time_ms=(time.time() - start_time) * 1000
         )
 
+    def _apply_after_agent_results(self, response, after_agent_input, results):
+        """Honour AFTER_AGENT hook outcomes on the final response.
+
+        A blocking hook/plugin (``GuardrailBlocked`` / ``PluginDecision.deny``)
+        replaces the returned value with a visible refusal marker, mirroring
+        the BEFORE_LLM block. Otherwise a string rewrite (e.g. PII redaction)
+        is adopted so it actually reaches the caller. The matching in-memory
+        assistant turn is updated too, so later turns and auto-memory never
+        see the withheld or unredacted text. Interaction callbacks, verbose
+        display and session/DB persistence run earlier and are not gated here.
+        """
+        if self._hook_runner.is_blocked(results):
+            reason = self._hook_runner.get_blocking_reason(results) or "Blocked by hook"
+            logging.warning(f"Agent {self.name} response blocked by AFTER_AGENT hook: {reason}")
+            final = f"[Response blocked by hook: {reason}]"
+        elif isinstance(response, str) and isinstance(after_agent_input.response, str):
+            final = after_agent_input.response
+        else:
+            return response
+        if final != response and isinstance(response, str):
+            # Prefer this turn's own messages (by identity) so a concurrent
+            # turn with identical text is never touched; fall back to a
+            # content match over shared history for untracked callers.
+            from .memory_mixin import _active_turn_owned
+            owned = _active_turn_owned.get() or []
+            with self._history_lock:
+                for candidates in (owned, self.chat_history):
+                    msg = next((m for m in reversed(candidates)
+                                if isinstance(m, dict) and m.get("role") == "assistant"
+                                and m.get("content") == response), None)
+                    if msg is not None:
+                        msg["content"] = final
+                        break
+        return final
+
     def _after_agent_side_effects(self, prompt, response, *, process_auto_memory=True):
         """Run loop-agnostic after-agent side effects (no ``await`` inside).
 
@@ -2125,7 +2160,8 @@ class ToolExecutionMixin:
         from ..hooks import HookEvent
         if self._hook_runner.registry.has_hooks(HookEvent.AFTER_AGENT):
             after_agent_input = self._build_after_agent_input(prompt, response, start_time, tools_used)
-            self._hook_runner.execute_sync(HookEvent.AFTER_AGENT, after_agent_input)
+            results = self._hook_runner.execute_sync(HookEvent.AFTER_AGENT, after_agent_input)
+            response = self._apply_after_agent_results(response, after_agent_input, results)
 
         self._after_agent_side_effects(prompt, response)
 
@@ -2183,7 +2219,8 @@ class ToolExecutionMixin:
         from ..hooks import HookEvent
         if self._hook_runner.registry.has_hooks(HookEvent.AFTER_AGENT):
             after_agent_input = self._build_after_agent_input(prompt, response, start_time, tools_used)
-            await self._hook_runner.execute(HookEvent.AFTER_AGENT, after_agent_input)
+            results = await self._hook_runner.execute(HookEvent.AFTER_AGENT, after_agent_input)
+            response = self._apply_after_agent_results(response, after_agent_input, results)
 
         if response and getattr(self, '_memory_instance', None):
             prompt_str = _memory_prompt_text(prompt)

src/praisonai-agents/praisonaiagents/plugins/manager.py

@@ -939,7 +939,9 @@ class PluginManager:
                         name=f"{name}:{event.value}",
                         # Delivery policies compose in order: apply a rewrite
                         # before the next policy, and stop the chain on denial.
-                        sequential=event == PluginHook.MESSAGE_SENDING,
+                        # AFTER_AGENT adapters mutate the shared response, so
+                        # stacked redactors must also run in order.
+                        sequential=event in (PluginHook.MESSAGE_SENDING, PluginHook.AFTER_AGENT),
                     )
                     hook_ids.append(hook_id)
                     count += 1
@@ -1171,8 +1173,14 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
 
     if _overrides("after_agent"):
         def after_agent_hook(data, _p=resolving_plugin):
-            new_resp = _p.after_agent(getattr(data, "response", ""),
-                                      {"agent_name": getattr(data, "agent_name", None)})
+            try:
+                new_resp = _p.after_agent(getattr(data, "response", ""),
+                                          {"agent_name": getattr(data, "agent_name", None)})
+            except GuardrailBlocked as e:
+                return HookResult.block(e.reason)
+            decision = _as_decision(new_resp)
+            if decision is not None:
+                return decision
             if isinstance(new_resp, str) and hasattr(data, "response"):
                 data.response = new_resp
             return HookResult.allow()

src/praisonai-agents/praisonaiagents/plugins/plugin.py

@@ -232,8 +232,18 @@ class Plugin(ABC):
         """
         return prompt
     
-    def after_agent(self, response: str, context: Dict[str, Any]) -> str:
-        """Called after agent execution. Can modify response."""
+    def after_agent(
+        self, response: str, context: Dict[str, Any]
+    ) -> Union[str, "PluginDecision", "HookResult", None]:
+        """Called after agent execution.
+
+        Return a modified ``response`` (rewrite), or a deny/block decision
+        (``PluginDecision.deny(reason)`` / a ``HookResult`` / raise
+        :class:`GuardrailBlocked`) to replace the value returned to the caller
+        (and the in-memory assistant turn) with a refusal marker, or ``None``
+        for no-op. Interaction callbacks, verbose display, streamed output and
+        session/DB persistence happen before this hook and are not gated by it.
+        """
         return response
     
     def before_tool(

src/praisonai-agents/tests/unit/plugins/test_after_agent_block.py

@@ -0,0 +1,173 @@
+"""
+Tests that AFTER_AGENT plugins can block or rewrite the final response (#5061).
+"""
+
+import asyncio
+import os
+import time
+
+from praisonaiagents.plugins.plugin import (
+    Plugin,
+    PluginInfo,
+    PluginHook,
+    PluginDecision,
+    GuardrailBlocked,
+)
+from praisonaiagents.plugins.manager import PluginManager
+from praisonaiagents.hooks.registry import HookRegistry
+from praisonaiagents.hooks.runner import HookRunner
+from praisonaiagents.hooks.types import HookEvent
+from praisonaiagents.hooks.events import AfterAgentInput
+
+SECRET = "123-45-6789"
+
+
+class RaiseBlockPlugin(Plugin):
+    @property
+    def info(self):
+        return PluginInfo(name="raiseblock", hooks=[PluginHook.AFTER_AGENT])
+
+    def after_agent(self, response, context):
+        if SECRET in response:
+            raise GuardrailBlocked("PII in response")
+        return response
+
+
+class DecisionBlockPlugin(Plugin):
+    @property
+    def info(self):
+        return PluginInfo(name="decisionblock", hooks=[PluginHook.AFTER_AGENT])
+
+    def after_agent(self, response, context):
+        if SECRET in response:
+            return PluginDecision.block("PII in response")
+        return response
+
+
+class RedactPlugin(Plugin):
+    @property
+    def info(self):
+        return PluginInfo(name="redact", hooks=[PluginHook.AFTER_AGENT])
+
+    def after_agent(self, response, context):
+        return response.replace(SECRET, "[REDACTED]")
+
+
+def _registry(plugin):
+    reg = HookRegistry()
+    mgr = PluginManager()
+    mgr.register(plugin)
+    mgr.set_plugin_options({plugin.info.name: {"allow_conversation": True}})
+    assert mgr.wire_into_hook_registry(reg) == 1
+    return reg
+
+
+def _input(response):
+    return AfterAgentInput(
+        session_id="s", cwd=os.getcwd(), event_name=HookEvent.AFTER_AGENT,
+        timestamp=str(time.time()), agent_name="a", response=response,
+    )
+
+
+def _agent(plugin):
+    from praisonaiagents import Agent
+    return Agent(name="a", role="r", goal="g", hooks=_registry(plugin))
+
+
+class TestAfterAgentAdapter:
+    def test_guardrail_blocked_blocks(self):
+        runner = HookRunner(registry=_registry(RaiseBlockPlugin()), cwd=os.getcwd())
+        results = runner.execute_sync(HookEvent.AFTER_AGENT, _input(f"SSN {SECRET}"))
+        assert runner.is_blocked(results) is True
+        assert runner.get_blocking_reason(results) == "PII in response"
+
+    def test_decision_blocks(self):
+        runner = HookRunner(registry=_registry(DecisionBlockPlugin()), cwd=os.getcwd())
+        results = runner.execute_sync(HookEvent.AFTER_AGENT, _input(f"SSN {SECRET}"))
+        assert runner.is_blocked(results) is True
+
+    def test_clean_response_allowed(self):
+        runner = HookRunner(registry=_registry(RaiseBlockPlugin()), cwd=os.getcwd())
+        data = _input("all good")
+        results = runner.execute_sync(HookEvent.AFTER_AGENT, data)
+        assert runner.is_blocked(results) is False
+        assert data.response == "all good"
+
+
+class TestAgentAfterAgentEnforcement:
+    def test_sync_block_withholds_response(self):
+        agent = _agent(RaiseBlockPlugin())
+        out = agent._trigger_after_agent_hook("q", f"SSN {SECRET}", time.time())
+        assert SECRET not in out
+        assert "blocked" in out.lower() and "PII in response" in out
+
+    def test_async_block_withholds_response(self):
+        agent = _agent(DecisionBlockPlugin())
+        out = asyncio.run(agent._atrigger_after_agent_hook("q", f"SSN {SECRET}", time.time()))
+        assert SECRET not in out
+        assert "PII in response" in out
+
+    def test_sync_rewrite_is_adopted(self):
+        agent = _agent(RedactPlugin())
+        out = agent._trigger_after_agent_hook("q", f"SSN {SECRET}", time.time())
+        assert out == "SSN [REDACTED]"
+
+    def test_clean_response_passes_through(self):
+        agent = _agent(RaiseBlockPlugin())
+        assert agent._trigger_after_agent_hook("q", "hello", time.time()) == "hello"
+
+    def test_block_scrubs_stored_assistant_turn(self):
+        agent = _agent(RaiseBlockPlugin())
+        agent.chat_history.append({"role": "assistant", "content": f"SSN {SECRET}"})
+        out = agent._trigger_after_agent_hook("q", f"SSN {SECRET}", time.time())
+        assert agent.chat_history[-1]["content"] == out
+        assert all(SECRET not in str(m.get("content")) for m in agent.chat_history)
+
+    def test_block_scrubs_turn_when_newer_turn_interleaved(self):
+        agent = _agent(RaiseBlockPlugin())
+        agent.chat_history.append({"role": "assistant", "content": f"SSN {SECRET}"})
+        agent.chat_history.append({"role": "user", "content": "other"})
+        agent.chat_history.append({"role": "assistant", "content": "concurrent reply"})
+        agent._trigger_after_agent_hook("q", f"SSN {SECRET}", time.time())
+        assert all(SECRET not in str(m.get("content")) for m in agent.chat_history)
+        assert agent.chat_history[-1]["content"] == "concurrent reply"
+
+    def test_block_scrubs_own_turn_not_identical_concurrent_turn(self):
+        agent = _agent(RaiseBlockPlugin())
+        own = {"role": "assistant", "content": f"SSN {SECRET}"}
+        # A concurrent (untracked here) turn appends an identical answer later.
+        other = {"role": "assistant", "content": f"SSN {SECRET}"}
+        token = agent._begin_turn_tracking()
+        try:
+            agent._append_to_chat_history(own)
+            agent.chat_history.append(other)
+            out = agent._trigger_after_agent_hook("q", f"SSN {SECRET}", time.time())
+        finally:
+            agent._end_turn_tracking(token)
+        assert own["content"] == out
+        assert other["content"] == f"SSN {SECRET}"
+
+
+class EmailRedactPlugin(Plugin):
+    @property
+    def info(self):
+        return PluginInfo(name="emailredact", hooks=[PluginHook.AFTER_AGENT])
+
+    def after_agent(self, response, context):
+        return response.replace("[email protected]", "[EMAIL]")
+
+
+def test_stacked_redactors_compose():
+    reg = HookRegistry()
+    mgr = PluginManager()
+    for p in (RedactPlugin(), EmailRedactPlugin()):
+        mgr.register(p)
+    mgr.set_plugin_options({
+        "redact": {"allow_conversation": True},
+        "emailredact": {"allow_conversation": True},
+    })
+    assert mgr.wire_into_hook_registry(reg) == 2
+    runner = HookRunner(registry=reg, cwd=os.getcwd())
+    data = _input(f"SSN {SECRET} mail [email protected]")
+    runner.execute_sync(HookEvent.AFTER_AGENT, data)
+    assert data.response == "SSN [REDACTED] mail [EMAIL]"