feat(beta): add ToolPolicyMiddleware.update_policy for runtime swaps

Address @Lancetnik follow-up on #2533: the dynamic-policy justification
needs a public API to be useful. Adds:

- ToolPolicyMiddleware.update_policy(config: ToolPolicyConfig)
  Atomically replaces the active policy. Subsequent tool calls see the
  new config; per-invocation instances already inside on_tool_execution
  keep their captured snapshot, so in-flight executions are never torn
  across a config change.

- ToolPolicyMiddleware.config property
  Read-only access to the active ToolPolicyConfig.

Intended hooks: rate-limit kill switches, per-tenant rule updates from
an ops dashboard, role changes during a long-running session -- all
without rebuilding the agent.

Thread-safety: _policy is assigned last so readers in __call__ never
observe a state where _config is new but _policy is stale (Python
attribute assignment is atomic for single references).

Tests: 19/19 pass, covering update-then-permit, in-flight-snapshot
isolation, restriction clearing, and the config property.

Author: amabito

autogen/beta/middleware/builtin/tool_policy.py

@@ -77,6 +77,11 @@ class ToolPolicyMiddleware(MiddlewareFactory):
     retained across invocations. To observe denied calls, pass an
     ``on_blocked`` callback when constructing the middleware.
 
+    The policy can be replaced at runtime via :meth:`update_policy` --
+    useful for rate-limit kill switches, per-tenant rules driven by an
+    ops dashboard, or role changes that land mid-session without
+    rebuilding the agent.
+
     Example::
 
         def audit(call: ToolCallEvent, reason: str) -> None:
@@ -89,6 +94,9 @@ def audit(call: ToolCallEvent, reason: str) -> None:
             on_blocked=audit,
         )
         agent = MyAgent(middleware=[mw])
+
+        # Later, swap the policy without rebuilding the agent:
+        mw.update_policy(ToolPolicyConfig(blocked_tools=["delete_all", "shell_exec"]))
     """
 
     def __init__(
@@ -104,6 +112,34 @@ def __init__(
         self._policy = _ToolPolicy(self._config)
         self._on_blocked = on_blocked
 
+    @property
+    def config(self) -> ToolPolicyConfig:
+        """Return the currently active policy configuration."""
+        return self._config
+
+    def update_policy(self, config: ToolPolicyConfig) -> None:
+        """Atomically replace the tool policy at runtime.
+
+        Subsequent tool calls use the new policy. Calls that have already
+        entered ``on_tool_execution`` continue with the policy snapshot
+        their per-invocation instance captured at construction time, so
+        in-flight executions are never torn across a configuration change.
+
+        This is the intended hook for dynamic policy sources: rate-limit
+        kill switches, per-tenant rule updates from an ops dashboard,
+        role changes during a long-running session, etc.
+
+        Args:
+            config: The new :class:`ToolPolicyConfig` to enforce. Passing
+                a fresh config with defaults (empty blocklist, no allowlist)
+                effectively clears all restrictions.
+        """
+        new_policy = _ToolPolicy(config)
+        # Assign _policy last so readers in __call__ never observe a state
+        # where _config is new but _policy is stale.
+        self._config = config
+        self._policy = new_policy
+
     def __call__(self, event: "BaseEvent", context: "Context") -> "BaseMiddleware":
         return _ToolPolicyInstance(event, context, self._policy, self._on_blocked)
 

test/beta/middleware/builtins/test_tool_policy_middleware.py

@@ -239,6 +239,71 @@ async def call_next(event: ToolCallEvent, context: mock.MagicMock) -> ToolResult
         assert recorded == []
         assert result is good_result
 
+    @pytest.mark.asyncio()
+    async def test_update_policy_takes_effect_on_next_call(self) -> None:
+        # Given a middleware that initially allows "search" only
+        factory = ToolPolicyMiddleware(allowed_tools=["search"])
+        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"))
+
+        # When "shell_exec" is attempted before the update -- denied
+        instance_before = factory(initial_event, ctx)
+        result_before = await instance_before.on_tool_execution(
+            call_next, ToolCallEvent(id="c1", name="shell_exec"), ctx
+        )
+        assert isinstance(result_before, ToolErrorEvent)
+
+        # And the policy is replaced at runtime to allow shell_exec
+        factory.update_policy(ToolPolicyConfig(allowed_tools=["search", "shell_exec"]))
+
+        # Then a new instance picks up the new policy and permits the call
+        instance_after = factory(initial_event, ctx)
+        result_after = await instance_after.on_tool_execution(call_next, ToolCallEvent(id="c2", name="shell_exec"), ctx)
+        assert not isinstance(result_after, ToolErrorEvent)
+        assert factory.config.allowed_tools == ("search", "shell_exec")
+
+    @pytest.mark.asyncio()
+    async def test_update_policy_in_flight_instance_uses_snapshot(self) -> None:
+        # Given a factory that allows "search" and an instance constructed under that policy
+        factory = ToolPolicyMiddleware(allowed_tools=["search"])
+        ctx = _make_context()
+        instance = factory(_make_event(), ctx)  # snapshot taken here
+
+        async def call_next(event: ToolCallEvent, context: mock.MagicMock) -> ToolResultEvent:
+            return ToolResultEvent(parent_id=event.id, name=event.name, result=ToolResult("ok"))
+
+        # When the factory policy is swapped to allow everything mid-flight
+        factory.update_policy(ToolPolicyConfig())
+
+        # Then the previously constructed instance still enforces the old snapshot
+        result = await instance.on_tool_execution(call_next, ToolCallEvent(id="c1", name="shell_exec"), ctx)
+        assert isinstance(result, ToolErrorEvent)
+
+    def test_update_policy_clears_restrictions(self) -> None:
+        # Given a middleware with both blocked and allowed lists
+        factory = ToolPolicyMiddleware(blocked_tools=["x"], allowed_tools=["y"])
+
+        # When the policy is replaced with an empty config
+        factory.update_policy(ToolPolicyConfig())
+
+        # Then check() permits any tool name
+        allowed, reason = factory._policy.check("anything")
+        assert allowed
+        assert reason == ""
+        assert factory.config.blocked_tools == ()
+        assert factory.config.allowed_tools is None
+
+    def test_config_property_returns_active_config(self) -> None:
+        # Given a middleware
+        factory = ToolPolicyMiddleware(blocked_tools=["delete_all"])
+
+        # Then config property exposes the active ToolPolicyConfig
+        assert isinstance(factory.config, ToolPolicyConfig)
+        assert factory.config.blocked_tools == ("delete_all",)
+
     @pytest.mark.asyncio()
     async def test_no_callback_means_no_error(self) -> None:
         # Given a middleware without an on_blocked callback