diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/AsyncRequestBroker.cs b/dotnet/src/Microsoft.Agents.AI/Harness/AsyncRequestBroker.cs new file mode 100644 index 00000000000..211de77f9b8 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI/Harness/AsyncRequestBroker.cs @@ -0,0 +1,142 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Concurrent; +using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Shared.DiagnosticIds; + +namespace Microsoft.Agents.AI; + +/// +/// A generic async 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 or , without polling or +/// blocking a thread. +/// +/// The type of the result delivered to the waiter. +[Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)] +public sealed class AsyncRequestBroker +{ + private const int DefaultSettledQueueSize = 256; + + private readonly ConcurrentDictionary> _pending = new(); + private readonly Queue _settledQueue = new(); + private readonly HashSet _settledSet = new(); + private readonly int _settledQueueSize; + private readonly object _settledLock = new(); + + /// + /// Initializes a new instance of the class. + /// + /// Maximum number of recently settled request IDs to track for late-response handling. + public AsyncRequestBroker(int settledQueueSize = DefaultSettledQueueSize) + { + if (settledQueueSize < 1) + { + throw new ArgumentOutOfRangeException(nameof(settledQueueSize), "Settled queue size must be at least 1."); + } + + this._settledQueueSize = settledQueueSize; + } + + /// Gets the number of requests currently awaiting a response. + public int PendingCount => this._pending.Count; + + /// Suspend until the request is resolved, rejected, or cancelled. + /// Unique identifier for this request. + /// Optional timeout. waits indefinitely. + /// Token to cancel the wait. + /// The result passed to . + public async Task RequestAsync(string requestId, TimeSpan? timeout = null, CancellationToken cancellationToken = default) + { + var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + lock (this._settledLock) + { + if (this._settledSet.Contains(requestId)) + { + throw new InvalidOperationException($"Request '{requestId}' was already settled."); + } + + if (!this._pending.TryAdd(requestId, tcs)) + { + throw new InvalidOperationException($"Request '{requestId}' is already pending."); + } + } + + using var cts = timeout.HasValue + ? CancellationTokenSource.CreateLinkedTokenSource(cancellationToken) + : null; + + if (cts is not null && timeout.HasValue) + { + cts.CancelAfter(timeout.Value); + } + + var effectiveToken = cts?.Token ?? cancellationToken; + using var registration = effectiveToken.Register(() => + { + // Atomic key+value removal: only cancel our own TCS, never a replacement registered under the same ID. + if (((ICollection>>)this._pending) + .Remove(new KeyValuePair>(requestId, tcs))) + { + this.TrackSettled(requestId); + tcs.TrySetCanceled(effectiveToken); + } + }); + + return await tcs.Task.ConfigureAwait(false); + } + + /// Deliver a successful result to a waiting call. + /// The request identifier. + /// The result to deliver. + /// if a waiter was notified; if the request was not found. + public bool Resolve(string requestId, TResult result) + { + if (!this._pending.TryRemove(requestId, out var tcs)) + { + return false; + } + + this.TrackSettled(requestId); + tcs.TrySetResult(result); + return true; + } + + /// Deliver an exception to a waiting call. + /// The request identifier. + /// The exception to deliver. + /// if a waiter was notified. + public bool Reject(string requestId, Exception exception) + { + if (!this._pending.TryRemove(requestId, out var tcs)) + { + return false; + } + + this.TrackSettled(requestId); + tcs.TrySetException(exception); + return true; + } + + private void TrackSettled(string requestId) + { + lock (this._settledLock) + { + if (this._settledQueue.Count >= this._settledQueueSize) + { + string oldest = this._settledQueue.Dequeue(); + this._settledSet.Remove(oldest); + } + + this._settledQueue.Enqueue(requestId); + this._settledSet.Add(requestId); + } + } +} diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net10.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net10.0/PublicAPI.Unshipped.txt index fada975c025..d37623ec1b5 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net10.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net10.0/PublicAPI.Unshipped.txt @@ -24,3 +24,9 @@ Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]virtual Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.FunctionInvocationContextExtensions [MAAI001]static Microsoft.Agents.AI.FunctionInvocationContextExtensions.WrapWithPendingMiddleware(this Microsoft.Extensions.AI.FunctionInvocationContext! context, Microsoft.Extensions.AI.AIFunction! function) -> Microsoft.Extensions.AI.AIFunction! +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.AsyncRequestBroker(int settledQueueSize = 256) -> void +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.PendingCount.get -> int +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.Reject(string! requestId, System.Exception! exception) -> bool +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.RequestAsync(string! requestId, System.TimeSpan? timeout = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.Resolve(string! requestId, TResult result) -> bool diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net472/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net472/PublicAPI.Unshipped.txt index fada975c025..d37623ec1b5 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net472/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net472/PublicAPI.Unshipped.txt @@ -24,3 +24,9 @@ Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]virtual Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.FunctionInvocationContextExtensions [MAAI001]static Microsoft.Agents.AI.FunctionInvocationContextExtensions.WrapWithPendingMiddleware(this Microsoft.Extensions.AI.FunctionInvocationContext! context, Microsoft.Extensions.AI.AIFunction! function) -> Microsoft.Extensions.AI.AIFunction! +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.AsyncRequestBroker(int settledQueueSize = 256) -> void +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.PendingCount.get -> int +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.Reject(string! requestId, System.Exception! exception) -> bool +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.RequestAsync(string! requestId, System.TimeSpan? timeout = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.Resolve(string! requestId, TResult result) -> bool diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net8.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net8.0/PublicAPI.Unshipped.txt index fada975c025..d37623ec1b5 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net8.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net8.0/PublicAPI.Unshipped.txt @@ -24,3 +24,9 @@ Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]virtual Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.FunctionInvocationContextExtensions [MAAI001]static Microsoft.Agents.AI.FunctionInvocationContextExtensions.WrapWithPendingMiddleware(this Microsoft.Extensions.AI.FunctionInvocationContext! context, Microsoft.Extensions.AI.AIFunction! function) -> Microsoft.Extensions.AI.AIFunction! +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.AsyncRequestBroker(int settledQueueSize = 256) -> void +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.PendingCount.get -> int +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.Reject(string! requestId, System.Exception! exception) -> bool +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.RequestAsync(string! requestId, System.TimeSpan? timeout = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.Resolve(string! requestId, TResult result) -> bool diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net9.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net9.0/PublicAPI.Unshipped.txt index fada975c025..d37623ec1b5 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net9.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net9.0/PublicAPI.Unshipped.txt @@ -24,3 +24,9 @@ Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]virtual Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.FunctionInvocationContextExtensions [MAAI001]static Microsoft.Agents.AI.FunctionInvocationContextExtensions.WrapWithPendingMiddleware(this Microsoft.Extensions.AI.FunctionInvocationContext! context, Microsoft.Extensions.AI.AIFunction! function) -> Microsoft.Extensions.AI.AIFunction! +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.AsyncRequestBroker(int settledQueueSize = 256) -> void +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.PendingCount.get -> int +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.Reject(string! requestId, System.Exception! exception) -> bool +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.RequestAsync(string! requestId, System.TimeSpan? timeout = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.Resolve(string! requestId, TResult result) -> bool diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt index fada975c025..d37623ec1b5 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt @@ -24,3 +24,9 @@ Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]virtual Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.FunctionInvocationContextExtensions [MAAI001]static Microsoft.Agents.AI.FunctionInvocationContextExtensions.WrapWithPendingMiddleware(this Microsoft.Extensions.AI.FunctionInvocationContext! context, Microsoft.Extensions.AI.AIFunction! function) -> Microsoft.Extensions.AI.AIFunction! +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.AsyncRequestBroker(int settledQueueSize = 256) -> void +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.PendingCount.get -> int +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.Reject(string! requestId, System.Exception! exception) -> bool +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.RequestAsync(string! requestId, System.TimeSpan? timeout = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! +[MAAI001]Microsoft.Agents.AI.AsyncRequestBroker.Resolve(string! requestId, TResult result) -> bool diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md index f9b01907f22..4d59b5ed6c8 100644 --- a/python/CHANGELOG.md +++ b/python/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- **agent-framework-core**: Add `AsyncRequestBroker` — a generic asyncio pause-and-wait utility for any server-side request/response pattern. Provides `request_async()`, `resolve()`, and `reject()` with timeout and race-condition handling. +- **agent-framework-tools**: Add `ShellAuditLog` and `ShellCommandRecord` to `LocalShellTool`; every command evaluation is recorded with policy decision, exit code, and timestamp. + ## [1.20.0] - 2026-10-02 ### Added 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