Skip to content

Python: Add per-tool rate limiting and admission control - #9293

Closed
Suhas Aggarwal (suhasagg) wants to merge 1 commit into
microsoft:mainfrom
suhasagg:feat-9239-per-tool-rate-limiting
Closed

Suhas Aggarwal (suhasagg) wants to merge 1 commit into
microsoft:mainfrom
suhasagg:feat-9239-per-tool-rate-limiting

Conversation

@suhasagg

Copy link
Copy Markdown
Contributor

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=False and 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 FunctionMiddleware architecture.

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_second and burst values.

  • Atomic admission decisions using per-tool asyncio.Lock instances.

  • 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:

from agent_framework import ToolRateLimit, ToolRateLimitMiddleware

rate_limiter = ToolRateLimitMiddleware(
{
"lookup_customer": ToolRateLimit(
requests_per_second=2,
burst=2,
),
"get_payment_status": ToolRateLimit(
requests_per_second=1,
burst=1,
),
"search_documents": ToolRateLimit(
requests_per_second=10,
burst=10,
),
}
)

Register with an agent using:

middleware=[rate_limiter]

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.py

The 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:

  1. Public API design: Whether ToolRateLimit and ToolRateLimitMiddleware are appropriate reusable abstractions, or whether a sample-first approach is preferred.

  2. Rejection semantics: Whether returning a recognizable result through FunctionInvocationContext.result is the preferred contract, or whether a structured error mechanism would better align with the framework.

  3. Concurrency and state lifetime: Whether per-tool locks and middleware-instance-local token buckets provide the appropriate initial scope.

  4. 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:

Validation Result
Existing middleware and feature tests 96 passed
New rate-limiting tests 19 passed
Ruff lint Passed
Ruff formatting Passed
MyPy on new implementation Passed
Standalone example Executed successfully
Git whitespace check Passed

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.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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",
),
@eavanvalkenburg

Copy link
Copy Markdown
Member

Suhas Aggarwal (@suhasagg) Thanks for the contribution. #9239 is a substantial feature request that has not completed triage or been marked ready-for-implementation, so I’m closing this PR while the issue planning process runs. Please continue from the issue once the direction is agreed and it is ready for implementation.

This branch was previously deployed

1 inactive deployment
github-app-auth — 7bf0b77f Deployed Oct 12, 2026 by suhasagg via team_check #6570
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Usage: [Issues, PRs], Target: documentation in the code base and learn docs python Usage: [Issues, PRs], Target: Python

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants