diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs index c3e91a9d1a..a472ae5355 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs @@ -186,7 +186,11 @@ public override async IAsyncEnumerable CreateAsync( // point below: AsyncLocal writes made in this streaming iterator are reverted across yield // boundaries, so a single up-front assignment would be lost before the toolbox/MCP calls run. var platformCallId = context.PlatformContext?.CallId; + var toolboxRequestScope = this._toolboxService?.CreateRequestScope() + ?? FoundryToolboxService.RequestToolboxScope.CreateInactive(); + await using var configuredToolboxRequestScope = toolboxRequestScope.ConfigureAwait(false); HostedCallContext.CallId = platformCallId; + HostedCallContext.ToolboxCacheScopeId = toolboxRequestScope?.Id; // Stamp/validate the hosted identity only when one was resolved. Locally (non-hosted) there is // no user identity, so there is nothing to partition or tamper-check and the session is shared. @@ -278,6 +282,7 @@ public override async IAsyncEnumerable CreateAsync( // Re-apply the call id: the EmitCreated/EmitInProgress yields above reverted the ambient // value, and the toolbox tools/list + consent egress below must carry it per request. HostedCallContext.CallId = platformCallId; + HostedCallContext.ToolboxCacheScopeId = toolboxRequestScope?.Id; var toolboxResolution = await this.ResolveToolboxToolsAsync( this._toolboxService, @@ -382,6 +387,7 @@ public override async IAsyncEnumerable CreateAsync( // below revert the ambient AsyncLocal, but the MCP tools/call egress that happens // inside MoveNextAsync must carry the platform call id on every request. HostedCallContext.CallId = platformCallId; + HostedCallContext.ToolboxCacheScopeId = toolboxRequestScope?.Id; bool shutdownDetected = false; McpConsentInfo? consentInfo = null; @@ -754,7 +760,7 @@ await toolboxService return new(Tools: [], Consents: pendingConsents, Error: null); } - List toolsToAdd = [.. toolboxService.Tools]; + List toolsToAdd = [.. toolboxService.Tools, .. toolboxService.GetCurrentRequestTools()]; var seen = new HashSet(StringComparer.OrdinalIgnoreCase); List? markerConsents = null; diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryToolboxService.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryToolboxService.cs index fa78008494..868bc752e8 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryToolboxService.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryToolboxService.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Collections.Concurrent; using System.Collections.Generic; using System.Diagnostics.CodeAnalysis; using System.Linq; @@ -52,9 +53,13 @@ public sealed class FoundryToolboxService : IHostedService, IAsyncDisposable private readonly ILogger _logger; private readonly Dictionary _toolboxes = new(StringComparer.OrdinalIgnoreCase); + private readonly ConcurrentDictionary> _requestToolboxes = + new(StringComparer.Ordinal); + private readonly ConcurrentDictionary>> _requestConsents = + new(StringComparer.Ordinal); + private readonly ConcurrentDictionary _requestOpenLocks = new(StringComparer.Ordinal); private readonly Dictionary> _pendingConsents = new(StringComparer.OrdinalIgnoreCase); private readonly HashSet _deferredToolboxNames = new(StringComparer.OrdinalIgnoreCase); - private readonly SemaphoreSlim _lazyOpenLock = new(1, 1); private string? _resolvedEndpoint; private string? _featuresHeader; @@ -264,7 +269,7 @@ private void RecomputeStatus() /// Retries enumeration for any pre-registered toolbox that was awaiting user OAuth consent at /// startup. Call this at the start of request handling: once the user has completed consent /// out of band, the proxy holds a valid token and tools/list now succeeds, so the - /// toolbox's tools become available and are appended to . + /// toolbox's tools become available within the current response scope. /// /// The request cancellation token. /// @@ -275,40 +280,41 @@ private void RecomputeStatus() /// internal async ValueTask> ResolvePendingConsentsAsync(CancellationToken cancellationToken) { - // Fast path: nothing awaiting consent. - if (this.ConsentRequiredToolboxNames.Count == 0) + var requestConsents = this.GetCurrentRequestConsents(); + + // Fast path: nothing awaiting consent from startup or the deferred retry in this request. + if (this.ConsentRequiredToolboxNames.Count == 0 && requestConsents.Count == 0) { return []; } - await this._lazyOpenLock.WaitAsync(cancellationToken).ConfigureAwait(false); + var requestOpenLock = this.GetCurrentRequestOpenLock(); + await requestOpenLock.WaitAsync(cancellationToken).ConfigureAwait(false); try { - if (this._pendingConsents.Count == 0) - { - return []; - } - - var stillPending = new List(); - var resolvedTools = new List(); + var stillPending = new List(requestConsents); foreach (var toolboxName in new List(this._pendingConsents.Keys)) { + if (this.HasCurrentRequestResolution(toolboxName)) + { + continue; + } + try { var result = await this.OpenToolboxAsync(toolboxName, version: null, cancellationToken).ConfigureAwait(false); if (result.Consents is { } consents) { - // Still gated: refresh the consent info (the URL may rotate) and surface it. - this._pendingConsents[toolboxName] = consents; + // Consent URLs can be user-specific. Keep the refreshed result inside this + // request instead of replacing the startup snapshot shared by the singleton. + this.CacheRequestConsents(toolboxName, consents); stillPending.AddRange(consents); continue; } var cached = result.Cached!; - this._toolboxes[toolboxName] = cached; - resolvedTools.AddRange(cached.Tools); - this._pendingConsents.Remove(toolboxName); + this.CacheRequestToolbox(toolboxName, cached); } catch (Exception ex) when (ex is not OperationCanceledException) { @@ -324,18 +330,11 @@ internal async ValueTask> ResolvePendingConsentsAs } } - if (resolvedTools.Count > 0) - { - this.Tools = [.. this.Tools, .. resolvedTools]; - } - - this.RecomputeStatus(); - return stillPending; } finally { - this._lazyOpenLock.Release(); + requestOpenLock.Release(); } } @@ -345,9 +344,9 @@ internal async ValueTask> ResolvePendingConsentsAs /// request handling, before : the request's egress carries /// the platform-injected per-user isolation key, so a toolbox that needs a delegated user identity /// (for example a Microsoft Graph / Agent365 connection) can now enumerate as that user. On success - /// the toolbox's tools are appended to ; if the proxy now reports the source needs - /// user OAuth consent, the toolbox is moved to the pending-consent set so the caller surfaces the - /// consent prompt; if it still fails, it stays deferred and is retried on a later request. + /// the toolbox's tools are retained only for the current response; if the proxy reports the source needs + /// user OAuth consent, that consent is also retained only for the current response; if it still fails, + /// the toolbox stays deferred and is retried on a later request. /// /// The request cancellation token. internal async ValueTask RetryDeferredToolboxesAsync(CancellationToken cancellationToken) @@ -358,7 +357,8 @@ internal async ValueTask RetryDeferredToolboxesAsync(CancellationToken cancellat return; } - await this._lazyOpenLock.WaitAsync(cancellationToken).ConfigureAwait(false); + var requestOpenLock = this.GetCurrentRequestOpenLock(); + await requestOpenLock.WaitAsync(cancellationToken).ConfigureAwait(false); try { if (this._deferredToolboxNames.Count == 0) @@ -366,27 +366,26 @@ internal async ValueTask RetryDeferredToolboxesAsync(CancellationToken cancellat return; } - var resolvedTools = new List(); - foreach (var toolboxName in new List(this._deferredToolboxNames)) { + if (this.HasCurrentRequestResolution(toolboxName)) + { + continue; + } + try { var result = await this.OpenToolboxAsync(toolboxName, version: null, cancellationToken).ConfigureAwait(false); if (result.Consents is { } consents) { - // With the per-user context now present, the proxy reports the tool source - // needs user OAuth consent. Move it to the pending-consent set (the handler - // surfaces the prompt) and drop it from the deferred set. - this._pendingConsents[toolboxName] = consents; - this._deferredToolboxNames.Remove(toolboxName); + // Consent state belongs to this request. A different request must enumerate + // independently because its caller context can produce a different result. + this.CacheRequestConsents(toolboxName, consents); continue; } var cached = result.Cached!; - this._toolboxes[toolboxName] = cached; - resolvedTools.AddRange(cached.Tools); - this._deferredToolboxNames.Remove(toolboxName); + this.CacheRequestToolbox(toolboxName, cached); } catch (Exception ex) when (ex is not OperationCanceledException) { @@ -400,23 +399,16 @@ internal async ValueTask RetryDeferredToolboxesAsync(CancellationToken cancellat } } } - - if (resolvedTools.Count > 0) - { - this.Tools = [.. this.Tools, .. resolvedTools]; - } - - this.RecomputeStatus(); } finally { - this._lazyOpenLock.Release(); + requestOpenLock.Release(); } } /// - /// Resolves the tools for a per-request toolbox marker. Returns cached tools when the - /// toolbox has already been opened; otherwise honors + /// Resolves the tools for a per-request toolbox marker. Returns startup-cached tools or tools + /// already opened in the current response; otherwise honors /// to either reject or lazily open it. /// /// The Foundry toolbox name from the marker. @@ -451,6 +443,11 @@ internal async ValueTask GetToolboxToolsAsync( return new ToolboxResolution(cached.Tools, []); } + if (this.TryGetCurrentRequestToolbox(toolboxName, out cached)) + { + return new ToolboxResolution(cached.Tools, []); + } + if (this._options.StrictMode && !this._options.ToolboxNames.Contains(toolboxName, StringComparer.OrdinalIgnoreCase)) { throw new InvalidOperationException( @@ -464,7 +461,8 @@ internal async ValueTask GetToolboxToolsAsync( $"Cannot resolve toolbox '{toolboxName}': FOUNDRY_PROJECT_ENDPOINT is not set."); } - await this._lazyOpenLock.WaitAsync(cancellationToken).ConfigureAwait(false); + var requestOpenLock = this.GetCurrentRequestOpenLock(); + await requestOpenLock.WaitAsync(cancellationToken).ConfigureAwait(false); try { // Double-check after acquiring the lock to avoid duplicate opens under concurrency. @@ -473,6 +471,11 @@ internal async ValueTask GetToolboxToolsAsync( return new ToolboxResolution(cached.Tools, []); } + if (this.TryGetCurrentRequestToolbox(toolboxName, out cached)) + { + return new ToolboxResolution(cached.Tools, []); + } + var result = await this.OpenToolboxAsync(toolboxName, version, cancellationToken).ConfigureAwait(false); if (result.Consents is { } pendingConsents) { @@ -485,12 +488,210 @@ internal async ValueTask GetToolboxToolsAsync( } cached = result.Cached!; - this._toolboxes[toolboxName] = cached; + this.CacheRequestToolbox(toolboxName, cached); return new ToolboxResolution(cached.Tools, []); } finally { - this._lazyOpenLock.Release(); + requestOpenLock.Release(); + } + } + + /// + /// Creates an async-disposable owner for toolbox clients opened while handling one response. + /// + internal RequestToolboxScope CreateRequestScope() + { + // Do not derive this key from a platform identifier. Concurrent recovery or duplicate + // delivery can reuse externally meaningful ids, but active response caches must never alias. + return RequestToolboxScope.Create(this, Guid.NewGuid().ToString("N")); + } + + /// + /// Gets tools opened inside the current response scope by deferred or consent retries. + /// + internal IReadOnlyList GetCurrentRequestTools() + { + var scopeId = GetCurrentRequestScopeId(); + if (scopeId is null || !this._requestToolboxes.TryGetValue(scopeId, out var toolboxes)) + { + return []; + } + + return [.. toolboxes.Values.SelectMany(static cached => cached.Tools)]; + } + + private static string? GetCurrentRequestScopeId() => + HostedCallContext.ToolboxCacheScopeId ?? HostedCallContext.CallId; + + private SemaphoreSlim GetCurrentRequestOpenLock() + { + var scopeId = GetCurrentRequestScopeId() + ?? throw new InvalidOperationException( + "A response-scoped toolbox cache identifier is required before resolving a toolbox during request handling."); + return this._requestOpenLocks.GetOrAdd(scopeId, static _ => new SemaphoreSlim(1, 1)); + } + + private bool HasCurrentRequestResolution(string toolboxName) => + this.TryGetCurrentRequestToolbox(toolboxName, out _) + || this.TryGetCurrentRequestConsents(toolboxName, out _); + + private bool TryGetCurrentRequestToolbox( + string toolboxName, + [NotNullWhen(true)] out CachedToolbox? cached) + { + var scopeId = GetCurrentRequestScopeId(); + if (scopeId is not null + && this._requestToolboxes.TryGetValue(scopeId, out var toolboxes) + && toolboxes.TryGetValue(toolboxName, out cached)) + { + return true; + } + + cached = null; + return false; + } + + private bool TryGetCurrentRequestConsents( + string toolboxName, + [NotNullWhen(true)] out IReadOnlyList? consents) + { + var scopeId = GetCurrentRequestScopeId(); + if (scopeId is not null + && this._requestConsents.TryGetValue(scopeId, out var toolboxConsents) + && toolboxConsents.TryGetValue(toolboxName, out consents)) + { + return true; + } + + consents = null; + return false; + } + + private IReadOnlyList GetCurrentRequestConsents() + { + var scopeId = GetCurrentRequestScopeId(); + if (scopeId is null || !this._requestConsents.TryGetValue(scopeId, out var toolboxConsents)) + { + return []; + } + + return [.. toolboxConsents.Values.SelectMany(static consents => consents)]; + } + + private void CacheRequestToolbox(string toolboxName, CachedToolbox cached) + { + var scopeId = GetCurrentRequestScopeId() + ?? throw new InvalidOperationException( + "A response-scoped toolbox cache identifier is required before opening a toolbox during request handling."); + var toolboxes = this._requestToolboxes.GetOrAdd( + scopeId, + static _ => new ConcurrentDictionary(StringComparer.OrdinalIgnoreCase)); + toolboxes[toolboxName] = cached; + + if (this._requestConsents.TryGetValue(scopeId, out var toolboxConsents)) + { + toolboxConsents.TryRemove(toolboxName, out _); + } + } + + private void CacheRequestConsents(string toolboxName, IReadOnlyList consents) + { + var scopeId = GetCurrentRequestScopeId() + ?? throw new InvalidOperationException( + "A response-scoped toolbox cache identifier is required before resolving toolbox consent during request handling."); + var toolboxConsents = this._requestConsents.GetOrAdd( + scopeId, + static _ => new ConcurrentDictionary>(StringComparer.OrdinalIgnoreCase)); + toolboxConsents[toolboxName] = consents; + } + + private async ValueTask ReleaseRequestScopeWithoutThrowAsync(string scopeId) + { + _ = await this.ReleaseRequestScopeAsync(scopeId).ConfigureAwait(false); + } + + private async ValueTask> ReleaseRequestScopeAsync(string scopeId) + { + List failures = []; + this._requestConsents.TryRemove(scopeId, out _); + + if (this._requestToolboxes.TryRemove(scopeId, out var toolboxes)) + { + foreach (var (toolboxName, cached) in toolboxes) + { + await this.DisposeCachedToolboxAsync( + cached, + toolboxName, + $"response scope '{scopeId}'", + failures).ConfigureAwait(false); + } + } + + if (this._requestOpenLocks.TryRemove(scopeId, out var requestOpenLock)) + { + try + { + requestOpenLock.Dispose(); + } + catch (Exception ex) + { + this.RecordDisposalFailure( + ex, + resource: "request open lock", + toolboxName: null, + owner: $"response scope '{scopeId}'", + failures); + } + } + + return failures; + } + + private async ValueTask DisposeCachedToolboxAsync( + CachedToolbox cached, + string toolboxName, + string owner, + List failures) + { + if (cached.Client is not null) + { + try + { + await cached.Client.DisposeAsync().ConfigureAwait(false); + } + catch (Exception ex) + { + this.RecordDisposalFailure(ex, "MCP client", toolboxName, owner, failures); + } + } + + try + { + cached.HttpClient.Dispose(); + } + catch (Exception ex) + { + this.RecordDisposalFailure(ex, "HTTP client", toolboxName, owner, failures); + } + } + + private void RecordDisposalFailure( + Exception exception, + string resource, + string? toolboxName, + string owner, + List failures) + { + failures.Add(exception); + if (this._logger.IsEnabled(LogLevel.Warning)) + { + this._logger.LogWarning( + exception, + "Failed to dispose {Resource} for toolbox {ToolboxName} owned by {Owner}.", + resource, + toolboxName ?? "(none)", + owner); } } @@ -706,18 +907,57 @@ private async Task OpenToolboxAsync( /// public async ValueTask DisposeAsync() { - foreach (var cached in this._toolboxes.Values) + List failures = []; + foreach (var (toolboxName, cached) in this._toolboxes) { - if (cached.Client is not null) - { - await cached.Client.DisposeAsync().ConfigureAwait(false); - } - - cached.HttpClient.Dispose(); + await this.DisposeCachedToolboxAsync( + cached, + toolboxName, + "service startup cache", + failures).ConfigureAwait(false); } this._toolboxes.Clear(); - this._lazyOpenLock.Dispose(); + var requestScopeIds = this._requestToolboxes.Keys + .Concat(this._requestConsents.Keys) + .Concat(this._requestOpenLocks.Keys) + .Distinct(StringComparer.Ordinal) + .ToArray(); + foreach (var scopeId in requestScopeIds) + { + failures.AddRange(await this.ReleaseRequestScopeAsync(scopeId).ConfigureAwait(false)); + } + + this._requestConsents.Clear(); + this._requestOpenLocks.Clear(); + + if (failures.Count > 0) + { + throw new AggregateException("One or more Foundry toolbox resources failed to dispose.", failures); + } + } + + internal sealed class RequestToolboxScope : IAsyncDisposable + { + private readonly FoundryToolboxService? _owner; + private int _disposed; + + private RequestToolboxScope(FoundryToolboxService? owner, string id) + { + this._owner = owner; + this.Id = id; + } + + internal string Id { get; } + + internal static RequestToolboxScope Create(FoundryToolboxService owner, string id) => new(owner, id); + + internal static RequestToolboxScope CreateInactive() => new(owner: null, Guid.NewGuid().ToString("N")); + + public ValueTask DisposeAsync() => + Interlocked.Exchange(ref this._disposed, 1) == 0 && this._owner is not null + ? this._owner.ReleaseRequestScopeWithoutThrowAsync(this.Id) + : ValueTask.CompletedTask; } internal sealed record CachedToolbox(McpClient? Client, HttpClient HttpClient, IReadOnlyList Tools); diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedCallContext.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedCallContext.cs index 52be548ffa..abcbefabda 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedCallContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedCallContext.cs @@ -17,6 +17,7 @@ namespace Microsoft.Agents.AI.Foundry.Hosting; internal static class HostedCallContext { private static readonly AsyncLocal s_callId = new(); + private static readonly AsyncLocal s_toolboxCacheScopeId = new(); /// Gets or sets the current request's call id, or when absent. public static string? CallId @@ -24,4 +25,14 @@ public static string? CallId get => s_callId.Value; set => s_callId.Value = value; } + + /// + /// Gets or sets the current response's toolbox cache scope, or when no + /// response-owned toolbox client may be opened. + /// + public static string? ToolboxCacheScopeId + { + get => s_toolboxCacheScopeId.Value; + set => s_toolboxCacheScopeId.Value = value; + } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryToolboxMarkerScopingTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryToolboxMarkerScopingTests.cs index f7e3d13bba..21b441e698 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryToolboxMarkerScopingTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryToolboxMarkerScopingTests.cs @@ -59,6 +59,7 @@ public async Task GetToolboxToolsAsync_MarkerConsent_IsRequestScopedAndDoesNotMu Consents: [new McpConsentInfo(name, $"{name}.tool", $"https://consent.example/{name}")]))); // Act: request A references marker-a and hits consent. + HostedCallContext.CallId = "request-a"; var resolutionA = await service.GetToolboxToolsAsync("marker-a", version: null, CancellationToken.None); // Assert: the consent is returned to THIS caller, with no tools. @@ -74,6 +75,7 @@ public async Task GetToolboxToolsAsync_MarkerConsent_IsRequestScopedAndDoesNotMu Assert.Empty(service.Tools); // Act: a different request references marker-b. Its consent must not accumulate globally. + HostedCallContext.CallId = "request-b"; var resolutionB = await service.GetToolboxToolsAsync("marker-b", version: null, CancellationToken.None); // Assert: still request-scoped, still no global mutation. @@ -101,6 +103,7 @@ public async Task GetToolboxToolsAsync_MarkerTools_AreReturnedToCallerNotInjecte Consents: [new McpConsentInfo(name, $"{name}.tool", $"https://consent.example/{name}")]))); // Act: request A references marker-a and resolves its tool. + HostedCallContext.CallId = "request-a"; var resolutionA = await service.GetToolboxToolsAsync("marker-a", version: null, CancellationToken.None); // Assert: the tool is returned to THIS caller only. diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryToolboxResponseScopeTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryToolboxResponseScopeTests.cs new file mode 100644 index 0000000000..2d6cd754c7 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryToolboxResponseScopeTests.cs @@ -0,0 +1,265 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Concurrent; +using System.Collections.Generic; +using System.Linq; +using System.Net.Http; +using System.Reflection; +using System.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Tasks; +using Azure.AI.AgentServer.Responses; +using Azure.AI.AgentServer.Responses.Models; +using Azure.Core; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Logging.Abstractions; +using Microsoft.Extensions.Options; +using Moq; + +namespace Microsoft.Agents.AI.Foundry.Hosting.UnitTests; + +[Collection(FoundryProjectEndpointEnvFixture.Name)] +public sealed class FoundryToolboxResponseScopeTests +{ + private static readonly string s_sharedResponseId = "resp_" + new string('7', 46); + + [Fact] + public async Task CreateAsync_ToolboxClientsAreIsolatedReusedAndDisposedPerResponseAsync() + { + // Arrange + HostedCallContext.CallId = null; + SetToolboxCacheScopeId(null); + + var openScopes = new ConcurrentBag(); + var agentScopes = new ConcurrentBag(); + var observedTools = new ConcurrentBag(); + var handlers = new ConcurrentBag(); + var bothConcurrentOpensEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var bothConcurrentRunsEnteredAgent = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var concurrentOpenCount = 0; + var chatCallCount = 0; + var openCount = 0; + + var options = new FoundryToolboxOptions + { + StrictMode = false, + EndpointOverride = "https://proj.example/api/projects/proj", + }; + options.ToolboxNames.Add("shared-toolbox"); + + await using var toolboxService = new FoundryToolboxService( + Options.Create(options), + Mock.Of()) + { + ToolboxOpener = async (_, _, cancellationToken) => + { + var scopeId = GetToolboxCacheScopeId(); + if (scopeId is null) + { + throw new InvalidOperationException("Startup has no response scope."); + } + + if (Interlocked.Increment(ref concurrentOpenCount) == 2) + { + bothConcurrentOpensEntered.TrySetResult(); + } + + await bothConcurrentOpensEntered.Task + .WaitAsync(TimeSpan.FromSeconds(5), cancellationToken); + openScopes.Add(scopeId); + var openNumber = Interlocked.Increment(ref openCount); + var handler = new TrackingHttpMessageHandler(throwOnDispose: openNumber == 1); + handlers.Add(handler); + AITool tool = AIFunctionFactory.Create(() => scopeId, name: $"scoped_tool_{openNumber}"); + return new FoundryToolboxService.ToolboxOpenResult( + new FoundryToolboxService.CachedToolbox( + Client: null, + new HttpClient(handler), + [tool]), + Consents: null); + }, + }; + await toolboxService.StartAsync(CancellationToken.None); + + var chatClient = new Mock(); + chatClient + .Setup(c => c.GetStreamingResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns((IEnumerable _, ChatOptions? chatOptions, CancellationToken cancellationToken) => + { + agentScopes.Add(GetToolboxCacheScopeId() ?? string.Empty); + observedTools.Add(Assert.Single(chatOptions?.Tools ?? [])); + + var callNumber = Interlocked.Increment(ref chatCallCount); + if (callNumber == 2) + { + bothConcurrentRunsEnteredAgent.TrySetResult(); + } + + return callNumber <= 2 + ? YieldAfterAsync(bothConcurrentRunsEnteredAgent.Task, cancellationToken) + : ThrowAsync(); + }); + + var agent = new ChatClientAgent(chatClient.Object); + var services = new ServiceCollection(); + services.AddSingleton(new InMemoryAgentSessionStore()); + services.AddSingleton(agent); + services.AddSingleton(new FakeHostedSessionIsolationKeyProvider()); + var handler = new AgentFrameworkResponseHandler( + services.BuildServiceProvider(), + NullLogger.Instance, + toolboxService); + + // Act: two responses deliberately share the same public response id. Their internally + // generated cache scopes must still be independent, while each response reuses its one + // deferred open when the same toolbox also appears as a request marker. + var first = RunAsync(handler, CreateRequest(), CreateContext(callId: "call-a")); + var second = RunAsync(handler, CreateRequest(), CreateContext(callId: "call-b")); + await Task.WhenAll(first, second); + + // Assert: both responses opened independently, carried the same opaque scope across the + // handler's initial lifecycle yields, and disposed their response-owned clients. + Assert.Equal(2, openCount); + Assert.Equal(2, openScopes.Distinct(StringComparer.Ordinal).Count()); + Assert.Equal( + openScopes.OrderBy(static value => value, StringComparer.Ordinal), + agentScopes.OrderBy(static value => value, StringComparer.Ordinal)); + Assert.Equal(2, observedTools.Distinct().Count()); + Assert.Equal(2, handlers.Count); + Assert.All(handlers, static item => Assert.True(item.IsDisposed)); + Assert.Equal(0, GetActiveRequestScopeCount(toolboxService)); + + // Act: an agent failure is an early terminal path. The handler may translate the exception + // to a failed response or propagate it, but the response-owned toolbox client must be released. + List? failedEvents = null; + var exception = await Record.ExceptionAsync( + async () => failedEvents = await RunAsync( + handler, + CreateRequest(), + CreateContext(callId: "call-failure"))); + + // Assert + Assert.True( + exception is InvalidOperationException + || failedEvents?.LastOrDefault() is ResponseFailedEvent); + Assert.Equal(3, openCount); + Assert.Equal(3, handlers.Count); + Assert.All(handlers, static item => Assert.True(item.IsDisposed)); + Assert.Equal(0, GetActiveRequestScopeCount(toolboxService)); + } + + private static CreateResponse CreateRequest() + { + var request = new CreateResponse { Model = "test" }; + request.Input = BinaryData.FromObjectAsJson(new[] + { + new + { + type = "message", + id = "msg_1", + status = "completed", + role = "user", + content = new[] { new { type = "input_text", text = "Hello" } }, + }, + }); + request.Tools.Add(new MCPTool("shared-toolbox") + { + ServerUrl = new Uri("foundry-toolbox://shared-toolbox"), + }); + return request; + } + + private static ResponseContext CreateContext(string callId) + { + var context = new Mock(s_sharedResponseId) { CallBase = true }; + context.Setup(x => x.PlatformContext).Returns(new PlatformContext("user", callId)); + context.Setup(x => x.GetHistoryAsync(It.IsAny())) + .ReturnsAsync(Array.Empty()); + context.Setup(x => x.GetInputItemsAsync(It.IsAny(), It.IsAny())) + .ReturnsAsync(Array.Empty()); + return context.Object; + } + + private static async Task> RunAsync( + AgentFrameworkResponseHandler handler, + CreateResponse request, + ResponseContext context) + { + var events = new List(); + await foreach (var responseEvent in handler.CreateAsync(request, context, CancellationToken.None)) + { + events.Add(responseEvent); + } + + return events; + } + + private static string? GetToolboxCacheScopeId() => + typeof(HostedCallContext) + .GetProperty("ToolboxCacheScopeId", BindingFlags.Public | BindingFlags.Static)? + .GetValue(null) as string; + + private static void SetToolboxCacheScopeId(string? value) => + typeof(HostedCallContext) + .GetProperty("ToolboxCacheScopeId", BindingFlags.Public | BindingFlags.Static)? + .SetValue(null, value); + + private static int GetActiveRequestScopeCount(FoundryToolboxService service) + { + var scopes = typeof(FoundryToolboxService) + .GetField("_requestToolboxes", BindingFlags.Instance | BindingFlags.NonPublic)? + .GetValue(service); + return scopes is null + ? -1 + : (int)scopes.GetType().GetProperty("Count")!.GetValue(scopes)!; + } + + private static async IAsyncEnumerable YieldAfterAsync( + Task gate, + [EnumeratorCancellation] CancellationToken cancellationToken) + { + await gate.WaitAsync(cancellationToken); + yield return new ChatResponseUpdate(ChatRole.Assistant, "ok") { MessageId = "resp_msg_1" }; + } + + private static async IAsyncEnumerable ThrowAsync() + { + await Task.Yield(); + throw new InvalidOperationException("Agent failed after toolbox resolution."); +#pragma warning disable CS0162 // Required to make this an async iterator. + yield break; +#pragma warning restore CS0162 + } + + private sealed class TrackingHttpMessageHandler : HttpMessageHandler + { + private readonly bool _throwOnDispose; + + internal TrackingHttpMessageHandler(bool throwOnDispose = false) + { + this._throwOnDispose = throwOnDispose; + } + + internal bool IsDisposed { get; private set; } + + protected override Task SendAsync( + HttpRequestMessage request, + CancellationToken cancellationToken) => + throw new NotSupportedException(); + + protected override void Dispose(bool disposing) + { + this.IsDisposed = true; + base.Dispose(disposing); + if (this._throwOnDispose) + { + throw new InvalidOperationException("Simulated HTTP client disposal failure."); + } + } + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryToolboxServiceTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryToolboxServiceTests.cs index 59b01ed9c4..d646a6fa97 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryToolboxServiceTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryToolboxServiceTests.cs @@ -1,11 +1,15 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Collections.Concurrent; +using System.Collections.Generic; +using System.Linq; using System.Net.Http; using System.Threading; using System.Threading.Tasks; using Azure.Core; using Microsoft.Extensions.AI; +using Microsoft.Extensions.Logging; using Microsoft.Extensions.Options; using Moq; @@ -235,12 +239,14 @@ public async Task RetryDeferredToolboxesAsync_StillUnreachable_StaysDeferredAndR }; options.ToolboxNames.Add("broken-toolbox"); - var service = new FoundryToolboxService( + await using var service = new FoundryToolboxService( Options.Create(options), Mock.Of()); await service.StartAsync(CancellationToken.None); Assert.Single(service.DeferredToolboxNames); + await using var scope = service.CreateRequestScope(); + HostedCallContext.ToolboxCacheScopeId = scope.Id; // Act: retry while the endpoint is still unreachable. await service.RetryDeferredToolboxesAsync(CancellationToken.None); @@ -350,6 +356,7 @@ public async Task GetToolboxToolsAsync_AllowsWellFormedSingleSegmentNameAsync(st Consents: null)), }; await service.StartAsync(CancellationToken.None); + HostedCallContext.CallId = $"request-{validName}"; // Act var resolution = await service.GetToolboxToolsAsync(validName, version: null, CancellationToken.None); @@ -359,4 +366,453 @@ public async Task GetToolboxToolsAsync_AllowsWellFormedSingleSegmentNameAsync(st Assert.Single(resolution.Tools); Assert.Same(tool, resolution.Tools[0]); } + + [Fact] + public async Task GetToolboxToolsAsync_DifferentRequestContexts_DoNotReuseCachedToolsAsync() + { + // Arrange + var options = new FoundryToolboxOptions + { + StrictMode = false, + EndpointOverride = "https://proj.example/api/projects/proj", + }; + var openCount = 0; + await using var service = new FoundryToolboxService( + Options.Create(options), + Mock.Of()) + { + ToolboxOpener = (_, _, _) => + { + openCount++; + AITool tool = AIFunctionFactory.Create(() => HostedCallContext.CallId, name: $"tool_{openCount}"); + return Task.FromResult( + new FoundryToolboxService.ToolboxOpenResult( + new FoundryToolboxService.CachedToolbox(Client: null, new HttpClient(), [tool]), + Consents: null)); + }, + }; + await service.StartAsync(CancellationToken.None); + + // Act + HostedCallContext.CallId = "request-a"; + var first = await service.GetToolboxToolsAsync("shared-toolbox", version: null, CancellationToken.None); + HostedCallContext.CallId = "request-b"; + var second = await service.GetToolboxToolsAsync("shared-toolbox", version: null, CancellationToken.None); + + // Assert + Assert.Equal(2, openCount); + Assert.NotSame(first.Tools[0], second.Tools[0]); + } + + [Fact] + public async Task GetToolboxToolsAsync_SameRequestContext_ReusesCachedToolsAsync() + { + // Arrange + var options = new FoundryToolboxOptions + { + StrictMode = false, + EndpointOverride = "https://proj.example/api/projects/proj", + }; + var openCount = 0; + await using var service = new FoundryToolboxService( + Options.Create(options), + Mock.Of()) + { + ToolboxOpener = (_, _, _) => + { + openCount++; + AITool tool = AIFunctionFactory.Create(() => "ok", name: "shared_tool"); + return Task.FromResult( + new FoundryToolboxService.ToolboxOpenResult( + new FoundryToolboxService.CachedToolbox(Client: null, new HttpClient(), [tool]), + Consents: null)); + }, + }; + await service.StartAsync(CancellationToken.None); + + // Act + HostedCallContext.CallId = "request-a"; + var first = await service.GetToolboxToolsAsync("shared-toolbox", version: null, CancellationToken.None); + var second = await service.GetToolboxToolsAsync("shared-toolbox", version: null, CancellationToken.None); + + // Assert + Assert.Equal(1, openCount); + Assert.Same(first.Tools[0], second.Tools[0]); + } + + [Fact] + public async Task RetryDeferredToolboxesAsync_DifferentRequestContexts_RetryIndependentlyAsync() + { + // Arrange + var options = new FoundryToolboxOptions + { + EndpointOverride = "https://proj.example/api/projects/proj", + }; + options.ToolboxNames.Add("shared-toolbox"); + var openCount = 0; + await using var service = new FoundryToolboxService( + Options.Create(options), + Mock.Of()) + { + ToolboxOpener = (_, _, _) => + { + openCount++; + if (HostedCallContext.CallId is null) + { + throw new InvalidOperationException("A request context is required."); + } + + AITool tool = AIFunctionFactory.Create(() => HostedCallContext.CallId, name: $"tool_{openCount}"); + return Task.FromResult( + new FoundryToolboxService.ToolboxOpenResult( + new FoundryToolboxService.CachedToolbox(Client: null, new HttpClient(), [tool]), + Consents: null)); + }, + }; + await service.StartAsync(CancellationToken.None); + + // Act + HostedCallContext.CallId = "request-a"; + await service.RetryDeferredToolboxesAsync(CancellationToken.None); + HostedCallContext.CallId = "request-b"; + await service.RetryDeferredToolboxesAsync(CancellationToken.None); + + // Assert + Assert.Equal(3, openCount); + } + + [Fact] + public async Task ResolvePendingConsentsAsync_DifferentRequestContexts_ResolveIndependentlyAsync() + { + // Arrange + var options = new FoundryToolboxOptions + { + EndpointOverride = "https://proj.example/api/projects/proj", + }; + options.ToolboxNames.Add("shared-toolbox"); + var openCount = 0; + await using var service = new FoundryToolboxService( + Options.Create(options), + Mock.Of()) + { + ToolboxOpener = (_, _, _) => + { + openCount++; + if (HostedCallContext.CallId is null) + { + IReadOnlyList consents = + [ + new("shared-toolbox", "send", "https://auth.example.com/consent"), + ]; + return Task.FromResult( + new FoundryToolboxService.ToolboxOpenResult(Cached: null, Consents: consents)); + } + + AITool tool = AIFunctionFactory.Create(() => HostedCallContext.CallId, name: $"tool_{openCount}"); + return Task.FromResult( + new FoundryToolboxService.ToolboxOpenResult( + new FoundryToolboxService.CachedToolbox(Client: null, new HttpClient(), [tool]), + Consents: null)); + }, + }; + await service.StartAsync(CancellationToken.None); + + // Act + HostedCallContext.CallId = "request-a"; + var first = await service.ResolvePendingConsentsAsync(CancellationToken.None); + HostedCallContext.CallId = "request-b"; + var second = await service.ResolvePendingConsentsAsync(CancellationToken.None); + + // Assert + Assert.Empty(first); + Assert.Empty(second); + Assert.Equal(3, openCount); + } + + [Fact] + public async Task GetToolboxToolsAsync_StartupOpenedToolbox_ReusesContainerCacheAcrossRequestsAsync() + { + // Arrange + var options = new FoundryToolboxOptions + { + EndpointOverride = "https://proj.example/api/projects/proj", + }; + options.ToolboxNames.Add("shared-toolbox"); + var openCount = 0; + await using var service = new FoundryToolboxService( + Options.Create(options), + Mock.Of()) + { + ToolboxOpener = (_, _, _) => + { + openCount++; + AITool tool = AIFunctionFactory.Create(() => "ok", name: "shared_tool"); + return Task.FromResult( + new FoundryToolboxService.ToolboxOpenResult( + new FoundryToolboxService.CachedToolbox(Client: null, new HttpClient(), [tool]), + Consents: null)); + }, + }; + await service.StartAsync(CancellationToken.None); + + // Act + HostedCallContext.CallId = "request-a"; + var first = await service.GetToolboxToolsAsync("shared-toolbox", version: null, CancellationToken.None); + HostedCallContext.CallId = "request-b"; + var second = await service.GetToolboxToolsAsync("shared-toolbox", version: null, CancellationToken.None); + + // Assert + Assert.Equal(1, openCount); + Assert.Same(first.Tools[0], second.Tools[0]); + } + + [Fact] + public async Task RequestToolboxScope_Dispose_ReleasesScopedClientAsync() + { + // Arrange + var options = new FoundryToolboxOptions + { + StrictMode = false, + EndpointOverride = "https://proj.example/api/projects/proj", + }; + var handler = new TrackingHttpMessageHandler(); + await using var service = new FoundryToolboxService( + Options.Create(options), + Mock.Of()) + { + ToolboxOpener = (_, _, _) => + { + AITool tool = AIFunctionFactory.Create(() => "ok", name: "shared_tool"); + return Task.FromResult( + new FoundryToolboxService.ToolboxOpenResult( + new FoundryToolboxService.CachedToolbox( + Client: null, + new HttpClient(handler), + [tool]), + Consents: null)); + }, + }; + await service.StartAsync(CancellationToken.None); + var scope = service.CreateRequestScope(); + HostedCallContext.ToolboxCacheScopeId = scope.Id; + _ = await service.GetToolboxToolsAsync("shared-toolbox", version: null, CancellationToken.None); + + // Act + await scope.DisposeAsync(); + + // Assert + Assert.True(handler.IsDisposed); + } + + [Fact] + public async Task GetToolboxToolsAsync_ConcurrentOpaqueScopes_DoNotCollideAsync() + { + // Arrange + var options = new FoundryToolboxOptions + { + StrictMode = false, + EndpointOverride = "https://proj.example/api/projects/proj", + }; + var openedScopes = new ConcurrentBag(); + var handlers = new ConcurrentBag(); + var bothOpenersEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var concurrentOpenCount = 0; + var openCount = 0; + await using var service = new FoundryToolboxService( + Options.Create(options), + Mock.Of()) + { + ToolboxOpener = async (_, _, cancellationToken) => + { + var scopeId = Assert.IsType(HostedCallContext.ToolboxCacheScopeId); + if (Interlocked.Increment(ref concurrentOpenCount) == 2) + { + bothOpenersEntered.TrySetResult(); + } + + await bothOpenersEntered.Task.WaitAsync(TimeSpan.FromSeconds(5), cancellationToken); + openedScopes.Add(scopeId); + var handler = new TrackingHttpMessageHandler(); + handlers.Add(handler); + var openNumber = Interlocked.Increment(ref openCount); + AITool tool = AIFunctionFactory.Create(() => scopeId, name: $"scoped_tool_{openNumber}"); + return new FoundryToolboxService.ToolboxOpenResult( + new FoundryToolboxService.CachedToolbox( + Client: null, + new HttpClient(handler), + [tool]), + Consents: null); + }, + }; + await service.StartAsync(CancellationToken.None); + + async Task ResolveInOwnScopeAsync() + { + await using var scope = service.CreateRequestScope(); + HostedCallContext.ToolboxCacheScopeId = scope.Id; + return await service.GetToolboxToolsAsync( + "shared-toolbox", + version: null, + CancellationToken.None); + } + + // Act + var resolutions = await Task.WhenAll( + Task.Run(ResolveInOwnScopeAsync), + Task.Run(ResolveInOwnScopeAsync)); + + // Assert + Assert.Equal(2, openCount); + Assert.Equal(2, openedScopes.Distinct(StringComparer.Ordinal).Count()); + Assert.NotSame(resolutions[0].Tools[0], resolutions[1].Tools[0]); + Assert.Equal(2, handlers.Count); + Assert.All(handlers, static item => Assert.True(item.IsDisposed)); + } + + [Fact] + public async Task RequestToolboxScope_DisposalFailure_CleansRemainingResourcesAndDoesNotThrowAsync() + { + // Arrange + var logger = new RecordingLogger(); + var handlers = new ConcurrentDictionary(StringComparer.Ordinal); + var options = new FoundryToolboxOptions + { + StrictMode = false, + EndpointOverride = "https://proj.example/api/projects/proj", + }; + var service = new FoundryToolboxService( + Options.Create(options), + Mock.Of(), + logger) + { + ToolboxOpener = (name, _, _) => + { + var handler = new TrackingHttpMessageHandler(throwOnDispose: name == "throws"); + handlers[name] = handler; + AITool tool = AIFunctionFactory.Create(() => "ok", name: $"{name}_tool"); + return Task.FromResult( + new FoundryToolboxService.ToolboxOpenResult( + new FoundryToolboxService.CachedToolbox( + Client: null, + new HttpClient(handler), + [tool]), + Consents: null)); + }, + }; + await service.StartAsync(CancellationToken.None); + var scope = service.CreateRequestScope(); + HostedCallContext.ToolboxCacheScopeId = scope.Id; + _ = await service.GetToolboxToolsAsync("throws", version: null, CancellationToken.None); + _ = await service.GetToolboxToolsAsync("continues", version: null, CancellationToken.None); + + // Act + var exception = await Record.ExceptionAsync(async () => await scope.DisposeAsync()); + + // Assert + Assert.Null(exception); + Assert.True(handlers["throws"].IsDisposed); + Assert.True(handlers["continues"].IsDisposed); + Assert.Single(logger.Exceptions); + Assert.Empty(service.GetCurrentRequestTools()); + + await service.DisposeAsync(); + } + + [Fact] + public async Task DisposeAsync_DisposalFailure_CleansRemainingResourcesAndThrowsAggregateAsync() + { + // Arrange + var logger = new RecordingLogger(); + var handlers = new ConcurrentDictionary(StringComparer.Ordinal); + var options = new FoundryToolboxOptions + { + StrictMode = false, + EndpointOverride = "https://proj.example/api/projects/proj", + }; + var service = new FoundryToolboxService( + Options.Create(options), + Mock.Of(), + logger) + { + ToolboxOpener = (name, _, _) => + { + var handler = new TrackingHttpMessageHandler(throwOnDispose: name == "throws"); + handlers[name] = handler; + AITool tool = AIFunctionFactory.Create(() => "ok", name: $"{name}_tool"); + return Task.FromResult( + new FoundryToolboxService.ToolboxOpenResult( + new FoundryToolboxService.CachedToolbox( + Client: null, + new HttpClient(handler), + [tool]), + Consents: null)); + }, + }; + await service.StartAsync(CancellationToken.None); + var scope = service.CreateRequestScope(); + HostedCallContext.ToolboxCacheScopeId = scope.Id; + _ = await service.GetToolboxToolsAsync("throws", version: null, CancellationToken.None); + _ = await service.GetToolboxToolsAsync("continues", version: null, CancellationToken.None); + + // Act + var exception = await Assert.ThrowsAsync( + async () => await service.DisposeAsync()); + + // Assert + Assert.Single(exception.InnerExceptions); + Assert.True(handlers["throws"].IsDisposed); + Assert.True(handlers["continues"].IsDisposed); + Assert.Single(logger.Exceptions); + } + + private sealed class TrackingHttpMessageHandler : HttpMessageHandler + { + private readonly bool _throwOnDispose; + + internal TrackingHttpMessageHandler(bool throwOnDispose = false) + { + this._throwOnDispose = throwOnDispose; + } + + internal bool IsDisposed { get; private set; } + + protected override Task SendAsync( + HttpRequestMessage request, + CancellationToken cancellationToken) => + throw new NotSupportedException(); + + protected override void Dispose(bool disposing) + { + this.IsDisposed = true; + base.Dispose(disposing); + if (this._throwOnDispose) + { + throw new InvalidOperationException("Simulated HTTP client disposal failure."); + } + } + } + + private sealed class RecordingLogger : ILogger + { + internal ConcurrentBag Exceptions { get; } = []; + + public IDisposable? BeginScope(TState state) + where TState : notnull => + null; + + public bool IsEnabled(LogLevel logLevel) => true; + + public void Log( + LogLevel logLevel, + EventId eventId, + TState state, + Exception? exception, + Func formatter) + { + if (exception is not null) + { + this.Exceptions.Add(exception); + } + } + } }