diff --git a/python/samples/02-agents/middleware/README.md b/python/samples/02-agents/middleware/README.md index 4b7d0bbdfb0..1c176ab2992 100644 --- a/python/samples/02-agents/middleware/README.md +++ b/python/samples/02-agents/middleware/README.md @@ -50,3 +50,45 @@ 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 execution deadlines + +The `per_tool_execution_deadlines.py` sample demonstrates independent +asynchronous deadlines using the existing `FunctionMiddleware` extension +point and Python 3.10-compatible `asyncio` APIs. + +Run from the `python` directory: + +```bash +python samples/02-agents/middleware/per_tool_execution_deadlines.py +python -m pytest samples/02-agents/middleware/test_per_tool_execution_deadlines.py -v +``` + +Each configured tool has an independent deadline. Unconfigured tools +execute normally. + +- **Read-only tools:** Deadline expiry produces a recoverable result. +- **Side-effecting tools:** Deadline expiry raises `MiddlewareFailure` + because side effects may already have occurred. +- **Tool-generated timeouts:** A tool's own `TimeoutError` propagates + separately from an enforced deadline. +- **Caller cancellation:** Cancellation propagates rather than being + converted into a normal timeout result. +- **Concurrency:** Different tools can have independent budgets. + +The sample introduces no new framework API and does not implement +automatic retries or distributed deadline enforcement. + +**Cancellation and cleanup limitations:** Async cancellation is +cooperative. Cleanup can take longer than the configured deadline. +Synchronous worker threads or remote operations may continue running. +Deadline expiry does not guarantee rollback, exactly-once execution, +or absence of side effects. + +Caller cancellation is propagated after requesting cancellation of the +underlying asynchronous task. That task may continue cooperative cleanup +after the caller has been cancelled. Any eventual task exception is +observed to avoid an unhandled-task warning. + +This sample does not provide hard termination, rollback, exactly-once +execution, or automatic retries. diff --git a/python/samples/02-agents/middleware/per_tool_execution_deadlines.py b/python/samples/02-agents/middleware/per_tool_execution_deadlines.py new file mode 100644 index 00000000000..8eb3e2987bc --- /dev/null +++ b/python/samples/02-agents/middleware/per_tool_execution_deadlines.py @@ -0,0 +1,130 @@ +# Copyright (c) Microsoft. All rights reserved. +"""Per-tool cooperative async deadlines using FunctionMiddleware. + +Run this sample with Python 3.10 or later. +""" + +import asyncio +import math +from collections.abc import Awaitable, Callable +from dataclasses import dataclass + +from agent_framework import FunctionInvocationContext, FunctionMiddleware, MiddlewareFailure + + +@dataclass(frozen=True) +class DeadlinePolicy: + """Timeout budget and side-effect classification for one tool.""" + + seconds: float + read_only: bool = False + + def __post_init__(self) -> None: + if not math.isfinite(self.seconds) or self.seconds <= 0: + raise ValueError("deadline must be a finite positive number") + + +class PerToolDeadlineMiddleware(FunctionMiddleware): + """Enforce independent deadlines for configured async tools. + + Cancellation is cooperative. The middleware waits for task cleanup; + therefore actual elapsed time may exceed the configured deadline. + """ + + def __init__(self, policies: dict[str, DeadlinePolicy]) -> None: + self._policies = dict(policies) + + @staticmethod + def _observe_task_completion(task: asyncio.Task[None]) -> None: + """Retrieve exceptions from a task finishing after cancellation.""" + if not task.cancelled(): + task.exception() + + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + policy = self._policies.get(context.function.name) + if policy is None: + await call_next() + return + + task = asyncio.create_task(call_next()) + deadline_cancel_requested = False + try: + done, _ = await asyncio.wait({task}, timeout=policy.seconds) + + if done: + # Propagates tool-raised TimeoutError and other exceptions. + await task + return + + # The deadline expired while the tool task was pending. + deadline_cancel_requested = True + task.cancel() + + # Wait for cooperative cleanup without suppressing + # cancellation of the calling middleware task. + try: + await asyncio.shield(task) + except asyncio.CancelledError: + if not task.done(): + # Caller cancellation during cleanup must propagate. + raise + if not task.cancelled(): + # A completed child may itself have raised cancellation. + raise + except MiddlewareFailure: + # Preserve fatal downstream middleware failures. + raise + except Exception: + # Ordinary cleanup errors do not establish that the + # timed-out operation completed successfully. + pass + + if policy.read_only: + context.result = f"Read-only tool '{context.function.name}' exceeded its deadline." + return + + raise MiddlewareFailure( + f"Tool '{context.function.name}' exceeded its deadline; remote side effects may have occurred." + ) + + except asyncio.CancelledError: + # Avoid cancelling the child a second time while it is + # already performing deadline-triggered cleanup. + if not deadline_cancel_requested: + task.cancel() + task.add_done_callback(self._observe_task_completion) + # Preserve caller cancellation. The tool may continue cleanup. + raise + + +async def main() -> None: + """Demonstrate recoverable and uncertain deadline expiry.""" + + from types import SimpleNamespace + + middleware = PerToolDeadlineMiddleware({ + "lookup": DeadlinePolicy(0.01, read_only=True), + "update": DeadlinePolicy(0.01, read_only=False), + }) + + async def slow_tool() -> None: + await asyncio.sleep(0.05) + + for name in ("lookup", "update"): + context = SimpleNamespace( + function=SimpleNamespace(name=name), + result=None, + ) + try: + await middleware.process(context, slow_tool) + print(f"{name}: {context.result}") + except MiddlewareFailure as exc: + print(f"{name}: {exc}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/python/samples/02-agents/middleware/test_per_tool_execution_deadlines.py b/python/samples/02-agents/middleware/test_per_tool_execution_deadlines.py new file mode 100644 index 00000000000..557bf91f110 --- /dev/null +++ b/python/samples/02-agents/middleware/test_per_tool_execution_deadlines.py @@ -0,0 +1,417 @@ +# Copyright (c) Microsoft. All rights reserved. +"""Offline tests for the per-tool deadline middleware sample.""" + +import asyncio +from types import SimpleNamespace +from typing import Any + +import pytest +from agent_framework import MiddlewareFailure +from per_tool_execution_deadlines import DeadlinePolicy, PerToolDeadlineMiddleware + + +def ctx(name: str = "tool") -> Any: + return SimpleNamespace(function=SimpleNamespace(name=name), result=None) + + +@pytest.mark.parametrize("seconds", [0, -1, float("nan"), float("inf"), -float("inf")]) +def test_invalid_deadline(seconds: float) -> None: + with pytest.raises(ValueError, match="finite positive"): + DeadlinePolicy(seconds) + + +@pytest.mark.asyncio +async def test_read_only_deadline_is_recoverable() -> None: + mw = PerToolDeadlineMiddleware({"tool": DeadlinePolicy(0.01, read_only=True)}) + context = ctx() + cleaned = asyncio.Event() + + async def slow() -> None: + try: + await asyncio.sleep(1) + finally: + cleaned.set() + + await mw.process(context, slow) + assert "exceeded its deadline" in context.result + assert cleaned.is_set() + + +@pytest.mark.asyncio +async def test_mutating_deadline_fails_closed() -> None: + mw = PerToolDeadlineMiddleware({"tool": DeadlinePolicy(0.01)}) + + async def slow() -> None: + await asyncio.sleep(1) + + with pytest.raises(MiddlewareFailure, match="side effects may have occurred"): + await mw.process(ctx(), slow) + + +@pytest.mark.asyncio +async def test_tool_timeout_is_not_deadline_expiry() -> None: + mw = PerToolDeadlineMiddleware({"tool": DeadlinePolicy(1, read_only=True)}) + + async def own_timeout() -> None: + raise TimeoutError("tool's own timeout") + + with pytest.raises(TimeoutError, match="own timeout"): + await mw.process(ctx(), own_timeout) + + +@pytest.mark.asyncio +async def test_caller_cancellation_propagates() -> None: + mw = PerToolDeadlineMiddleware({"tool": DeadlinePolicy(1)}) + started = asyncio.Event() + cleaned = asyncio.Event() + + async def slow() -> None: + started.set() + try: + await asyncio.sleep(1) + finally: + cleaned.set() + + outer = asyncio.create_task(mw.process(ctx(), slow)) + await started.wait() + outer.cancel() + + with pytest.raises(asyncio.CancelledError): + await outer + + assert cleaned.is_set() + + +@pytest.mark.asyncio +async def test_independent_deadlines_concurrently() -> None: + mw = PerToolDeadlineMiddleware({ + "fast": DeadlinePolicy(0.01, read_only=True), + "slow": DeadlinePolicy(1, read_only=True), + }) + fast, slow = ctx("fast"), ctx("slow") + + async def wait() -> None: + await asyncio.sleep(0.03) + + await asyncio.gather( + mw.process(fast, wait), + mw.process(slow, wait), + ) + assert "exceeded" in fast.result + assert slow.result is None + + +@pytest.mark.asyncio +async def test_unconfigured_tool_passes_through() -> None: + mw = PerToolDeadlineMiddleware({}) + called = False + + async def run() -> None: + nonlocal called + called = True + + await mw.process(ctx("unconfigured"), run) + assert called + + +@pytest.mark.asyncio +async def test_normal_completion() -> None: + mw = PerToolDeadlineMiddleware({"tool": DeadlinePolicy(1, read_only=True)}) + context = ctx() + + async def run() -> None: + context.result = "ok" + + await mw.process(context, run) + assert context.result == "ok" + + +@pytest.mark.asyncio +async def test_tool_exception_propagates() -> None: + mw = PerToolDeadlineMiddleware({"tool": DeadlinePolicy(1, read_only=True)}) + + async def fail() -> None: + raise RuntimeError("tool failed") + + with pytest.raises(RuntimeError, match="tool failed"): + await mw.process(ctx(), fail) + + +@pytest.mark.asyncio +async def test_mutating_timeout_does_not_cancel_sibling() -> None: + mw = PerToolDeadlineMiddleware({ + "write": DeadlinePolicy(0.01), + "read": DeadlinePolicy(1, read_only=True), + }) + sibling_finished = asyncio.Event() + + async def write() -> None: + await asyncio.sleep(1) + + async def read() -> None: + await asyncio.sleep(0.03) + sibling_finished.set() + + results = await asyncio.gather( + mw.process(ctx("write"), write), + mw.process(ctx("read"), read), + return_exceptions=True, + ) + + assert isinstance(results[0], MiddlewareFailure) + assert results[1] is None + assert sibling_finished.is_set() + + +@pytest.mark.asyncio +async def test_caller_cancellation_with_delayed_cleanup() -> None: + mw = PerToolDeadlineMiddleware({"tool": DeadlinePolicy(1)}) + started = asyncio.Event() + cleanup_started = asyncio.Event() + cleanup_finished = asyncio.Event() + + async def slow() -> None: + started.set() + try: + await asyncio.sleep(10) + finally: + cleanup_started.set() + await asyncio.sleep(0.02) + cleanup_finished.set() + + outer = asyncio.create_task(mw.process(ctx(), slow)) + await started.wait() + outer.cancel() + + with pytest.raises(asyncio.CancelledError): + await outer + + await asyncio.wait_for(cleanup_finished.wait(), timeout=1) + assert cleanup_started.is_set() + + +@pytest.mark.asyncio +async def test_timeout_does_not_retry_tool() -> None: + mw = PerToolDeadlineMiddleware({"tool": DeadlinePolicy(0.01, read_only=True)}) + calls = 0 + + async def slow() -> None: + nonlocal calls + calls += 1 + await asyncio.sleep(1) + + await mw.process(ctx(), slow) + assert calls == 1 + + +@pytest.mark.asyncio +async def test_timeout_cleanup_can_exceed_deadline() -> None: + mw = PerToolDeadlineMiddleware({"tool": DeadlinePolicy(0.01, read_only=True)}) + cleaned = asyncio.Event() + + async def slow() -> None: + try: + await asyncio.sleep(1) + finally: + await asyncio.sleep(0.02) + cleaned.set() + + await mw.process(ctx(), slow) + assert cleaned.is_set() + + +@pytest.mark.asyncio +async def test_deadline_through_real_function_middleware_pipeline() -> None: + """Exercise the framework pipeline rather than calling process directly.""" + from agent_framework._middleware import ( + FunctionInvocationContext, + FunctionMiddlewarePipeline, + ) + from agent_framework._tools import FunctionTool + + middleware = PerToolDeadlineMiddleware({"slow_tool": DeadlinePolicy(0.01, read_only=True)}) + pipeline = FunctionMiddlewarePipeline(middleware) + function = FunctionTool( + name="slow_tool", + func=lambda: "unused", + approval_mode="never_require", + ) + context = FunctionInvocationContext(function=function, arguments={}) + + calls = 0 + + async def handler(ctx: FunctionInvocationContext) -> str: + nonlocal calls + calls += 1 + await asyncio.sleep(1) + return "late result" + + result = await pipeline.execute(context, handler) + + assert result == "Read-only tool 'slow_tool' exceeded its deadline." + assert calls == 1 + + +@pytest.mark.asyncio +async def test_mutating_deadline_through_real_function_middleware_pipeline() -> None: + """A side-effecting timeout must escape the pipeline as MiddlewareFailure.""" + from agent_framework._middleware import ( + FunctionInvocationContext, + FunctionMiddlewarePipeline, + ) + from agent_framework._tools import FunctionTool + + middleware = PerToolDeadlineMiddleware({"write_tool": DeadlinePolicy(0.01, read_only=False)}) + pipeline = FunctionMiddlewarePipeline(middleware) + function = FunctionTool( + name="write_tool", + func=lambda: "unused", + approval_mode="never_require", + ) + context = FunctionInvocationContext(function=function, arguments={}) + + async def handler(ctx: FunctionInvocationContext) -> str: + await asyncio.sleep(1) + return "unexpected" + + with pytest.raises(MiddlewareFailure, match="side effects may have occurred"): + await pipeline.execute(context, handler) + + +@pytest.mark.asyncio +async def test_mutating_deadline_aborts_real_agent_concurrent_batch() -> None: + """A mutating deadline aborts the batch and cancels an in-flight sibling.""" + import sys + + # Load the framework's existing mock-client fixture for this test. + from agent_framework import Agent, ChatResponse, Content, Message + from agent_framework._tools import FunctionTool + + tests_root = "packages/core/tests" + if tests_root not in sys.path: + sys.path.insert(0, tests_root) + + from core.conftest import MockBaseChatClient + + chat_client_base = MockBaseChatClient() + + sibling_started = asyncio.Event() + sibling_cancelled = asyncio.Event() + write_calls = 0 + + async def write_impl(location: str) -> str: + nonlocal write_calls + write_calls += 1 + await sibling_started.wait() + await asyncio.Event().wait() + return "never" + + async def sibling_impl(location: str) -> str: + sibling_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + sibling_cancelled.set() + raise + return "never" + + write_tool = FunctionTool( + name="write_tool", + func=write_impl, + approval_mode="never_require", + ) + sibling_tool = FunctionTool( + name="sibling_tool", + func=sibling_impl, + approval_mode="never_require", + ) + + chat_client_base.run_responses = [ + ChatResponse( + messages=[ + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="call_1", + name="write_tool", + arguments='{"location": "Seattle"}', + ), + Content.from_function_call( + call_id="call_2", + name="sibling_tool", + arguments='{"location": "Seattle"}', + ), + ], + ) + ] + ) + ] + + middleware = PerToolDeadlineMiddleware({"write_tool": DeadlinePolicy(0.05, read_only=False)}) + + agent = Agent( + client=chat_client_base, + middleware=[middleware], + tools=[write_tool, sibling_tool], + ) + + with pytest.raises( + MiddlewareFailure, + match="side effects may have occurred", + ): + await asyncio.wait_for(agent.run("run both tools"), timeout=2) + + assert write_calls == 1 + assert sibling_cancelled.is_set() + assert chat_client_base.call_count == 1 + + +@pytest.mark.asyncio +async def test_cleanup_middleware_failure_propagates() -> None: + """A fatal downstream failure during cleanup must not be swallowed.""" + mw = PerToolDeadlineMiddleware({ + "tool": DeadlinePolicy(0.01, read_only=True), + }) + + async def downstream() -> None: + try: + await asyncio.sleep(10) + finally: + raise MiddlewareFailure("fatal downstream failure") + + with pytest.raises(MiddlewareFailure, match="fatal downstream failure"): + await mw.process(ctx(), downstream) + + +@pytest.mark.asyncio +async def test_caller_cancellation_during_deadline_cleanup_propagates() -> None: + """Caller cancellation during cleanup must not become a timeout.""" + mw = PerToolDeadlineMiddleware({ + "tool": DeadlinePolicy(0.01, read_only=True), + }) + cleanup_started = asyncio.Event() + cleanup_release = asyncio.Event() + cleanup_finished = asyncio.Event() + + async def slow() -> None: + try: + await asyncio.Event().wait() + finally: + cleanup_started.set() + await cleanup_release.wait() + cleanup_finished.set() + + outer = asyncio.create_task(mw.process(ctx(), slow)) + + try: + await asyncio.wait_for(cleanup_started.wait(), timeout=1) + + # Cancel only after the deadline has fired and cleanup has begun. + outer.cancel() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(outer, timeout=1) + finally: + cleanup_release.set() + await asyncio.wait_for(cleanup_finished.wait(), timeout=1)