feat(beta): add ToolPolicyMiddleware for tool access control

[amabito]

Why are these changes needed?

Pulls the tool allow/block list logic out of GovernanceMiddleware (#2501,
closed) into its own middleware. Sits on on_tool_execution, checks the
tool name against blocked and allowed lists, returns ToolErrorEvent if denied.

blocked_tools overrides allowed_tools. allowed_tools=None means no
restriction; allowed_tools=[] means deny all. Config is frozen after
construction.

Related issue number

Replaces the tool policy portion of #2501 (closed).
Related: #2531 (BudgetMiddleware), #2532 (CircuitBreakerMiddleware).

Checks

• I've included any doc changes needed for https://docs.ag2.ai/. See https://docs.ag2.ai/latest/docs/contributor-guide/documentation/ to build and test documentation locally.
• I've added tests (if relevant) corresponding to the changes introduced in this PR.
• I've made sure all auto checks have passed.

[Lancetnik]

@amabito hi! Thank you for the middlewares contribution, but I don't see a sense of the current one

1. If you don't wont to let an Agent execute a tool, you can just do not pass it to context. Provide models with information about tool and forbid its execution looks strange to me
2. Global counters _total_tool_calls and _total_blocked also could slowly lead to memory leaks because we have no options to drop them. Also, I don't see a reason to have such counters

[amabito]

Hey @Lancetnik, thanks for the pushback -- both points are worth addressing directly.

"Just don't pass the tool to context"

Static exclusion at registration time works when the policy is fixed for the session, and for that case you're right that a middleware adds nothing. Where it stops working:

1. Audit of denied attempts. If a tool was advertised to the model earlier in the session (cached context, multi-turn history, a prior planning step) and policy changes mid-session, the model may still try to call it. Static exclusion at registration doesn't intercept that call or leave any record. The middleware catches the attempt, returns a structured ToolErrorEvent via _make_tool_error (tool_policy.py L37-L47), and surfaces it on the event stream where observability hooks can pick it up. "The model tried to call delete_all and was blocked" is a signal you want to see, not silence.

2. Dynamic policy updates. Policies driven by external state -- rate limiting, user role, per-tenant rules, kill-switches triggered by an ops dashboard -- can't be expressed by omitting the tool at construction time, because the agent is already built. The middleware holds _policy as a swappable reference, so a supervisor can update it without rebuilding the agent. I haven't exposed a setter yet, but the shape is there.

3. Deny-by-default sandboxes. An allowlist-based agent (allowed_tools=["search", "calc"]) is a cleaner primitive than "register every tool except these", especially when the tool set is assembled dynamically from plugins or MCP servers. The check in _ToolPolicy.check (L59-L68) enforces deny-beats-allow precedence, which is the usual policy semantics.

If even the audit case feels out of scope for AG2's middleware layer and you'd rather see this as a dedicated capability or event hook rather than a middleware, I'm open to that -- say the word and I'll repackage it. But I don't think "don't pass the tool" covers these cases.

Counters and the leak concern

The counters at L90-L91 are instance attributes on ToolPolicyMiddleware, not module-level globals, so they don't persist across unrelated factories. What they do is grow unbounded over the lifetime of a long-lived factory with no reset path, which is a real API design smell even if it's not strictly a leak.

I'd rather just drop the counters. They were a convenience for testing and a half-hearted observability hook, and neither justifies stateful middleware. Concrete change:

• Remove _total_tool_calls, _total_blocked, _lock, and the two properties (L90-L104).
• Add an optional on_blocked: Callable[[ToolCallEvent, str], None] | None parameter. Users who want metrics wire it to their own Prometheus counter / log / whatever. The middleware stays stateless.
• Tests that check counter values get rewritten against a mock callback.

Landing it in the next revision alongside the constructor simplification from #2531 / #2532.

[Lancetnik]

• Dynamic policy updates. Policies driven by external state -- rate limiting, user role, per-tenant rules, kill-switches triggered by an ops dashboard -- can't be expressed by omitting the tool at construction time, because the agent is already built. The middleware holds _policy as a swappable reference, so a supervisor can update it without rebuilding the agent. I haven't exposed a setter yet, but the shape is there.

Fair, it seems the commonest case for such feature - but I don't see a public API covers it. Can we introduce a some?

[Lancetnik]

I prefer end-to-end test here using public API (or at least do not use non-public objects). Such tests are not so implementation-tied and not so fragile. Also, they cover real usecases and could be used as a some kind of documentation

Please, try to implement each test using TestClient / TracingClient. Reference: https://github.com/ag2ai/ag2/blob/main/test/beta/middleware/test_tool_execution.py

[amabito]

