Repository navigation
Python: Add per-tool rate limiting and admission control - #9293
Closed
Suhas Aggarwal (suhasagg) wants to merge 1 commit into
Closed
Suhas Aggarwal (suhasagg) wants to merge 1 commit into
Suhas Aggarwal (suhasagg) wants to merge 1 commit into
Conversation
Suhas Aggarwal (suhasagg)
requested review from
SergeyMenshykh,
Tao Chen (TaoChenOSU),
Eduard van Valkenburg (eavanvalkenburg),
Giles Odigwe (giles17),
Jose Alvarez (jpalvarezl),
Evan Mattson (moonbox3),
Roger Barreto (rogerbarreto),
sophia-ramsey and
westey (westey-m)
as code owners
October 12, 2026 10:21
Suhas Aggarwal (suhasagg)
temporarily deployed
to
github-app-auth
October 12, 2026 10:21 — with
GitHub Actions
Inactive
Suhas Aggarwal (suhasagg)
temporarily deployed
to
github-app-auth
October 12, 2026 10:21 — with
GitHub Actions
Inactive
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Event-loop reuse can fail, and the public typing stub is not synchronized with the new exports.
2 open findings
What changed in this PR
Adds opt-in, per-tool token-bucket rate limiting to Python function middleware.
Changes:
- Adds rate-limit configuration and middleware.
- Exposes the API and provides a runnable sample.
- Adds concurrency, validation, cancellation, and integration tests.
| File | Description |
|---|---|
python/packages/core/agent_framework/_rate_limit.py |
Implements token-bucket admission control. |
python/packages/core/agent_framework/__init__.py |
Adds public runtime exports. |
python/packages/core/tests/core/test_rate_limit.py |
Tests rate-limiting behavior and integration. |
python/samples/02-agents/middleware/per_tool_rate_limiting.py |
Demonstrates independent tool quotas. |
python/samples/02-agents/middleware/README.md |
Lists the new sample. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.
| class _Bucket: | ||
| tokens: float | ||
| updated_at: float | ||
| lock: asyncio.Lock = field(default_factory=asyncio.Lock) |
Comment on lines
+197
to
+200
| "._rate_limit": ( | ||
| "ToolRateLimit", | ||
| "ToolRateLimitMiddleware", | ||
| ), |
Member
|
Suhas Aggarwal (@suhasagg) Thanks for the contribution. #9239 is a substantial feature request that has not completed triage or been marked |
This branch was previously deployed
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
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.


Motivation & Context
This PR introduces opt-in, concurrency-safe per-tool request-rate limiting and admission control for Microsoft Agent Framework's Python SDK.
Production AI agents frequently invoke external services with different request quotas, including enterprise APIs, CRM systems, payment services, databases, and retrieval systems.
When multiple tool calls execute concurrently, requests to the same downstream service may exceed its permitted request rate, resulting in HTTP 429 responses, unnecessary failures, and unpredictable agent execution.
Existing controls such as
allow_concurrent_invocation=Falseand function invocation limits do not enforce requests-per-second quotas.This contribution addresses the capability proposed in #9239 by introducing a reusable token-bucket middleware integrated with the existing
FunctionMiddlewarearchitecture.The implementation is opt-in, backward compatible, and does not modify the existing function invocation scheduler.
Description & Review Guide
What are the major changes?
1. Per-tool token-bucket admission control
Introduces two public Python classes:
ToolRateLimit— immutable configuration for request rate and burst capacity.ToolRateLimitMiddleware— concurrency-safe middleware enforcing admission policies before tool execution.The implementation supports:
Independent request-rate limits for individual named tools.
Configurable
requests_per_secondandburstvalues.Atomic admission decisions using per-tool
asyncio.Lockinstances.Monotonic-clock-based token replenishment.
Defensive handling of backward clock movement.
Immediate rejection when capacity is exhausted.
No automatic retries, waiting, or replay of rejected operations.
Unrestricted execution for tools without configured policies.
Cancellation propagation without holding the admission lock during tool execution.
Validation of invalid configuration values.
2. Public API integration
Both classes are exported through
agent_framework.Example configuration:
3. Regression and integration tests
Adds 19 tests covering:
Admission within configured quotas.
Burst exhaustion and pre-execution rejection.
Independent per-tool quotas.
Token replenishment using a deterministic clock.
Concurrent admission safety with 100 simultaneous invocations.
Cancellation propagation and accounting.
Backward-clock recovery.
Invalid configuration rejection.
Integration with
FunctionMiddlewarePipeline.Integration with automatic function invocation, including preservation of function call IDs.
4. Runnable sample and documentation
Adds:
python/samples/02-agents/middleware/per_tool_rate_limiting.pyThe example demonstrates two tools with independent burst capacities and fail-fast rejection. It runs without external APIs, cloud credentials, or network access.
The middleware samples README is updated accordingly.
What is the impact of these changes?
Prevents excessive locally invoked tool calls from reaching downstream services.
Reduces avoidable quota violations under concurrent execution.
Preserves concurrency between independently limited tools.
Provides reusable admission control instead of requiring applications to implement their own synchronization and token accounting.
Preserves existing behavior unless the middleware is explicitly registered.
Introduces no additional framework-level dependencies.
Implementation boundaries
Rate limits are local to each middleware instance.
Quotas are not coordinated across processes, replicas, or distributed workers.
Rejected calls return an identifiable rate-limit message through
FunctionInvocationContext.result.Distributed rate limiting, bounded waiting, automatic retries, and provider-hosted tool quotas are outside the scope of this PR.
What do you want reviewers to focus on?
I would particularly appreciate maintainer feedback on:
Public API design: Whether
ToolRateLimitandToolRateLimitMiddlewareare appropriate reusable abstractions, or whether a sample-first approach is preferred.Rejection semantics: Whether returning a recognizable result through
FunctionInvocationContext.resultis the preferred contract, or whether a structured error mechanism would better align with the framework.Concurrency and state lifetime: Whether per-tool locks and middleware-instance-local token buckets provide the appropriate initial scope.
Framework integration: Whether the implementation follows the expected middleware registration and function invocation patterns.
The implementation intentionally avoids modifying the core function invocation scheduler and can be refined based on maintainer feedback.
Validation
The following checks were completed locally:
No external services or paid API credentials were required for testing.
Related Issue
Related to #9239.
This PR implements the proposed initial middleware approach. I recognize that the issue has not yet received explicit maintainer approval for the implementation direction, and I welcome feedback on scope or API design before merging.
AI Assistance
No material AI assistance was used.
This is an AI-assisted contribution. I reviewed, understood, and verified all submitted content and accept responsibility for it.
AI assistance details:
AI assistance was used for implementation guidance, test development . I reviewed the changes, executed the local tests and validation commands, and take responsibility for the submitted implementation.
Contribution Checklist
The code builds clean without any errors or warnings.
All relevant local unit tests pass, and I have added new tests where possible.
The PR follows the Contribution Guidelines to the best of my knowledge.
This PR links to an agreed issue with no competing open PR, or the Related Issue section documents an applicable exception.
This is not a breaking change.