From d7f2a7a8019d13f5f10f3bd72135d6c6b1cc28c1 Mon Sep 17 00:00:00 2001 From: Marco Minerva Date: Mon, 19 Oct 2026 08:52:23 +0000 Subject: [PATCH 1/3] .NET: Honor hosted agent lifetimes and add request-time agent selection for OpenAI Responses hosting Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- dotnet/eng/MSBuild/Shared.props | 3 + ...ndpointRouteBuilderExtensions.Responses.cs | 109 +++++++++-- .../Microsoft.Agents.AI.Hosting.OpenAI.csproj | 1 + .../Responses/AIAgentResponseExecutor.cs | 9 +- .../Responses/IResponsesService.cs | 36 ++++ .../Responses/InMemoryResponsesService.cs | 100 ++++++++-- .../Responses/ResponsesHttpHandler.cs | 97 ++++++++-- .../HostedAgents/HostedAgentResolution.cs | 52 +++++ .../InMemoryResponsesServiceTests.cs | 32 +++ .../ResponseSessionRegressionTests.cs | 183 ++++++++++++++++++ 10 files changed, 572 insertions(+), 50 deletions(-) create mode 100644 dotnet/src/Shared/HostedAgents/HostedAgentResolution.cs diff --git a/dotnet/eng/MSBuild/Shared.props b/dotnet/eng/MSBuild/Shared.props index 10db0a47ef4..3c8f9670284 100644 --- a/dotnet/eng/MSBuild/Shared.props +++ b/dotnet/eng/MSBuild/Shared.props @@ -38,6 +38,9 @@ + + + diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/EndpointRouteBuilderExtensions.Responses.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/EndpointRouteBuilderExtensions.Responses.cs index 0403d1eb172..63c66399bc6 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/EndpointRouteBuilderExtensions.Responses.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/EndpointRouteBuilderExtensions.Responses.cs @@ -38,6 +38,9 @@ public static IEndpointConventionBuilder MapOpenAIResponses(this IEndpointRouteB /// Custom route path for the OpenAI Responses endpoint. /// Optional options controlling how incoming requests are mapped onto the agent run. /// + /// Singleton registrations are resolved once when the endpoint is mapped. Scoped and transient registrations are + /// resolved for each validation or execution operation. The keyed session store is always resolved for each operation, + /// and persisted sessions are partitioned by the registration name. /// See for endpoint authorization /// and caller isolation requirements. /// @@ -46,21 +49,24 @@ public static IEndpointConventionBuilder MapOpenAIResponses(this IEndpointRouteB ArgumentNullException.ThrowIfNull(endpoints); ArgumentNullException.ThrowIfNull(agentBuilder); - ValidateAgentName(agentBuilder.Name); - path ??= $"/{agentBuilder.Name}/v1/responses"; + string agentName = agentBuilder.Name; + ValidateAgentName(agentName); + path ??= $"/{agentName}/v1/responses"; - // Defer agent and session-store resolution so their registered lifetimes are owned by - // each validation or execution operation rather than by the application root. var scopeFactory = endpoints.ServiceProvider.GetRequiredService(); - var executor = new AIAgentResponseExecutor( - agentBuilder.Name, - scopeFactory, - mapOptions); + var executor = agentBuilder.Lifetime == ServiceLifetime.Singleton + ? new AIAgentResponseExecutor( + endpoints.ServiceProvider.GetRequiredKeyedService(agentName), + agentName, + scopeFactory, + mapOptions, + sessionStorageIdentity: agentName) + : new AIAgentResponseExecutor(agentName, scopeFactory, mapOptions); return MapOpenAIResponses( endpoints, executor, path, - agentBuilder.Name); + agentName); } /// @@ -114,11 +120,72 @@ public static IEndpointConventionBuilder MapOpenAIResponses( agent.Name); } + /// + /// Maps OpenAI Responses API endpoints that select the for each request. + /// + /// The to add the OpenAI Responses endpoints to. + /// The route path for the responses endpoint. + /// + /// A delegate invoked once per request that returns the agent to serve, or when + /// no agent matches the request. The delegate can use route values, claims, headers, and + /// to select or resolve the agent. + /// + /// Optional options controlling how incoming requests are mapped onto the agent run. + /// An for further endpoint configuration. + /// + /// + /// The selector applies to every operation: creation, retrieval, cancellation, deletion, and input item listing. + /// When returns , the operation responds with + /// 404 Not Found before any response, session, or conversation state is loaded. + /// + /// + /// The selected agent's is its stable logical identity: it is used to resolve the + /// keyed , to partition persisted sessions, and to bind stored responses to the + /// agent that created them, so a response can only be retrieved, canceled, or deleted through requests that select + /// the same agent name. Conversations continue across requests even when the selector returns a new agent instance + /// each time. The selected agent must therefore have a non-empty name. Background responses keep running the + /// selected agent after the HTTP request completes, so agents used for background responses must not depend on + /// request-scoped services. + /// + /// + /// Selection is not an authorization boundary. See + /// for endpoint authorization and caller isolation requirements. + /// + /// + public static IEndpointConventionBuilder MapOpenAIResponses( + this IEndpointRouteBuilder endpoints, + [StringSyntax("Route")] string responsesPath, + Func agentSelector, + OpenAIResponsesMapOptions? mapOptions = null) + { + ArgumentNullException.ThrowIfNull(endpoints); + ArgumentException.ThrowIfNullOrWhiteSpace(responsesPath); + ArgumentNullException.ThrowIfNull(agentSelector); + + var scopeFactory = endpoints.ServiceProvider.GetRequiredService(); + return MapOpenAIResponses( + endpoints, + executor: null, + responsesPath, + endpointAgentName: null, + context => + { + if (agentSelector(context) is not { } agent) + { + return null; + } + + string agentName = HostedAgentResolution.GetSelectedAgentName(agent); + return (agentName, new AIAgentResponseExecutor(agent, agentName, scopeFactory, mapOptions, sessionStorageIdentity: agentName)); + }); + } + private static RouteGroupBuilder MapOpenAIResponses( IEndpointRouteBuilder endpoints, - IResponseExecutor executor, + IResponseExecutor? executor, string responsesPath, - string endpointAgentName) + string? endpointAgentName, + Func? agentSelector = null) { // Resolve the response storage settings and optional conversation storage. var storageOptions = endpoints.ServiceProvider.GetService() ?? new InMemoryStorageOptions(); @@ -132,38 +199,46 @@ private static RouteGroupBuilder MapOpenAIResponses( // Create the response service so response and conversation operations are scoped by the caller's isolation key. var responsesService = new InMemoryResponsesService(executor, storageOptions, conversationStorage, isolationKeyResolver); - var handlers = new ResponsesHttpHandler(responsesService); + var handlers = new ResponsesHttpHandler(responsesService, agentSelector is null ? null : SelectAgent); var group = endpoints.MapGroup(responsesPath); // Create response endpoint group.MapPost("/", handlers.CreateResponseAsync) - .WithName(endpointAgentName + "/CreateResponse") + .WithEndpointName(endpointAgentName, "CreateResponse") .WithSummary("Creates a model response for the given input"); // Get response endpoint group.MapGet("{responseId}", handlers.GetResponseAsync) - .WithName(endpointAgentName + "/GetResponse") + .WithEndpointName(endpointAgentName, "GetResponse") .WithSummary("Retrieves a response by ID"); // Cancel response endpoint group.MapPost("{responseId}/cancel", handlers.CancelResponseAsync) - .WithName(endpointAgentName + "/CancelResponse") + .WithEndpointName(endpointAgentName, "CancelResponse") .WithSummary("Cancels an in-progress response"); // Delete response endpoint group.MapDelete("{responseId}", handlers.DeleteResponseAsync) - .WithName(endpointAgentName + "/DeleteResponse") + .WithEndpointName(endpointAgentName, "DeleteResponse") .WithSummary("Deletes a response"); // List response input items endpoint group.MapGet("{responseId}/input_items", handlers.ListResponseInputItemsAsync) - .WithName(endpointAgentName + "/ListResponseInputItems") + .WithEndpointName(endpointAgentName, "ListResponseInputItems") .WithSummary("Lists the input items for a response"); MarkFeatureUsed(); return group; + + (IResponsesService, IResponseExecutor)? SelectAgent(HttpContext context) => + agentSelector(context) is { } selection + ? (responsesService.ForAgent(selection.AgentName), selection.Executor) + : null; } + private static RouteHandlerBuilder WithEndpointName(this RouteHandlerBuilder builder, string? endpointAgentName, string operationName) + => endpointAgentName is null ? builder : builder.WithName(endpointAgentName + "/" + operationName); + /// /// Maps OpenAI Responses API endpoints to the specified . /// diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Microsoft.Agents.AI.Hosting.OpenAI.csproj b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Microsoft.Agents.AI.Hosting.OpenAI.csproj index baf2152635c..b99905c72e1 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Microsoft.Agents.AI.Hosting.OpenAI.csproj +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Microsoft.Agents.AI.Hosting.OpenAI.csproj @@ -14,6 +14,7 @@ true true + true diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/AIAgentResponseExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/AIAgentResponseExecutor.cs index 84e95d7f0d7..92b943ac2d6 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/AIAgentResponseExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/AIAgentResponseExecutor.cs @@ -19,6 +19,7 @@ internal sealed class AIAgentResponseExecutor : IResponseExecutor { private readonly AIAgent? _agent; private readonly string _registrationKey; + private readonly string? _sessionStorageIdentity; private readonly IServiceScopeFactory _scopeFactory; private readonly OpenAIResponsesMapOptions _mapOptions; @@ -26,7 +27,8 @@ public AIAgentResponseExecutor( AIAgent agent, string registrationKey, IServiceScopeFactory scopeFactory, - OpenAIResponsesMapOptions? mapOptions = null) + OpenAIResponsesMapOptions? mapOptions = null, + string? sessionStorageIdentity = null) { ArgumentNullException.ThrowIfNull(agent); ArgumentException.ThrowIfNullOrEmpty(registrationKey); @@ -34,6 +36,7 @@ public AIAgentResponseExecutor( this._agent = agent; this._registrationKey = registrationKey; + this._sessionStorageIdentity = sessionStorageIdentity; this._scopeFactory = scopeFactory; this._mapOptions = mapOptions ?? new OpenAIResponsesMapOptions(); } @@ -47,6 +50,7 @@ public AIAgentResponseExecutor( ArgumentNullException.ThrowIfNull(scopeFactory); this._registrationKey = registrationKey; + this._sessionStorageIdentity = registrationKey; this._scopeFactory = scopeFactory; this._mapOptions = mapOptions ?? new OpenAIResponsesMapOptions(); } @@ -123,7 +127,6 @@ private AIAgent ResolveExecutionAgent(IServiceProvider services) return agent; } - string? sessionStorageIdentity = this._agent is null ? this._registrationKey : null; - return new AIHostAgent(agent, sessionStore, sessionStorageIdentity); + return new AIHostAgent(agent, sessionStore, this._sessionStorageIdentity); } } diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/IResponsesService.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/IResponsesService.cs index b1676ac99cc..5c5b34192ba 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/IResponsesService.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/IResponsesService.cs @@ -49,6 +49,42 @@ IAsyncEnumerable CreateResponseStreamingAsync( CreateResponse request, CancellationToken cancellationToken = default); + /// + /// Validates a create response request that will run with the specified executor. + /// + /// The create response request to validate. + /// The executor selected for this request. + /// Cancellation token. + /// A ResponseError if validation fails, null if validation succeeds. + ValueTask ValidateRequestAsync( + CreateResponse request, + IResponseExecutor executor, + CancellationToken cancellationToken = default); + + /// + /// Creates a model response for the given input using the specified executor. + /// + /// The create response request. + /// The executor selected for this request. + /// Cancellation token. + /// The created response. + Task CreateResponseAsync( + CreateResponse request, + IResponseExecutor executor, + CancellationToken cancellationToken = default); + + /// + /// Creates a streaming model response for the given input using the specified executor. + /// + /// The create response request. + /// The executor selected for this request. + /// Cancellation token. + /// An async enumerable of streaming response events. + IAsyncEnumerable CreateResponseStreamingAsync( + CreateResponse request, + IResponseExecutor executor, + CancellationToken cancellationToken = default); + /// /// Retrieves a response by ID. /// diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/InMemoryResponsesService.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/InMemoryResponsesService.cs index c696c406817..4b7981c0996 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/InMemoryResponsesService.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/InMemoryResponsesService.cs @@ -2,6 +2,7 @@ using System; using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; using System.Linq; using System.Runtime.CompilerServices; using System.Threading; @@ -19,11 +20,12 @@ namespace Microsoft.Agents.AI.Hosting.OpenAI.Responses; /// internal sealed class InMemoryResponsesService : IResponsesService, IDisposable { - private readonly IResponseExecutor _executor; private readonly MemoryCache _cache; + private readonly bool _ownsCache; private readonly InMemoryStorageOptions _options; private readonly Conversations.IConversationStorage? _conversationStorage; private readonly IsolationKeyResolver? _isolationKeyResolver; + private readonly string? _agentName; private sealed class ResponseState { @@ -35,6 +37,7 @@ private sealed class ResponseState public CreateResponse? Request { get; set; } public string? ConversationStorageId { get; set; } public string? IsolationKey { get; set; } + public IResponseExecutor? Executor { get; set; } public List StreamingUpdates { get; } = []; public Task? CompletionTask { get; set; } public CancellationTokenSource? CancellationTokenSource { get; set; } @@ -145,25 +148,58 @@ public InMemoryResponsesService(IResponseExecutor executor, InMemoryStorageOptio { } + // A null executor requires every create operation to supply the executor selected for its request. public InMemoryResponsesService( - IResponseExecutor executor, + IResponseExecutor? executor, InMemoryStorageOptions options, Conversations.IConversationStorage? conversationStorage, IsolationKeyResolver? isolationKeyResolver) { - ArgumentNullException.ThrowIfNull(executor); ArgumentNullException.ThrowIfNull(options); - this._executor = executor; + this.DefaultExecutor = executor; this._options = options; this._cache = new MemoryCache(options.ToMemoryCacheOptions()); + this._ownsCache = true; this._conversationStorage = conversationStorage; this._isolationKeyResolver = isolationKeyResolver; } + private InMemoryResponsesService(InMemoryResponsesService owner, string agentName) + { + this._options = owner._options; + this._cache = owner._cache; + this._conversationStorage = owner._conversationStorage; + this._isolationKeyResolver = owner._isolationKeyResolver; + this._agentName = agentName; + } + + /// + /// Creates a view that shares this service's storage but only exposes the responses created through it. + /// + /// The stable name of the agent that owns the responses. + /// A responses service bound to . + public InMemoryResponsesService ForAgent(string agentName) + { + ArgumentException.ThrowIfNullOrEmpty(agentName); + return new InMemoryResponsesService(this, agentName); + } + + [AllowNull] + private IResponseExecutor DefaultExecutor => + field ?? throw new InvalidOperationException("This responses service has no default executor. Supply the executor selected for the request."); + + public ValueTask ValidateRequestAsync( + CreateResponse request, + CancellationToken cancellationToken = default) + => this.ValidateRequestAsync(request, this.DefaultExecutor, cancellationToken); + public async ValueTask ValidateRequestAsync( CreateResponse request, + IResponseExecutor executor, CancellationToken cancellationToken = default) { + ArgumentNullException.ThrowIfNull(executor); + if (request.Conversation is not null && !string.IsNullOrEmpty(request.Conversation.Id) && !string.IsNullOrEmpty(request.PreviousResponseId)) { @@ -192,13 +228,21 @@ public InMemoryResponsesService( } } - return await this._executor.ValidateRequestAsync(request, cancellationToken).ConfigureAwait(false); + return await executor.ValidateRequestAsync(request, cancellationToken).ConfigureAwait(false); } + public Task CreateResponseAsync( + CreateResponse request, + CancellationToken cancellationToken = default) + => this.CreateResponseAsync(request, this.DefaultExecutor, cancellationToken); + public async Task CreateResponseAsync( CreateResponse request, + IResponseExecutor executor, CancellationToken cancellationToken = default) { + ArgumentNullException.ThrowIfNull(executor); + if (request.Stream == true) { throw new InvalidOperationException("Cannot create a streaming response using CreateResponseAsync. Use CreateResponseStreamingAsync instead."); @@ -209,7 +253,7 @@ public async Task CreateResponseAsync( // Resolve ambient caller identity before a background response detaches from the HTTP request. string? isolationKey = await this.GetIsolationKeyAsync(cancellationToken).ConfigureAwait(false); - string responseStorageId = IsolationKeyResolver.ScopeId(responseId, isolationKey); + string responseStorageId = this.ScopeToAgent(IsolationKeyResolver.ScopeId(responseId, isolationKey)); var conversationStorageId = request.Conversation?.Id is { } conversationId ? IsolationKeyResolver.ScopeId(conversationId, isolationKey) @@ -222,7 +266,7 @@ public async Task CreateResponseAsync( }; var state = this.InitializeResponse( - responseId, responseStorageId, conversationStorageId, isolationKey, request); + responseId, responseStorageId, conversationStorageId, isolationKey, request, executor); state.CompletionTask = this.ExecuteResponseAsync(responseId, state, ct); // For background responses, start execution and return immediately @@ -236,10 +280,18 @@ public async Task CreateResponseAsync( return state.Response!; } + public IAsyncEnumerable CreateResponseStreamingAsync( + CreateResponse request, + CancellationToken cancellationToken = default) + => this.CreateResponseStreamingAsync(request, this.DefaultExecutor, cancellationToken); + public async IAsyncEnumerable CreateResponseStreamingAsync( CreateResponse request, + IResponseExecutor executor, [EnumeratorCancellation] CancellationToken cancellationToken = default) { + ArgumentNullException.ThrowIfNull(executor); + if (request.Stream == false) { throw new InvalidOperationException("Cannot create a non-streaming response using CreateResponseStreamingAsync. Use CreateResponseAsync instead."); @@ -250,7 +302,7 @@ public async IAsyncEnumerable CreateResponseStreamingAsy // Streaming execution is also detached from the request cancellation token, so capture identity now. string? isolationKey = await this.GetIsolationKeyAsync(cancellationToken).ConfigureAwait(false); - string responseStorageId = IsolationKeyResolver.ScopeId(responseId, isolationKey); + string responseStorageId = this.ScopeToAgent(IsolationKeyResolver.ScopeId(responseId, isolationKey)); var conversationStorageId = request.Conversation?.Id is { } conversationId ? IsolationKeyResolver.ScopeId(conversationId, isolationKey) @@ -258,7 +310,7 @@ public async IAsyncEnumerable CreateResponseStreamingAsy // Start execution var state = this.InitializeResponse( - responseId, responseStorageId, conversationStorageId, isolationKey, request); + responseId, responseStorageId, conversationStorageId, isolationKey, request, executor); state.CompletionTask = this.ExecuteResponseAsync(responseId, state, CancellationToken.None); // Stream updates as they become available @@ -270,7 +322,7 @@ public async IAsyncEnumerable CreateResponseStreamingAsy public async Task GetResponseAsync(string responseId, CancellationToken cancellationToken = default) { - var responseStorageId = await this.GetStorageIdAsync(responseId, cancellationToken).ConfigureAwait(false); + var responseStorageId = await this.GetResponseStorageIdAsync(responseId, cancellationToken).ConfigureAwait(false); this._cache.TryGetValue(responseStorageId, out ResponseState? state); return state?.Response; } @@ -280,7 +332,7 @@ public async IAsyncEnumerable GetResponseStreamingAsync( int? startingAfter = null, [EnumeratorCancellation] CancellationToken cancellationToken = default) { - var responseStorageId = await this.GetStorageIdAsync(responseId, cancellationToken).ConfigureAwait(false); + var responseStorageId = await this.GetResponseStorageIdAsync(responseId, cancellationToken).ConfigureAwait(false); if (!this._cache.TryGetValue(responseStorageId, out ResponseState? state) || state is null) { yield break; @@ -295,7 +347,7 @@ public async IAsyncEnumerable GetResponseStreamingAsync( public async Task CancelResponseAsync(string responseId, CancellationToken cancellationToken = default) { - var responseStorageId = await this.GetStorageIdAsync(responseId, cancellationToken).ConfigureAwait(false); + var responseStorageId = await this.GetResponseStorageIdAsync(responseId, cancellationToken).ConfigureAwait(false); if (!this._cache.TryGetValue(responseStorageId, out ResponseState? state) || state is null) { throw new InvalidOperationException($"Response '{responseId}' not found."); @@ -324,7 +376,7 @@ public async Task CancelResponseAsync(string responseId, CancellationT public async Task DeleteResponseAsync(string responseId, CancellationToken cancellationToken = default) { - var responseStorageId = await this.GetStorageIdAsync(responseId, cancellationToken).ConfigureAwait(false); + var responseStorageId = await this.GetResponseStorageIdAsync(responseId, cancellationToken).ConfigureAwait(false); if (!this._cache.TryGetValue(responseStorageId, out ResponseState? state)) { return false; @@ -349,7 +401,7 @@ public async Task> ListResponseInputItemsAsync( int effectiveLimit = Math.Clamp(limit ?? IResponsesService.DefaultListLimit, 1, 100); SortOrder effectiveOrder = order ?? SortOrder.Descending; - var responseStorageId = await this.GetStorageIdAsync(responseId, cancellationToken).ConfigureAwait(false); + var responseStorageId = await this.GetResponseStorageIdAsync(responseId, cancellationToken).ConfigureAwait(false); if (!this._cache.TryGetValue(responseStorageId, out ResponseState? state)) { throw new InvalidOperationException($"Response '{responseId}' not found."); @@ -409,6 +461,12 @@ public async Task> ListResponseInputItemsAsync( }; } + private async ValueTask GetResponseStorageIdAsync(string responseId, CancellationToken cancellationToken) + => this.ScopeToAgent(await this.GetStorageIdAsync(responseId, cancellationToken).ConfigureAwait(false)); + + // Reuses the isolation key escaping so agent names cannot collide with isolation-scoped identifiers. + private string ScopeToAgent(string storageId) => IsolationKeyResolver.ScopeId(storageId, this._agentName); + private ValueTask GetStorageIdAsync(string id, CancellationToken cancellationToken) { if (this._isolationKeyResolver is null) @@ -429,7 +487,8 @@ private ResponseState InitializeResponse( string responseStorageId, string? conversationStorageId, string? isolationKey, - CreateResponse request) + CreateResponse request, + IResponseExecutor executor) { var metadata = request.Metadata ?? []; @@ -480,6 +539,7 @@ private ResponseState InitializeResponse( Request = request, ConversationStorageId = conversationStorageId, IsolationKey = isolationKey, + Executor = executor, CancellationTokenSource = new CancellationTokenSource() }; @@ -531,7 +591,7 @@ private async Task ExecuteResponseAsync(string responseId, ResponseState state, List outputItems = []; // Execute using the injected executor - await foreach (var streamingEvent in this._executor.ExecuteAsync(context, request, conversationHistory, linkedCts.Token).ConfigureAwait(false)) + await foreach (var streamingEvent in state.Executor!.ExecuteAsync(context, request, conversationHistory, linkedCts.Token).ConfigureAwait(false)) { state.AddStreamingEvent(streamingEvent); @@ -616,6 +676,9 @@ private async Task ExecuteResponseAsync(string responseId, ResponseState state, } finally { + // Stored responses outlive execution, so release the executor and the agent graph it references. + state.Executor = null; + // Signal one final time to unblock any waiting consumers state.SignalUpdate(); } @@ -640,6 +703,9 @@ private static List GetInputItems(string responseId, ResponseState public void Dispose() { - this._cache.Dispose(); + if (this._ownsCache) + { + this._cache.Dispose(); + } } } diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/ResponsesHttpHandler.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/ResponsesHttpHandler.cs index 2c5f705a563..266835e5fc3 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/ResponsesHttpHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/ResponsesHttpHandler.cs @@ -16,26 +16,44 @@ namespace Microsoft.Agents.AI.Hosting.OpenAI.Responses; internal sealed class ResponsesHttpHandler { private readonly IResponsesService _responsesService; + private readonly Func? _agentSelector; /// /// Initializes a new instance of the class. /// /// The responses service. - public ResponsesHttpHandler(IResponsesService responsesService) + /// + /// An optional delegate that selects, for every operation, the responses service bound to the selected agent and + /// the executor that runs it, returning when no agent matches the request. When omitted, + /// and its executor are used. + /// + public ResponsesHttpHandler( + IResponsesService responsesService, + Func? agentSelector = null) { this._responsesService = responsesService ?? throw new ArgumentNullException(nameof(responsesService)); + this._agentSelector = agentSelector; } /// /// Creates a model response for the given input. /// public async Task CreateResponseAsync( + HttpContext httpContext, [FromBody] CreateResponse request, [FromQuery] bool? stream, CancellationToken cancellationToken) { + // Select the agent before validation so an unmatched request never loads session or conversation state. + if (!this.TrySelectAgent(httpContext, out IResponsesService responsesService, out IResponseExecutor? executor)) + { + return AgentNotFound(); + } + // Validate the request first - ResponseError? validationError = await this._responsesService.ValidateRequestAsync(request, cancellationToken).ConfigureAwait(false); + ResponseError? validationError = executor is null + ? await responsesService.ValidateRequestAsync(request, cancellationToken).ConfigureAwait(false) + : await responsesService.ValidateRequestAsync(request, executor, cancellationToken).ConfigureAwait(false); if (validationError is not null) { var (statusCode, wireCode) = ResponseErrorCodes.MapValidationError(validationError.Code); @@ -62,9 +80,9 @@ public async Task CreateResponseAsync( if (shouldStream) { - var streamingResponse = this._responsesService.CreateResponseStreamingAsync( - request, - cancellationToken: cancellationToken); + var streamingResponse = executor is null + ? responsesService.CreateResponseStreamingAsync(request, cancellationToken: cancellationToken) + : responsesService.CreateResponseStreamingAsync(request, executor, cancellationToken); return new SseJsonResult( streamingResponse, @@ -72,9 +90,9 @@ public async Task CreateResponseAsync( OpenAIHostingJsonContext.Default.StreamingResponseEvent); } - var response = await this._responsesService.CreateResponseAsync( - request, - cancellationToken: cancellationToken).ConfigureAwait(false); + var response = executor is null + ? await responsesService.CreateResponseAsync(request, cancellationToken: cancellationToken).ConfigureAwait(false) + : await responsesService.CreateResponseAsync(request, executor, cancellationToken).ConfigureAwait(false); return response.Status switch { @@ -101,16 +119,22 @@ public async Task CreateResponseAsync( /// Retrieves a response by ID. /// public async Task GetResponseAsync( + HttpContext httpContext, string responseId, [FromQuery] string[]? include, [FromQuery] bool? stream, [FromQuery] int? starting_after, CancellationToken cancellationToken) { + if (!this.TrySelectAgent(httpContext, out IResponsesService responsesService, out _)) + { + return AgentNotFound(); + } + // If streaming is requested, return SSE stream if (stream == true) { - var streamingResponse = this._responsesService.GetResponseStreamingAsync( + var streamingResponse = responsesService.GetResponseStreamingAsync( responseId, startingAfter: starting_after, cancellationToken: cancellationToken); @@ -122,7 +146,7 @@ public async Task GetResponseAsync( } // Non-streaming: return the response object - var response = await this._responsesService.GetResponseAsync(responseId, cancellationToken).ConfigureAwait(false); + var response = await responsesService.GetResponseAsync(responseId, cancellationToken).ConfigureAwait(false); return response is not null ? Results.Ok(response) : Results.NotFound(new ErrorResponse @@ -139,12 +163,18 @@ public async Task GetResponseAsync( /// Cancels an in-progress response. /// public async Task CancelResponseAsync( + HttpContext httpContext, string responseId, CancellationToken cancellationToken) { + if (!this.TrySelectAgent(httpContext, out IResponsesService responsesService, out _)) + { + return AgentNotFound(); + } + try { - var response = await this._responsesService.CancelResponseAsync(responseId, cancellationToken).ConfigureAwait(false); + var response = await responsesService.CancelResponseAsync(responseId, cancellationToken).ConfigureAwait(false); return Results.Ok(response); } catch (InvalidOperationException ex) @@ -164,10 +194,16 @@ public async Task CancelResponseAsync( /// Deletes a response. /// public async Task DeleteResponseAsync( + HttpContext httpContext, string responseId, CancellationToken cancellationToken) { - var deleted = await this._responsesService.DeleteResponseAsync(responseId, cancellationToken).ConfigureAwait(false); + if (!this.TrySelectAgent(httpContext, out IResponsesService responsesService, out _)) + { + return AgentNotFound(); + } + + var deleted = await responsesService.DeleteResponseAsync(responseId, cancellationToken).ConfigureAwait(false); return deleted ? Results.Ok(new DeleteResponse { Id = responseId, Object = "response", Deleted = true }) : Results.NotFound(new ErrorResponse @@ -184,6 +220,7 @@ public async Task DeleteResponseAsync( /// Lists the input items for a response. /// public async Task ListResponseInputItemsAsync( + HttpContext httpContext, string responseId, [FromQuery] int? limit, [FromQuery] string? order, @@ -191,6 +228,11 @@ public async Task ListResponseInputItemsAsync( [FromQuery] string? before, CancellationToken cancellationToken) { + if (!this.TrySelectAgent(httpContext, out IResponsesService responsesService, out _)) + { + return AgentNotFound(); + } + try { // Convert string order to SortOrder enum @@ -202,7 +244,7 @@ string s when s.Equals("desc", StringComparison.OrdinalIgnoreCase) => SortOrder. _ => throw new InvalidOperationException($"Invalid order value: {order}. Must be 'asc' or 'desc'.") }; - var result = await this._responsesService.ListResponseInputItemsAsync( + var result = await responsesService.ListResponseInputItemsAsync( responseId, limit, sortOrder, @@ -224,4 +266,33 @@ string s when s.Equals("desc", StringComparison.OrdinalIgnoreCase) => SortOrder. }); } } + + private bool TrySelectAgent(HttpContext httpContext, out IResponsesService responsesService, out IResponseExecutor? executor) + { + responsesService = this._responsesService; + executor = null; + if (this._agentSelector is null) + { + return true; + } + + if (this._agentSelector(httpContext) is not { } selection) + { + return false; + } + + (responsesService, executor) = selection; + return true; + } + + private static IResult AgentNotFound() => + Results.NotFound(new ErrorResponse + { + Error = new ErrorDetails + { + Message = "No agent matches the request.", + Type = "invalid_request_error", + Code = "agent_not_found" + } + }); } diff --git a/dotnet/src/Shared/HostedAgents/HostedAgentResolution.cs b/dotnet/src/Shared/HostedAgents/HostedAgentResolution.cs new file mode 100644 index 00000000000..42d9638d61a --- /dev/null +++ b/dotnet/src/Shared/HostedAgents/HostedAgentResolution.cs @@ -0,0 +1,52 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using Microsoft.Extensions.DependencyInjection; + +namespace Microsoft.Agents.AI.Hosting; + +/// +/// Provides validation shared by protocol hosting integrations that resolve agents for each request. +/// +internal static class HostedAgentResolution +{ + /// + /// Verifies that a keyed registration exists without resolving it. + /// + /// The application service provider. + /// The name of the keyed agent registration. + /// No agent is registered with . + /// + /// Resolving the agent from the application service provider would capture scoped or transient registrations, + /// so the check is skipped when the provider cannot report keyed registrations. + /// + public static void EnsureAgentRegistered(IServiceProvider services, string agentName) + { + if (services.GetService() is { } isKeyedService && + !isKeyedService.IsKeyedService(typeof(AIAgent), agentName)) + { + throw new InvalidOperationException( + $"No {nameof(AIAgent)} is registered with the name '{agentName}'. " + + $"Register it with AddAIAgent(\"{agentName}\", ...) before mapping the endpoint."); + } + } + + /// + /// Gets the stable logical identity of an agent returned by a request-time agent selector. + /// + /// The selected agent. + /// The agent name, which identifies the agent's hosted services and persisted sessions. + /// The agent does not have a name. + public static string GetSelectedAgentName(AIAgent agent) + { + string? agentName = agent.Name; + if (string.IsNullOrWhiteSpace(agentName)) + { + throw new InvalidOperationException( + $"The agent selector returned an {nameof(AIAgent)} without a {nameof(AIAgent.Name)}. " + + "Selected agents must have a stable name, which identifies their hosted services and persisted sessions across requests."); + } + + return agentName!; + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/InMemoryResponsesServiceTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/InMemoryResponsesServiceTests.cs index 1827d852330..e25d1a855f5 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/InMemoryResponsesServiceTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/InMemoryResponsesServiceTests.cs @@ -138,6 +138,38 @@ public async Task ListResponseInputItemsAsync_ReturnsTheItemsTheCursorsSelectAsy }; } + [Fact] + public async Task CreateResponseAsync_ReleasesExecutorAfterExecutionCompletesAsync() + { + // Arrange + using var service = new InMemoryResponsesService( + executor: null, new InMemoryStorageOptions(), conversationStorage: null, isolationKeyResolver: null); + + // Act + (WeakReference executor, string responseId) = await CreateResponseWithCollectableExecutorAsync(service); + + // Leave the completion stack so the finished execution state machine has released its locals. + await Task.Yield(); + GC.Collect(); + GC.WaitForPendingFinalizers(); + GC.Collect(); + + // Assert + Assert.False(executor.IsAlive); + Assert.NotNull(await service.GetResponseAsync(responseId)); + } + + [MethodImpl(MethodImplOptions.NoInlining)] + private static async Task<(WeakReference Executor, string ResponseId)> CreateResponseWithCollectableExecutorAsync( + InMemoryResponsesService service) + { + var executor = new StubResponseExecutor(); + Response response = await service.CreateResponseAsync( + new CreateResponse { Input = ResponseInput.FromText("hello") }, + executor); + return (new WeakReference(executor), response.Id); + } + private sealed class StubResponseExecutor : IResponseExecutor { public ValueTask ValidateRequestAsync(CreateResponse request, CancellationToken cancellationToken = default) diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/ResponseSessionRegressionTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/ResponseSessionRegressionTests.cs index d68adfe50af..ad09d8dfa5c 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/ResponseSessionRegressionTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/ResponseSessionRegressionTests.cs @@ -11,6 +11,7 @@ using System.Threading; using System.Threading.Tasks; using Microsoft.AspNetCore.Builder; +using Microsoft.AspNetCore.Routing; using Microsoft.AspNetCore.TestHost; using Microsoft.Extensions.AI; using Microsoft.Extensions.DependencyInjection; @@ -76,6 +77,188 @@ public async Task ApprovalContinuation_TransientAgent_UsesStableRegistrationIden Assert.Equal(2, modelCalls); } + [Fact] + public async Task ApprovalContinuation_AgentSelectorReturningNewInstances_UsesStableAgentNameAsync() + { + // Arrange + int modelCalls = 0; + int toolCalls = 0; + int selectedAgents = 0; + AIFunction function = new ApprovalRequiredAIFunction(AIFunctionFactory.Create( + () => + { + Interlocked.Increment(ref toolCalls); + return "Sunny"; + }, + ToolName)); + Mock chatClient = CreateApprovalChatClient(() => Interlocked.Increment(ref modelCalls)); + + WebApplicationBuilder builder = CreateBuilder(validateScopes: true); + builder.Services.AddKeyedSingleton(AgentName, new InMemoryAgentSessionStore()); + builder.AddOpenAIResponses(); + + await using WebApplication app = builder.Build(); + app.MapOpenAIResponses("/agents/{agentId}/v1/responses", context => + { + Interlocked.Increment(ref selectedAgents); + return context.GetRouteValue("agentId") as string == AgentName + ? new ChatClientAgent(chatClient.Object, name: AgentName, tools: [function]) + : null; + }); + await app.StartAsync(); + using HttpClient client = app.GetTestClient(); + string path = $"/agents/{AgentName}/v1/responses"; + + (string responseId, JsonElement approvalEvent) = await CreatePendingApprovalAsync( + client, + path, + includeAgentName: false); + using StringContent approvalContent = JsonContent(CreateApprovalResponseJson( + responseId, + approvalEvent, + includeAgentName: false)); + + // Act + using HttpResponseMessage response = await client.PostAsync( + new Uri(path, UriKind.Relative), + approvalContent); + string responseBody = await response.Content.ReadAsStringAsync(); + + // Assert + Assert.True(response.IsSuccessStatusCode, responseBody); + Assert.Equal(2, selectedAgents); + Assert.Equal(1, toolCalls); + Assert.Equal(2, modelCalls); + } + + [Fact] + public async Task AgentSelector_ReturningNull_ReturnsNotFoundAsync() + { + // Arrange + WebApplicationBuilder builder = CreateBuilder(); + builder.AddOpenAIResponses(); + + await using WebApplication app = builder.Build(); + app.MapOpenAIResponses("/agents/{agentId}/v1/responses", _ => null); + await app.StartAsync(); + using HttpClient client = app.GetTestClient(); + + // Act + using HttpResponseMessage response = await PostJsonAsync( + client, + "/agents/unknown/v1/responses", + """{"input":"hello"}"""); + string responseBody = await response.Content.ReadAsStringAsync(); + + // Assert + Assert.Equal(System.Net.HttpStatusCode.NotFound, response.StatusCode); + Assert.Contains("agent_not_found", responseBody, StringComparison.Ordinal); + } + + [Theory] + [InlineData("GET", "/agents/unknown/v1/responses/resp_1")] + [InlineData("POST", "/agents/unknown/v1/responses/resp_1/cancel")] + [InlineData("DELETE", "/agents/unknown/v1/responses/resp_1")] + [InlineData("GET", "/agents/unknown/v1/responses/resp_1/input_items")] + public async Task AgentSelector_ReturningNull_ReturnsNotFoundForStoredResponseOperationsAsync(string method, string path) + { + // Arrange + int selections = 0; + WebApplicationBuilder builder = CreateBuilder(); + builder.AddOpenAIResponses(); + + await using WebApplication app = builder.Build(); + app.MapOpenAIResponses("/agents/{agentId}/v1/responses", _ => + { + Interlocked.Increment(ref selections); + return null; + }); + await app.StartAsync(); + using HttpClient client = app.GetTestClient(); + + // Act + using HttpRequestMessage request = new(new HttpMethod(method), new Uri(path, UriKind.Relative)); + using HttpResponseMessage response = await client.SendAsync(request); + string responseBody = await response.Content.ReadAsStringAsync(); + + // Assert + Assert.Equal(System.Net.HttpStatusCode.NotFound, response.StatusCode); + Assert.Contains("agent_not_found", responseBody, StringComparison.Ordinal); + Assert.Equal(1, selections); + } + + [Fact] + public async Task AgentSelector_BindsStoredResponsesToSelectedAgentAsync() + { + // Arrange + Dictionary agents = new(StringComparer.Ordinal) + { + ["alpha"] = new ChatClientAgent(new TestHelpers.SimpleMockChatClient("Alpha response"), name: "alpha"), + ["beta"] = new ChatClientAgent(new TestHelpers.SimpleMockChatClient("Beta response"), name: "beta"), + }; + WebApplicationBuilder builder = CreateBuilder(); + builder.AddOpenAIResponses(); + + await using WebApplication app = builder.Build(); + app.MapOpenAIResponses( + "/agents/{agentId}/v1/responses", + context => context.GetRouteValue("agentId") is string agentId ? agents.GetValueOrDefault(agentId) : null); + await app.StartAsync(); + using HttpClient client = app.GetTestClient(); + + using HttpResponseMessage createResponse = await PostJsonAsync( + client, + "/agents/alpha/v1/responses", + """{"input":"hello"}"""); + string createBody = await createResponse.Content.ReadAsStringAsync(); + Assert.True(createResponse.IsSuccessStatusCode, createBody); + string responseId = JsonDocument.Parse(createBody).RootElement.GetProperty("id").GetString()!; + + // Act + using HttpResponseMessage otherAgentGet = await client.GetAsync(new Uri($"/agents/beta/v1/responses/{responseId}", UriKind.Relative)); + using HttpResponseMessage otherAgentInputItems = await client.GetAsync(new Uri($"/agents/beta/v1/responses/{responseId}/input_items", UriKind.Relative)); + using HttpResponseMessage otherAgentDelete = await client.DeleteAsync(new Uri($"/agents/beta/v1/responses/{responseId}", UriKind.Relative)); + using HttpResponseMessage ownerGet = await client.GetAsync(new Uri($"/agents/alpha/v1/responses/{responseId}", UriKind.Relative)); + string ownerBody = await ownerGet.Content.ReadAsStringAsync(); + + // Assert + Assert.Equal(System.Net.HttpStatusCode.NotFound, otherAgentGet.StatusCode); + Assert.Equal(System.Net.HttpStatusCode.NotFound, otherAgentInputItems.StatusCode); + Assert.Equal(System.Net.HttpStatusCode.NotFound, otherAgentDelete.StatusCode); + Assert.True(ownerGet.IsSuccessStatusCode, ownerBody); + Assert.Contains("Alpha response", ownerBody, StringComparison.Ordinal); + } + + [Fact] + public async Task AgentSelector_ResolvesScopedAgentFromRequestServicesAsync() + { + // Arrange + WebApplicationBuilder builder = CreateBuilder(validateScopes: true); + builder.AddAIAgent( + AgentName, + (_, name) => new ChatClientAgent(new TestHelpers.SimpleMockChatClient("Scoped response"), name: name), + ServiceLifetime.Scoped); + builder.AddOpenAIResponses(); + + await using WebApplication app = builder.Build(); + app.MapOpenAIResponses( + "/agents/{agentId}/v1/responses", + context => context.RequestServices.GetKeyedService(context.GetRouteValue("agentId"))); + await app.StartAsync(); + using HttpClient client = app.GetTestClient(); + + // Act + using HttpResponseMessage response = await PostJsonAsync( + client, + $"/agents/{AgentName}/v1/responses", + """{"input":"hello"}"""); + string responseBody = await response.Content.ReadAsStringAsync(); + + // Assert + Assert.True(response.IsSuccessStatusCode, responseBody); + Assert.Contains("Scoped response", responseBody, StringComparison.Ordinal); + } + [Theory] [InlineData(false)] [InlineData(true)] From e091559db0e36e6a030d20cb02eb6b3619bc306b Mon Sep 17 00:00:00 2001 From: Marco Minerva Date: Mon, 19 Oct 2026 13:31:02 +0000 Subject: [PATCH 2/3] .NET: Resolve the OpenAI Responses selector session store from the request Agents selected for OpenAI Responses can be resolved from HttpContext.RequestServices, but their keyed session store was resolved from separate per-operation scopes, so a store depending on request-initialized state received a different, uninitialized dependency. Validation and responses that are not run in the background now resolve the store from the request services of the selected agent, while background responses keep using a scope they own because they outlive the request. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...ndpointRouteBuilderExtensions.Responses.cs | 14 +- .../Responses/AIAgentResponseExecutor.cs | 84 +++++-- .../ResponseSessionRegressionTests.cs | 213 ++++++++++++++++++ 3 files changed, 293 insertions(+), 18 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/EndpointRouteBuilderExtensions.Responses.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/EndpointRouteBuilderExtensions.Responses.cs index 63c66399bc6..a95b3f6a99f 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/EndpointRouteBuilderExtensions.Responses.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/EndpointRouteBuilderExtensions.Responses.cs @@ -143,9 +143,14 @@ public static IEndpointConventionBuilder MapOpenAIResponses( /// keyed , to partition persisted sessions, and to bind stored responses to the /// agent that created them, so a response can only be retrieved, canceled, or deleted through requests that select /// the same agent name. Conversations continue across requests even when the selector returns a new agent instance - /// each time. The selected agent must therefore have a non-empty name. Background responses keep running the - /// selected agent after the HTTP request completes, so agents used for background responses must not depend on - /// request-scoped services. + /// each time. The selected agent must therefore have a non-empty name. + /// + /// + /// The keyed is resolved from for + /// validation and for responses that are not run in the background, so it shares the request-scoped dependencies of + /// an agent resolved from the same services. Background responses keep running the selected agent after the HTTP + /// request completes and resolve the session store from a separate service scope that they own, so agents used for + /// background responses must not depend on request-scoped services. /// /// /// Selection is not an authorization boundary. See @@ -176,7 +181,8 @@ public static IEndpointConventionBuilder MapOpenAIResponses( } string agentName = HostedAgentResolution.GetSelectedAgentName(agent); - return (agentName, new AIAgentResponseExecutor(agent, agentName, scopeFactory, mapOptions, sessionStorageIdentity: agentName)); + return (agentName, new AIAgentResponseExecutor( + agent, agentName, scopeFactory, mapOptions, sessionStorageIdentity: agentName, requestServices: context.RequestServices)); }); } diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/AIAgentResponseExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/AIAgentResponseExecutor.cs index 92b943ac2d6..4e1c52a7e7f 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/AIAgentResponseExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/AIAgentResponseExecutor.cs @@ -21,14 +21,31 @@ internal sealed class AIAgentResponseExecutor : IResponseExecutor private readonly string _registrationKey; private readonly string? _sessionStorageIdentity; private readonly IServiceScopeFactory _scopeFactory; + private readonly IServiceProvider? _requestServices; private readonly OpenAIResponsesMapOptions _mapOptions; + /// + /// Initializes a new instance of the class for a fixed agent instance. + /// + /// The agent that executes the responses. + /// The name that keys the agent's session store. + /// The factory used to create a scope for operations that are not bound to a request. + /// Optional options controlling how incoming requests are mapped onto the agent run. + /// + /// The stable logical identity that partitions persisted sessions, or to use the agent instance identity. + /// + /// + /// The services of the request that selected , or when the agent is not + /// bound to a request. When specified, validation and responses that are not run in the background resolve the + /// session store from these services, so the store shares the agent's request-scoped dependencies. + /// public AIAgentResponseExecutor( AIAgent agent, string registrationKey, IServiceScopeFactory scopeFactory, OpenAIResponsesMapOptions? mapOptions = null, - string? sessionStorageIdentity = null) + string? sessionStorageIdentity = null, + IServiceProvider? requestServices = null) { ArgumentNullException.ThrowIfNull(agent); ArgumentException.ThrowIfNullOrEmpty(registrationKey); @@ -38,6 +55,7 @@ public AIAgentResponseExecutor( this._registrationKey = registrationKey; this._sessionStorageIdentity = sessionStorageIdentity; this._scopeFactory = scopeFactory; + this._requestServices = requestServices; this._mapOptions = mapOptions ?? new OpenAIResponsesMapOptions(); } @@ -59,19 +77,33 @@ public AIAgentResponseExecutor( CreateResponse request, CancellationToken cancellationToken = default) { + // Validation always runs within the request. + if (this._requestServices is { } requestServices) + { + return await this.ValidateRequestAsync(request, requestServices, cancellationToken).ConfigureAwait(false); + } + AsyncServiceScope scope = this._scopeFactory.CreateAsyncScope(); await using (scope.ConfigureAwait(false)) { - AIAgent executionAgent = this.ResolveExecutionAgent(scope.ServiceProvider); - ResponseError? sessionError = - AgentResponseExecution.ValidateSessionRequirements(request, executionAgent is AIHostAgent); - return sessionError - ?? this.ValidateRunOptions(request) - ?? await AgentResponseExecution.ValidatePendingApprovalResponsesAsync( - executionAgent, request, cancellationToken).ConfigureAwait(false); + return await this.ValidateRequestAsync(request, scope.ServiceProvider, cancellationToken).ConfigureAwait(false); } } + private async ValueTask ValidateRequestAsync( + CreateResponse request, + IServiceProvider services, + CancellationToken cancellationToken) + { + AIAgent executionAgent = this.ResolveExecutionAgent(services); + ResponseError? sessionError = + AgentResponseExecution.ValidateSessionRequirements(request, executionAgent is AIHostAgent); + return sessionError + ?? this.ValidateRunOptions(request) + ?? await AgentResponseExecution.ValidatePendingApprovalResponsesAsync( + executionAgent, request, cancellationToken).ConfigureAwait(false); + } + internal ResponseError? ValidateRunOptions(CreateResponse request) { try @@ -97,23 +129,47 @@ public async IAsyncEnumerable ExecuteAsync( IReadOnlyList? conversationHistory = null, [EnumeratorCancellation] CancellationToken cancellationToken = default) { + // Responses that are not run in the background are owned by the request that selected the agent, so their + // session store shares the agent's request-scoped dependencies. + if (this._requestServices is { } requestServices && request.Background != true) + { + await foreach (StreamingResponseEvent streamingEvent in this.ExecuteAsync( + requestServices, context, request, conversationHistory, cancellationToken).ConfigureAwait(false)) + { + yield return streamingEvent; + } + + yield break; + } + // The scope surrounds the full enumeration so background execution retains scoped // agents and stores through session persistence after the HTTP request has returned. AsyncServiceScope scope = this._scopeFactory.CreateAsyncScope(); await using (scope.ConfigureAwait(false)) { - AIAgent executionAgent = this.ResolveExecutionAgent(scope.ServiceProvider); - - // Agent selection is fixed by the endpoint. The remaining response execution behavior is - // shared with request-routed hosted agents. - await foreach (StreamingResponseEvent streamingEvent in AgentResponseExecution.ExecuteAsync( - executionAgent, this._mapOptions, context, request, conversationHistory, cancellationToken).ConfigureAwait(false)) + await foreach (StreamingResponseEvent streamingEvent in this.ExecuteAsync( + scope.ServiceProvider, context, request, conversationHistory, cancellationToken).ConfigureAwait(false)) { yield return streamingEvent; } } } + private IAsyncEnumerable ExecuteAsync( + IServiceProvider services, + AgentInvocationContext context, + CreateResponse request, + IReadOnlyList? conversationHistory, + CancellationToken cancellationToken) + { + AIAgent executionAgent = this.ResolveExecutionAgent(services); + + // Agent selection is fixed by the endpoint. The remaining response execution behavior is + // shared with request-routed hosted agents. + return AgentResponseExecution.ExecuteAsync( + executionAgent, this._mapOptions, context, request, conversationHistory, cancellationToken); + } + private AIAgent ResolveExecutionAgent(IServiceProvider services) { AIAgent agent = this._agent ?? diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/ResponseSessionRegressionTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/ResponseSessionRegressionTests.cs index ad09d8dfa5c..092d4b6634f 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/ResponseSessionRegressionTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/ResponseSessionRegressionTests.cs @@ -476,6 +476,84 @@ public async Task StreamingApproval_RequiresTerminalEventBeforeContinuationAsync } } + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task AgentSelector_ResolvesSessionStoreFromRequestServicesAsync(bool stream) + { + // Arrange + TenantObservations observations = new(); + WebApplicationBuilder builder = CreateBuilder(validateScopes: true); + builder.Services.AddScoped(); + builder.Services.AddKeyedScoped( + TenantAgent.AgentName, + (services, _) => new TenantSessionStore(services.GetRequiredService(), observations)); + builder.AddOpenAIResponses(); + + await using WebApplication app = builder.Build(); + app.Use((context, next) => + { + context.RequestServices.GetRequiredService().Tenant = context.Request.Headers["X-Tenant"].ToString(); + return next(context); + }); + app.MapOpenAIResponses( + "/v1/responses", + context => new TenantAgent(context.RequestServices.GetRequiredService(), observations)); + await app.StartAsync(); + using HttpClient client = app.GetTestClient(); + client.DefaultRequestHeaders.Add("X-Tenant", "contoso"); + + // Act + using HttpResponseMessage response = await PostJsonAsync( + client, + "/v1/responses", + $$"""{"input":"hello","stream":{{(stream ? "true" : "false")}}}"""); + string responseBody = await response.Content.ReadAsStringAsync(); + + // Assert + Assert.True(response.IsSuccessStatusCode, responseBody); + TenantState agentState = Assert.Single(observations.AgentStates); + Assert.Equal("contoso", agentState.Tenant); + Assert.NotEmpty(observations.StoreStates); + Assert.All(observations.StoreStates, storeState => Assert.Same(agentState, storeState)); + Assert.Equal(["contoso", "contoso"], observations.StoreOperationTenants); + } + + [Fact] + public async Task AgentSelector_BackgroundResponse_ResolvesSessionStoreFromOwnScopeAsync() + { + // Arrange + TenantObservations observations = new(); + WebApplicationBuilder builder = CreateBuilder(validateScopes: true); + builder.Services.AddScoped(); + builder.Services.AddKeyedScoped( + TenantAgent.AgentName, + (services, _) => new TenantSessionStore(services.GetRequiredService(), observations)); + builder.AddOpenAIResponses(); + + await using WebApplication app = builder.Build(); + app.MapOpenAIResponses( + "/v1/responses", + context => new TenantAgent(context.RequestServices.GetRequiredService(), observations)); + await app.StartAsync(); + using HttpClient client = app.GetTestClient(); + + // Act + using HttpResponseMessage response = await PostJsonAsync(client, "/v1/responses", """{"input":"hello","background":true}"""); + string responseBody = await response.Content.ReadAsStringAsync(); + Assert.True(response.IsSuccessStatusCode, responseBody); + using JsonDocument document = JsonDocument.Parse(responseBody); + await WaitForResponseCompletionAsync(client, "/v1/responses", document.RootElement.GetProperty("id").GetString()!); + + // Assert + // The selector also runs for the polling requests, so the first selection is the one that created the response. + TenantState agentState = observations.AgentStates[0]; + TenantState executionStoreState = observations.StoreStates[^1]; + Assert.DoesNotContain(executionStoreState, observations.AgentStates); + Assert.NotSame(agentState, executionStoreState); + Assert.Equal(2, observations.StoreOperationTenants.Count); + } + private static WebApplicationBuilder CreateBuilder(bool validateScopes = false) { var options = new WebApplicationOptions @@ -781,4 +859,139 @@ public void Dispose() this.Disposed.TrySetResult(); } } + + private sealed class TenantState + { + public string? Tenant { get; set; } + } + + private sealed class TenantObservations + { + private readonly object _gate = new(); + private readonly List _agentStates = []; + private readonly List _storeStates = []; + private readonly List _storeOperationTenants = []; + + public IReadOnlyList AgentStates { get { lock (this._gate) { return [.. this._agentStates]; } } } + + public IReadOnlyList StoreStates { get { lock (this._gate) { return [.. this._storeStates]; } } } + + public IReadOnlyList StoreOperationTenants { get { lock (this._gate) { return [.. this._storeOperationTenants]; } } } + + public void RecordAgent(TenantState state) + { + lock (this._gate) + { + this._agentStates.Add(state); + } + } + + public void RecordStore(TenantState state) + { + lock (this._gate) + { + this._storeStates.Add(state); + } + } + + public void RecordStoreOperation(TenantState state) + { + lock (this._gate) + { + this._storeOperationTenants.Add(state.Tenant); + } + } + } + + private sealed class TenantSessionStore : AgentSessionStore + { + private readonly InMemoryAgentSessionStore _inner = new(); + private readonly TenantState _tenant; + private readonly TenantObservations _observations; + + public TenantSessionStore(TenantState tenant, TenantObservations observations) + { + this._tenant = tenant; + this._observations = observations; + observations.RecordStore(tenant); + } + + public override ValueTask GetSessionAsync( + AIAgent agent, + AgentSessionStoreKey key, + CancellationToken cancellationToken = default) + { + this._observations.RecordStoreOperation(this._tenant); + return this._inner.GetSessionAsync(agent, key, cancellationToken); + } + + public override ValueTask SaveSessionAsync( + AIAgent agent, + AgentSessionStoreKey key, + AgentSession session, + CancellationToken cancellationToken = default) + { + this._observations.RecordStoreOperation(this._tenant); + return this._inner.SaveSessionAsync(agent, key, session, cancellationToken); + } + } + + private sealed class TenantAgent : AIAgent + { + public const string AgentName = "tenant-agent"; + + private readonly TenantState _tenant; + + public TenantAgent(TenantState tenant, TenantObservations observations) + { + this._tenant = tenant; + observations.RecordAgent(tenant); + } + + public override string Name => AgentName; + + protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) => + new(new TenantSession()); + + protected override ValueTask SerializeSessionCoreAsync( + AgentSession session, + JsonSerializerOptions? jsonSerializerOptions = null, + CancellationToken cancellationToken = default) => + new(session.StateBag.Serialize()); + + protected override ValueTask DeserializeSessionCoreAsync( + JsonElement serializedState, + JsonSerializerOptions? jsonSerializerOptions = null, + CancellationToken cancellationToken = default) => + new(new TenantSession(AgentSessionStateBag.Deserialize(serializedState))); + + protected override Task RunCoreAsync( + IEnumerable messages, + AgentSession? session = null, + AgentRunOptions? options = null, + CancellationToken cancellationToken = default) => + throw new NotSupportedException("This regression test uses streaming execution."); + + protected override async IAsyncEnumerable RunCoreStreamingAsync( + IEnumerable messages, + AgentSession? session = null, + AgentRunOptions? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + await Task.Yield(); + yield return new AgentResponseUpdate(ChatRole.Assistant, $"Hello {this._tenant.Tenant}"); + } + + private sealed class TenantSession : AgentSession + { + public TenantSession() + { + } + + public TenantSession(AgentSessionStateBag stateBag) + : base(stateBag) + { + } + } + } } From 1907d1b8c55e05c857b88ff2ec0108e3fcb1336e Mon Sep 17 00:00:00 2001 From: Marco Minerva Date: Mon, 19 Oct 2026 14:00:41 +0000 Subject: [PATCH 3/3] .NET: Keep OpenAI Responses selector request services alive until the run completes Responses for selected agents that are not run in the background use services owned by the HTTP request, but streaming execution kept running after the client disconnected while the request ended, and a canceled non-streaming request ended while its execution was still unwinding. Executors now report whether an execution is bound to the request, and the responses service waits for such executions before the request ends: streaming runs are not canceled and complete normally, while canceled non-streaming runs are awaited until they stop. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...ndpointRouteBuilderExtensions.Responses.cs | 8 +- .../Responses/AIAgentResponseExecutor.cs | 10 +- .../Responses/IResponseExecutor.cs | 11 + .../Responses/InMemoryResponsesService.cs | 28 ++- .../ResponseSessionRegressionTests.cs | 226 ++++++++++++++++-- 5 files changed, 258 insertions(+), 25 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/EndpointRouteBuilderExtensions.Responses.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/EndpointRouteBuilderExtensions.Responses.cs index a95b3f6a99f..696bba2c1c8 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/EndpointRouteBuilderExtensions.Responses.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/EndpointRouteBuilderExtensions.Responses.cs @@ -148,9 +148,11 @@ public static IEndpointConventionBuilder MapOpenAIResponses( /// /// The keyed is resolved from for /// validation and for responses that are not run in the background, so it shares the request-scoped dependencies of - /// an agent resolved from the same services. Background responses keep running the selected agent after the HTTP - /// request completes and resolve the session store from a separate service scope that they own, so agents used for - /// background responses must not depend on request-scoped services. + /// an agent resolved from the same services. These responses are bound to the request: a streaming response keeps + /// running when the client disconnects, and the request ends only after the run completes, while a non-streaming + /// response is canceled with the request and the request ends after the run stops. Background responses keep + /// running the selected agent after the HTTP request completes and resolve the session store from a separate service + /// scope that they own, so agents used for background responses must not depend on request-scoped services. /// /// /// Selection is not an authorization boundary. See diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/AIAgentResponseExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/AIAgentResponseExecutor.cs index 4e1c52a7e7f..b515506ffd0 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/AIAgentResponseExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/AIAgentResponseExecutor.cs @@ -131,7 +131,7 @@ public async IAsyncEnumerable ExecuteAsync( { // Responses that are not run in the background are owned by the request that selected the agent, so their // session store shares the agent's request-scoped dependencies. - if (this._requestServices is { } requestServices && request.Background != true) + if (this._requestServices is { } requestServices && this.IsExecutionBoundToRequest(request)) { await foreach (StreamingResponseEvent streamingEvent in this.ExecuteAsync( requestServices, context, request, conversationHistory, cancellationToken).ConfigureAwait(false)) @@ -170,6 +170,14 @@ private IAsyncEnumerable ExecuteAsync( executionAgent, this._mapOptions, context, request, conversationHistory, cancellationToken); } + /// + /// + /// Executions for an agent selected by a request use that request's services unless they run in the background, + /// which outlives the request and therefore owns a separate service scope. + /// + public bool IsExecutionBoundToRequest(CreateResponse request) => + this._requestServices is not null && request.Background != true; + private AIAgent ResolveExecutionAgent(IServiceProvider services) { AIAgent agent = this._agent ?? diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/IResponseExecutor.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/IResponseExecutor.cs index 84f47af3ed1..46414d3f570 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/IResponseExecutor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/IResponseExecutor.cs @@ -37,4 +37,15 @@ IAsyncEnumerable ExecuteAsync( CreateResponse request, IReadOnlyList? conversationHistory = null, CancellationToken cancellationToken = default); + + /// + /// Gets a value indicating whether the execution of uses services owned by the HTTP request + /// that creates the response. + /// + /// The create response request. + /// + /// when the request must not end before the execution completes; when + /// the execution owns the services it uses and can outlive the request. + /// + bool IsExecutionBoundToRequest(CreateResponse request) => false; } diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/InMemoryResponsesService.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/InMemoryResponsesService.cs index 4b7981c0996..4d20ef15112 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/InMemoryResponsesService.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Responses/InMemoryResponsesService.cs @@ -276,7 +276,17 @@ public async Task CreateResponseAsync( } // For non-background responses, wait for completion - await state.CompletionTask!.WaitAsync(cancellationToken).ConfigureAwait(false); + try + { + await state.CompletionTask!.WaitAsync(cancellationToken).ConfigureAwait(false); + } + catch (OperationCanceledException) when (executor.IsExecutionBoundToRequest(request)) + { + // The canceled execution still uses services owned by this request, so let it finish unwinding first. + await state.CompletionTask!.ConfigureAwait(false); + throw; + } + return state.Response!; } @@ -314,9 +324,21 @@ public async IAsyncEnumerable CreateResponseStreamingAsy state.CompletionTask = this.ExecuteResponseAsync(responseId, state, CancellationToken.None); // Stream updates as they become available - await foreach (StreamingResponseEvent update in state.StreamUpdatesAsync(cancellationToken: cancellationToken).ConfigureAwait(false)) + try { - yield return update; + await foreach (StreamingResponseEvent update in state.StreamUpdatesAsync(cancellationToken: cancellationToken).ConfigureAwait(false)) + { + yield return update; + } + } + finally + { + // The execution keeps running when the stream ends early, for example because the client disconnects. + // When it uses services owned by this request, the request must not end before the execution completes. + if (executor.IsExecutionBoundToRequest(request)) + { + await state.CompletionTask.ConfigureAwait(false); + } } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/ResponseSessionRegressionTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/ResponseSessionRegressionTests.cs index 092d4b6634f..2deb0220e60 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/ResponseSessionRegressionTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/ResponseSessionRegressionTests.cs @@ -28,6 +28,8 @@ public sealed class ResponseSessionRegressionTests private const string AgentName = "session-regression-agent"; private const string ToolName = "get_weather"; + private static readonly TimeSpan s_timeout = TimeSpan.FromSeconds(10); + [Fact] public async Task ApprovalContinuation_TransientAgent_UsesStableRegistrationIdentityAsync() { @@ -554,6 +556,103 @@ public async Task AgentSelector_BackgroundResponse_ResolvesSessionStoreFromOwnSc Assert.Equal(2, observations.StoreOperationTenants.Count); } + [Fact] + public async Task AgentSelector_StreamingDisconnect_KeepsRequestServicesUntilRunCompletesAsync() + { + // Arrange + RequestBoundRunProbe probe = new(); + await using WebApplication app = await StartRequestBoundAppAsync(probe); + using HttpClient client = app.GetTestClient(); + using HttpRequestMessage request = new(HttpMethod.Post, new Uri("/v1/responses", UriKind.Relative)) + { + Content = JsonContent("""{"input":"hello","stream":true}""") + }; + + // Act +#pragma warning disable CA2025 // The request is awaited below, before the request message is disposed. + Task sending = client.SendAsync(request, HttpCompletionOption.ResponseHeadersRead); +#pragma warning restore CA2025 + await probe.AgentStarted.Task.WaitAsync(s_timeout); + HttpResponseMessage response = await sending; + + // Disposing the response before reading it to the end aborts the request on the server, like a client disconnect. + response.Dispose(); + bool requestCompletedBeforeRun = await CompletesWithinAsync(probe.RequestCompleted.Task, TimeSpan.FromMilliseconds(500)); + probe.ReleaseAgent.SetResult(); + await probe.SessionSaved.Task.WaitAsync(s_timeout); + await probe.RequestCompleted.Task.WaitAsync(s_timeout); + + // Assert + Assert.False(requestCompletedBeforeRun); + Assert.False(probe.DependencyDisposedWhenRunFinished); + Assert.False(probe.StoreDisposedWhenSessionSaved); + } + + [Fact] + public async Task AgentSelector_NonStreamingDisconnect_KeepsRequestServicesUntilCanceledRunCompletesAsync() + { + // Arrange + RequestBoundRunProbe probe = new() { CleanupDelay = TimeSpan.FromMilliseconds(300) }; + await using WebApplication app = await StartRequestBoundAppAsync(probe); + using HttpClient client = app.GetTestClient(); + using CancellationTokenSource disconnect = new(); + + // Act +#pragma warning disable CA2025 // The request is awaited below, before the token source is disposed. + Task sending = PostJsonAsync(client, "/v1/responses", """{"input":"hello"}""", disconnect.Token); +#pragma warning restore CA2025 + await probe.AgentStarted.Task.WaitAsync(s_timeout); + disconnect.Cancel(); + bool requestCanceled = false; + try + { + (await sending).Dispose(); + } + catch (OperationCanceledException) + { + requestCanceled = true; + } + + await probe.RunFinished.Task.WaitAsync(s_timeout); + await probe.RequestCompleted.Task.WaitAsync(s_timeout); + + // Assert + Assert.True(requestCanceled); + Assert.True(probe.RunCanceled); + Assert.False(probe.DependencyDisposedWhenRunFinished); + } + + private static async Task StartRequestBoundAppAsync(RequestBoundRunProbe probe) + { + WebApplicationBuilder builder = CreateBuilder(validateScopes: true); + builder.Services.AddScoped(); + builder.Services.AddKeyedScoped( + BlockingAgent.AgentName, + (_, _) => new DisposalTrackingSessionStore(probe)); + builder.AddOpenAIResponses(); + + WebApplication app = builder.Build(); + app.Use(async (context, next) => + { + try + { + await next(context); + } + finally + { + probe.RequestCompleted.TrySetResult(); + } + }); + app.MapOpenAIResponses( + "/v1/responses", + context => new BlockingAgent(context.RequestServices.GetRequiredService(), probe)); + await app.StartAsync(); + return app; + } + + private static async Task CompletesWithinAsync(Task task, TimeSpan timeout) => + await Task.WhenAny(task, Task.Delay(timeout)) == task; + private static WebApplicationBuilder CreateBuilder(bool validateScopes = false) { var options = new WebApplicationOptions @@ -688,10 +787,11 @@ private static StringContent JsonContent(string json) => private static async Task PostJsonAsync( HttpClient client, string path, - string json) + string json, + CancellationToken cancellationToken = default) { using StringContent content = JsonContent(json); - return await client.PostAsync(new Uri(path, UriKind.Relative), content); + return await client.PostAsync(new Uri(path, UriKind.Relative), content, cancellationToken); } private static async Task WaitForResponseCompletionAsync( @@ -936,7 +1036,7 @@ public override ValueTask SaveSessionAsync( } } - private sealed class TenantAgent : AIAgent + private sealed class TenantAgent : StreamingTestAgent { public const string AgentName = "tenant-agent"; @@ -950,8 +1050,108 @@ public TenantAgent(TenantState tenant, TenantObservations observations) public override string Name => AgentName; + protected override async IAsyncEnumerable RunCoreStreamingAsync( + IEnumerable messages, + AgentSession? session = null, + AgentRunOptions? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + await Task.Yield(); + yield return new AgentResponseUpdate(ChatRole.Assistant, $"Hello {this._tenant.Tenant}"); + } + } + + private sealed class RequestBoundRunProbe + { + public TimeSpan CleanupDelay { get; init; } + + public TaskCompletionSource AgentStarted { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public TaskCompletionSource ReleaseAgent { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public TaskCompletionSource RunFinished { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public TaskCompletionSource SessionSaved { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public TaskCompletionSource RequestCompleted { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public bool RunCanceled { get; set; } + + public bool? DependencyDisposedWhenRunFinished { get; set; } + + public bool? StoreDisposedWhenSessionSaved { get; set; } + } + + private sealed class DisposableRequestDependency : IDisposable + { + public bool IsDisposed { get; private set; } + + public void Dispose() => this.IsDisposed = true; + } + + private sealed class DisposalTrackingSessionStore(RequestBoundRunProbe probe) : AgentSessionStore, IDisposable + { + private readonly InMemoryAgentSessionStore _inner = new(); + private bool _disposed; + + public override ValueTask GetSessionAsync( + AIAgent agent, + AgentSessionStoreKey key, + CancellationToken cancellationToken = default) => + this._inner.GetSessionAsync(agent, key, cancellationToken); + + public override async ValueTask SaveSessionAsync( + AIAgent agent, + AgentSessionStoreKey key, + AgentSession session, + CancellationToken cancellationToken = default) + { + probe.StoreDisposedWhenSessionSaved = this._disposed; + await this._inner.SaveSessionAsync(agent, key, session, cancellationToken); + probe.SessionSaved.TrySetResult(); + } + + public void Dispose() => this._disposed = true; + } + + private sealed class BlockingAgent(DisposableRequestDependency dependency, RequestBoundRunProbe probe) : StreamingTestAgent + { + public const string AgentName = "blocking-agent"; + + public override string Name => AgentName; + + protected override async IAsyncEnumerable RunCoreStreamingAsync( + IEnumerable messages, + AgentSession? session = null, + AgentRunOptions? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + probe.AgentStarted.TrySetResult(); + try + { + await probe.ReleaseAgent.Task.WaitAsync(cancellationToken); + } + finally + { + probe.RunCanceled = cancellationToken.IsCancellationRequested; + + // Keep using the request-scoped dependency while the run unwinds. + await Task.Delay(probe.CleanupDelay, CancellationToken.None); + probe.DependencyDisposedWhenRunFinished = dependency.IsDisposed; + probe.RunFinished.TrySetResult(); + } + + yield return new AgentResponseUpdate(ChatRole.Assistant, "done"); + } + } + + /// + /// An agent that persists its session state bag and only supports streaming runs. + /// + private abstract class StreamingTestAgent : AIAgent + { protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) => - new(new TenantSession()); + new(new TestAgentSession()); protected override ValueTask SerializeSessionCoreAsync( AgentSession session, @@ -963,7 +1163,7 @@ protected override ValueTask DeserializeSessionCoreAsync( JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) => - new(new TenantSession(AgentSessionStateBag.Deserialize(serializedState))); + new(new TestAgentSession(AgentSessionStateBag.Deserialize(serializedState))); protected override Task RunCoreAsync( IEnumerable messages, @@ -972,23 +1172,13 @@ protected override Task RunCoreAsync( CancellationToken cancellationToken = default) => throw new NotSupportedException("This regression test uses streaming execution."); - protected override async IAsyncEnumerable RunCoreStreamingAsync( - IEnumerable messages, - AgentSession? session = null, - AgentRunOptions? options = null, - [EnumeratorCancellation] CancellationToken cancellationToken = default) - { - await Task.Yield(); - yield return new AgentResponseUpdate(ChatRole.Assistant, $"Hello {this._tenant.Tenant}"); - } - - private sealed class TenantSession : AgentSession + private sealed class TestAgentSession : AgentSession { - public TenantSession() + public TestAgentSession() { } - public TenantSession(AgentSessionStateBag stateBag) + public TestAgentSession(AgentSessionStateBag stateBag) : base(stateBag) { }