refactor(beta): simplify ToolPolicyMiddleware per review

Address review feedback from @Lancetnik on ag2ai#2533:

- Flatten constructor API: pass blocked_tools/allowed_tools/on_blocked
  directly instead of requiring callers to build a ToolPolicyConfig
- Drop stateful counters (_total_tool_calls, _total_blocked, _lock)
  and their accessor properties -- the middleware is now stateless
- Add optional on_blocked callback so callers can wire their own
  observability (metrics, audit log) without the library retaining
  any per-factory state

The ToolPolicyConfig dataclass is kept as an internal struct for
immutability (tuple freezing in __post_init__).

Author: amabito

autogen/beta/middleware/builtin/tool_policy.py

@@ -4,14 +4,16 @@
 
 """Tool policy middleware -- blocks disallowed tool calls before execution."""
 
-import threading
+from collections.abc import Callable, Sequence
 from dataclasses import dataclass, field
 
 from autogen.beta.annotations import Context
 from autogen.beta.events import BaseEvent, ToolCallEvent, ToolErrorEvent
 from autogen.beta.middleware.base import BaseMiddleware, MiddlewareFactory, ToolExecution, ToolResultType
 from autogen.beta.tools import ToolResult
 
+OnBlockedCallback = Callable[[ToolCallEvent, str], None]
+
 
 @dataclass
 class ToolPolicyConfig:
@@ -71,40 +73,39 @@ def check(self, tool_name: str) -> tuple[bool, str]:
 class ToolPolicyMiddleware(MiddlewareFactory):
     """Factory that creates per-invocation tool policy middleware instances.
 
-    A single ToolPolicyMiddleware shares counters across all instances it
-    creates, so statistics accumulate over the lifetime of the factory.
+    The middleware is stateless: no counters or shared mutable state are
+    retained across invocations. To observe denied calls, pass an
+    ``on_blocked`` callback when constructing the middleware.
 
     Example::
 
-        config = ToolPolicyConfig(
+        def audit(call: ToolCallEvent, reason: str) -> None:
+            logger.info("blocked %s: %s", call.name, reason)
+
+
+        mw = ToolPolicyMiddleware(
             blocked_tools=["delete_all"],
             allowed_tools=["search", "calc"],
+            on_blocked=audit,
         )
-        mw = ToolPolicyMiddleware(config)
         agent = MyAgent(middleware=[mw])
     """
 
-    def __init__(self, config: ToolPolicyConfig | None = None) -> None:
-        self._config = config or ToolPolicyConfig()
+    def __init__(
+        self,
+        blocked_tools: Sequence[str] | None = None,
+        allowed_tools: Sequence[str] | None = None,
+        on_blocked: OnBlockedCallback | None = None,
+    ) -> None:
+        self._config = ToolPolicyConfig(
+            blocked_tools=list(blocked_tools) if blocked_tools else [],
+            allowed_tools=list(allowed_tools) if allowed_tools is not None else None,
+        )
         self._policy = _ToolPolicy(self._config)
-        self._total_tool_calls: int = 0
-        self._total_blocked: int = 0
-        self._lock = threading.Lock()
-
-    @property
-    def total_tool_calls(self) -> int:
-        """Number of tool calls that passed the policy check and were forwarded."""
-        with self._lock:
-            return self._total_tool_calls
-
-    @property
-    def total_blocked(self) -> int:
-        """Number of tool calls that were blocked."""
-        with self._lock:
-            return self._total_blocked
+        self._on_blocked = on_blocked
 
     def __call__(self, event: "BaseEvent", context: "Context") -> "BaseMiddleware":
-        return _ToolPolicyInstance(event, context, self._policy, self)
+        return _ToolPolicyInstance(event, context, self._policy, self._on_blocked)
 
 
 class _ToolPolicyInstance(BaseMiddleware):
@@ -115,11 +116,11 @@ def __init__(
         event: "BaseEvent",
         context: "Context",
         policy: _ToolPolicy,
-        factory: ToolPolicyMiddleware,
+        on_blocked: OnBlockedCallback | None,
     ) -> None:
         super().__init__(event, context)
         self._policy = policy
