From 3d34b8c2f65ca805bbb930b5d2a60ddcd661df79 Mon Sep 17 00:00:00 2001 From: suhasagg Date: Wed, 28 Oct 2026 12:55:09 +0530 Subject: [PATCH 1/2] docs(python): add concurrency-safe per-tool rate limiting sample --- python/samples/02-agents/middleware/README.md | 12 ++ .../middleware/per_tool_rate_limiting.py | 83 ++++++++ .../middleware/test_per_tool_rate_limiting.py | 199 ++++++++++++++++++ 3 files changed, 294 insertions(+) create mode 100644 python/samples/02-agents/middleware/per_tool_rate_limiting.py create mode 100644 python/samples/02-agents/middleware/test_per_tool_rate_limiting.py diff --git a/python/samples/02-agents/middleware/README.md b/python/samples/02-agents/middleware/README.md index 4b7d0bbdfb0..5defd5ee573 100644 --- a/python/samples/02-agents/middleware/README.md +++ b/python/samples/02-agents/middleware/README.md @@ -6,6 +6,7 @@ This folder contains focused middleware samples for `Agent`, chat clients, tools | File | Description | |------|-------------| +| [`per_tool_rate_limiting.py`](./per_tool_rate_limiting.py) | In-process per-tool token-bucket middleware; offline tests in `test_per_tool_rate_limiting.py`. | | [`agent_and_run_level_middleware.py`](./agent_and_run_level_middleware.py) | Demonstrates combining agent-level and run-level middleware. | | [`agent_loop_middleware_refinement.py`](./agent_loop_middleware_refinement.py) | Demonstrates `AgentLoopMiddleware` with a `should_continue` predicate: a completion-marker refinement loop with feedback tracking and `fresh_context`. | | [`agent_loop_middleware_todos.py`](./agent_loop_middleware_todos.py) | Demonstrates `AgentLoopMiddleware` with a `should_continue` predicate built from a `TodoProvider` via `todos_remaining`, so the agent keeps working while open todos remain. | @@ -50,3 +51,14 @@ agent's latest response to a second, external judge chat client on every iterati or malicious judge endpoint could exfiltrate that data, or return a manipulated verdict/gap analysis that gets fed back into the loop as feedback — a form of indirect prompt injection. Only configure a judge client that points at a service you trust as much as the primary model. + +## Per-tool rate limiting + +Run `python per_tool_rate_limiting.py` to see two admitted calls and one rejected call. Quotas are per middleware instance and tool name, in-process only. Rejections are tool-visible results before execution, not `MiddlewareTermination`; no retries or distributed guarantees. Thread-safe accounting is shared across event loops. +## Per-tool rate limiting + +Run `python per_tool_rate_limiting.py` for a self-contained example. The implementation uses the existing `FunctionMiddleware` extension point and one token bucket per configured tool name. Unconfigured tools are admitted. All buckets belong to **one middleware instance**: sharing an instance across agents shares its quotas; separate processes/instances do not coordinate. Accounting uses a monotonic clock and a short `threading.Lock`, held only for token calculations and never while awaiting a tool. + +An exhausted bucket short-circuits before downstream middleware and the tool execute. It sets `context.result` to a tool-visible error string rather than raising `MiddlewareTermination` (which can end the tool loop). A rejected tool call is not retried automatically. The agent may choose to request the tool again, so configure appropriate policies and prompts. Tokens are charged **on admission**, including for subsequently cancelled or failed executions; there is no refund. No waiting queue is used, so cancellation while awaiting quota admission does not arise. The sample does **not** implement distributed limits, remote-operation rollback, or idempotency. + +For agent integration, pass a shared instance in `Agent(..., middleware=[PerToolRateLimitMiddleware({"lookup": RatePolicy(2, 1.0)})])` alongside your actual client and tools. See `test_per_tool_rate_limiting.py` for an offline integration test using the framework's real `FunctionInvocationContext` and `FunctionMiddlewarePipeline`. 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..1330213f4a9 --- /dev/null +++ b/python/samples/02-agents/middleware/per_tool_rate_limiting.py @@ -0,0 +1,83 @@ +# Copyright (c) Microsoft. All rights reserved. +"""Local per-tool token-bucket admission control using FunctionMiddleware. + +A quota belongs to this middleware instance and is keyed by the tool name. +It is not a distributed limit and it does not make tool execution idempotent. +A rejected invocation returns a tool-visible result without invoking the tool. +""" + +import asyncio +import math +import threading +import time +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any + +from agent_framework import FunctionInvocationContext, FunctionMiddleware + + +@dataclass(frozen=True) +class RatePolicy: + capacity: int + refill_per_second: float + + def __post_init__(self) -> None: + if ( + isinstance(self.capacity, bool) + or not isinstance(self.capacity, int) + or self.capacity <= 0 + or not math.isfinite(self.refill_per_second) + or self.refill_per_second <= 0 + ): + raise ValueError("capacity and refill_per_second must be finite and positive") + + +class PerToolRateLimitMiddleware(FunctionMiddleware): + """Share one quota per tool name across concurrent calls and event loops.""" + + def __init__(self, policies: Mapping[str, RatePolicy]) -> None: + self._policies = dict(policies) + self._buckets: dict[str, tuple[float, float]] = {} + # No asyncio primitives: one instance can be reused across event loops. + self._lock = threading.Lock() + + def admit(self, name: str, *, now: float | None = None) -> bool: + """Atomically spend one token, or reject without spending a token.""" + policy = self._policies.get(name) + if policy is None: + return True + timestamp = time.monotonic() if now is None else now + with self._lock: + tokens, last = self._buckets.get(name, (float(policy.capacity), timestamp)) + elapsed = max(0.0, timestamp - last) + tokens = min(float(policy.capacity), tokens + elapsed * policy.refill_per_second) + admitted = tokens >= 1.0 + self._buckets[name] = (tokens - 1.0 if admitted else tokens, max(last, timestamp)) + return admitted + + async def process(self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None: + name = context.function.name + if not self.admit(name): + # Do not use MiddlewareTermination: it can terminate the tool loop. + context.result = f"Tool '{name}' is rate limited. Retry later if appropriate." + return + # Never hold the lock during execution. Do not swallow cancellation. + await call_next() + + +async def main() -> None: + limiter = PerToolRateLimitMiddleware({"lookup": RatePolicy(2, 0.01)}) + for _ in range(3): + context: Any = SimpleNamespace(function=SimpleNamespace(name="lookup"), result=None) + + async def invoke(current: Any = context) -> None: + current.result = "tool executed" + + await limiter.process(context, invoke) + print(context.result) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/python/samples/02-agents/middleware/test_per_tool_rate_limiting.py b/python/samples/02-agents/middleware/test_per_tool_rate_limiting.py new file mode 100644 index 00000000000..67ef6840aa0 --- /dev/null +++ b/python/samples/02-agents/middleware/test_per_tool_rate_limiting.py @@ -0,0 +1,199 @@ +# Copyright (c) Microsoft. All rights reserved. +"""Offline regression tests for the per-tool rate-limit sample.""" + +import asyncio +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace +from typing import Any + +import pytest +from per_tool_rate_limiting import PerToolRateLimitMiddleware, RatePolicy + + +def context(name: str = "lookup") -> Any: + return SimpleNamespace(function=SimpleNamespace(name=name), result=None) + + +@pytest.mark.parametrize( + "capacity,refill", [(0, 1), (-1, 1), (1, 0), (1, -1), (1, float("inf")), (1, float("nan")), (True, 1)] +) +def test_invalid_policy(capacity: int, refill: float) -> None: + with pytest.raises(ValueError): + RatePolicy(capacity, refill) + + +def test_refill_and_independent_tool_quotas() -> None: + mw = PerToolRateLimitMiddleware({"a": RatePolicy(1, 1), "b": RatePolicy(2, 1)}) + assert mw.admit("a", now=10) + assert not mw.admit("a", now=10) + assert mw.admit("b", now=10) + assert mw.admit("b", now=10) + assert not mw.admit("b", now=10) + assert mw.admit("a", now=11) + assert mw.admit("unconfigured", now=11) + + +def test_partial_refill_and_clock_rollback() -> None: + mw = PerToolRateLimitMiddleware({"a": RatePolicy(1, 0.5)}) + assert mw.admit("a", now=10) + assert not mw.admit("a", now=11) + assert not mw.admit("a", now=9) + assert mw.admit("a", now=12) + + +def test_thread_safe_admission() -> None: + mw = PerToolRateLimitMiddleware({"a": RatePolicy(7, 0.000001)}) + with ThreadPoolExecutor(max_workers=24) as pool: + results = list(pool.map(lambda _: mw.admit("a", now=100), range(200))) + assert sum(results) == 7 + + +@pytest.mark.asyncio +async def test_concurrent_invocations_never_exceed_capacity() -> None: + mw = PerToolRateLimitMiddleware({"a": RatePolicy(3, 0.000001)}) + called = 0 + + async def one() -> str | None: + nonlocal called + ctx = context("a") + + async def execute() -> None: + nonlocal called + called += 1 + await asyncio.sleep(0) + ctx.result = "ok" + + await mw.process(ctx, execute) + return ctx.result + + results = await asyncio.gather(*(one() for _ in range(50))) + assert called == 3 + assert results.count("ok") == 3 + assert sum("rate limited" in r for r in results if r is not None) == 47 + + +@pytest.mark.asyncio +async def test_rejection_before_side_effects() -> None: + mw = PerToolRateLimitMiddleware({"a": RatePolicy(1, 0.000001)}) + called = 0 + + async def execute() -> None: + nonlocal called + called += 1 + + first, second = context("a"), context("a") + await mw.process(first, execute) + await mw.process(second, execute) + assert called == 1 + assert "rate limited" in second.result + + +def test_reuse_across_separate_event_loops() -> None: + mw = PerToolRateLimitMiddleware({"a": RatePolicy(1, 0.000001)}) + + async def run() -> str | None: + ctx = context("a") + + async def execute() -> None: + ctx.result = "ok" + + await mw.process(ctx, execute) + return ctx.result + + assert asyncio.run(run()) == "ok" + assert "rate limited" in asyncio.run(run()) + + +@pytest.mark.asyncio +async def test_caller_cancellation_is_not_swallowed() -> None: + mw = PerToolRateLimitMiddleware({"a": RatePolicy(1, 1)}) + started = asyncio.Event() + + async def execute() -> None: + started.set() + await asyncio.sleep(3600) + + task = asyncio.create_task(mw.process(context("a"), execute)) + await started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + +@pytest.mark.asyncio +async def test_tool_exception_is_not_swallowed() -> None: + mw = PerToolRateLimitMiddleware({"a": RatePolicy(1, 1)}) + + async def execute() -> None: + raise RuntimeError("tool failure") + + with pytest.raises(RuntimeError, match="tool failure"): + await mw.process(context("a"), execute) + + +@pytest.mark.asyncio +async def test_real_framework_context_and_pipeline() -> None: + """Exercise the framework's actual context, tool and middleware pipeline.""" + from agent_framework import FunctionInvocationContext, tool + from agent_framework._middleware import FunctionMiddlewarePipeline + + calls: list[str] = [] + + @tool(approval_mode="never_require") + def lookup(query: str) -> str: + calls.append(query) + return f"found:{query}" + + limiter = PerToolRateLimitMiddleware({"lookup": RatePolicy(1, 0.000001)}) + pipeline = FunctionMiddlewarePipeline(limiter) + + async def invoke(query: str) -> Any: + ctx = FunctionInvocationContext(function=lookup, arguments={"query": query}) + + async def execute(current: FunctionInvocationContext) -> Any: + return await current.function.invoke(arguments=current.arguments) + + await pipeline.execute(ctx, execute) + return ctx.result + + admitted = await invoke("first") + rejected = await invoke("second") + assert calls == ["first"] + assert admitted is not None + assert isinstance(rejected, str) + assert "rate limited" in rejected + + +@pytest.mark.asyncio +async def test_rejected_call_does_not_run_other_middleware_or_tool() -> None: + from agent_framework import FunctionInvocationContext, FunctionMiddleware, tool + from agent_framework._middleware import FunctionMiddlewarePipeline + + events: list[str] = [] + + @tool(approval_mode="never_require") + def lookup(query: str) -> str: + events.append("tool") + return query + + class Downstream(FunctionMiddleware): + async def process(self, context: FunctionInvocationContext, call_next: Any) -> None: + events.append("downstream") + await call_next() + + pipeline = FunctionMiddlewarePipeline(PerToolRateLimitMiddleware({"lookup": RatePolicy(1, 0.000001)}), Downstream()) + + async def invoke() -> Any: + ctx = FunctionInvocationContext(function=lookup, arguments={"query": "x"}) + + async def execute(current: FunctionInvocationContext) -> Any: + return await current.function.invoke(arguments=current.arguments) + + await pipeline.execute(ctx, execute) + return ctx.result + + await invoke() + assert events == ["downstream", "tool"] + rejected = await invoke() + assert events == ["downstream", "tool"] + assert "rate limited" in rejected From 105054fd1c81db0781a1aaf018e621374f0f682a Mon Sep 17 00:00:00 2001 From: suhasagg Date: Wed, 28 Oct 2026 15:49:56 +0530 Subject: [PATCH 2/2] fix(samples): address Copilot rate limiting review findings --- .../test_per_tool_rate_limiting_sample.py} | 57 ++++++++++++++++++- .../middleware/per_tool_rate_limiting.py | 2 +- 2 files changed, 57 insertions(+), 2 deletions(-) rename python/{samples/02-agents/middleware/test_per_tool_rate_limiting.py => packages/core/tests/core/test_per_tool_rate_limiting_sample.py} (79%) diff --git a/python/samples/02-agents/middleware/test_per_tool_rate_limiting.py b/python/packages/core/tests/core/test_per_tool_rate_limiting_sample.py similarity index 79% rename from python/samples/02-agents/middleware/test_per_tool_rate_limiting.py rename to python/packages/core/tests/core/test_per_tool_rate_limiting_sample.py index 67ef6840aa0..7b65e3d0260 100644 --- a/python/samples/02-agents/middleware/test_per_tool_rate_limiting.py +++ b/python/packages/core/tests/core/test_per_tool_rate_limiting_sample.py @@ -2,12 +2,26 @@ """Offline regression tests for the per-tool rate-limit sample.""" import asyncio +import importlib.util +import sys from concurrent.futures import ThreadPoolExecutor +from pathlib import Path from types import SimpleNamespace from typing import Any import pytest -from per_tool_rate_limiting import PerToolRateLimitMiddleware, RatePolicy + +_SAMPLE = Path(__file__).resolve().parents[4] / "samples/02-agents/middleware/per_tool_rate_limiting.py" + +_SPEC = importlib.util.spec_from_file_location("_per_tool_rate_limiting_sample", _SAMPLE) +assert _SPEC is not None and _SPEC.loader is not None + +_MODULE = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _MODULE +_SPEC.loader.exec_module(_MODULE) + +PerToolRateLimitMiddleware = _MODULE.PerToolRateLimitMiddleware +RatePolicy = _MODULE.RatePolicy def context(name: str = "lookup") -> Any: @@ -197,3 +211,44 @@ async def execute(current: FunctionInvocationContext) -> Any: rejected = await invoke() assert events == ["downstream", "tool"] assert "rate limited" in rejected + + +def test_clock_sampled_only_after_acquiring_lock( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A contending caller must sample time after acquiring the lock.""" + import threading + import time + + limiter = PerToolRateLimitMiddleware({"lookup": RatePolicy(1, 1.0)}) + + # Start with one available token at t=100. + assert limiter.admit("lookup", now=100) + assert limiter.admit("lookup", now=101) + + clock_called = threading.Event() + result: list[bool] = [] + + def fake_monotonic() -> float: + clock_called.set() + return 102.0 + + monkeypatch.setattr(time, "monotonic", fake_monotonic) + + def attempt() -> None: + result.append(limiter.admit("lookup")) + + with limiter._lock: + worker = threading.Thread(target=attempt) + worker.start() + + # While the lock is held, the worker must not sample time. + assert not clock_called.wait(timeout=0.1) + + worker.join(timeout=2) + assert not worker.is_alive() + assert result == [True] + assert clock_called.is_set() + + # No second admission at the same timestamp. + assert not limiter.admit("lookup", now=102) diff --git a/python/samples/02-agents/middleware/per_tool_rate_limiting.py b/python/samples/02-agents/middleware/per_tool_rate_limiting.py index 1330213f4a9..33157558aba 100644 --- a/python/samples/02-agents/middleware/per_tool_rate_limiting.py +++ b/python/samples/02-agents/middleware/per_tool_rate_limiting.py @@ -48,8 +48,8 @@ def admit(self, name: str, *, now: float | None = None) -> bool: policy = self._policies.get(name) if policy is None: return True - timestamp = time.monotonic() if now is None else now with self._lock: + timestamp = time.monotonic() if now is None else now tokens, last = self._buckets.get(name, (float(policy.capacity), timestamp)) elapsed = max(0.0, timestamp - last) tokens = min(float(policy.capacity), tokens + elapsed * policy.refill_per_second)