feat(beta): add governance middleware (budget, circuit breaker, tool policy) - #2501
Closed
amabito wants to merge 18 commits into
Closed
feat(beta): add governance middleware (budget, circuit breaker, tool policy)#2501amabito wants to merge 18 commits into
amabito wants to merge 18 commits into
Conversation
amabito
force-pushed
the
feat/governance-middleware-v2
branch
from
March 20, 2026 05:45
474c174 to
df043a4
Compare
…policy) Adds GovernanceMiddleware to autogen/beta/middleware/builtin/ with: - Budget enforcement via on_llm_call (tracks ModelResponse.usage) - Circuit breaker (CLOSED/OPEN/HALF_OPEN with atomic probe claim) - Tool policy (block/allowlist) via on_tool_execution - Degradation (log + disable expensive tools near budget limit) - Thread-safe (Lock on budget, circuit breaker, stats) 33 tests, all passing.
- BudgetState: NaN/Inf/negative validation, __setattr__ write guard, record_tokens() with lock, string sentinel key - Budget middleware: .get() + isinstance check (no KeyError crash), getattr fallback for event.name - Helpers: generic error content instead of repr(error) - Policy: fail-closed on non-dict JSON args, set copy on init, generic PolicyViolationError message, 64KB args size guard - Redaction: max_content_bytes ReDoS guard, tool argument redaction, field-preserving event copy, tool-arg redaction in on_tool_call - Tests: refactored to pytest-asyncio, 12 new adversarial tests
- Redaction: fail-closed on copy failure (raise instead of forward unredacted) - Policy: fail-closed on JSON parse error (DENY instead of empty args) - BudgetState: try_consume_tool_call returns (bool, str) matching llm API - BudgetState: _COUNTER_FIELDS as class-level frozenset (avoid per-call rebuild) - Policy: document intentional size cap difference (64KB args vs 100KB content) - Tests: update for try_consume_tool_call return type and fail-closed policy
- BudgetState: guard limit fields (max_tokens etc.) against post-init mutation - BudgetState: reject non-int max_llm_calls/max_tool_calls (float boundary bug) - Redaction: validate max_content_bytes >= 0 at init - Policy: fail-closed on non-string arguments (bytes/bytearray)
- policy.py: use getattr(event, 'name', None) instead of event.name to
prevent AttributeError when event has no name attribute (fail-closed
path already handles non-str names including None)
- redaction.py: add isinstance(int)/bool guard on max_content_bytes to
prevent float('nan') from bypassing the < 0 validation (IEEE 754:
nan < 0 is always False)
- test_adversarial.py: add 3 tests covering the above fixes
- test_redaction_max_content_bytes_rejects_float
- test_redaction_max_content_bytes_rejects_bool
- test_policy_event_missing_name_attribute_is_denied (async)
30/30 tests pass. Rounds 6+7 both CLEAN.
- ToolError.name: use generic "<denied>" to prevent tool name enumeration - Redaction on_tool_call: emit ToolError instead of raising RuntimeError - Policy: reject empty string tool names (fail-closed) - BudgetState: reject bool values for limits and record_tokens
- Redaction on_llm_call: catch RuntimeError from _copy_event_with_content, emit ModelResponse denial instead of crashing middleware chain - Policy: catch RecursionError/MemoryError from json.loads (fail-closed) - Demo: update assertions to use isinstance(ToolError) after name anonymization
Codex CLI (gpt-5.4) 3-body x 6-round review-fix loop findings: Round 1: event type names (ToolCall/ToolError, not ToolCallEvent), CancelledError probe leak, fallback_model documented as observation-only Round 2: split except Exception/BaseException for probe management, failure_threshold NaN/Inf validation (isinstance int check) Round 5: Gemini token key support (prompt_token_count/candidates_token_count) Also: GovernanceConfig.__post_init__ validates all numeric fields, _safe_float() coerces string/None/NaN token values, disable_tools_on_degrade frozen to tuple, record() type widened to dict[str, object], budget soft-limit race documented in docstring, total_llm_calls/total_tool_calls properties exposed. Tests: 33 -> 47 (+14). All new tests cover review findings.
Pre-commit check fix: ruff formatter adjustments (line wrapping, blank line removal, parametrize formatting).
The v2 branch events/__init__.py exports ToolCallEvent/ToolErrorEvent (renamed in ag2ai#2491), not ToolCall/ToolError. Local tests passed due to stale pyc cache. CI correctly caught the import failure.
The standalone ag2_governance_middleware/ PoC was superseded by autogen/beta/middleware/builtin/governance.py. Remove to avoid pre-commit failures on files not part of this PR.
Covers the except BaseException path in on_llm_call -- verifies that asyncio.CancelledError releases the HALF_OPEN probe slot without incrementing the failure counter.
Adds test for the probe_claimed=False path in except BaseException handler (line 409->411). governance.py now at 100% statement + branch coverage.
Align with TelemetryMiddleware pattern (PR ag2ai#2488): - Import GovernanceConfig + GovernanceMiddleware from .governance - Add both to __all__ tuple - No optional dependency guard needed (stdlib only)
- Freeze blocked_tools in GovernanceConfig.__post_init__ (consistency with disable_tools_on_degrade) - Add test_governance_importable_from_builtin_init (verify __init__.py export) - Add test_zero_budget_allows_calls (budget=0 full async path) 59 tests passing.
Consistent with disable_tools_on_degrade: all list config fields are frozen to tuple in __post_init__ to prevent post-construction mutation. _ToolPolicy already uses frozenset internally, but the config object itself was mutable. 2 new tests: blocked_tools and allowed_tools are tuple after init. 61 tests passing.
amabito
force-pushed
the
feat/governance-middleware-v2
branch
from
March 21, 2026 06:21
368b2b2 to
bf44ef8
Compare
amabito
marked this pull request as ready for review
March 22, 2026 12:13
3 tasks
Contributor
Author
|
Splitting this into three single-responsibility middlewares to match the v2 design pattern. New PRs incoming shortly. |
Codecov Report❌ Patch coverage is
... and 13 files with indirect coverage changes 🚀 New features to boost your workflow:
|
Contributor
Author
|
Split into three single-responsibility middlewares:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Governance middleware for AG2 Beta: budget tracking, circuit breaker, tool policy.
Uses the
on_llm_call/on_tool_executionhooks from #2439.Context
With #2498 adding Arcade.dev (8,000+ tools), there's no built-in way to cap
LLM spend or restrict which tools an agent can call. This fills that gap.
What's in here
autogen/beta/middleware/builtin/governance.py-- 448 lines, single file.ModelResponse.usage(OpenAI + Gemini key formats). Blocks calls at limit. Soft cap -- concurrent calls can overshoot slightly.61 tests, 893 lines. Budget, CB, policy, concurrency, adversarial, config validation.
CC: @Lancetnik @marklysze