From e04335937bfc8bc20e6e528668ff1f7011e80887 Mon Sep 17 00:00:00 2001 From: Mark Licata Date: Wed, 7 Oct 2026 08:33:52 -0700 Subject: [PATCH 1/5] Python: add AsyncRequestBroker and shell command audit log Co-Authored-By: Claude Opus 5.5 --- .../packages/core/agent_framework/__init__.py | 2 + .../core/agent_framework/__init__.pyi | 2 + .../_harness/_request_broker.py | 121 ++++++++++++++++++ .../core/tests/core/test_request_broker.py | 59 +++++++++ .../agent_framework_tools/shell/__init__.py | 3 + .../agent_framework_tools/shell/_audit.py | 74 +++++++++++ .../agent_framework_tools/shell/_tool.py | 19 ++- .../tools/tests/test_local_shell_tool.py | 36 ++++++ 8 files changed, 314 insertions(+), 2 deletions(-) create mode 100644 python/packages/core/agent_framework/_harness/_request_broker.py create mode 100644 python/packages/core/tests/core/test_request_broker.py create mode 100644 python/packages/tools/agent_framework_tools/shell/_audit.py diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 3953187578e..9720b557024 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -148,6 +148,7 @@ "MemoryTopicRecord", ), "._harness._mode": ("DEFAULT_MODE_SOURCE_ID", "AgentModeProvider", "get_agent_mode", "set_agent_mode"), + "._harness._request_broker": ("AsyncRequestBroker",), "._harness._todo": ( "DEFAULT_TODO_SOURCE_ID", "TodoFileStore", @@ -451,6 +452,7 @@ "AgentSession", "AggregatingSkillsSource", "Annotation", + "AsyncRequestBroker", "BackgroundAgentsProvider", "BackgroundTaskInfo", "BackgroundTaskStatus", diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index e8d1bc6e411..265ee28a071 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/python/packages/core/agent_framework/__init__.pyi @@ -106,6 +106,7 @@ from ._harness._memory import ( MemoryTopicRecord, ) from ._harness._mode import DEFAULT_MODE_SOURCE_ID, AgentModeProvider, get_agent_mode, set_agent_mode +from ._harness._request_broker import AsyncRequestBroker from ._harness._todo import ( DEFAULT_TODO_SOURCE_ID, TodoFileStore, @@ -406,6 +407,7 @@ __all__ = [ "AgentSession", "AggregatingSkillsSource", "Annotation", + "AsyncRequestBroker", "BackgroundAgentsProvider", "BackgroundTaskInfo", "BackgroundTaskStatus", diff --git a/python/packages/core/agent_framework/_harness/_request_broker.py b/python/packages/core/agent_framework/_harness/_request_broker.py new file mode 100644 index 00000000000..1d9488b04e5 --- /dev/null +++ b/python/packages/core/agent_framework/_harness/_request_broker.py @@ -0,0 +1,121 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""AsyncRequestBroker: a generic asyncio pause-and-wait primitive. + +Provides :class:`AsyncRequestBroker`, which lets any server-side code suspend +execution on a request ID and wait asynchronously for a client to deliver the +matching result via :meth:`AsyncRequestBroker.resolve` or +:meth:`AsyncRequestBroker.reject`. + +The broker tracks recently settled request IDs in a bounded deque (default +256 entries) so that re-using a settled ID is reported instead of silently +hanging. Late responses for unknown IDs return ``False`` rather than raising. + +Usage:: + + broker = AsyncRequestBroker() + + # Server side - suspend until answered (or timeout) + result = await broker.request_async("req-123", timeout=30.0) + + # Client side - deliver the answer + broker.resolve("req-123", decision) +""" + +from __future__ import annotations + +import asyncio +import collections +from typing import Generic, TypeVar + +from .._feature_stage import ExperimentalFeature, experimental + +__all__ = ["AsyncRequestBroker"] + +T = TypeVar("T") + + +@experimental(feature_id=ExperimentalFeature.HARNESS) +class AsyncRequestBroker(Generic[T]): + """Generic asyncio pause-and-wait broker for request/response patterns. + + Allows server-side code to suspend on a request ID and await a client + response delivered via :meth:`resolve` or :meth:`reject`, without + polling or blocking a thread. + + Thread-safety: designed for single-threaded asyncio use. The broker must + be used within a single event loop. + """ + + DEFAULT_SETTLED_DEQUE_SIZE = 256 + + def __init__(self, settled_deque_size: int = DEFAULT_SETTLED_DEQUE_SIZE) -> None: + self._pending: dict[str, asyncio.Future[T]] = {} + self._settled: collections.deque[str] = collections.deque(maxlen=settled_deque_size) + + async def request_async(self, request_id: str, *, timeout: float | None = None) -> T: + """Suspend until the request is resolved or rejected. + + Args: + request_id: Unique identifier for this request. + + Keyword Args: + timeout: Maximum seconds to wait. ``None`` waits indefinitely. + + Returns: + The result passed to :meth:`resolve`. + + Raises: + ValueError: If the ID is already pending or was recently settled. + asyncio.TimeoutError: If ``timeout`` elapses before a response arrives. + asyncio.CancelledError: If the waiting task is cancelled. + Exception: Any exception passed to :meth:`reject`. + """ + if request_id in self._settled: + raise ValueError(f"Request '{request_id}' was already settled.") + if request_id in self._pending: + raise ValueError(f"Request '{request_id}' is already pending.") + loop = asyncio.get_running_loop() + future: asyncio.Future[T] = loop.create_future() + self._pending[request_id] = future + try: + return await asyncio.wait_for(asyncio.shield(future), timeout=timeout) + except BaseException: + if self._pending.get(request_id) is future: + del self._pending[request_id] + future.cancel() + raise + + def resolve(self, request_id: str, result: T) -> bool: + """Deliver a successful result to a waiting :meth:`request_async` call. + + Returns: + ``True`` if a waiter was notified, ``False`` if the request was not found + (already settled or never registered). + """ + future = self._pending.pop(request_id, None) + if future is None: + return False + self._settled.append(request_id) + if not future.done(): + future.set_result(result) + return True + + def reject(self, request_id: str, exception: BaseException) -> bool: + """Deliver an exception to a waiting :meth:`request_async` call. + + Returns: + ``True`` if a waiter was notified. + """ + future = self._pending.pop(request_id, None) + if future is None: + return False + self._settled.append(request_id) + if not future.done(): + future.set_exception(exception) + return True + + @property + def pending_count(self) -> int: + """Number of outstanding requests currently awaiting a response.""" + return len(self._pending) diff --git a/python/packages/core/tests/core/test_request_broker.py b/python/packages/core/tests/core/test_request_broker.py new file mode 100644 index 00000000000..07f2d4c0ef0 --- /dev/null +++ b/python/packages/core/tests/core/test_request_broker.py @@ -0,0 +1,59 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio + +import pytest + +from agent_framework import AsyncRequestBroker + + +async def test_resolve_delivers_result() -> None: + broker: AsyncRequestBroker[str] = AsyncRequestBroker() + task = asyncio.create_task(broker.request_async("r1")) + await asyncio.sleep(0) + assert broker.resolve("r1", "value") is True + assert await task == "value" + + +async def test_reject_delivers_exception() -> None: + broker: AsyncRequestBroker[str] = AsyncRequestBroker() + task = asyncio.create_task(broker.request_async("r1")) + await asyncio.sleep(0) + assert broker.reject("r1", RuntimeError("boom")) is True + with pytest.raises(RuntimeError, match="boom"): + await task + + +async def test_timeout_raises() -> None: + broker: AsyncRequestBroker[str] = AsyncRequestBroker() + with pytest.raises(asyncio.TimeoutError): + await broker.request_async("r1", timeout=0.001) + assert broker.pending_count == 0 + + +async def test_resolve_returns_false_when_not_found() -> None: + broker: AsyncRequestBroker[str] = AsyncRequestBroker() + assert broker.resolve("missing", "x") is False + assert broker.reject("missing", RuntimeError("x")) is False + + +async def test_pending_count() -> None: + broker: AsyncRequestBroker[int] = AsyncRequestBroker() + task = asyncio.create_task(broker.request_async("r1")) + await asyncio.sleep(0) + assert broker.pending_count == 1 + broker.resolve("r1", 1) + await task + assert broker.pending_count == 0 + + +async def test_duplicate_request_raises() -> None: + broker: AsyncRequestBroker[int] = AsyncRequestBroker() + task = asyncio.create_task(broker.request_async("r1")) + await asyncio.sleep(0) + with pytest.raises(ValueError, match="already pending"): + await broker.request_async("r1") + broker.resolve("r1", 1) + await task + with pytest.raises(ValueError, match="already settled"): + await broker.request_async("r1") diff --git a/python/packages/tools/agent_framework_tools/shell/__init__.py b/python/packages/tools/agent_framework_tools/shell/__init__.py index 06517a96c22..0be0e982aed 100644 --- a/python/packages/tools/agent_framework_tools/shell/__init__.py +++ b/python/packages/tools/agent_framework_tools/shell/__init__.py @@ -4,6 +4,7 @@ from __future__ import annotations +from ._audit import ShellAuditLog, ShellCommandRecord from ._docker import ( DEFAULT_IMAGE as DOCKER_DEFAULT_IMAGE, ) @@ -35,7 +36,9 @@ "DockerNotAvailableError", "DockerShellTool", "LocalShellTool", + "ShellAuditLog", "ShellCommandError", + "ShellCommandRecord", "ShellDecision", "ShellEnvironmentProvider", "ShellEnvironmentProviderOptions", diff --git a/python/packages/tools/agent_framework_tools/shell/_audit.py b/python/packages/tools/agent_framework_tools/shell/_audit.py new file mode 100644 index 00000000000..d6a51ab7851 --- /dev/null +++ b/python/packages/tools/agent_framework_tools/shell/_audit.py @@ -0,0 +1,74 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Shell command audit log for compliance and observability.""" + +from __future__ import annotations + +import datetime +from dataclasses import dataclass, field + +__all__ = ["ShellAuditLog", "ShellCommandRecord"] + + +@dataclass +class ShellCommandRecord: + """A record of a single shell command execution attempt. + + Stored in the tool's audit log for compliance and observability. + All timestamps are UTC ISO 8601. + """ + + command: str + """The command string as submitted.""" + + policy_decision: str + """The :class:`~agent_framework_tools.shell.ShellDecision` outcome: ``"allow"`` or ``"deny"``.""" + + timestamp_utc: str = field(default_factory=lambda: datetime.datetime.now(datetime.timezone.utc).isoformat()) + """UTC ISO 8601 timestamp when the command was evaluated.""" + + exit_code: int | None = None + """Process exit code, set after execution. ``None`` if the command was denied or not yet run.""" + + denial_reason: str = "" + """Reason from :class:`~agent_framework_tools.shell.ShellDecision` when denied.""" + + +class ShellAuditLog: + """An in-memory audit log of shell command decisions and outcomes. + + Records every command evaluated by :class:`~agent_framework_tools.shell.LocalShellTool`, + including the policy decision, execution outcome, and exit code. + + Access via ``LocalShellTool.audit_log``. + + Examples: + .. code-block:: python + + tool = LocalShellTool() + await tool.run("echo hello") + + for record in tool.audit_log.records: + print(record.command, record.policy_decision, record.exit_code) + """ + + def __init__(self) -> None: + self._records: list[ShellCommandRecord] = [] + + def record(self, command: str, policy_decision: str, denial_reason: str = "") -> ShellCommandRecord: + """Add a new record at evaluation time (before execution).""" + rec = ShellCommandRecord( + command=command, + policy_decision=policy_decision, + denial_reason=denial_reason, + ) + self._records.append(rec) + return rec + + @property + def records(self) -> list[ShellCommandRecord]: + """All recorded command attempts, in chronological order.""" + return list(self._records) + + def __len__(self) -> int: + return len(self._records) diff --git a/python/packages/tools/agent_framework_tools/shell/_tool.py b/python/packages/tools/agent_framework_tools/shell/_tool.py index d85beeebd03..397b9a58b47 100644 --- a/python/packages/tools/agent_framework_tools/shell/_tool.py +++ b/python/packages/tools/agent_framework_tools/shell/_tool.py @@ -15,6 +15,7 @@ from agent_framework._tools import SHELL_TOOL_KIND_VALUE from .._feature_usage import FeatureIndex +from ._audit import ShellAuditLog from ._executor import run_stateless from ._policy import ShellPolicy, ShellRequest from ._resolve import is_powershell, resolve_shell @@ -143,6 +144,10 @@ class LocalShellTool: will execute any command the model emits. on_command: Optional audit hook called with the command string for every command that passes policy. Use for logging / telemetry. + + Attributes: + audit_log: :class:`ShellAuditLog` recording every evaluated command + (allowed or denied) with its policy decision and exit code. """ def __init__( @@ -181,6 +186,7 @@ def __init__( self._max_output_bytes = max_output_bytes self._approval_mode: Literal["always_require", "never_require"] = approval_mode self._on_command = on_command + self.audit_log: ShellAuditLog = ShellAuditLog() merged_env: dict[str, str] | None if env is None and not clean_env: @@ -256,6 +262,11 @@ async def run(self, command: str, *, timeout: float | None = None) -> ShellResul mark_feature_used(FeatureIndex.TOOLS_SHELL) request = ShellRequest(command=command, workdir=self._workdir) decision = self._policy.evaluate(request) + record = self.audit_log.record( + command=command, + policy_decision=decision.decision, + denial_reason=decision.reason if decision.decision == "deny" else "", + ) if decision.decision == "deny": raise ShellCommandError(f"Command rejected by policy: {decision.reason}") if self._on_command is not None: @@ -272,9 +283,11 @@ async def run(self, command: str, *, timeout: float | None = None) -> ShellResul if self._session is None: raise RuntimeError("LocalShellTool session failed to start") effective = self._maybe_reanchor(command) - return await self._session.run(effective, timeout=effective_timeout) + result = await self._session.run(effective, timeout=effective_timeout) + record.exit_code = result.exit_code + return result - return await run_stateless( + result = await run_stateless( self._stateless_argv, command, workdir=self._workdir, @@ -282,6 +295,8 @@ async def run(self, command: str, *, timeout: float | None = None) -> ShellResul timeout=effective_timeout, max_output_bytes=self._max_output_bytes, ) + record.exit_code = result.exit_code + return result # ------------------------------------------------------------------ AF wiring diff --git a/python/packages/tools/tests/test_local_shell_tool.py b/python/packages/tools/tests/test_local_shell_tool.py index 0f2dd7f3483..0aa3f5daec1 100644 --- a/python/packages/tools/tests/test_local_shell_tool.py +++ b/python/packages/tools/tests/test_local_shell_tool.py @@ -577,3 +577,39 @@ async def test_persistent_confines_workdir_by_default_powershell(tmp_path: os.Pa await tool.run(f"Set-Location -LiteralPath '{subdir}'") pwd = await tool.run("(Get-Location).Path") assert os.path.realpath(pwd.stdout.strip()) == os.path.realpath(str(tmp_path)) + + +async def test_audit_log_records_allowed_command() -> None: + tool = LocalShellTool(mode="stateless", approval_mode="never_require", acknowledge_unsafe=True) + cmd = "Write-Output hello" if sys.platform == "win32" else "echo hello" + await tool.run(cmd) + + record = tool.audit_log.records[0] + assert record.command == cmd + assert record.policy_decision == "allow" + assert record.exit_code == 0 + assert record.timestamp_utc + + +async def test_audit_log_records_denied_command() -> None: + tool = LocalShellTool( + mode="stateless", + policy=ShellPolicy(denylist=[r"forbidden"]), + approval_mode="never_require", + acknowledge_unsafe=True, + ) + with pytest.raises(ShellCommandError): + await tool.run("echo forbidden") + + record = tool.audit_log.records[0] + assert record.policy_decision == "deny" + assert record.exit_code is None + assert record.denial_reason + + +async def test_audit_log_accumulates_multiple() -> None: + tool = LocalShellTool(mode="stateless", approval_mode="never_require", acknowledge_unsafe=True) + cmd = "Write-Output hello" if sys.platform == "win32" else "echo hello" + for _ in range(3): + await tool.run(cmd) + assert len(tool.audit_log) == 3 From ba5c367e8f093c7678aa72eeee1fcff66557224b Mon Sep 17 00:00:00 2001 From: Mark Licata Date: Wed, 7 Oct 2026 09:01:23 -0700 Subject: [PATCH 2/5] Address review feedback: broker settled-ID tracking, immutable audit snapshots Co-Authored-By: Claude Opus 5.5 --- .../_harness/_request_broker.py | 26 +++++++--- .../core/tests/core/test_request_broker.py | 50 +++++++++++++++++++ .../agent_framework_tools/shell/_audit.py | 10 ++-- .../tools/tests/test_local_shell_tool.py | 18 +++++++ 4 files changed, 93 insertions(+), 11 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_request_broker.py b/python/packages/core/agent_framework/_harness/_request_broker.py index 1d9488b04e5..98c913e274f 100644 --- a/python/packages/core/agent_framework/_harness/_request_broker.py +++ b/python/packages/core/agent_framework/_harness/_request_broker.py @@ -8,8 +8,11 @@ :meth:`AsyncRequestBroker.reject`. The broker tracks recently settled request IDs in a bounded deque (default -256 entries) so that re-using a settled ID is reported instead of silently -hanging. Late responses for unknown IDs return ``False`` rather than raising. +256 entries). A request counts as settled once it is resolved, rejected, +timed out, or cancelled. Re-using a settled ID raises ``ValueError`` instead of +creating a new waiter that a late response could complete with stale data. +Late responses for recently settled IDs return ``True`` (and are ignored); +responses for unknown or evicted IDs return ``False`` rather than raising. Usage:: @@ -66,7 +69,8 @@ async def request_async(self, request_id: str, *, timeout: float | None = None) The result passed to :meth:`resolve`. Raises: - ValueError: If the ID is already pending or was recently settled. + ValueError: If the ID is already pending or was recently settled + (resolved, rejected, timed out, or cancelled). asyncio.TimeoutError: If ``timeout`` elapses before a response arrives. asyncio.CancelledError: If the waiting task is cancelled. Exception: Any exception passed to :meth:`reject`. @@ -83,6 +87,9 @@ async def request_async(self, request_id: str, *, timeout: float | None = None) except BaseException: if self._pending.get(request_id) is future: del self._pending[request_id] + # Remember the ID so a late response cannot complete a new waiter + # that re-uses it, and so re-use is reported as already settled. + self._settled.append(request_id) future.cancel() raise @@ -90,12 +97,13 @@ def resolve(self, request_id: str, result: T) -> bool: """Deliver a successful result to a waiting :meth:`request_async` call. Returns: - ``True`` if a waiter was notified, ``False`` if the request was not found - (already settled or never registered). + ``True`` if a waiter was notified or the request was recently settled + (the late result is ignored). ``False`` if the request ID is unknown + (never registered, or evicted from the settled history). """ future = self._pending.pop(request_id, None) if future is None: - return False + return request_id in self._settled self._settled.append(request_id) if not future.done(): future.set_result(result) @@ -105,11 +113,13 @@ def reject(self, request_id: str, exception: BaseException) -> bool: """Deliver an exception to a waiting :meth:`request_async` call. Returns: - ``True`` if a waiter was notified. + ``True`` if a waiter was notified or the request was recently settled + (the late exception is ignored). ``False`` if the request ID is unknown + (never registered, or evicted from the settled history). """ future = self._pending.pop(request_id, None) if future is None: - return False + return request_id in self._settled self._settled.append(request_id) if not future.done(): future.set_exception(exception) diff --git a/python/packages/core/tests/core/test_request_broker.py b/python/packages/core/tests/core/test_request_broker.py index 07f2d4c0ef0..9dd1abd2bd1 100644 --- a/python/packages/core/tests/core/test_request_broker.py +++ b/python/packages/core/tests/core/test_request_broker.py @@ -57,3 +57,53 @@ async def test_duplicate_request_raises() -> None: await task with pytest.raises(ValueError, match="already settled"): await broker.request_async("r1") + + +async def test_resolve_and_reject_return_true_for_recently_settled_id() -> None: + broker: AsyncRequestBroker[int] = AsyncRequestBroker() + task = asyncio.create_task(broker.request_async("r1")) + await asyncio.sleep(0) + assert broker.resolve("r1", 1) is True + assert await task == 1 + # Late duplicate responses for a recently settled ID are acknowledged. + assert broker.resolve("r1", 2) is True + assert broker.reject("r1", RuntimeError("late")) is True + # Unknown IDs are still reported as not found. + assert broker.resolve("unknown", 1) is False + assert broker.reject("unknown", RuntimeError("x")) is False + + +async def test_evicted_settled_id_returns_false() -> None: + broker: AsyncRequestBroker[int] = AsyncRequestBroker(settled_deque_size=1) + for request_id in ("r1", "r2"): + task = asyncio.create_task(broker.request_async(request_id)) + await asyncio.sleep(0) + broker.resolve(request_id, 0) + await task + assert broker.resolve("r1", 1) is False + assert broker.resolve("r2", 1) is True + + +async def test_timed_out_id_is_settled_and_cannot_be_reused() -> None: + broker: AsyncRequestBroker[str] = AsyncRequestBroker() + with pytest.raises(asyncio.TimeoutError): + await broker.request_async("r1", timeout=0.001) + # A late response for the expired request is acknowledged but not delivered. + assert broker.resolve("r1", "stale") is True + with pytest.raises(ValueError, match="already settled"): + await broker.request_async("r1") + assert broker.pending_count == 0 + + +async def test_cancelled_id_is_settled_and_cannot_be_reused() -> None: + broker: AsyncRequestBroker[str] = AsyncRequestBroker() + task = asyncio.create_task(broker.request_async("r1")) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert broker.pending_count == 0 + assert broker.resolve("r1", "stale") is True + with pytest.raises(ValueError, match="already settled"): + await broker.request_async("r1") + assert broker.pending_count == 0 diff --git a/python/packages/tools/agent_framework_tools/shell/_audit.py b/python/packages/tools/agent_framework_tools/shell/_audit.py index d6a51ab7851..a4cad847d3f 100644 --- a/python/packages/tools/agent_framework_tools/shell/_audit.py +++ b/python/packages/tools/agent_framework_tools/shell/_audit.py @@ -5,7 +5,7 @@ from __future__ import annotations import datetime -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace __all__ = ["ShellAuditLog", "ShellCommandRecord"] @@ -67,8 +67,12 @@ def record(self, command: str, policy_decision: str, denial_reason: str = "") -> @property def records(self) -> list[ShellCommandRecord]: - """All recorded command attempts, in chronological order.""" - return list(self._records) + """Snapshots of all recorded command attempts, in chronological order. + + Each call returns detached copies, so mutating a returned record does + not alter the stored audit history. + """ + return [replace(rec) for rec in self._records] def __len__(self) -> int: return len(self._records) diff --git a/python/packages/tools/tests/test_local_shell_tool.py b/python/packages/tools/tests/test_local_shell_tool.py index 0aa3f5daec1..cd174362953 100644 --- a/python/packages/tools/tests/test_local_shell_tool.py +++ b/python/packages/tools/tests/test_local_shell_tool.py @@ -613,3 +613,21 @@ async def test_audit_log_accumulates_multiple() -> None: for _ in range(3): await tool.run(cmd) assert len(tool.audit_log) == 3 + + +async def test_audit_log_records_are_detached_snapshots() -> None: + tool = LocalShellTool(mode="stateless", approval_mode="never_require", acknowledge_unsafe=True) + cmd = "Write-Output hello" if sys.platform == "win32" else "echo hello" + await tool.run(cmd) + + snapshot = tool.audit_log.records[0] + snapshot.command = "tampered" + snapshot.policy_decision = "deny" + snapshot.exit_code = 99 + tool.audit_log.records.clear() + + record = tool.audit_log.records[0] + assert len(tool.audit_log) == 1 + assert record.command == cmd + assert record.policy_decision == "allow" + assert record.exit_code == 0 From 3d5879397cc93db8321173ca34c1e5ff1eb45ae6 Mon Sep 17 00:00:00 2001 From: Mark Licata Date: Tue, 13 Oct 2026 08:43:33 -0700 Subject: [PATCH 3/5] Address review feedback: rename request_async, record workdir in shell audit Co-Authored-By: Claude Opus 5.5 --- .../_harness/_request_broker.py | 8 ++-- .../core/tests/core/test_request_broker.py | 26 ++++++------ .../agent_framework_tools/shell/_audit.py | 17 +++++++- .../agent_framework_tools/shell/_tool.py | 1 + .../tools/tests/test_local_shell_tool.py | 40 +++++++++++++++++++ 5 files changed, 74 insertions(+), 18 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_request_broker.py b/python/packages/core/agent_framework/_harness/_request_broker.py index 98c913e274f..abb0e8afc0a 100644 --- a/python/packages/core/agent_framework/_harness/_request_broker.py +++ b/python/packages/core/agent_framework/_harness/_request_broker.py @@ -19,7 +19,7 @@ broker = AsyncRequestBroker() # Server side - suspend until answered (or timeout) - result = await broker.request_async("req-123", timeout=30.0) + result = await broker.request("req-123", timeout=30.0) # Client side - deliver the answer broker.resolve("req-123", decision) @@ -56,7 +56,7 @@ def __init__(self, settled_deque_size: int = DEFAULT_SETTLED_DEQUE_SIZE) -> None self._pending: dict[str, asyncio.Future[T]] = {} self._settled: collections.deque[str] = collections.deque(maxlen=settled_deque_size) - async def request_async(self, request_id: str, *, timeout: float | None = None) -> T: + async def request(self, request_id: str, *, timeout: float | None = None) -> T: """Suspend until the request is resolved or rejected. Args: @@ -94,7 +94,7 @@ async def request_async(self, request_id: str, *, timeout: float | None = None) raise def resolve(self, request_id: str, result: T) -> bool: - """Deliver a successful result to a waiting :meth:`request_async` call. + """Deliver a successful result to a waiting :meth:`request` call. Returns: ``True`` if a waiter was notified or the request was recently settled @@ -110,7 +110,7 @@ def resolve(self, request_id: str, result: T) -> bool: return True def reject(self, request_id: str, exception: BaseException) -> bool: - """Deliver an exception to a waiting :meth:`request_async` call. + """Deliver an exception to a waiting :meth:`request` call. Returns: ``True`` if a waiter was notified or the request was recently settled diff --git a/python/packages/core/tests/core/test_request_broker.py b/python/packages/core/tests/core/test_request_broker.py index 9dd1abd2bd1..1ba13a72d73 100644 --- a/python/packages/core/tests/core/test_request_broker.py +++ b/python/packages/core/tests/core/test_request_broker.py @@ -9,7 +9,7 @@ async def test_resolve_delivers_result() -> None: broker: AsyncRequestBroker[str] = AsyncRequestBroker() - task = asyncio.create_task(broker.request_async("r1")) + task = asyncio.create_task(broker.request("r1")) await asyncio.sleep(0) assert broker.resolve("r1", "value") is True assert await task == "value" @@ -17,7 +17,7 @@ async def test_resolve_delivers_result() -> None: async def test_reject_delivers_exception() -> None: broker: AsyncRequestBroker[str] = AsyncRequestBroker() - task = asyncio.create_task(broker.request_async("r1")) + task = asyncio.create_task(broker.request("r1")) await asyncio.sleep(0) assert broker.reject("r1", RuntimeError("boom")) is True with pytest.raises(RuntimeError, match="boom"): @@ -27,7 +27,7 @@ async def test_reject_delivers_exception() -> None: async def test_timeout_raises() -> None: broker: AsyncRequestBroker[str] = AsyncRequestBroker() with pytest.raises(asyncio.TimeoutError): - await broker.request_async("r1", timeout=0.001) + await broker.request("r1", timeout=0.001) assert broker.pending_count == 0 @@ -39,7 +39,7 @@ async def test_resolve_returns_false_when_not_found() -> None: async def test_pending_count() -> None: broker: AsyncRequestBroker[int] = AsyncRequestBroker() - task = asyncio.create_task(broker.request_async("r1")) + task = asyncio.create_task(broker.request("r1")) await asyncio.sleep(0) assert broker.pending_count == 1 broker.resolve("r1", 1) @@ -49,19 +49,19 @@ async def test_pending_count() -> None: async def test_duplicate_request_raises() -> None: broker: AsyncRequestBroker[int] = AsyncRequestBroker() - task = asyncio.create_task(broker.request_async("r1")) + task = asyncio.create_task(broker.request("r1")) await asyncio.sleep(0) with pytest.raises(ValueError, match="already pending"): - await broker.request_async("r1") + await broker.request("r1") broker.resolve("r1", 1) await task with pytest.raises(ValueError, match="already settled"): - await broker.request_async("r1") + await broker.request("r1") async def test_resolve_and_reject_return_true_for_recently_settled_id() -> None: broker: AsyncRequestBroker[int] = AsyncRequestBroker() - task = asyncio.create_task(broker.request_async("r1")) + task = asyncio.create_task(broker.request("r1")) await asyncio.sleep(0) assert broker.resolve("r1", 1) is True assert await task == 1 @@ -76,7 +76,7 @@ async def test_resolve_and_reject_return_true_for_recently_settled_id() -> None: async def test_evicted_settled_id_returns_false() -> None: broker: AsyncRequestBroker[int] = AsyncRequestBroker(settled_deque_size=1) for request_id in ("r1", "r2"): - task = asyncio.create_task(broker.request_async(request_id)) + task = asyncio.create_task(broker.request(request_id)) await asyncio.sleep(0) broker.resolve(request_id, 0) await task @@ -87,17 +87,17 @@ async def test_evicted_settled_id_returns_false() -> None: async def test_timed_out_id_is_settled_and_cannot_be_reused() -> None: broker: AsyncRequestBroker[str] = AsyncRequestBroker() with pytest.raises(asyncio.TimeoutError): - await broker.request_async("r1", timeout=0.001) + await broker.request("r1", timeout=0.001) # A late response for the expired request is acknowledged but not delivered. assert broker.resolve("r1", "stale") is True with pytest.raises(ValueError, match="already settled"): - await broker.request_async("r1") + await broker.request("r1") assert broker.pending_count == 0 async def test_cancelled_id_is_settled_and_cannot_be_reused() -> None: broker: AsyncRequestBroker[str] = AsyncRequestBroker() - task = asyncio.create_task(broker.request_async("r1")) + task = asyncio.create_task(broker.request("r1")) await asyncio.sleep(0) task.cancel() with pytest.raises(asyncio.CancelledError): @@ -105,5 +105,5 @@ async def test_cancelled_id_is_settled_and_cannot_be_reused() -> None: assert broker.pending_count == 0 assert broker.resolve("r1", "stale") is True with pytest.raises(ValueError, match="already settled"): - await broker.request_async("r1") + await broker.request("r1") assert broker.pending_count == 0 diff --git a/python/packages/tools/agent_framework_tools/shell/_audit.py b/python/packages/tools/agent_framework_tools/shell/_audit.py index a4cad847d3f..005b76f862b 100644 --- a/python/packages/tools/agent_framework_tools/shell/_audit.py +++ b/python/packages/tools/agent_framework_tools/shell/_audit.py @@ -24,6 +24,13 @@ class ShellCommandRecord: policy_decision: str """The :class:`~agent_framework_tools.shell.ShellDecision` outcome: ``"allow"`` or ``"deny"``.""" + workdir: str | None = None + """Working directory used for policy evaluation and execution. + + ``None`` means no working directory was configured and the command ran in + the host process's current working directory. + """ + timestamp_utc: str = field(default_factory=lambda: datetime.datetime.now(datetime.timezone.utc).isoformat()) """UTC ISO 8601 timestamp when the command was evaluated.""" @@ -55,11 +62,19 @@ class ShellAuditLog: def __init__(self) -> None: self._records: list[ShellCommandRecord] = [] - def record(self, command: str, policy_decision: str, denial_reason: str = "") -> ShellCommandRecord: + def record( + self, + command: str, + policy_decision: str, + denial_reason: str = "", + *, + workdir: str | None = None, + ) -> ShellCommandRecord: """Add a new record at evaluation time (before execution).""" rec = ShellCommandRecord( command=command, policy_decision=policy_decision, + workdir=workdir, denial_reason=denial_reason, ) self._records.append(rec) diff --git a/python/packages/tools/agent_framework_tools/shell/_tool.py b/python/packages/tools/agent_framework_tools/shell/_tool.py index 397b9a58b47..56cb202f8a7 100644 --- a/python/packages/tools/agent_framework_tools/shell/_tool.py +++ b/python/packages/tools/agent_framework_tools/shell/_tool.py @@ -266,6 +266,7 @@ async def run(self, command: str, *, timeout: float | None = None) -> ShellResul command=command, policy_decision=decision.decision, denial_reason=decision.reason if decision.decision == "deny" else "", + workdir=request.workdir, ) if decision.decision == "deny": raise ShellCommandError(f"Command rejected by policy: {decision.reason}") diff --git a/python/packages/tools/tests/test_local_shell_tool.py b/python/packages/tools/tests/test_local_shell_tool.py index cd174362953..72dcceb417f 100644 --- a/python/packages/tools/tests/test_local_shell_tool.py +++ b/python/packages/tools/tests/test_local_shell_tool.py @@ -588,6 +588,7 @@ async def test_audit_log_records_allowed_command() -> None: assert record.command == cmd assert record.policy_decision == "allow" assert record.exit_code == 0 + assert record.workdir is None assert record.timestamp_utc @@ -605,6 +606,43 @@ async def test_audit_log_records_denied_command() -> None: assert record.policy_decision == "deny" assert record.exit_code is None assert record.denial_reason + assert record.workdir is None + + +async def test_audit_log_records_workdir_for_denied_command(tmp_path: Any) -> None: + tool = LocalShellTool( + mode="stateless", + policy=ShellPolicy(denylist=[r"forbidden"]), + approval_mode="never_require", + acknowledge_unsafe=True, + workdir=tmp_path, + ) + with pytest.raises(ShellCommandError): + await tool.run("echo forbidden") + + record = tool.audit_log.records[0] + assert record.policy_decision == "deny" + assert record.workdir == str(tmp_path) + + +async def test_audit_log_records_distinct_workdirs_for_same_command(tmp_path: Any) -> None: + dir_a = tmp_path / "a" + dir_b = tmp_path / "b" + dir_a.mkdir() + dir_b.mkdir() + cmd = "Write-Output hello" if sys.platform == "win32" else "echo hello" + + tool_a = LocalShellTool(mode="stateless", approval_mode="never_require", acknowledge_unsafe=True, workdir=dir_a) + tool_b = LocalShellTool(mode="stateless", approval_mode="never_require", acknowledge_unsafe=True, workdir=dir_b) + await tool_a.run(cmd) + await tool_b.run(cmd) + + record_a = tool_a.audit_log.records[0] + record_b = tool_b.audit_log.records[0] + assert record_a.command == record_b.command == cmd + assert record_a.workdir == str(dir_a) + assert record_b.workdir == str(dir_b) + assert record_a.workdir != record_b.workdir async def test_audit_log_accumulates_multiple() -> None: @@ -624,6 +662,7 @@ async def test_audit_log_records_are_detached_snapshots() -> None: snapshot.command = "tampered" snapshot.policy_decision = "deny" snapshot.exit_code = 99 + snapshot.workdir = "tampered" tool.audit_log.records.clear() record = tool.audit_log.records[0] @@ -631,3 +670,4 @@ async def test_audit_log_records_are_detached_snapshots() -> None: assert record.command == cmd assert record.policy_decision == "allow" assert record.exit_code == 0 + assert record.workdir is None From e5cd62da7f2b23474cbd629a38e808650121e1ce Mon Sep 17 00:00:00 2001 From: Mark Licata Date: Tue, 13 Oct 2026 08:45:50 -0700 Subject: [PATCH 4/5] Add AsyncRequestBroker sample Co-Authored-By: Claude Opus 5.5 --- python/samples/02-agents/harness/README.md | 1 + .../02-agents/harness/async_request_broker.py | 110 ++++++++++++++++++ 2 files changed, 111 insertions(+) create mode 100644 python/samples/02-agents/harness/async_request_broker.py diff --git a/python/samples/02-agents/harness/README.md b/python/samples/02-agents/harness/README.md index 20ccc49780d..a949cd9403d 100644 --- a/python/samples/02-agents/harness/README.md +++ b/python/samples/02-agents/harness/README.md @@ -30,6 +30,7 @@ Each feature can be disabled or customized via keyword arguments. |------|-------------| | `harness_research.py` | Interactive research assistant with web search, a plan/execute workflow, and an execute-mode loop that re-invokes the agent until every todo is complete | | `harness_data_processing.py` | Data-processing assistant over a folder of CSV files, demonstrating file-access tools and tool approval | +| `async_request_broker.py` | `AsyncRequestBroker` basics with no LLM required: server-side code awaits a request that a client resolves, rejects, or never answers (timeout) | | [`build_your_own_claw/`](./build_your_own_claw/README.md) | *Build your own claw* blog series — a personal finance assistant built step by step | ## Running diff --git a/python/samples/02-agents/harness/async_request_broker.py b/python/samples/02-agents/harness/async_request_broker.py new file mode 100644 index 00000000000..53ca49075aa --- /dev/null +++ b/python/samples/02-agents/harness/async_request_broker.py @@ -0,0 +1,110 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import warnings + +from agent_framework import AsyncRequestBroker + +"""AsyncRequestBroker: pause server-side work until a client responds. + +This sample demonstrates how to use `AsyncRequestBroker` to let server-side +code (for example a running agent or tool) stop partway through, ask a +connected client for a decision, and continue once the client answers. + +A typical case is a tool that needs human approval before it runs. The server +calls `await broker.request(request_id)` and sends the request ID to the UI +over whatever transport you use (websocket, SSE, HTTP). When the user answers, +the client handler calls `broker.resolve(request_id, value)` or +`broker.reject(request_id, exception)`, and the awaiting coroutine continues. + +The sample covers three cases: +1. The server awaits a request and the client resolves it. +2. The server awaits a request, the client rejects it, and the server handles + the error. +3. The server awaits a request with a timeout that nobody answers. + +No LLM or credentials are needed. The "client" is simulated with +`asyncio.create_task` in the same event loop. +""" + +# AsyncRequestBroker is experimental; silence the warning to keep output tidy. +warnings.filterwarnings("ignore", message=r"\[HARNESS\]") + + +class ApprovalDenied(Exception): + """Raised on the server when the client denies an approval request.""" + + +async def server_delete_files(broker: AsyncRequestBroker[bool], request_id: str, path: str) -> None: + """Server-side work that pauses for client approval before continuing.""" + print(f"[server] Requesting approval to delete '{path}' (id={request_id})") + try: + # Suspends here, without blocking the event loop, until the client answers. + approved = await broker.request(request_id) + except ApprovalDenied as ex: + print(f"[server] Request {request_id} was rejected: {ex}. Skipping.") + return + print(f"[server] Request {request_id} answered: approved={approved}. Deleting '{path}'.") + + +async def client_respond(broker: AsyncRequestBroker[bool], request_id: str, *, approve: bool) -> None: + """Simulate a client (e.g. a UI) answering the request after a short delay.""" + await asyncio.sleep(0.1) # The user thinks about it... + if approve: + print(f"[client] Approving {request_id}") + broker.resolve(request_id, True) + else: + print(f"[client] Denying {request_id}") + broker.reject(request_id, ApprovalDenied("user said no")) + + +async def main() -> None: + broker: AsyncRequestBroker[bool] = AsyncRequestBroker() + + # 1. Resolve: the client approves and the server continues with the result. + print("=== Resolve ===") + await asyncio.gather( + server_delete_files(broker, "req-1", "./tmp/cache"), + client_respond(broker, "req-1", approve=True), + ) + + # 2. Reject: the client denies and the server gets the exception. + print("\n=== Reject ===") + await asyncio.gather( + server_delete_files(broker, "req-2", "./important"), + client_respond(broker, "req-2", approve=False), + ) + + # 3. Timeout: nobody answers, so the server stops waiting. + print("\n=== Timeout ===") + try: + await broker.request("req-3", timeout=0.2) + except asyncio.TimeoutError: + print("[server] Request req-3 timed out with no response.") + + # A late reply for a settled request is ignored. resolve() returns True to + # show the ID is known but already settled. + print(f"[client] Late reply to req-3 accepted as already settled: {broker.resolve('req-3', True)}") + print(f"Pending requests: {broker.pending_count}") + + +if __name__ == "__main__": + asyncio.run(main()) + +""" +Sample output: +=== Resolve === +[server] Requesting approval to delete './tmp/cache' (id=req-1) +[client] Approving req-1 +[server] Request req-1 answered: approved=True. Deleting './tmp/cache'. + +=== Reject === +[server] Requesting approval to delete './important' (id=req-2) +[client] Denying req-2 +[server] Request req-2 was rejected: user said no. Skipping. + +=== Timeout === +[server] Request req-3 timed out with no response. +[client] Late reply to req-3 accepted as already settled: True +Pending requests: 0 +""" From 36a9e5658cbd7fb6e30f76f68f54cea01e70cd6c Mon Sep 17 00:00:00 2001 From: Mark Licata Date: Tue, 3 Nov 2026 11:48:13 -0800 Subject: [PATCH 5/5] Python: remove AsyncRequestBroker from core; replace ShellAuditLog with Python logging - Remove AsyncRequestBroker from agent_framework public API and core harness; any in-process broker belongs in the application, not the framework - Remove ShellAuditLog and ShellCommandRecord public types; shell tool now emits structured events through logging.getLogger('agent_framework.shell.audit') with shell.policy and shell.execution records - Fix effective workdir: record os.getcwd() for stateless mode when no workdir configured, instead of None - Fix outcome tracking: try/finally ensures exit_code and outcome are always logged even on exception or cancellation - Rewrite async_request_broker.py sample to show MAF pause-and-resume approval contract with a locally-defined transport, not a framework import - Addresses eavanvalkenburg feedback on #8861 Generated with Amplifier Co-Authored-By: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com> --- .../packages/core/agent_framework/__init__.py | 2 - .../core/agent_framework/__init__.pyi | 2 - .../_harness/_request_broker.py | 131 -------- .../core/tests/core/test_request_broker.py | 109 ------- .../agent_framework_tools/shell/__init__.py | 3 - .../agent_framework_tools/shell/_audit.py | 93 ------ .../agent_framework_tools/shell/_tool.py | 89 ++++-- .../tools/tests/test_local_shell_tool.py | 122 +++----- python/samples/02-agents/harness/README.md | 2 +- .../02-agents/harness/async_request_broker.py | 282 +++++++++++++----- 10 files changed, 308 insertions(+), 527 deletions(-) delete mode 100644 python/packages/core/agent_framework/_harness/_request_broker.py delete mode 100644 python/packages/core/tests/core/test_request_broker.py delete mode 100644 python/packages/tools/agent_framework_tools/shell/_audit.py diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 9720b557024..3953187578e 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -148,7 +148,6 @@ "MemoryTopicRecord", ), "._harness._mode": ("DEFAULT_MODE_SOURCE_ID", "AgentModeProvider", "get_agent_mode", "set_agent_mode"), - "._harness._request_broker": ("AsyncRequestBroker",), "._harness._todo": ( "DEFAULT_TODO_SOURCE_ID", "TodoFileStore", @@ -452,7 +451,6 @@ "AgentSession", "AggregatingSkillsSource", "Annotation", - "AsyncRequestBroker", "BackgroundAgentsProvider", "BackgroundTaskInfo", "BackgroundTaskStatus", diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index 265ee28a071..e8d1bc6e411 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/python/packages/core/agent_framework/__init__.pyi @@ -106,7 +106,6 @@ from ._harness._memory import ( MemoryTopicRecord, ) from ._harness._mode import DEFAULT_MODE_SOURCE_ID, AgentModeProvider, get_agent_mode, set_agent_mode -from ._harness._request_broker import AsyncRequestBroker from ._harness._todo import ( DEFAULT_TODO_SOURCE_ID, TodoFileStore, @@ -407,7 +406,6 @@ __all__ = [ "AgentSession", "AggregatingSkillsSource", "Annotation", - "AsyncRequestBroker", "BackgroundAgentsProvider", "BackgroundTaskInfo", "BackgroundTaskStatus", diff --git a/python/packages/core/agent_framework/_harness/_request_broker.py b/python/packages/core/agent_framework/_harness/_request_broker.py deleted file mode 100644 index abb0e8afc0a..00000000000 --- a/python/packages/core/agent_framework/_harness/_request_broker.py +++ /dev/null @@ -1,131 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -"""AsyncRequestBroker: a generic asyncio pause-and-wait primitive. - -Provides :class:`AsyncRequestBroker`, which lets any server-side code suspend -execution on a request ID and wait asynchronously for a client to deliver the -matching result via :meth:`AsyncRequestBroker.resolve` or -:meth:`AsyncRequestBroker.reject`. - -The broker tracks recently settled request IDs in a bounded deque (default -256 entries). A request counts as settled once it is resolved, rejected, -timed out, or cancelled. Re-using a settled ID raises ``ValueError`` instead of -creating a new waiter that a late response could complete with stale data. -Late responses for recently settled IDs return ``True`` (and are ignored); -responses for unknown or evicted IDs return ``False`` rather than raising. - -Usage:: - - broker = AsyncRequestBroker() - - # Server side - suspend until answered (or timeout) - result = await broker.request("req-123", timeout=30.0) - - # Client side - deliver the answer - broker.resolve("req-123", decision) -""" - -from __future__ import annotations - -import asyncio -import collections -from typing import Generic, TypeVar - -from .._feature_stage import ExperimentalFeature, experimental - -__all__ = ["AsyncRequestBroker"] - -T = TypeVar("T") - - -@experimental(feature_id=ExperimentalFeature.HARNESS) -class AsyncRequestBroker(Generic[T]): - """Generic asyncio pause-and-wait broker for request/response patterns. - - Allows server-side code to suspend on a request ID and await a client - response delivered via :meth:`resolve` or :meth:`reject`, without - polling or blocking a thread. - - Thread-safety: designed for single-threaded asyncio use. The broker must - be used within a single event loop. - """ - - DEFAULT_SETTLED_DEQUE_SIZE = 256 - - def __init__(self, settled_deque_size: int = DEFAULT_SETTLED_DEQUE_SIZE) -> None: - self._pending: dict[str, asyncio.Future[T]] = {} - self._settled: collections.deque[str] = collections.deque(maxlen=settled_deque_size) - - async def request(self, request_id: str, *, timeout: float | None = None) -> T: - """Suspend until the request is resolved or rejected. - - Args: - request_id: Unique identifier for this request. - - Keyword Args: - timeout: Maximum seconds to wait. ``None`` waits indefinitely. - - Returns: - The result passed to :meth:`resolve`. - - Raises: - ValueError: If the ID is already pending or was recently settled - (resolved, rejected, timed out, or cancelled). - asyncio.TimeoutError: If ``timeout`` elapses before a response arrives. - asyncio.CancelledError: If the waiting task is cancelled. - Exception: Any exception passed to :meth:`reject`. - """ - if request_id in self._settled: - raise ValueError(f"Request '{request_id}' was already settled.") - if request_id in self._pending: - raise ValueError(f"Request '{request_id}' is already pending.") - loop = asyncio.get_running_loop() - future: asyncio.Future[T] = loop.create_future() - self._pending[request_id] = future - try: - return await asyncio.wait_for(asyncio.shield(future), timeout=timeout) - except BaseException: - if self._pending.get(request_id) is future: - del self._pending[request_id] - # Remember the ID so a late response cannot complete a new waiter - # that re-uses it, and so re-use is reported as already settled. - self._settled.append(request_id) - future.cancel() - raise - - def resolve(self, request_id: str, result: T) -> bool: - """Deliver a successful result to a waiting :meth:`request` call. - - Returns: - ``True`` if a waiter was notified or the request was recently settled - (the late result is ignored). ``False`` if the request ID is unknown - (never registered, or evicted from the settled history). - """ - future = self._pending.pop(request_id, None) - if future is None: - return request_id in self._settled - self._settled.append(request_id) - if not future.done(): - future.set_result(result) - return True - - def reject(self, request_id: str, exception: BaseException) -> bool: - """Deliver an exception to a waiting :meth:`request` call. - - Returns: - ``True`` if a waiter was notified or the request was recently settled - (the late exception is ignored). ``False`` if the request ID is unknown - (never registered, or evicted from the settled history). - """ - future = self._pending.pop(request_id, None) - if future is None: - return request_id in self._settled - self._settled.append(request_id) - if not future.done(): - future.set_exception(exception) - return True - - @property - def pending_count(self) -> int: - """Number of outstanding requests currently awaiting a response.""" - return len(self._pending) diff --git a/python/packages/core/tests/core/test_request_broker.py b/python/packages/core/tests/core/test_request_broker.py deleted file mode 100644 index 1ba13a72d73..00000000000 --- a/python/packages/core/tests/core/test_request_broker.py +++ /dev/null @@ -1,109 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import asyncio - -import pytest - -from agent_framework import AsyncRequestBroker - - -async def test_resolve_delivers_result() -> None: - broker: AsyncRequestBroker[str] = AsyncRequestBroker() - task = asyncio.create_task(broker.request("r1")) - await asyncio.sleep(0) - assert broker.resolve("r1", "value") is True - assert await task == "value" - - -async def test_reject_delivers_exception() -> None: - broker: AsyncRequestBroker[str] = AsyncRequestBroker() - task = asyncio.create_task(broker.request("r1")) - await asyncio.sleep(0) - assert broker.reject("r1", RuntimeError("boom")) is True - with pytest.raises(RuntimeError, match="boom"): - await task - - -async def test_timeout_raises() -> None: - broker: AsyncRequestBroker[str] = AsyncRequestBroker() - with pytest.raises(asyncio.TimeoutError): - await broker.request("r1", timeout=0.001) - assert broker.pending_count == 0 - - -async def test_resolve_returns_false_when_not_found() -> None: - broker: AsyncRequestBroker[str] = AsyncRequestBroker() - assert broker.resolve("missing", "x") is False - assert broker.reject("missing", RuntimeError("x")) is False - - -async def test_pending_count() -> None: - broker: AsyncRequestBroker[int] = AsyncRequestBroker() - task = asyncio.create_task(broker.request("r1")) - await asyncio.sleep(0) - assert broker.pending_count == 1 - broker.resolve("r1", 1) - await task - assert broker.pending_count == 0 - - -async def test_duplicate_request_raises() -> None: - broker: AsyncRequestBroker[int] = AsyncRequestBroker() - task = asyncio.create_task(broker.request("r1")) - await asyncio.sleep(0) - with pytest.raises(ValueError, match="already pending"): - await broker.request("r1") - broker.resolve("r1", 1) - await task - with pytest.raises(ValueError, match="already settled"): - await broker.request("r1") - - -async def test_resolve_and_reject_return_true_for_recently_settled_id() -> None: - broker: AsyncRequestBroker[int] = AsyncRequestBroker() - task = asyncio.create_task(broker.request("r1")) - await asyncio.sleep(0) - assert broker.resolve("r1", 1) is True - assert await task == 1 - # Late duplicate responses for a recently settled ID are acknowledged. - assert broker.resolve("r1", 2) is True - assert broker.reject("r1", RuntimeError("late")) is True - # Unknown IDs are still reported as not found. - assert broker.resolve("unknown", 1) is False - assert broker.reject("unknown", RuntimeError("x")) is False - - -async def test_evicted_settled_id_returns_false() -> None: - broker: AsyncRequestBroker[int] = AsyncRequestBroker(settled_deque_size=1) - for request_id in ("r1", "r2"): - task = asyncio.create_task(broker.request(request_id)) - await asyncio.sleep(0) - broker.resolve(request_id, 0) - await task - assert broker.resolve("r1", 1) is False - assert broker.resolve("r2", 1) is True - - -async def test_timed_out_id_is_settled_and_cannot_be_reused() -> None: - broker: AsyncRequestBroker[str] = AsyncRequestBroker() - with pytest.raises(asyncio.TimeoutError): - await broker.request("r1", timeout=0.001) - # A late response for the expired request is acknowledged but not delivered. - assert broker.resolve("r1", "stale") is True - with pytest.raises(ValueError, match="already settled"): - await broker.request("r1") - assert broker.pending_count == 0 - - -async def test_cancelled_id_is_settled_and_cannot_be_reused() -> None: - broker: AsyncRequestBroker[str] = AsyncRequestBroker() - task = asyncio.create_task(broker.request("r1")) - await asyncio.sleep(0) - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - assert broker.pending_count == 0 - assert broker.resolve("r1", "stale") is True - with pytest.raises(ValueError, match="already settled"): - await broker.request("r1") - assert broker.pending_count == 0 diff --git a/python/packages/tools/agent_framework_tools/shell/__init__.py b/python/packages/tools/agent_framework_tools/shell/__init__.py index 0be0e982aed..06517a96c22 100644 --- a/python/packages/tools/agent_framework_tools/shell/__init__.py +++ b/python/packages/tools/agent_framework_tools/shell/__init__.py @@ -4,7 +4,6 @@ from __future__ import annotations -from ._audit import ShellAuditLog, ShellCommandRecord from ._docker import ( DEFAULT_IMAGE as DOCKER_DEFAULT_IMAGE, ) @@ -36,9 +35,7 @@ "DockerNotAvailableError", "DockerShellTool", "LocalShellTool", - "ShellAuditLog", "ShellCommandError", - "ShellCommandRecord", "ShellDecision", "ShellEnvironmentProvider", "ShellEnvironmentProviderOptions", diff --git a/python/packages/tools/agent_framework_tools/shell/_audit.py b/python/packages/tools/agent_framework_tools/shell/_audit.py deleted file mode 100644 index 005b76f862b..00000000000 --- a/python/packages/tools/agent_framework_tools/shell/_audit.py +++ /dev/null @@ -1,93 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -"""Shell command audit log for compliance and observability.""" - -from __future__ import annotations - -import datetime -from dataclasses import dataclass, field, replace - -__all__ = ["ShellAuditLog", "ShellCommandRecord"] - - -@dataclass -class ShellCommandRecord: - """A record of a single shell command execution attempt. - - Stored in the tool's audit log for compliance and observability. - All timestamps are UTC ISO 8601. - """ - - command: str - """The command string as submitted.""" - - policy_decision: str - """The :class:`~agent_framework_tools.shell.ShellDecision` outcome: ``"allow"`` or ``"deny"``.""" - - workdir: str | None = None - """Working directory used for policy evaluation and execution. - - ``None`` means no working directory was configured and the command ran in - the host process's current working directory. - """ - - timestamp_utc: str = field(default_factory=lambda: datetime.datetime.now(datetime.timezone.utc).isoformat()) - """UTC ISO 8601 timestamp when the command was evaluated.""" - - exit_code: int | None = None - """Process exit code, set after execution. ``None`` if the command was denied or not yet run.""" - - denial_reason: str = "" - """Reason from :class:`~agent_framework_tools.shell.ShellDecision` when denied.""" - - -class ShellAuditLog: - """An in-memory audit log of shell command decisions and outcomes. - - Records every command evaluated by :class:`~agent_framework_tools.shell.LocalShellTool`, - including the policy decision, execution outcome, and exit code. - - Access via ``LocalShellTool.audit_log``. - - Examples: - .. code-block:: python - - tool = LocalShellTool() - await tool.run("echo hello") - - for record in tool.audit_log.records: - print(record.command, record.policy_decision, record.exit_code) - """ - - def __init__(self) -> None: - self._records: list[ShellCommandRecord] = [] - - def record( - self, - command: str, - policy_decision: str, - denial_reason: str = "", - *, - workdir: str | None = None, - ) -> ShellCommandRecord: - """Add a new record at evaluation time (before execution).""" - rec = ShellCommandRecord( - command=command, - policy_decision=policy_decision, - workdir=workdir, - denial_reason=denial_reason, - ) - self._records.append(rec) - return rec - - @property - def records(self) -> list[ShellCommandRecord]: - """Snapshots of all recorded command attempts, in chronological order. - - Each call returns detached copies, so mutating a returned record does - not alter the stored audit history. - """ - return [replace(rec) for rec in self._records] - - def __len__(self) -> int: - return len(self._records) diff --git a/python/packages/tools/agent_framework_tools/shell/_tool.py b/python/packages/tools/agent_framework_tools/shell/_tool.py index 56cb202f8a7..7f95dee259e 100644 --- a/python/packages/tools/agent_framework_tools/shell/_tool.py +++ b/python/packages/tools/agent_framework_tools/shell/_tool.py @@ -15,7 +15,6 @@ from agent_framework._tools import SHELL_TOOL_KIND_VALUE from .._feature_usage import FeatureIndex -from ._audit import ShellAuditLog from ._executor import run_stateless from ._policy import ShellPolicy, ShellRequest from ._resolve import is_powershell, resolve_shell @@ -29,6 +28,7 @@ ) logger = logging.getLogger(__name__) +_audit_logger = logging.getLogger("agent_framework.shell.audit") def _quote_posix(value: str) -> str: @@ -145,9 +145,14 @@ class LocalShellTool: on_command: Optional audit hook called with the command string for every command that passes policy. Use for logging / telemetry. - Attributes: - audit_log: :class:`ShellAuditLog` recording every evaluated command - (allowed or denied) with its policy decision and exit code. + Note: + Shell execution events are emitted through + ``logging.getLogger("agent_framework.shell.audit")`` with ``shell.policy`` + (policy decisions) and ``shell.execution`` (outcomes) messages. + Configure a handler on that logger to collect, store, or forward events. + Commands are not included in log records by default because they may + contain secrets; use the ``on_command`` hook if you need them in a + controlled environment. """ def __init__( @@ -186,7 +191,6 @@ def __init__( self._max_output_bytes = max_output_bytes self._approval_mode: Literal["always_require", "never_require"] = approval_mode self._on_command = on_command - self.audit_log: ShellAuditLog = ShellAuditLog() merged_env: dict[str, str] | None if env is None and not clean_env: @@ -242,6 +246,9 @@ async def __aexit__(self, *_exc: object) -> None: # ------------------------------------------------------------------ core run + def _effective_workdir(self) -> str: + return os.path.abspath(self._workdir) if self._workdir is not None else os.getcwd() + async def run(self, command: str, *, timeout: float | None = None) -> ShellResult: """Execute ``command`` directly and return its :class:`ShellResult`. @@ -262,11 +269,18 @@ async def run(self, command: str, *, timeout: float | None = None) -> ShellResul mark_feature_used(FeatureIndex.TOOLS_SHELL) request = ShellRequest(command=command, workdir=self._workdir) decision = self._policy.evaluate(request) - record = self.audit_log.record( - command=command, - policy_decision=decision.decision, - denial_reason=decision.reason if decision.decision == "deny" else "", - workdir=request.workdir, + # Resolve the effective cwd so records never report ``None``. + # In persistent mode the shell's real cwd may differ if + # ``confine_workdir=False``; this reports the configured location. + effective_workdir = self._effective_workdir() + _audit_logger.info( + "shell.policy", + extra={ + "policy_decision": decision.decision, + "denial_reason": decision.reason if decision.decision == "deny" else "", + "workdir": effective_workdir, + # command is intentionally omitted -- it may contain secrets. + }, ) if decision.decision == "deny": raise ShellCommandError(f"Command rejected by policy: {decision.reason}") @@ -278,26 +292,43 @@ async def run(self, command: str, *, timeout: float | None = None) -> ShellResul effective_timeout = self._timeout if timeout is None else timeout - if self._mode == "persistent": - if self._session is None: - await self.start() - if self._session is None: - raise RuntimeError("LocalShellTool session failed to start") - effective = self._maybe_reanchor(command) - result = await self._session.run(effective, timeout=effective_timeout) - record.exit_code = result.exit_code + exit_code: int | None = None + timed_out = False + outcome = "unknown" + try: + if self._mode == "persistent": + if self._session is None: + await self.start() + if self._session is None: + raise RuntimeError("LocalShellTool session failed to start") + effective = self._maybe_reanchor(command) + result = await self._session.run(effective, timeout=effective_timeout) + else: + result = await run_stateless( + self._stateless_argv, + command, + workdir=self._workdir, + env=self._env, + timeout=effective_timeout, + max_output_bytes=self._max_output_bytes, + ) + exit_code = result.exit_code + timed_out = bool(getattr(result, "timed_out", False)) + outcome = "timeout" if timed_out else ("success" if exit_code == 0 else "failure") return result - - result = await run_stateless( - self._stateless_argv, - command, - workdir=self._workdir, - env=self._env, - timeout=effective_timeout, - max_output_bytes=self._max_output_bytes, - ) - record.exit_code = result.exit_code - return result + except BaseException as exc: + outcome = "cancelled" if isinstance(exc, asyncio.CancelledError) else "exception" + raise + finally: + _audit_logger.info( + "shell.execution", + extra={ + "exit_code": exit_code, + "timed_out": timed_out, + "outcome": outcome, + "workdir": effective_workdir, + }, + ) # ------------------------------------------------------------------ AF wiring diff --git a/python/packages/tools/tests/test_local_shell_tool.py b/python/packages/tools/tests/test_local_shell_tool.py index 34ac7662945..0e89dd75996 100644 --- a/python/packages/tools/tests/test_local_shell_tool.py +++ b/python/packages/tools/tests/test_local_shell_tool.py @@ -2,6 +2,7 @@ import asyncio import json +import logging import os import shlex import shutil @@ -848,95 +849,66 @@ async def test_persistent_confines_workdir_by_default_powershell(tmp_path: os.Pa assert os.path.realpath(pwd.stdout.strip()) == os.path.realpath(str(tmp_path)) -async def test_audit_log_records_allowed_command() -> None: - tool = LocalShellTool(mode="stateless", approval_mode="never_require", acknowledge_unsafe=True) - cmd = "Write-Output hello" if sys.platform == "win32" else "echo hello" - await tool.run(cmd) +def _audit_records(caplog: pytest.LogCaptureFixture, message: str) -> list[Any]: + return [r for r in caplog.records if r.name == "agent_framework.shell.audit" and r.message == message] - record = tool.audit_log.records[0] - assert record.command == cmd - assert record.policy_decision == "allow" - assert record.exit_code == 0 - assert record.workdir is None - assert record.timestamp_utc +async def test_shell_logging_records_policy_decision_for_allowed_command( + caplog: pytest.LogCaptureFixture, +) -> None: + tool = LocalShellTool(mode="stateless", approval_mode="never_require", acknowledge_unsafe=True) + cmd = "Write-Output hi" if sys.platform == "win32" else "echo hi" + with caplog.at_level(logging.INFO, logger="agent_framework.shell.audit"): + await tool.run(cmd) + records = _audit_records(caplog, "shell.policy") + assert len(records) == 1 + assert records[0].__dict__["policy_decision"] == "allow" + assert records[0].__dict__["denial_reason"] == "" + assert "command" not in records[0].__dict__ -async def test_audit_log_records_denied_command() -> None: - tool = LocalShellTool( - mode="stateless", - policy=ShellPolicy(denylist=[r"forbidden"]), - approval_mode="never_require", - acknowledge_unsafe=True, - ) - with pytest.raises(ShellCommandError): - await tool.run("echo forbidden") - record = tool.audit_log.records[0] - assert record.policy_decision == "deny" - assert record.exit_code is None - assert record.denial_reason - assert record.workdir is None +async def test_shell_logging_records_execution_outcome(caplog: pytest.LogCaptureFixture) -> None: + tool = LocalShellTool(mode="stateless", approval_mode="never_require", acknowledge_unsafe=True) + cmd = "Write-Output hi" if sys.platform == "win32" else "echo hi" + with caplog.at_level(logging.INFO, logger="agent_framework.shell.audit"): + await tool.run(cmd) + records = _audit_records(caplog, "shell.execution") + assert len(records) == 1 + assert records[0].__dict__["exit_code"] == 0 + assert records[0].__dict__["timed_out"] is False + assert records[0].__dict__["outcome"] == "success" -async def test_audit_log_records_workdir_for_denied_command(tmp_path: Any) -> None: +async def test_shell_logging_records_denial(caplog: pytest.LogCaptureFixture) -> None: tool = LocalShellTool( mode="stateless", - policy=ShellPolicy(denylist=[r"forbidden"]), approval_mode="never_require", acknowledge_unsafe=True, - workdir=tmp_path, + policy=ShellPolicy(denylist=[r".*"]), ) - with pytest.raises(ShellCommandError): - await tool.run("echo forbidden") - - record = tool.audit_log.records[0] - assert record.policy_decision == "deny" - assert record.workdir == str(tmp_path) - - -async def test_audit_log_records_distinct_workdirs_for_same_command(tmp_path: Any) -> None: - dir_a = tmp_path / "a" - dir_b = tmp_path / "b" - dir_a.mkdir() - dir_b.mkdir() - cmd = "Write-Output hello" if sys.platform == "win32" else "echo hello" - - tool_a = LocalShellTool(mode="stateless", approval_mode="never_require", acknowledge_unsafe=True, workdir=dir_a) - tool_b = LocalShellTool(mode="stateless", approval_mode="never_require", acknowledge_unsafe=True, workdir=dir_b) - await tool_a.run(cmd) - await tool_b.run(cmd) - - record_a = tool_a.audit_log.records[0] - record_b = tool_b.audit_log.records[0] - assert record_a.command == record_b.command == cmd - assert record_a.workdir == str(dir_a) - assert record_b.workdir == str(dir_b) - assert record_a.workdir != record_b.workdir + with caplog.at_level(logging.INFO, logger="agent_framework.shell.audit"), pytest.raises(ShellCommandError): + await tool.run("echo hi") + policy_records = _audit_records(caplog, "shell.policy") + assert len(policy_records) == 1 + assert policy_records[0].__dict__["policy_decision"] == "deny" + assert policy_records[0].__dict__["denial_reason"] + assert _audit_records(caplog, "shell.execution") == [] -async def test_audit_log_accumulates_multiple() -> None: +async def test_shell_logging_effective_workdir_resolves_none(caplog: pytest.LogCaptureFixture) -> None: + """With no workdir configured, the logged workdir is the process cwd, not None.""" tool = LocalShellTool(mode="stateless", approval_mode="never_require", acknowledge_unsafe=True) - cmd = "Write-Output hello" if sys.platform == "win32" else "echo hello" - for _ in range(3): + cmd = "Write-Output hi" if sys.platform == "win32" else "echo hi" + with caplog.at_level(logging.INFO, logger="agent_framework.shell.audit"): await tool.run(cmd) - assert len(tool.audit_log) == 3 + records = _audit_records(caplog, "shell.policy") + assert records[0].__dict__["workdir"] == os.getcwd() -async def test_audit_log_records_are_detached_snapshots() -> None: - tool = LocalShellTool(mode="stateless", approval_mode="never_require", acknowledge_unsafe=True) - cmd = "Write-Output hello" if sys.platform == "win32" else "echo hello" - await tool.run(cmd) - - snapshot = tool.audit_log.records[0] - snapshot.command = "tampered" - snapshot.policy_decision = "deny" - snapshot.exit_code = 99 - snapshot.workdir = "tampered" - tool.audit_log.records.clear() - - record = tool.audit_log.records[0] - assert len(tool.audit_log) == 1 - assert record.command == cmd - assert record.policy_decision == "allow" - assert record.exit_code == 0 - assert record.workdir is None +async def test_shell_logging_records_configured_workdir(caplog: pytest.LogCaptureFixture, tmp_path: Any) -> None: + tool = LocalShellTool(mode="stateless", approval_mode="never_require", acknowledge_unsafe=True, workdir=tmp_path) + cmd = "Write-Output hi" if sys.platform == "win32" else "echo hi" + with caplog.at_level(logging.INFO, logger="agent_framework.shell.audit"): + await tool.run(cmd) + assert _audit_records(caplog, "shell.policy")[0].__dict__["workdir"] == os.path.abspath(tmp_path) + assert _audit_records(caplog, "shell.execution")[0].__dict__["workdir"] == os.path.abspath(tmp_path) diff --git a/python/samples/02-agents/harness/README.md b/python/samples/02-agents/harness/README.md index a949cd9403d..474f1945d90 100644 --- a/python/samples/02-agents/harness/README.md +++ b/python/samples/02-agents/harness/README.md @@ -30,7 +30,7 @@ Each feature can be disabled or customized via keyword arguments. |------|-------------| | `harness_research.py` | Interactive research assistant with web search, a plan/execute workflow, and an execute-mode loop that re-invokes the agent until every todo is complete | | `harness_data_processing.py` | Data-processing assistant over a folder of CSV files, demonstrating file-access tools and tool approval | -| `async_request_broker.py` | `AsyncRequestBroker` basics with no LLM required: server-side code awaits a request that a client resolves, rejects, or never answers (timeout) | +| `async_request_broker.py` | MAF approval pause-and-resume: agent pauses for shell-command approval, application delivers the request to a simulated client, and resumes the agent with the decision | | [`build_your_own_claw/`](./build_your_own_claw/README.md) | *Build your own claw* blog series — a personal finance assistant built step by step | ## Running diff --git a/python/samples/02-agents/harness/async_request_broker.py b/python/samples/02-agents/harness/async_request_broker.py index 53ca49075aa..ff35906771d 100644 --- a/python/samples/02-agents/harness/async_request_broker.py +++ b/python/samples/02-agents/harness/async_request_broker.py @@ -1,110 +1,228 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio -import warnings +import logging +import sys +from collections.abc import Awaitable, Mapping, Sequence +from typing import Any + +from agent_framework import ( + Agent, + AgentSession, + BaseChatClient, + ChatResponse, + Content, + FunctionInvocationLayer, + Message, +) +from agent_framework_tools.shell import LocalShellTool + +"""MAF approval pause-and-resume for shell commands. + +Agent Framework's approval contract is pause-and-resume, not "block and wait": + +1. A tool created with ``approval_mode="always_require"`` does NOT run. The + agent returns early and ``response.user_input_requests`` contains a + function-approval request. +2. Your application delivers that request to whoever decides (a UI, a chat + channel, a reviewer) over a transport that YOU own. +3. When a decision arrives, call ``agent.run(...)`` again with the SAME + ``AgentSession`` and an approval response built with + ``request.to_function_approval_response(approved)``. + +The framework deliberately ships no request broker: where pending requests +live, how they reach a client, and how long they wait are application concerns. + +This sample shows that split. ``_PendingApprovals`` below is a minimal +``asyncio.Future``-based mailbox that stands in for your transport. It is +process-lifetime only: if the process exits, pending requests are lost. It is +NOT a substitute for a durable queue or workflow engine; use one of those when +approvals must survive restarts or span multiple workers. + +No LLM or credentials are needed. A scripted chat client emits a shell tool +call, and the real ``LocalShellTool`` runs it only after approval. Shell +policy/outcome events are emitted on the ``agent_framework.shell.audit`` +logger; the sample enables it so you can see them. +""" -from agent_framework import AsyncRequestBroker -"""AsyncRequestBroker: pause server-side work until a client responds. +class _AuditPrinter(logging.Handler): + """Print shell audit events (policy decision and outcome) with their structured fields.""" -This sample demonstrates how to use `AsyncRequestBroker` to let server-side -code (for example a running agent or tool) stop partway through, ask a -connected client for a decision, and continue once the client answers. + _FIELDS = ("policy_decision", "outcome", "exit_code", "timed_out") -A typical case is a tool that needs human approval before it runs. The server -calls `await broker.request(request_id)` and sends the request ID to the UI -over whatever transport you use (websocket, SSE, HTTP). When the user answers, -the client handler calls `broker.resolve(request_id, value)` or -`broker.reject(request_id, exception)`, and the awaiting coroutine continues. + def emit(self, record: logging.LogRecord) -> None: + fields = {k: record.__dict__[k] for k in self._FIELDS if k in record.__dict__} + print(f" [audit] {record.getMessage()} {fields}") -The sample covers three cases: -1. The server awaits a request and the client resolves it. -2. The server awaits a request, the client rejects it, and the server handles - the error. -3. The server awaits a request with a timeout that nobody answers. -No LLM or credentials are needed. The "client" is simulated with -`asyncio.create_task` in the same event loop. -""" +_audit_logger = logging.getLogger("agent_framework.shell.audit") +_audit_logger.setLevel(logging.INFO) +_audit_logger.addHandler(_AuditPrinter()) -# AsyncRequestBroker is experimental; silence the warning to keep output tidy. -warnings.filterwarnings("ignore", message=r"\[HARNESS\]") +class _PendingApprovals: + """Application-owned mailbox: request id -> Future[bool]. Process-lifetime only.""" -class ApprovalDenied(Exception): - """Raised on the server when the client denies an approval request.""" + def __init__(self) -> None: + self._futures: dict[str, asyncio.Future[bool]] = {} + def submit(self, request_id: str) -> asyncio.Future[bool]: + future: asyncio.Future[bool] = asyncio.get_running_loop().create_future() + self._futures[request_id] = future + return future -async def server_delete_files(broker: AsyncRequestBroker[bool], request_id: str, path: str) -> None: - """Server-side work that pauses for client approval before continuing.""" - print(f"[server] Requesting approval to delete '{path}' (id={request_id})") - try: - # Suspends here, without blocking the event loop, until the client answers. - approved = await broker.request(request_id) - except ApprovalDenied as ex: - print(f"[server] Request {request_id} was rejected: {ex}. Skipping.") - return - print(f"[server] Request {request_id} answered: approved={approved}. Deleting '{path}'.") + def answer(self, request_id: str, approved: bool) -> None: + future = self._futures.pop(request_id, None) + if future is not None and not future.done(): + future.set_result(approved) + def discard(self, request_id: str) -> None: + self._futures.pop(request_id, None) -async def client_respond(broker: AsyncRequestBroker[bool], request_id: str, *, approve: bool) -> None: - """Simulate a client (e.g. a UI) answering the request after a short delay.""" - await asyncio.sleep(0.1) # The user thinks about it... - if approve: - print(f"[client] Approving {request_id}") - broker.resolve(request_id, True) - else: - print(f"[client] Denying {request_id}") - broker.reject(request_id, ApprovalDenied("user said no")) + @property + def pending_count(self) -> int: + return len(self._futures) -async def main() -> None: - broker: AsyncRequestBroker[bool] = AsyncRequestBroker() +class _ScriptedShellClient(FunctionInvocationLayer, BaseChatClient): + """Stands in for an LLM: asks to run one shell command, then summarizes.""" - # 1. Resolve: the client approves and the server continues with the result. - print("=== Resolve ===") - await asyncio.gather( - server_delete_files(broker, "req-1", "./tmp/cache"), - client_respond(broker, "req-1", approve=True), - ) + def __init__(self, command: str) -> None: + super().__init__() + self._command = command + self._asked = False - # 2. Reject: the client denies and the server gets the exception. - print("\n=== Reject ===") - await asyncio.gather( - server_delete_files(broker, "req-2", "./important"), - client_respond(broker, "req-2", approve=False), - ) + def _inner_get_response( + self, + *, + messages: Sequence[Message], + stream: bool, + options: Mapping[str, Any], + **kwargs: Any, + ) -> Awaitable[ChatResponse]: + async def _get() -> ChatResponse: + if not self._asked: + self._asked = True + call = Content.from_function_call( + call_id="call-1", name="run_shell", arguments={"command": self._command} + ) + return ChatResponse(messages=Message(role="assistant", contents=[call])) + return ChatResponse(messages=Message(role="assistant", contents=["The command finished."])) + + return _get() + + +def _make_agent(command: str) -> Agent[Any]: + shell = LocalShellTool(mode="stateless", approval_mode="always_require") + return Agent(client=_ScriptedShellClient(command), tools=[shell.as_function()]) + + +async def simulated_client(pending: _PendingApprovals, request_id: str, *, approve: bool, delay: float) -> None: + """Simulate a remote reviewer answering after a delay (or never, if delay is None-like).""" + await asyncio.sleep(delay) + print(f"[client] {'Approving' if approve else 'Denying'} {request_id}") + pending.answer(request_id, approve) + + +async def run_with_approval( + agent: Agent[Any], + session: AgentSession, + pending: _PendingApprovals, + prompt: str, + *, + answer_timeout: float, + client_task_factory: Any, +) -> str: + """Drive pause-and-resume. ``client_task_factory(request_id)`` starts the simulated client.""" + response = await agent.run(prompt, session=session) + if not response.user_input_requests: + return response.text - # 3. Timeout: nobody answers, so the server stops waiting. - print("\n=== Timeout ===") + request = response.user_input_requests[0] + assert request.function_call is not None + request_id = request.id or request.function_call.call_id + print(f"[server] Paused: approval needed for {request.function_call.name}({request.function_call.arguments})") + + # Deliver the request to the client over *your* transport, then wait for a decision. + decision = pending.submit(request_id) + client_task = client_task_factory(request_id) try: - await broker.request("req-3", timeout=0.2) + approved = await asyncio.wait_for(decision, timeout=answer_timeout) except asyncio.TimeoutError: - print("[server] Request req-3 timed out with no response.") + # Nobody answered. The agent is NOT blocked -- it already returned -- but this + # conversation is stuck in a "pending approval" state until you act. Decide a + # policy: keep waiting, escalate, or resume with a denial. + print(f"[server] No answer within {answer_timeout}s; request {request_id} is still unresolved.") + pending.discard(request_id) + client_task.cancel() + return "STUCK: approval was never answered (session left paused)" + + # Resume: same session, plus the approval response. + print(f"[server] Resuming with approved={approved}") + resumed = await agent.run( + [ + Message("assistant", [request]), + Message("user", [request.to_function_approval_response(approved)]), + ], + session=session, + ) + return resumed.text + + +async def main() -> None: + command = "Write-Output hello" if sys.platform == "win32" else "echo hello from approved shell" + pending = _PendingApprovals() + + print("=== 1. Approved: the command runs after the client says yes ===") + agent, session = _make_agent(command), AgentSession() + print( + "[agent]", + await run_with_approval( + agent, + session, + pending, + "Run the greeting command", + answer_timeout=2.0, + client_task_factory=lambda rid: asyncio.create_task( + simulated_client(pending, rid, approve=True, delay=0.1) + ), + ), + ) - # A late reply for a settled request is ignored. resolve() returns True to - # show the ID is known but already settled. - print(f"[client] Late reply to req-3 accepted as already settled: {broker.resolve('req-3', True)}") - print(f"Pending requests: {broker.pending_count}") + print("\n=== 2. Denied: the command never runs ===") + agent, session = _make_agent(command), AgentSession() + print( + "[agent]", + await run_with_approval( + agent, + session, + pending, + "Run the greeting command", + answer_timeout=2.0, + client_task_factory=lambda rid: asyncio.create_task( + simulated_client(pending, rid, approve=False, delay=0.1) + ), + ), + ) + + print("\n=== 3. Nobody answers: the session is stuck pending ===") + agent, session = _make_agent(command), AgentSession() + print( + "[agent]", + await run_with_approval( + agent, + session, + pending, + "Run the greeting command", + answer_timeout=0.3, + client_task_factory=lambda rid: asyncio.create_task( + simulated_client(pending, rid, approve=True, delay=30.0) + ), + ), + ) + print(f"Pending requests: {pending.pending_count}") if __name__ == "__main__": asyncio.run(main()) - -""" -Sample output: -=== Resolve === -[server] Requesting approval to delete './tmp/cache' (id=req-1) -[client] Approving req-1 -[server] Request req-1 answered: approved=True. Deleting './tmp/cache'. - -=== Reject === -[server] Requesting approval to delete './important' (id=req-2) -[client] Denying req-2 -[server] Request req-2 was rejected: user said no. Skipping. - -=== Timeout === -[server] Request req-3 timed out with no response. -[client] Late reply to req-3 accepted as already settled: True -Pending requests: 0 -"""