@Lancetnik Thanks -- fair point, "it's in the shape" is not a public API. Added one in 5984625.

ToolPolicyMiddleware.update_policy(config)

mw = ToolPolicyMiddleware(allowed_tools=["search", "calc"])
agent = MyAgent(middleware=[mw])

# Later, from an ops callback, tenant-context hook, rate-limit watchdog, etc.:
mw.update_policy(ToolPolicyConfig(
    blocked_tools=["delete_all", "shell_exec"],
    allowed_tools=["search", "calc"],
))

The swap is atomic with a specific contract:

• Subsequent calls to __call__ (i.e. new events arriving at the agent) pick up the new policy.
• Per-invocation _ToolPolicyInstance objects that are already inside on_tool_execution keep the policy snapshot they captured at construction time. In-flight tool executions are never torn across a config change -- you won't see a call half-evaluated against the old rules and half against the new.
• _policy is assigned last (after _config) so a concurrent reader in __call__ never observes a state where _config is new but _policy is stale. Single-reference attribute assignment is atomic in CPython, so no lock is needed for the happy path.

Also added mw.config as a read-only property for supervisors that want to log or diff the active configuration.

Tests cover:

1. test_update_policy_takes_effect_on_next_call -- deny shell_exec, swap policy, next call permits it.
2. test_update_policy_in_flight_instance_uses_snapshot -- construct instance, swap to "no restrictions", confirm the in-flight instance still enforces the old allowlist.
3. test_update_policy_clears_restrictions -- passing a fresh ToolPolicyConfig() drops all rules.
4. test_config_property_returns_active_config.

Branch is also rebased onto current main -- the CONFLICTING state is resolved.

Other changes carried from the earlier reply (already in the branch via 1d96b4a): _total_tool_calls / _total_blocked / _lock and their properties are gone; observation goes through the on_blocked callback.

19/19 tests pass, pre-commit clean.

[amabito]

Sorry for the rapid-fire -- ran the middleware through another review pass and wanted to send the consolidated state.

Since the last update:

• update_policy() contract. Docstring now states clearly that self._policy is the sole authority for enforcement in __call__; self._config is updated for the .config property. Under CPython the two STORE_ATTR are not a single transaction, so a .config reader may briefly see the new config while __call__ uses the previous policy for one final instantiation. test_concurrent_update_policy_consistent_decisions exercises 20 parallel on_tool_execution calls while a background writer alternates policies and confirms every result is a valid decision against some snapshot.
• Callback isolation. on_blocked is now wrapped in contextlib.suppress(Exception), so a failing audit callback does not prevent the ToolErrorEvent from being returned.
• Export. ToolPolicyMiddleware and ToolPolicyConfig are re-exported from both autogen.beta.middleware and autogen.beta.middleware.builtin. Regression test test_tool_policy_importable_from_middleware_top_level covers the top-level path. Note: denials currently produce ToolErrorEvent(PermissionError(reason)), not a dedicated ToolDeniedError -- happy to add a named exception in a follow-up if that matters for callers.
• Edge-case tests added. test_empty_tool_name_denied_by_allowlist and test_empty_allowlist_blocks_all_via_public_interface cover corner cases through the public middleware interface.

25 tests, all passing. Pre-commit clean.

[marklysze]

CI is failing due to API staleness. ToolResult no longer accepts a content keyword argument — the constructor signature is:

ToolResult(*inputs: SendableMessage | Input, parts: Iterable[...] = (), final: bool = False, metadata: dict | None = None)

Any ToolResult(content=..., ...) calls in the PR or its tests need to be updated to pass the content as a positional argument or via parts=[...].

There are also failures in unrelated beta tests (test_watch.py, test_network/) — those appear to be pre-existing flakiness, not caused by this PR.

The ToolPolicyMiddleware design itself (allow/block list, on_blocked callback, update_policy() for mutable runtime config) looks solid once the API call sites are fixed.

[amabito]

@marklysze Thanks for catching this. Pushed the API fix in 13eeb3e748a:

• _make_tool_error simplified to one line: return ToolErrorEvent.from_call(event, PermissionError(reason)) -- matching the upstream idiom.
• 1 ToolResult(content=...) site fixed (plus removed the now-unused ToolResult import from tool_policy.py).
• 2 test assertions rewritten: .content checks replaced with isinstance(error_event.error, PermissionError) + str(error_event.error) == reason, since ToolErrorEvent no longer carries a .content property.
• Also rebased onto upstream/main to pick up the new ToolResult definition.
• 25/25 tests passing locally. Pre-commit clean (ruff, mypy, codespell).

CI should be green this round. Happy to address any remaining concerns on the middleware design.

[Lancetnik]

