Repository navigation
Python: Python + .NET: add AsyncRequestBroker and shell command audit log #8995
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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; | ||
|
|
||
| /// <summary> | ||
| /// A generic async pause-and-wait broker for request/response patterns. | ||
| /// </summary> | ||
| /// <remarks> | ||
| /// Allows server-side code to suspend on a request ID and await a client response | ||
| /// delivered via <see cref="Resolve"/> or <see cref="Reject"/>, without polling or | ||
| /// blocking a thread. | ||
| /// </remarks> | ||
| /// <typeparam name="TResult">The type of the result delivered to the waiter.</typeparam> | ||
| [Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)] | ||
| public sealed class AsyncRequestBroker<TResult> | ||
| { | ||
| private const int DefaultSettledQueueSize = 256; | ||
|
|
||
| private readonly ConcurrentDictionary<string, TaskCompletionSource<TResult>> _pending = new(); | ||
| private readonly Queue<string> _settledQueue = new(); | ||
| private readonly HashSet<string> _settledSet = new(); | ||
| private readonly int _settledQueueSize; | ||
| private readonly object _settledLock = new(); | ||
|
|
||
| /// <summary> | ||
| /// Initializes a new instance of the <see cref="AsyncRequestBroker{TResult}"/> class. | ||
| /// </summary> | ||
| /// <param name="settledQueueSize">Maximum number of recently settled request IDs to track for late-response handling.</param> | ||
| public AsyncRequestBroker(int settledQueueSize = DefaultSettledQueueSize) | ||
| { | ||
| if (settledQueueSize < 1) | ||
| { | ||
| throw new ArgumentOutOfRangeException(nameof(settledQueueSize), "Settled queue size must be at least 1."); | ||
| } | ||
|
|
||
| this._settledQueueSize = settledQueueSize; | ||
| } | ||
|
|
||
| /// <summary>Gets the number of requests currently awaiting a response.</summary> | ||
| public int PendingCount => this._pending.Count; | ||
|
|
||
| /// <summary>Suspend until the request is resolved, rejected, or cancelled.</summary> | ||
| /// <param name="requestId">Unique identifier for this request.</param> | ||
| /// <param name="timeout">Optional timeout. <see langword="null"/> waits indefinitely.</param> | ||
| /// <param name="cancellationToken">Token to cancel the wait.</param> | ||
| /// <returns>The result passed to <see cref="Resolve"/>.</returns> | ||
| public async Task<TResult> RequestAsync(string requestId, TimeSpan? timeout = null, CancellationToken cancellationToken = default) | ||
| { | ||
| var tcs = new TaskCompletionSource<TResult>(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<KeyValuePair<string, TaskCompletionSource<TResult>>>)this._pending) | ||
| .Remove(new KeyValuePair<string, TaskCompletionSource<TResult>>(requestId, tcs))) | ||
| { | ||
| this.TrackSettled(requestId); | ||
| tcs.TrySetCanceled(effectiveToken); | ||
| } | ||
| }); | ||
|
|
||
| return await tcs.Task.ConfigureAwait(false); | ||
| } | ||
|
|
||
| /// <summary>Deliver a successful result to a waiting <see cref="RequestAsync"/> call.</summary> | ||
| /// <param name="requestId">The request identifier.</param> | ||
| /// <param name="result">The result to deliver.</param> | ||
| /// <returns><see langword="true"/> if a waiter was notified; <see langword="false"/> if the request was not found.</returns> | ||
| 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; | ||
| } | ||
|
|
||
| /// <summary>Deliver an exception to a waiting <see cref="RequestAsync"/> call.</summary> | ||
| /// <param name="requestId">The request identifier.</param> | ||
| /// <param name="exception">The exception to deliver.</param> | ||
| /// <returns><see langword="true"/> if a waiter was notified.</returns> | ||
| 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); | ||
| } | ||
| } | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
| @@ -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: | ||||||
|
There was a problem hiding this comment. Across python MAF we never use
Suggested change
|
||||||
| """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 | ||||||
|
marklicata marked this conversation as resolved.
|
||||||
| 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) | ||||||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
please remove this, we add all relevant changes in version update PR's