diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 3953187578e..c1263e5bcca 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -194,6 +194,10 @@ "chat_middleware", "function_middleware", ), + "._rate_limit": ( + "ToolRateLimit", + "ToolRateLimitMiddleware", + ), "._sessions": ( "AgentSession", "ContextProvider", @@ -638,6 +642,8 @@ "ToolApprovalRuleCallback", "ToolApprovalState", "ToolMode", + "ToolRateLimit", + "ToolRateLimitMiddleware", "ToolResultCompactionStrategy", "ToolTypes", "TruncationStrategy", diff --git a/python/packages/core/agent_framework/_rate_limit.py b/python/packages/core/agent_framework/_rate_limit.py new file mode 100644 index 00000000000..ec744e9638e --- /dev/null +++ b/python/packages/core/agent_framework/_rate_limit.py @@ -0,0 +1,98 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Opt-in, per-tool token-bucket admission control for function middleware.""" + +import asyncio +import math +import time +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass, field + +from ._middleware import FunctionInvocationContext, FunctionMiddleware + + +@dataclass(frozen=True, slots=True) +class ToolRateLimit: + """Maximum average requests per second and available burst tokens.""" + + requests_per_second: float + burst: int = 1 + + def __post_init__(self) -> None: + if isinstance(self.requests_per_second, bool): + raise ValueError("requests_per_second must be finite and positive") + if not math.isfinite(self.requests_per_second) or self.requests_per_second <= 0: + raise ValueError("requests_per_second must be finite and positive") + if isinstance(self.burst, bool) or not isinstance(self.burst, int) or self.burst < 1: + raise ValueError("burst must be a positive integer") + + +@dataclass(slots=True) +class _Bucket: + tokens: float + updated_at: float + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + + +class ToolRateLimitMiddleware(FunctionMiddleware): + """Apply instance-local, fail-fast token-bucket limits to named tools. + + Each configured tool has its own bucket. The initial bucket is full. + Rejected calls do not execute the tool, sleep, or retry. + + Limits are local to this middleware instance and are not shared across + processes, replicas, or distributed workers. Configure a shared external + rate limiter when deployment-wide quotas are required. + """ + + def __init__( + self, + limits: Mapping[str, ToolRateLimit], + *, + clock: Callable[[], float] = time.monotonic, + ) -> None: + self._limits = dict(limits) + if any(not isinstance(name, str) or not name for name in self._limits): + raise ValueError("tool names must be non-empty strings") + if any(not isinstance(limit, ToolRateLimit) for limit in self._limits.values()): + raise TypeError("limits must contain ToolRateLimit values") + self._clock = clock + self._buckets: dict[str, _Bucket] = {} + + # Initialize buckets eagerly so concurrent calls never race to create one. + for name, limit in self._limits.items(): + self._buckets[name] = _Bucket( + tokens=float(limit.burst), + updated_at=self._clock(), + ) + + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + name = context.function.name + limit = self._limits.get(name) + + if limit is None: + await call_next() + return + + bucket = self._buckets[name] + + async with bucket.lock: + now = self._clock() + elapsed = max(0.0, now - bucket.updated_at) + bucket.tokens = min( + float(limit.burst), + bucket.tokens + elapsed * limit.requests_per_second, + ) + bucket.updated_at = max(bucket.updated_at, now) + + if bucket.tokens < 1.0: + context.result = f"Tool rate limit exceeded for '{name}'; request rejected before execution." + return + + bucket.tokens -= 1.0 + + await call_next() diff --git a/python/packages/core/tests/core/test_rate_limit.py b/python/packages/core/tests/core/test_rate_limit.py new file mode 100644 index 00000000000..a608b56ba59 --- /dev/null +++ b/python/packages/core/tests/core/test_rate_limit.py @@ -0,0 +1,384 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Tests for per-tool token-bucket admission control.""" + +from collections.abc import Awaitable, Callable + +import pytest + +from agent_framework._middleware import FunctionInvocationContext +from agent_framework._rate_limit import ToolRateLimit, ToolRateLimitMiddleware +from agent_framework._tools import FunctionTool + + +class FakeClock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +def make_context(name: str) -> FunctionInvocationContext: + return FunctionInvocationContext( + function=FunctionTool( + name=name, + description="Test tool", + func=lambda: None, + ), + arguments={}, + ) + + +async def invoke( + middleware: ToolRateLimitMiddleware, + name: str, + execute: Callable[[], Awaitable[None]], +) -> FunctionInvocationContext: + context = make_context(name) + await middleware.process(context, execute) + return context + + +@pytest.mark.asyncio +async def test_first_request_is_admitted() -> None: + clock = FakeClock() + limiter = ToolRateLimitMiddleware( + {"search": ToolRateLimit(requests_per_second=1)}, + clock=clock, + ) + executed = 0 + + async def execute() -> None: + nonlocal executed + executed += 1 + + context = await invoke(limiter, "search", execute) + + assert executed == 1 + assert context.result is None + + +@pytest.mark.asyncio +async def test_burst_exhaustion_rejects_without_execution() -> None: + limiter = ToolRateLimitMiddleware( + {"search": ToolRateLimit(requests_per_second=1, burst=2)}, + clock=FakeClock(), + ) + executed = 0 + + async def execute() -> None: + nonlocal executed + executed += 1 + + await invoke(limiter, "search", execute) + await invoke(limiter, "search", execute) + rejected = await invoke(limiter, "search", execute) + + assert executed == 2 + assert "rate limit exceeded" in rejected.result.lower() + + +@pytest.mark.asyncio +async def test_tokens_refill_using_injected_clock() -> None: + clock = FakeClock() + limiter = ToolRateLimitMiddleware( + {"search": ToolRateLimit(requests_per_second=2)}, + clock=clock, + ) + executed = 0 + + async def execute() -> None: + nonlocal executed + executed += 1 + + await invoke(limiter, "search", execute) + rejected = await invoke(limiter, "search", execute) + assert "rate limit exceeded" in rejected.result.lower() + + clock.advance(0.5) + admitted = await invoke(limiter, "search", execute) + + assert executed == 2 + assert admitted.result is None + + +@pytest.mark.asyncio +async def test_unconfigured_tool_is_unrestricted() -> None: + limiter = ToolRateLimitMiddleware( + {"search": ToolRateLimit(requests_per_second=1)}, + clock=FakeClock(), + ) + executed = 0 + + async def execute() -> None: + nonlocal executed + executed += 1 + + for _ in range(5): + await invoke(limiter, "other_tool", execute) + + assert executed == 5 + + +@pytest.mark.parametrize( + ("rate", "burst"), + [ + (0, 1), + (-1, 1), + (float("nan"), 1), + (float("inf"), 1), + (1, 0), + (1, -1), + (1, True), + ], +) +def test_invalid_configuration_is_rejected(rate: float, burst: int) -> None: + with pytest.raises(ValueError): + ToolRateLimit(requests_per_second=rate, burst=burst) + + +@pytest.mark.asyncio +async def test_concurrent_calls_cannot_exceed_burst() -> None: + """Concurrent calls must reserve tokens atomically.""" + import asyncio + + clock = FakeClock() + limiter = ToolRateLimitMiddleware( + {"search": ToolRateLimit(requests_per_second=1, burst=3)}, + clock=clock, + ) + + executed = 0 + + async def execute() -> None: + nonlocal executed + executed += 1 + await asyncio.sleep(0) + + results = await asyncio.gather(*(invoke(limiter, "search", execute) for _ in range(100))) + + admitted = sum(context.result is None for context in results) + rejected = sum( + isinstance(context.result, str) and "rate limit exceeded" in context.result.lower() for context in results + ) + + assert admitted == 3 + assert rejected == 97 + assert executed == 3 + + +@pytest.mark.asyncio +async def test_independent_tool_quotas() -> None: + """Exhausting one tool's quota must not affect another tool.""" + clock = FakeClock() + limiter = ToolRateLimitMiddleware( + { + "search": ToolRateLimit(requests_per_second=1, burst=1), + "database": ToolRateLimit(requests_per_second=2, burst=2), + }, + clock=clock, + ) + + executions: list[str] = [] + + async def execute_search() -> None: + executions.append("search") + + async def execute_database() -> None: + executions.append("database") + + first_search = await invoke(limiter, "search", execute_search) + rejected_search = await invoke(limiter, "search", execute_search) + + first_database = await invoke(limiter, "database", execute_database) + second_database = await invoke(limiter, "database", execute_database) + rejected_database = await invoke(limiter, "database", execute_database) + + assert first_search.result is None + assert "rate limit exceeded" in rejected_search.result.lower() + + assert first_database.result is None + assert second_database.result is None + assert "rate limit exceeded" in rejected_database.result.lower() + + assert executions == ["search", "database", "database"] + + +@pytest.mark.asyncio +async def test_cancellation_propagates_and_preserves_bucket_state() -> None: + """Cancellation after admission propagates without corrupting the bucket.""" + import asyncio + + clock = FakeClock() + limiter = ToolRateLimitMiddleware( + {"search": ToolRateLimit(requests_per_second=1, burst=1)}, + clock=clock, + ) + + started = asyncio.Event() + + async def cancelled_tool() -> None: + started.set() + await asyncio.Event().wait() + + task = asyncio.create_task(invoke(limiter, "search", cancelled_tool)) + await asyncio.wait_for(started.wait(), timeout=1) + + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + executions = 0 + + async def execute() -> None: + nonlocal executions + executions += 1 + + rejected = await invoke(limiter, "search", execute) + + assert executions == 0 + assert "rate limit exceeded" in rejected.result.lower() + + clock.advance(1.0) + + admitted = await invoke(limiter, "search", execute) + + assert admitted.result is None + assert executions == 1 + + +@pytest.mark.asyncio +async def test_backward_clock_does_not_create_tokens() -> None: + """A backward clock jump must not replenish exhausted capacity.""" + clock = FakeClock() + limiter = ToolRateLimitMiddleware( + {"search": ToolRateLimit(requests_per_second=1, burst=1)}, + clock=clock, + ) + + executions = 0 + + async def execute() -> None: + nonlocal executions + executions += 1 + + await invoke(limiter, "search", execute) + + clock.advance(-10.0) + rejected = await invoke(limiter, "search", execute) + + assert "rate limit exceeded" in rejected.result.lower() + assert executions == 1 + + clock.advance(10.0) + still_rejected = await invoke(limiter, "search", execute) + assert "rate limit exceeded" in still_rejected.result.lower() + assert executions == 1 + + clock.advance(1.0) + admitted = await invoke(limiter, "search", execute) + + assert admitted.result is None + assert executions == 2 + + +@pytest.mark.asyncio +async def test_rate_limit_integrates_with_function_middleware_pipeline() -> None: + """Rejected calls must not reach the pipeline's final handler.""" + from agent_framework._middleware import FunctionMiddlewarePipeline + + limiter = ToolRateLimitMiddleware( + {"search": ToolRateLimit(requests_per_second=1, burst=1)}, + clock=FakeClock(), + ) + pipeline = FunctionMiddlewarePipeline(limiter) + + executions = 0 + + async def final_handler(context: FunctionInvocationContext) -> str: + nonlocal executions + executions += 1 + return "tool executed" + + first = make_context("search") + first_result = await pipeline.execute(first, final_handler) + + second = make_context("search") + second_result = await pipeline.execute(second, final_handler) + + assert first_result == "tool executed" + assert isinstance(second_result, str) + assert "rate limit exceeded" in second_result.lower() + assert executions == 1 + + +@pytest.mark.asyncio +async def test_rate_limit_auto_function_invocation() -> None: + """Rate limiting must short-circuit real function invocation.""" + from agent_framework import Content, tool + from agent_framework._middleware import FunctionMiddlewarePipeline + from agent_framework._tools import ( + _auto_invoke_function, + normalize_function_invocation_configuration, + ) + + executions = 0 + + @tool(name="limited_search") + def limited_search() -> str: + nonlocal executions + executions += 1 + return "search completed" + + limiter = ToolRateLimitMiddleware( + {"limited_search": ToolRateLimit(requests_per_second=1, burst=1)}, + clock=FakeClock(), + ) + pipeline = FunctionMiddlewarePipeline(limiter) + config = normalize_function_invocation_configuration(None) + tool_map = {limited_search.name: limited_search} + + first = await _auto_invoke_function( + Content.from_function_call( + call_id="rate-limit-first", + name=limited_search.name, + arguments={}, + ), + config=config, + tool_map=tool_map, + middleware_pipeline=pipeline, + ) + + second = await _auto_invoke_function( + Content.from_function_call( + call_id="rate-limit-second", + name=limited_search.name, + arguments={}, + ), + config=config, + tool_map=tool_map, + middleware_pipeline=pipeline, + ) + + assert first.type == "function_result" + assert first.call_id == "rate-limit-first" + assert first.result == "search completed" + + assert second.type == "function_result" + assert second.call_id == "rate-limit-second" + assert isinstance(second.result, str) + assert "rate limit exceeded" in second.result.lower() + + assert executions == 1 + + +@pytest.mark.parametrize("rate", [True, False]) +def test_boolean_request_rates_are_rejected(rate: bool) -> None: + """Boolean values must not be accepted as request rates.""" + with pytest.raises(ValueError, match="requests_per_second"): + ToolRateLimit(requests_per_second=rate) diff --git a/python/samples/02-agents/middleware/README.md b/python/samples/02-agents/middleware/README.md index 4b7d0bbdfb0..e3067c83880 100644 --- a/python/samples/02-agents/middleware/README.md +++ b/python/samples/02-agents/middleware/README.md @@ -19,6 +19,7 @@ This folder contains focused middleware samples for `Agent`, chat clients, tools | [`function_based_middleware.py`](./function_based_middleware.py) | Shows function-based agent and function middleware. | | [`middleware_termination.py`](./middleware_termination.py) | Demonstrates stopping a middleware pipeline early. | | [`message_injection_middleware.py`](./message_injection_middleware.py) | Demonstrates `MessageInjectionMiddleware` with a real Foundry chat client: enqueueing a follow-up message into the active session while a long-running async tool is awaiting. | +| [`per_tool_rate_limiting.py`](./per_tool_rate_limiting.py) | Demonstrates opt-in, per-tool token-bucket admission control, independent burst quotas, and fail-fast rejection without external services. | | [`override_result_with_middleware.py`](./override_result_with_middleware.py) | Shows how middleware registers result transforms and buffered re-derivation on its context before execution, then post-processes regular and streaming responses. | | [`runtime_context_delegation.py`](./runtime_context_delegation.py) | Demonstrates delegating arguments with runtime context data. | | [`session_behavior_middleware.py`](./session_behavior_middleware.py) | Shows how middleware interacts with session-backed runs. | diff --git a/python/samples/02-agents/middleware/per_tool_rate_limiting.py b/python/samples/02-agents/middleware/per_tool_rate_limiting.py new file mode 100644 index 00000000000..edc21086949 --- /dev/null +++ b/python/samples/02-agents/middleware/per_tool_rate_limiting.py @@ -0,0 +1,56 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Demonstrate per-tool request-rate limiting without external services.""" + +import asyncio + +from agent_framework import ( + FunctionInvocationContext, + ToolRateLimit, + ToolRateLimitMiddleware, + tool, +) + + +@tool +def search() -> str: + """Search for information.""" + return "Search completed" + + +@tool +def lookup() -> str: + """Look up a record.""" + return "Lookup completed" + + +async def main() -> None: + """Demonstrate independent quotas and fail-fast rejection.""" + limiter = ToolRateLimitMiddleware({ + "search": ToolRateLimit(requests_per_second=1, burst=2), + "lookup": ToolRateLimit(requests_per_second=1, burst=1), + }) + + async def run_tool(function) -> None: + context = FunctionInvocationContext( + function=function, + arguments={}, + ) + + async def execute() -> None: + context.result = function.func() + + await limiter.process(context, execute) + print(f"{function.name}: {context.result}") + + print("Search: burst capacity of 2") + for _ in range(3): + await run_tool(search) + + print("\nLookup: independent burst capacity of 1") + for _ in range(2): + await run_tool(lookup) + + +if __name__ == "__main__": + asyncio.run(main())