-        self._factory = factory
+        self._on_blocked = on_blocked
 
     async def on_tool_execution(
         self,
@@ -129,10 +130,8 @@ async def on_tool_execution(
     ) -> ToolResultType:
         allowed, reason = self._policy.check(event.name)
         if not allowed:
-            with self._factory._lock:
-                self._factory._total_blocked += 1
+            if self._on_blocked is not None:
+                self._on_blocked(event, reason)
             return _make_tool_error(event, reason)
 
-        with self._factory._lock:
-            self._factory._total_tool_calls += 1
         return await call_next(event, context)

test/beta/middleware/builtins/test_tool_policy_middleware.py

@@ -162,36 +162,98 @@ def test_tool_error_content_matches_reason(self) -> None:
 
 
 # ---------------------------------------------------------------------------
-# TestToolPolicyStats
+# TestToolPolicyMiddleware
 # ---------------------------------------------------------------------------
 
 
-class TestToolPolicyStats:
+class TestToolPolicyMiddleware:
+    def test_flat_constructor(self) -> None:
+        # Given arguments passed directly (no explicit config object)
+        mw = ToolPolicyMiddleware(
+            blocked_tools=["delete_all"],
+            allowed_tools=["search"],
+        )
+
+        # Then the internal config is built from those arguments
+        assert mw._config.blocked_tools == ("delete_all",)
+        assert mw._config.allowed_tools == ("search",)
+
+    def test_default_constructor(self) -> None:
+        # Given no arguments
+        mw = ToolPolicyMiddleware()
+
+        # Then defaults are empty blocklist and None allowlist (no restriction)
+        assert mw._config.blocked_tools == ()
+        assert mw._config.allowed_tools is None
+
     @pytest.mark.asyncio()
-    async def test_call_and_blocked_counters(self) -> None:
-        # Given a middleware factory with a blocklist
-        config = ToolPolicyConfig(blocked_tools=["bad_tool"])
-        factory = ToolPolicyMiddleware(config)
+    async def test_on_blocked_callback_fires_on_denial(self) -> None:
+        # Given a middleware with a blocklist and an on_blocked callback
+        recorded: list[tuple[str, str]] = []
+
+        def audit(call: ToolCallEvent, reason: str) -> None:
+            recorded.append((call.name, reason))
+
+        factory = ToolPolicyMiddleware(blocked_tools=["bad_tool"], on_blocked=audit)
 
         ctx = _make_context()
         initial_event = _make_event()
 
-        # Prepare a real ToolResultEvent to return from call_next
+        async def call_next(event: ToolCallEvent, context: mock.MagicMock) -> ToolResultEvent:
+            return ToolResultEvent(parent_id=event.id, name=event.name, result=ToolResult("ok"))
+
+        # When a blocked call is denied
+        blocked_call = ToolCallEvent(id="c1", name="bad_tool")
+        instance = factory(initial_event, ctx)
+        result = await instance.on_tool_execution(call_next, blocked_call, ctx)
+
+        # Then the callback was invoked with the call and reason, and a ToolErrorEvent was returned
+        assert len(recorded) == 1
+        assert recorded[0][0] == "bad_tool"
+        assert "blocked" in recorded[0][1]
+        assert isinstance(result, ToolErrorEvent)
+
+    @pytest.mark.asyncio()
+    async def test_on_blocked_not_called_when_allowed(self) -> None:
+        # Given a middleware with a callback and an allowed call
+        recorded: list[tuple[str, str]] = []
+
+        def audit(call: ToolCallEvent, reason: str) -> None:
+            recorded.append((call.name, reason))
+
+        factory = ToolPolicyMiddleware(allowed_tools=["good_tool"], on_blocked=audit)
+
+        ctx = _make_context()
+        initial_event = _make_event()
         good_call = ToolCallEvent(id="c1", name="good_tool")
         good_result = ToolResultEvent(parent_id="c1", name="good_tool", result=ToolResult("ok"))
 
         async def call_next(event: ToolCallEvent, context: mock.MagicMock) -> ToolResultEvent:
             return good_result
 
-        # When processing one allowed call and one blocked call
-        instance_allow = factory(initial_event, ctx)
-        await instance_allow.on_tool_execution(call_next, good_call, ctx)
+        # When the call passes the policy
+        instance = factory(initial_event, ctx)
+        result = await instance.on_tool_execution(call_next, good_call, ctx)
+
+        # Then the callback was not invoked and the downstream result was returned
+        assert recorded == []
+        assert result is good_result
+
+    @pytest.mark.asyncio()
+    async def test_no_callback_means_no_error(self) -> None:
+        # Given a middleware without an on_blocked callback
+        factory = ToolPolicyMiddleware(blocked_tools=["bad_tool"])
+
+        ctx = _make_context()
+        initial_event = _make_event()
+
+        async def call_next(event: ToolCallEvent, context: mock.MagicMock) -> ToolResultEvent:
+            return ToolResultEvent(parent_id=event.id, name=event.name, result=ToolResult("ok"))
 
-        blocked_call = ToolCallEvent(id="c2", name="bad_tool")
-        instance_block = factory(initial_event, ctx)
-        result = await instance_block.on_tool_execution(call_next, blocked_call, ctx)
+        # When a blocked call is denied without a callback
+        blocked_call = ToolCallEvent(id="c1", name="bad_tool")
+        instance = factory(initial_event, ctx)
+        result = await instance.on_tool_execution(call_next, blocked_call, ctx)
 
-        # Then counters reflect the activity
-        assert factory.total_tool_calls == 1
-        assert factory.total_blocked == 1
+        # Then denial still works, just without observation
         assert isinstance(result, ToolErrorEvent)