@amabito sorry for waiting so long... There were too many things we changed and contributed our side last month. Now we have a stable enough API and I'll be happy to revisit your PRs next week

[chengduman]

Good separation of concerns — pulling tool policy out of GovernanceMiddleware makes sense. A couple of production observations:

Policy ordering matters

If both ToolPolicyMiddleware and the broader governance hooks run on the same agent, the execution order determines behavior. In practice:

# Recommended order: cheap checks first, expensive checks last
interceptors = [
    ToolPolicyMiddleware(blocked_tools=["rm", "drop_db"]),  # O(1) regex
    RateLimiter(max_per_minute=10),                         # O(1) counter
    CostEstimator(max_budget=0.50),                        # O(n) estimation
    AuditLogger(),                                          # O(n) write
]

Fail-fast on tool policy before spending compute on cost estimation.

blocked_tools overrides allowed_tools — edge case

If someone sets blocked_tools=["exec_bash", "rm"] AND allowed_tools=["exec_bash"] simultaneously (misconfig), the blocked list wins. Worth a log warning.

Cost estimation tie-in

ToolPolicyMiddleware and CostGovernor (#2942) are complementary: policy checks WHAT the agent can call, cost checks HOW MUCH each call costs. A combined dashboard that surfaces both in one place would be useful for operators.

References:

• GovernanceMiddleware (feat(beta): add governance middleware (budget, circuit breaker, tool policy) #2501)
• Cost governance discussion ([Feature Request]: governance interceptor for tool calls and inter-agent messages (TealTiger integration) #2942)

[chengduman]

Glad this was useful. If anyone is looking for help designing a production-grade interceptor chain with cost budgeting, we do lightweight architecture reviews (from $2) — details at https://chengduman.github.io/zhongpu-consulting-advisory/services.html

[marklysze]

@amabito, we have unfortunately been unable to review your PR just yet. We are hoping to get V1 over the line in the next couple of weeks and once that's done, we'll jump on it. Thank you for your patience.

[Lancetnik]

Hi @amabito — this PR now has merge conflicts with main and has been moved to status:changes-requested.

Please rebase onto the current main and resolve the conflicts. Once you push the updated branch, the PR will return to the review queue automatically.

⏳ If there's no update within 30 days, this PR will be marked stale and closed 7 days after that.

[genisis0x]

Clean design overall. A few things I like: blocked_tools taking precedence over allowed_tools is the safe default and it's documented; keeping the None (no allowlist) vs [] (allow nothing) distinction through __post_init__ is exactly right and easy to get wrong; freezing to frozenset and creating a per-invocation instance from a policy snapshot means an update_policy mid-flight can't tear an in-progress check. The update_policy docstring is honest about the two non-atomic STORE_ATTRs, and since __call__ only ever reads self._policy (one atomic attribute read) enforcement is always consistent with a single policy snapshot — the only visible skew is the config property vs enforcement, which is cosmetic. Good.

Three things worth pinning down before this is relied on as a security control, since for an access gate the interesting cases are the bypasses:

1. Matching is exact string membership (tool_name in self._blocked). Is event.name guaranteed to be the canonical registered tool name at the point on_tool_execution runs? If a tool can ever reach here under a differently-cased or whitespace-variant name (or an alias) than what an operator put in blocked_tools, an exact-match blocklist silently lets it through. For allowed_tools the failure is safe (unknown variant is denied), but for blocked_tools it's the dangerous direction. If names are already normalized upstream this is moot — worth a one-line confirmation, and if there's any doubt, normalizing both sides (and matching on the canonical name) closes it.

2. Coverage: is on_tool_execution invoked for every path a tool can execute through — direct calls, handoffs/sub-agent tools, any nested or framework-internal execution? A policy that wraps most invocations but not all isn't a policy an operator can trust for "delete_all is blocked." A test that asserts a blocked tool stays blocked when reached via a sub-agent/handoff (not just a direct call) would document that the gate has no side doors.

3. The on_blocked callback is wrapped in contextlib.suppress(Exception) — right instinct that enforcement must not depend on the observer. But on_blocked is the audit hook, and silently swallowing its failure means you can block a dangerous call and lose the record that you did. For a security audit sink that's the one failure you'd want to see. Consider logging the suppressed exception (still non-fatal) so a broken audit path is visible rather than silent.

None of these are blockers on the mechanism — the allow/deny logic itself is correct. They're the questions that decide whether it holds up as a security boundary vs a best-effort filter.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
autogen/beta/middleware/__init__.py	100.00% <ø> (ø)
autogen/beta/middleware/builtin/__init__.py	83.33% <100.00%> (-16.67%)	⬇️
autogen/beta/middleware/builtin/tool_policy.py	100.00% <100.00%> (ø)

... and 66 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.