diff --git a/dotnet/samples/04-hosting/FoundryHostedAgents/responses/Hosted-ChatClientAgent/README.md b/dotnet/samples/04-hosting/FoundryHostedAgents/responses/Hosted-ChatClientAgent/README.md index eada65ad288..e8d6096d7ad 100644 --- a/dotnet/samples/04-hosting/FoundryHostedAgents/responses/Hosted-ChatClientAgent/README.md +++ b/dotnet/samples/04-hosting/FoundryHostedAgents/responses/Hosted-ChatClientAgent/README.md @@ -69,6 +69,38 @@ AZURE_TOKEN_CREDENTIALS=dev > credentials (Azure CLI, Visual Studio, `azd`) and skips the managed-identity probe. This variable > is only for local runs; the deployed agent in Foundry uses the platform-injected managed identity. +## Agent lifetime + +`builder.AddFoundryResponses(agent)` registers one agent object that serves every request for the +life of the process. When the agent needs services that should not live that long, for example a +database context, register a factory with a DI lifetime instead: + +```csharp +builder.AddFoundryResponses( + agentName, + sp => new MyAgent(sp.GetRequiredService()), + ServiceLifetime.Scoped); +``` + +Each response run (streaming, background, or resumed after a restart) gets its own DI scope, which is +disposed when that run's event stream ends. A `Scoped` agent is therefore created once per response and +disposed afterwards, even when a background response keeps running after the HTTP request has returned. +Conversation state is not kept on the agent object; each turn loads the saved session, and sessions are +stored under `agentName`, so switching from the instance overload keeps existing conversations. + +The session store has its own registration and lifetime, independent of the agent. Without one, every +agent uses the default store, a single shared instance. A store passed to `AddFoundryResponses(agent, store)` +is used only by that agent. Each store keeps its own sessions, so moving an agent to a different store +means the sessions it saved in the old store are no longer found: + +```csharp +// For every agent that has no store of its own. A store object you pass in is never disposed by DI. +builder.AddFoundryAgentSessionStore(sp => new MySessionStore(sp.GetRequiredService()), ServiceLifetime.Scoped); + +// Or for one agent only. +builder.AddFoundryAgentSessionStore(mySessionStore, agentName); +``` + ## Run and test locally Local runs use two terminals: one hosts the agent, the other is a code-first client that talks to it diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs index 715d0819091..0c5c1fa0277 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs @@ -42,6 +42,12 @@ internal sealed partial class AgentFrameworkResponseHandler : ResponseHandler private readonly FoundryResponsesServerOptions _options; private readonly FoundryToolboxService? _toolboxService; + /// + /// The factory agent that answers requests without a matching agent name when no agent is registered + /// without a key. See . + /// + private readonly string? _defaultAgentName; + /// /// Initializes a new instance of the class /// that resolves agents from keyed DI services. @@ -82,6 +88,7 @@ public AgentFrameworkResponseHandler( this._logger = logger; this._options = foundryResponsesOptions; this._toolboxService = toolboxService; + this._defaultAgentName = serviceProvider.GetService()?.Name; } /// @@ -90,13 +97,11 @@ public override async IAsyncEnumerable CreateAsync( ResponseContext context, [EnumeratorCancellation] CancellationToken cancellationToken) { - // 1. Resolve agent - var agent = this.ResolveAgent(request); - var sessionStore = this.ResolveSessionStore(request); - - // 2. Resolve the per-request hosted user identity, so the session can be + // 1. Resolve the per-request hosted user identity, so the session can be // loaded from a per-user partition. Fresh sessions are tagged once; resumed sessions are // validated against the live request to detect cross-user session leaks and in-process tampering. + // This runs before any agent or store is created, so a rejected request never invokes a + // caller-supplied agent or session store factory. string? resolvedUserId = context.PlatformContext?.UserIdKey; if (string.IsNullOrWhiteSpace(resolvedUserId)) { @@ -114,6 +119,16 @@ public override async IAsyncEnumerable CreateAsync( "Ensure the x-agent-user-id header is present."); } + // 2. Resolve agent and session store from a DI scope owned by this response run. The scope + // covers the whole enumeration, not the HTTP request: background runs continue after the + // HTTP response is sent, and recovery runs have no HTTP request at all. Scoped and transient + // registrations are therefore created per run and disposed when the event stream ends. + AsyncServiceScope serviceScope = this._serviceProvider.CreateAsyncScope(); + await using var configuredServiceScope = serviceScope.ConfigureAwait(false); + IServiceProvider services = serviceScope.ServiceProvider; + var (agent, sessionStoreKey) = this.ResolveAgent(request, services); + var sessionStore = this.ResolveSessionStore(sessionStoreKey, services); + // When resolvedUserId is null here the container is NOT hosted by Foundry (local // development: docker run / dotnet run outside the platform, so no x-agent-user-id header). // Per-user isolation simply does not apply in that case: no user partition is added to the @@ -878,45 +893,71 @@ private bool ShouldPersistForResilience(CreateResponse request) /// /// Resolves an from the request. /// Tries agent.name first, then falls back to metadata["entity_id"]. - /// If neither is present, attempts to resolve a default (non-keyed) . + /// If neither matches a keyed agent, attempts to resolve a default (non-keyed) , and + /// then the agent registered through the factory overload of AddFoundryResponses as the default. /// - private AIAgent ResolveAgent(CreateResponse request) + /// The incoming request. + /// The service provider of the DI scope owned by the current response run. + /// The prepared agent and the name to look its session store up with. + private (AIAgent Agent, string? SessionStoreKey) ResolveAgent(CreateResponse request, IServiceProvider services) { var agentName = GetAgentName(request); + var defaultAgent = services.GetService(); if (!string.IsNullOrEmpty(agentName)) { - var agent = this._serviceProvider.GetKeyedService(agentName); + var agent = services.GetKeyedService(agentName); if (agent is not null) { - string storageIdentity = FoundryHostingAgent.ResolveSessionStorageIdentity( - agent, - agentName, - this._serviceProvider.GetService()); - return this.PrepareResolvedAgent(agent, storageIdentity); + // The factory default has no non-keyed registration to compare against, so it is + // recognized by name. It keeps the location AddFoundryResponses(agent) uses for the + // default agent, unless a non-keyed agent took the default role. + string storageIdentity = defaultAgent is null + && string.Equals(agentName, this._defaultAgentName, StringComparison.Ordinal) + ? DefaultAgentStorageIdentity(agentName) + : FoundryHostingAgent.ResolveSessionStorageIdentity(agent, agentName, defaultAgent); + return (this.PrepareResolvedAgent(agent, storageIdentity), agentName); } LogKeyedAgentNotFound(this._logger, agentName); } - // Try non-keyed default - var defaultAgent = this._serviceProvider.GetService(); if (defaultAgent is not null) { string storageIdentity = FoundryHostingAgent.ResolveSessionStorageIdentity( defaultAgent, registrationKey: null, defaultAgent); - return this.PrepareResolvedAgent(defaultAgent, storageIdentity); + + // The default agent uses the store registered under its own name, the same store a request + // naming that agent uses. + string? sessionStoreKey = string.IsNullOrWhiteSpace(defaultAgent.Name) ? null : defaultAgent.Name; + return (this.PrepareResolvedAgent(defaultAgent, storageIdentity), sessionStoreKey); + } + + // The factory overload registers its agent only under its name, so that DI owns each created + // object once. Requests without a matching agent name reach it through the recorded default name. + if (this._defaultAgentName is { } defaultAgentName + && services.GetKeyedService(defaultAgentName) is { } factoryDefaultAgent) + { + return ( + this.PrepareResolvedAgent(factoryDefaultAgent, DefaultAgentStorageIdentity(defaultAgentName)), + defaultAgentName); } var errorMessage = string.IsNullOrEmpty(agentName) ? "No agent name specified in the request (via agent.name or metadata[\"entity_id\"]) and no default AIAgent is registered." - : $"Agent '{agentName}' not found. Ensure it is registered via builder.AddFoundryResponses(agent) or services.AddKeyedSingleton(\"{agentName}\", ...)."; + : $"Agent '{agentName}' not found. Ensure it is registered via builder.AddFoundryResponses(agent), builder.AddFoundryResponses(\"{agentName}\", agentFactory, lifetime), or a keyed AIAgent service registration such as services.AddKeyedSingleton(\"{agentName}\", ...)."; throw new InvalidOperationException(errorMessage); } + /// + /// The storage identity AddFoundryResponses(agent) gives the default agent, used for the factory + /// default whatever its lifetime. + /// + private static string DefaultAgentStorageIdentity(string agentName) => $"name:{agentName}"; + private AIAgent PrepareResolvedAgent(AIAgent agent, string sessionStorageIdentity) { ServiceCollectionExtensions.TryApplyUserAgent(agent); @@ -930,17 +971,16 @@ private AIAgent PrepareResolvedAgent(AIAgent agent, string sessionStorageIdentit } /// - /// Resolves an from the request. - /// Tries agent.name first, then falls back to metadata["entity_id"]. - /// If neither is present, attempts to resolve a default (non-keyed) . + /// Resolves the registered under , or the + /// store registered for all agents (non-keyed) when there is none. /// - private AgentSessionStore ResolveSessionStore(CreateResponse request) + /// The key returned by . + /// The service provider of the DI scope owned by the current response run. + private AgentSessionStore ResolveSessionStore(string? agentName, IServiceProvider services) { - var agentName = GetAgentName(request); - if (!string.IsNullOrEmpty(agentName)) { - var sessionStore = this._serviceProvider.GetKeyedService(agentName); + var sessionStore = services.GetKeyedService(agentName); if (sessionStore is not null) { return sessionStore; @@ -949,8 +989,8 @@ private AgentSessionStore ResolveSessionStore(CreateResponse request) LogKeyedSessionStoreNotFound(this._logger, agentName); } - // Try non-keyed default - var defaultSessionStore = this._serviceProvider.GetService(); + // An agent without a store of its own uses the store registered for all agents. + var defaultSessionStore = services.GetService(); if (defaultSessionStore is not null) { return defaultSessionStore; @@ -958,7 +998,7 @@ private AgentSessionStore ResolveSessionStore(CreateResponse request) var errorMessage = string.IsNullOrEmpty(agentName) ? "No agent name specified in the request (via agent.name or metadata[\"entity_id\"]) and no default AgentSessionStore is registered." - : $"AgentSessionStore for agent '{agentName}' not found. Ensure it is registered via builder.AddFoundryResponses(agent, agentSessionStore) or services.AddKeyedSingleton(\"{agentName}\", ...)."; + : $"AgentSessionStore for agent '{agentName}' not found. Ensure it is registered via services.AddFoundryAgentSessionStore(agentSessionStore, \"{agentName}\"), services.AddFoundryAgentSessionStore(agentSessionStore) for all agents, or builder.AddFoundryResponses(agent, agentSessionStore)."; throw new InvalidOperationException(errorMessage); } @@ -1034,6 +1074,6 @@ private static partial void LogConsentLinkRejected( [LoggerMessage(LogLevel.Warning, "Agent '{AgentName}' not found in keyed services. Attempting default resolution.", EventId = 0)] private static partial void LogKeyedAgentNotFound(ILogger logger, string agentName); - [LoggerMessage(LogLevel.Warning, "SessionStore for agent '{AgentName}' not found in keyed services. Attempting default resolution.", EventId = 0)] + [LoggerMessage(LogLevel.Debug, "No session store registered for agent '{AgentName}'. Using the store registered for all agents.", EventId = 0)] private static partial void LogKeyedSessionStoreNotFound(ILogger logger, string agentName); } diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryDefaultAgentName.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryDefaultAgentName.cs new file mode 100644 index 00000000000..50015608272 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryDefaultAgentName.cs @@ -0,0 +1,34 @@ +// Copyright (c) Microsoft. All rights reserved. + +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Agents.AI.Foundry.Hosting; + +/// +/// The name of the agent registered through the factory overload of AddFoundryResponses that answers +/// requests without a matching agent name. +/// +/// +/// +/// The factory overload registers its agent only under the agent name. It does not add a non-keyed +/// registration that forwards to the keyed one, because Microsoft DI disposes the object a factory returns +/// once per registration that created or returned it, so a forwarding registration would dispose the same +/// object twice. This value is registered instead, only when no default agent existed at registration time, +/// so the response handler can route requests without an agent name to that agent. +/// +/// +/// The response handler ignores it when an is registered without a key, because that +/// agent receives requests without an agent name. Sessions of the agent named here are stored in the location +/// AddFoundryResponses(agent) uses for the default agent, name:{Name}, whatever its lifetime. +/// +/// +internal sealed class FoundryDefaultAgentName +{ + public FoundryDefaultAgentName(string name) + { + this.Name = Throw.IfNullOrWhitespace(name); + } + + /// Gets the key the default factory agent is registered under. + public string Name { get; } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostApplicationBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostApplicationBuilderExtensions.cs index 35aed0370db..e6969ce8721 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostApplicationBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostApplicationBuilderExtensions.cs @@ -1,5 +1,6 @@ // Copyright (c) Microsoft. All rights reserved. +using System; using Microsoft.Extensions.Configuration; using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.DependencyInjection.Extensions; @@ -21,6 +22,15 @@ public static class HostApplicationBuilderExtensions /// Use this overload for applications that can access their host builder. In a Foundry hosted /// environment it binds the response-storage and resilient-task-storage identity and endpoint /// from the same ResponsesServer configuration section, as required by AgentServer. + /// Keyed agent registrations may use any DI lifetime. Each response run resolves its agent and + /// session store from a DI scope that is disposed when the run's event stream ends, so scoped and + /// transient agents are created per run, including background runs that outlive the HTTP request. + /// For each run the handler uses the registered under the agent's name, + /// then the store registered for all agents (see + /// ), then the + /// default store. A request answered by the default agent uses that agent's store. Each store keeps its own + /// sessions, so changing which store an agent uses means sessions saved in the previous store are no + /// longer found. /// /// builder.Services.AddKeyedSingleton<AIAgent>("my-agent", myAgent); /// builder.AddFoundryResponses(); @@ -58,13 +68,26 @@ public static IHostApplicationBuilder AddFoundryResponses( /// application builder. /// /// + /// + /// A store passed in agentSessionStore is registered only under the agent's name. Every request this + /// agent answers uses it, including requests without a matching agent name, and other agents do not use it. + /// When the agent has no name, the store has no key to be registered under and is used for all agents + /// instead. When no store is passed, the agent uses the store registered for all agents, with + /// or a + /// non-keyed registration, and otherwise the default store. Each store keeps + /// its own sessions, so changing which store an agent uses means sessions saved in the previous store are no + /// longer found. + /// /// /// builder.AddFoundryResponses(myAgent); /// /// /// The host application builder. /// The agent instance to register. - /// The agent session store to use, or for the default. + /// + /// The session store for this agent, used only by this agent. If null, the agent uses the store registered + /// for all agents, or the default store. + /// /// Optional hosting behavior. Required inputs belong on the extension method parameters. /// The host application builder for chaining. public static IHostApplicationBuilder AddFoundryResponses( @@ -89,19 +112,140 @@ public static IHostApplicationBuilder AddFoundryResponses( configuredOptions, includeServerOptions: serverAdded, applyOptions: serverAdded || options is not null); - agentSessionStore ??= ServiceCollectionExtensions.CreateDefaultAgentSessionStore(); + ServiceCollectionExtensions.RegisterAgentInstance(host.Services, agent, agentSessionStore); + ServiceCollectionExtensions.RegisterResponseHandler(host.Services); + ServiceCollectionExtensions.MarkFeatureUsed(); + return host; + } - if (!string.IsNullOrWhiteSpace(agent.Name)) - { - host.Services.TryAddKeyedSingleton(agent.Name, agent); - host.Services.TryAddKeyedSingleton(agent.Name, agentSessionStore); - } + /// + /// Registers the Azure AI Responses Server SDK and an created by a factory + /// with the specified DI lifetime on a host application builder. + /// + /// + /// + /// Use this overload when the agent depends on services that are not singletons, for example a + /// database context, or when it should not live for the whole process. The agent is registered only + /// under . When no default agent is registered yet, requests without a + /// matching agent name are also routed to this agent. An existing registration under + /// is kept, the same as for AddFoundryResponses(agent). + /// + /// + /// Every response run (streaming, non-streaming, background, or a run resumed after a restart) gets + /// its own DI scope, created when the run starts and disposed when its event stream ends. That is + /// later than the end of the HTTP request for a background run. The lifetime applies to the agent: + /// + /// : one instance for the process, disposed at shutdown. + /// : one instance per response run, disposed when the run ends. + /// : a new instance on every lookup, disposed when the run ends. + /// + /// This scope is not the ASP.NET Core request scope, so scoped services that middleware fills in for an + /// HTTP request are new, empty instances here. Do not derive the caller's identity from such services in + /// the factory; the per-request user identity is supplied by the platform and applied by the response + /// handler when it loads and saves sessions. + /// + /// + /// The session store is registered separately with + /// , + /// with its own lifetime. Without such a registration for , the agent uses the + /// store registered for all agents, or the default store. Sessions are + /// stored in the same location AddFoundryResponses(agent) uses for an agent with that name + /// registered in the same order, whatever the lifetime, so switching between the two overloads or + /// between lifetimes keeps existing sessions readable. + /// + /// + /// builder.AddFoundryResponses( + /// "my-agent", + /// sp => new MyAgent(sp.GetRequiredService<MyDbContext>()), + /// ServiceLifetime.Scoped); + /// + /// + /// The host application builder. + /// The name the agent is registered under and the name requests use to select it. + /// Creates the agent from the service provider of the current DI scope. + /// The DI lifetime of the agent. Defaults to . + /// Optional hosting behavior. Required inputs belong on the extension method parameters. + /// The host application builder for chaining. + public static IHostApplicationBuilder AddFoundryResponses( + this IHostApplicationBuilder host, + string agentName, + Func agentFactory, + ServiceLifetime lifetime = ServiceLifetime.Singleton, + FoundryResponsesServerOptions? options = null) + { + _ = Throw.IfNull(host); + ServiceCollectionExtensions.RegisterAgentFactory(host.Services, agentName, agentFactory, lifetime); + return host.AddFoundryResponses(options); + } - host.Services.TryAddSingleton(agent); - host.Services.TryAddSingleton(agentSessionStore); + /// + /// Registers the that Foundry Responses hosting uses for one agent, or for + /// every agent that has no store of its own, on a host application builder. + /// + /// + /// The response handler uses the store registered under the agent's name, then the store registered for + /// all agents, then the default store, for agents registered with any AddFoundryResponses overload + /// or directly in DI. The DI container never disposes . + /// Calling this method again for the same , or again for all agents, replaces + /// the earlier store, including one passed to AddFoundryResponses(agent, agentSessionStore) or + /// registered directly in DI, in whatever order it is called relative to AddFoundryResponses. Each + /// store keeps its own sessions, so changing which store an agent uses means sessions saved in the previous + /// store are no longer found. + /// + /// builder.AddFoundryAgentSessionStore(mySessionStore); + /// + /// + /// The host application builder. + /// The session store to use. + /// + /// The name of the agent the store is for, or to use it for every agent that has no + /// store of its own. + /// + /// The host application builder for chaining. + public static IHostApplicationBuilder AddFoundryAgentSessionStore( + this IHostApplicationBuilder host, + AgentSessionStore agentSessionStore, + string? agentName = null) + { + _ = Throw.IfNull(host); + host.Services.AddFoundryAgentSessionStore(agentSessionStore, agentName); + return host; + } - ServiceCollectionExtensions.RegisterResponseHandler(host.Services); - ServiceCollectionExtensions.MarkFeatureUsed(); + /// + /// Registers a factory for the that Foundry Responses hosting uses for one + /// agent, or for every agent that has no store of its own, with the specified DI lifetime on a host + /// application builder. + /// + /// + /// The response handler looks the store up once per response run, from the DI scope of that run, for + /// agents registered with any AddFoundryResponses overload or directly in DI. The store's lifetime is + /// independent of the agent's lifetime. Calling this method again for the same , + /// or again for all agents, replaces the earlier store, including one passed to + /// AddFoundryResponses(agent, agentSessionStore) or registered directly in DI, in whatever order it is + /// called relative to AddFoundryResponses. + /// + /// builder.AddFoundryAgentSessionStore( + /// sp => new MySessionStore(sp.GetRequiredService<MyDbContext>()), + /// ServiceLifetime.Scoped); + /// + /// + /// The host application builder. + /// Creates the session store from the service provider of the current DI scope. + /// The DI lifetime of the session store. Defaults to . + /// + /// The name of the agent the store is for, or to use it for every agent that has no + /// store of its own. + /// + /// The host application builder for chaining. + public static IHostApplicationBuilder AddFoundryAgentSessionStore( + this IHostApplicationBuilder host, + Func agentSessionStoreFactory, + ServiceLifetime lifetime = ServiceLifetime.Singleton, + string? agentName = null) + { + _ = Throw.IfNull(host); + host.Services.AddFoundryAgentSessionStore(agentSessionStoreFactory, lifetime, agentName); return host; } diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedStoredOutputHealthCheck.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedStoredOutputHealthCheck.cs index 127770bcb69..a1ed328413b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedStoredOutputHealthCheck.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedStoredOutputHealthCheck.cs @@ -5,6 +5,7 @@ using System.Globalization; using System.Threading; using System.Threading.Tasks; +using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Diagnostics.HealthChecks; using Microsoft.Extensions.Logging; using Microsoft.Shared.Diagnostics; @@ -64,19 +65,25 @@ public async Task CheckHealthAsync(HealthCheckContext context List storingAgents = []; var checkedAgents = 0; - foreach (var agent in this.ResolveAgents()) + // Agents may be registered as scoped or transient, so they are resolved from a scope that is + // disposed once the check is done, the same way the response handler resolves them per run. + AsyncServiceScope serviceScope = this._serviceProvider.CreateAsyncScope(); + await using (serviceScope.ConfigureAwait(false)) { - if (agent.GetService() is not { } chatClientAgent) + foreach (var agent in ServiceCollectionExtensions.ResolveRegisteredAgents(serviceScope.ServiceProvider)) { - // Hosting only reaches the store setting through ChatClientAgent's chat options, so any - // other agent runs untouched and there is nothing to report. - continue; - } - - checkedAgents++; - if (await this.StoresItsOwnResponsesAsync(agent, cancellationToken).ConfigureAwait(false)) - { - storingAgents.Add(agent.Name ?? agent.Id); + if (agent.GetService() is not { } chatClientAgent) + { + // Hosting only reaches the store setting through ChatClientAgent's chat options, so any + // other agent runs untouched and there is nothing to report. + continue; + } + + checkedAgents++; + if (await this.StoresItsOwnResponsesAsync(agent, cancellationToken).ConfigureAwait(false)) + { + storingAgents.Add(agent.Name ?? agent.Id); + } } } @@ -154,11 +161,6 @@ private async Task StoresItsOwnResponsesAsync(AIAgent agent, CancellationT return probe.StoredOutputEnabled is true; } - /// - /// Every agent this container can serve: the ones registered under a name, plus the default. - /// - private List ResolveAgents() => ServiceCollectionExtensions.ResolveRegisteredAgents(this._serviceProvider); - [LoggerMessage(LogLevel.Debug, "Could not probe the stored output setting for agent '{AgentName}'.", EventId = 0)] private static partial void LogStoredOutputProbeFailed( ILogger logger, diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedWorkflowCheckpointingHealthCheck.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedWorkflowCheckpointingHealthCheck.cs index ec823f1b473..103f31ac1dd 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedWorkflowCheckpointingHealthCheck.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/HostedWorkflowCheckpointingHealthCheck.cs @@ -49,7 +49,7 @@ public HostedWorkflowCheckpointingHealthCheck(IServiceProvider serviceProvider) /// internal bool IsHosted { get; set; } = FoundryEnvironment.IsHosted; - public Task CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default) + public async Task CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default) { _ = Throw.IfNull(context); @@ -57,42 +57,48 @@ public Task CheckHealthAsync(HealthCheckContext context, Canc { // Nothing redirects checkpoints outside a Foundry container, so an agent that brings its // own checkpoint manager is not competing with anything. - return Task.FromResult(HealthCheckResult.Healthy( - "Workflow checkpointing: not running in a Foundry container, so workflow checkpoints are left where each agent puts them.")); + return HealthCheckResult.Healthy( + "Workflow checkpointing: not running in a Foundry container, so workflow checkpoints are left where each agent puts them."); } List incompatibleAgents = []; var checkedAgents = 0; - foreach (var agent in ServiceCollectionExtensions.ResolveRegisteredAgents(this._serviceProvider)) + // Agents may be registered as scoped or transient, so they are resolved from a scope that is + // disposed once the check is done, the same way the response handler resolves them per run. + AsyncServiceScope serviceScope = this._serviceProvider.CreateAsyncScope(); + await using (serviceScope.ConfigureAwait(false)) { - if (agent.GetService() is not { } metadata) + foreach (var agent in ServiceCollectionExtensions.ResolveRegisteredAgents(serviceScope.ServiceProvider)) { - continue; - } + if (agent.GetService() is not { } metadata) + { + continue; + } - checkedAgents++; - AIAgent redirected = ServiceCollectionExtensions.ApplyWorkflowCheckpointing( - agent, - this._serviceProvider.GetService()); + checkedAgents++; + AIAgent redirected = ServiceCollectionExtensions.ApplyWorkflowCheckpointing( + agent, + this._serviceProvider.GetService()); - if (metadata.UsesOwnCheckpointStorage || ReferenceEquals(redirected, agent)) - { - incompatibleAgents.Add(agent.Name ?? agent.Id); + if (metadata.UsesOwnCheckpointStorage || ReferenceEquals(redirected, agent)) + { + incompatibleAgents.Add(agent.Name ?? agent.Id); + } } } if (incompatibleAgents.Count > 0) { - return Task.FromResult(new HealthCheckResult( + return new HealthCheckResult( status: context.Registration.FailureStatus, description: string.Create( CultureInfo.InvariantCulture, $"Workflow checkpointing: {incompatibleAgents.Count} registered workflow agent(s) cannot use the checkpoint store supplied by hosting. Remove a caller-configured checkpoint manager and register the workflow agent directly rather than behind middleware."), - data: new Dictionary(StringComparer.Ordinal) { ["incompatibleAgents"] = incompatibleAgents })); + data: new Dictionary(StringComparer.Ordinal) { ["incompatibleAgents"] = incompatibleAgents }); } - return Task.FromResult(HealthCheckResult.Healthy( - string.Create(CultureInfo.InvariantCulture, $"Workflow checkpointing: {checkedAgents} workflow agent(s) checked, all leaving their checkpoint storage to hosting."))); + return HealthCheckResult.Healthy( + string.Create(CultureInfo.InvariantCulture, $"Workflow checkpointing: {checkedAgents} workflow agent(s) checked, all leaving their checkpoint storage to hosting.")); } } diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/ServiceCollectionExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/ServiceCollectionExtensions.cs index bac9b10b17d..6268421227e 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/ServiceCollectionExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/ServiceCollectionExtensions.cs @@ -52,6 +52,20 @@ public static class ServiceCollectionExtensions /// a caller-provided identity for session isolation is unsupported. /// /// + /// Keyed agent registrations may use any DI lifetime. Each response run resolves its agent and + /// session store from a DI scope that is disposed when the run's event stream ends, so scoped and + /// transient agents are created per run, including background runs that outlive the HTTP request. + /// + /// + /// For each response run the handler uses the registered under the agent's + /// name, for example with AddFoundryAgentSessionStore(store, "my-agent"). Otherwise it uses the store + /// registered for all agents, with AddFoundryAgentSessionStore(store) or a non-keyed + /// registration, and otherwise the default store, which uses the AgentServer + /// SDK's local state-store fallback. A request answered by the default agent uses that agent's store. Each + /// store keeps its own sessions, so changing which store an agent uses means sessions saved in the previous + /// store are no longer found. + /// + /// /// Example: /// /// builder.Services.AddKeyedSingleton<AIAgent>("my-agent", myAgent); @@ -109,6 +123,16 @@ public static IServiceCollection AddFoundryResponses(this IServiceCollection ser /// a caller-provided identity for session isolation is unsupported. /// /// + /// A store passed in agentSessionStore is registered only under the agent's name. Every request this + /// agent answers uses it, including requests without a matching agent name, and other agents do not use it. + /// When the agent has no name, the store has no key to be registered under and is used for all agents + /// instead. When no store is passed, the agent uses the store registered for all agents, with + /// or a non-keyed + /// registration, and otherwise the default store. Each store keeps its own + /// sessions, so changing which store an agent uses means sessions saved in the previous store are no + /// longer found. + /// + /// /// Example: /// /// builder.Services.AddFoundryResponses(myAgent); @@ -120,7 +144,10 @@ public static IServiceCollection AddFoundryResponses(this IServiceCollection ser /// /// The service collection. /// The agent instance to register. - /// The agent session store to use for managing agent sessions server-side. If null, the default store uses the AgentServer SDK's local state-store fallback. + /// + /// The session store for this agent, used only by this agent. If null, the agent uses the store registered + /// for all agents, or the default store, which uses the AgentServer SDK's local state-store fallback. + /// /// /// Optional , for example to allow the /// agent's own service to store the responses it produces, or to opt in to durable long-running @@ -149,22 +176,324 @@ public static IServiceCollection AddFoundryResponses( configuredOptions, includeServerOptions: serverAdded, applyOptions: serverAdded || options is not null); - agentSessionStore ??= CreateDefaultAgentSessionStore(); + RegisterAgentInstance(services, agent, agentSessionStore); + RegisterResponseHandler(services); + MarkFeatureUsed(); + return services; + } + + /// + /// Registers the Azure AI Responses Server SDK and an created by a factory + /// with the specified DI lifetime, as the handler for all incoming requests. + /// + /// + /// + /// Use this overload when the agent depends on services that are not singletons, for example a + /// database context, or when it should not live for the whole process. The agent is registered only + /// under . When no default agent is registered yet, requests without a + /// matching agent name are also routed to this agent. An existing registration under + /// is kept, the same as for AddFoundryResponses(agent). + /// + /// + /// Every response run (streaming, non-streaming, background, or a run resumed after a restart) gets + /// its own DI scope, created when the run starts and disposed when its event stream ends. That is + /// later than the end of the HTTP request for a background run. The lifetime applies to the agent: + /// + /// : one instance for the process, disposed at shutdown. + /// : one instance per response run, disposed when the run ends. + /// : a new instance on every lookup, disposed when the run ends. + /// + /// This scope is not the ASP.NET Core request scope, so scoped services that middleware fills in for an + /// HTTP request are new, empty instances here. Do not derive the caller's identity from such services in + /// the factory; the per-request user identity is supplied by the platform and applied by the response + /// handler when it loads and saves sessions. + /// + /// + /// The session store is registered separately with + /// , + /// with its own lifetime, because the response handler uses the store directly rather than through the + /// agent. Without such a registration for , the agent uses the store registered + /// for all agents, or the default store. Conversation state is not kept on + /// the agent object: each turn loads the saved session from the store's backing storage. Sessions are + /// stored in the same location AddFoundryResponses(agent) uses for an agent with that name + /// registered in the same order, whatever the lifetime, so switching between the two overloads or + /// between lifetimes keeps existing sessions readable. + /// + /// + /// This overload is for local or non-hosted scenarios. Foundry hosted applications must use + /// the overload. + /// + /// + /// Example: + /// + /// builder.Services.AddFoundryResponses( + /// "my-agent", + /// sp => new MyAgent(sp.GetRequiredService<MyDbContext>()), + /// ServiceLifetime.Scoped); + /// + /// + /// + /// The service collection. + /// The name the agent is registered under and the name requests use to select it. + /// Creates the agent from the service provider of the current DI scope. + /// The DI lifetime of the agent. Defaults to . + /// + /// Optional , for example to allow the + /// agent's own service to store the responses it produces, or to opt in to durable long-running + /// (resilient) background responses via . + /// + /// The service collection for chaining. + public static IServiceCollection AddFoundryResponses( + this IServiceCollection services, + string agentName, + Func agentFactory, + ServiceLifetime lifetime = ServiceLifetime.Singleton, + FoundryResponsesServerOptions? options = null) + { + _ = Throw.IfNull(services); + RegisterAgentFactory(services, agentName, agentFactory, lifetime); + return services.AddFoundryResponses(options); + } + + /// + /// Registers the that Foundry Responses hosting uses for one agent, or for + /// every agent that has no store of its own. + /// + /// + /// + /// For each response run the response handler uses the store registered under the agent's name, then the + /// store registered for all agents, then the default store, which uses the AgentServer SDK's local + /// state-store fallback. This applies to agents registered with any AddFoundryResponses overload or + /// directly in DI, and a request answered by the default agent uses that agent's store. The store receives + /// the agent on every call and partitions sessions by agent and by user, so one store can serve several + /// agents. Each store keeps its own sessions, so changing which store an agent uses means sessions saved in + /// the previous store are no longer found. + /// + /// + /// The DI container never disposes , because the caller created it. + /// Calling this method again for the same , or again for all agents, replaces + /// the earlier store, including one passed to AddFoundryResponses(agent, agentSessionStore) or + /// registered directly as a keyed or non-keyed , in whatever order the + /// methods are called. + /// + /// + /// Example: + /// + /// builder.Services.AddFoundryAgentSessionStore(mySessionStore); + /// + /// + /// + /// The service collection. + /// The session store to use. + /// + /// The name of the agent the store is for, or to use it for every agent that has no + /// store of its own. + /// + /// The service collection for chaining. + public static IServiceCollection AddFoundryAgentSessionStore( + this IServiceCollection services, + AgentSessionStore agentSessionStore, + string? agentName = null) + { + _ = Throw.IfNull(services); + _ = Throw.IfNull(agentSessionStore); + ValidateSessionStoreAgentName(agentName); + RemoveAgentSessionStoreRegistrations(services, agentName); + services.Add(agentName is null + ? ServiceDescriptor.Singleton(agentSessionStore) + : ServiceDescriptor.KeyedSingleton(agentName, agentSessionStore)); + return services; + } + + /// + /// Registers a factory for the that Foundry Responses hosting uses for one + /// agent, or for every agent that has no store of its own, with the specified DI lifetime. + /// + /// + /// + /// For each response run the response handler looks the store up once, from the DI scope of that run, + /// using the store registered under the agent's name, then the store registered for all agents, then the + /// default store. This applies to agents registered with any AddFoundryResponses overload or + /// directly in DI. The store's lifetime is independent of the agent's lifetime: + /// + /// : one instance for the process, disposed at shutdown. + /// : one instance per response run, disposed when the run ends. + /// : a new instance on every lookup, disposed when the run ends. + /// + /// The DI container disposes what the factory returns according to that lifetime, even when the factory + /// returns an object it did not create. To share an object you dispose yourself, use + /// . + /// + /// + /// Calling this method again for the same , or again for all agents, replaces + /// the earlier store, including one passed to AddFoundryResponses(agent, agentSessionStore) or + /// registered directly as a keyed or non-keyed , in whatever order the + /// methods are called. Each store keeps its own sessions, so changing which store an agent uses means + /// sessions saved in the previous store are no longer found. + /// + /// + /// Example: + /// + /// builder.Services.AddFoundryAgentSessionStore( + /// sp => new MySessionStore(sp.GetRequiredService<MyDbContext>()), + /// ServiceLifetime.Scoped); + /// + /// + /// + /// The service collection. + /// Creates the session store from the service provider of the current DI scope. + /// The DI lifetime of the session store. Defaults to . + /// + /// The name of the agent the store is for, or to use it for every agent that has no + /// store of its own. + /// + /// The service collection for chaining. + public static IServiceCollection AddFoundryAgentSessionStore( + this IServiceCollection services, + Func agentSessionStoreFactory, + ServiceLifetime lifetime = ServiceLifetime.Singleton, + string? agentName = null) + { + _ = Throw.IfNull(services); + _ = Throw.IfNull(agentSessionStoreFactory); + ValidateLifetime(lifetime); + ValidateSessionStoreAgentName(agentName); + + string target = agentName is null ? "all agents" : $"agent '{agentName}'"; + AgentSessionStore CreateStore(IServiceProvider serviceProvider) => + agentSessionStoreFactory(serviceProvider) + ?? throw new InvalidOperationException($"The agent session store factory registered for {target} returned null."); + + RemoveAgentSessionStoreRegistrations(services, agentName); + services.Add(agentName is null + ? new ServiceDescriptor(typeof(AgentSessionStore), CreateStore, lifetime) + : new ServiceDescriptor(typeof(AgentSessionStore), agentName, (serviceProvider, _) => CreateStore(serviceProvider), lifetime)); + return services; + } + + /// + /// Registers an agent factory under before the shared Responses registration runs. + /// + /// + /// Only a keyed registration is added, so each created agent has exactly one registration that owns its + /// disposal. When no default agent exists yet, a tells the response + /// handler to route requests without a matching agent name to this agent. An existing keyed registration + /// wins, the same as for AddFoundryResponses(agent); this agent is then not recorded as the default, + /// so the existing registration keeps its routing and storage location. + /// + internal static void RegisterAgentFactory( + IServiceCollection services, + string agentName, + Func agentFactory, + ServiceLifetime lifetime) + { + _ = Throw.IfNullOrWhitespace(agentName); + _ = Throw.IfNull(agentFactory); + ValidateLifetime(lifetime); + + if (services.Any(descriptor => + descriptor.ServiceType == typeof(AIAgent) + && descriptor.IsKeyedService + && Equals(descriptor.ServiceKey, agentName))) + { + return; + } + + services.Add(new ServiceDescriptor( + typeof(AIAgent), + agentName, + (serviceProvider, _) => agentFactory(serviceProvider) + ?? throw new InvalidOperationException($"The agent factory registered for '{agentName}' returned null."), + lifetime)); + + if (!HasDefaultAgentRegistration(services)) + { + services.AddSingleton(new FoundryDefaultAgentName(agentName)); + } + } + + /// + /// Registers an agent passed to AddFoundryResponses(agent, agentSessionStore), and its session store. + /// + /// + /// The store is registered only under the agent's name, so other agents do not use it. An agent without a + /// name has no key to register its store under, so its store is registered for all agents. The default + /// store is registered for all agents unless another store already is. + /// + internal static void RegisterAgentInstance(IServiceCollection services, AIAgent agent, AgentSessionStore? agentSessionStore) + { if (!string.IsNullOrWhiteSpace(agent.Name)) { services.TryAddKeyedSingleton(agent.Name, agent); - services.TryAddKeyedSingleton(agent.Name, agentSessionStore); + if (agentSessionStore is not null) + { + services.TryAddKeyedSingleton(agent.Name, agentSessionStore); + } + } + else if (agentSessionStore is not null) + { + services.TryAddSingleton(agentSessionStore); } // Also register as the default (non-keyed) agent so requests // without an agent name can resolve it (e.g., local dev tooling). - services.TryAddSingleton(agent); - services.TryAddSingleton(agentSessionStore); + TryAddDefaultAgent(services, agent); + services.TryAddSingleton(_ => CreateDefaultAgentSessionStore()); + } - RegisterResponseHandler(services); - MarkFeatureUsed(); - return services; + /// + /// Registers as the default (non-keyed) agent unless a default already exists, + /// including a factory registration that receives requests without an agent name. + /// + internal static void TryAddDefaultAgent(IServiceCollection services, AIAgent agent) + { + if (!HasDefaultAgentRegistration(services)) + { + services.AddSingleton(agent); + } + } + + /// + /// Whether a registration already receives requests without an agent name: a non-keyed + /// or a factory agent recorded as the default. + /// + private static bool HasDefaultAgentRegistration(IServiceCollection services) => + services.Any(descriptor => + (descriptor.ServiceType == typeof(AIAgent) && !descriptor.IsKeyedService) + || descriptor.ServiceType == typeof(FoundryDefaultAgentName)); + + private static void ValidateLifetime(ServiceLifetime lifetime) + { + if (lifetime is not (ServiceLifetime.Singleton or ServiceLifetime.Scoped or ServiceLifetime.Transient)) + { + Throw.ArgumentOutOfRangeException(nameof(lifetime), lifetime, "Unsupported service lifetime."); + } + } + + private static void ValidateSessionStoreAgentName(string? agentName) + { + if (agentName is not null) + { + _ = Throw.IfNullOrWhitespace(agentName); + } + } + + /// + /// Removes the session store registrations for , or the non-keyed ones when it + /// is , so the store registered next replaces them. + /// + private static void RemoveAgentSessionStoreRegistrations(IServiceCollection services, string? agentName) + { + for (int i = services.Count - 1; i >= 0; i--) + { + ServiceDescriptor descriptor = services[i]; + if (descriptor.ServiceType == typeof(AgentSessionStore) + && Equals(descriptor.ServiceKey, agentName)) + { + services.RemoveAt(i); + } + } } /// diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryAgentLifetimeTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryAgentLifetimeTests.cs new file mode 100644 index 00000000000..f94b8ba178b --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/FoundryAgentLifetimeTests.cs @@ -0,0 +1,1067 @@ +// 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.Runtime.CompilerServices; +using System.Text; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Azure.AI.AgentServer.Responses; +using Azure.AI.AgentServer.Responses.Models; +using Microsoft.AspNetCore.Builder; +using Microsoft.AspNetCore.Hosting.Server; +using Microsoft.AspNetCore.TestHost; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Hosting; +using Microsoft.Extensions.Logging.Abstractions; +using Moq; +using MeaiTextContent = Microsoft.Extensions.AI.TextContent; + +namespace Microsoft.Agents.AI.Foundry.Hosting.UnitTests; + +/// +/// Covers how the response handler resolves agents and session stores per response run, the +/// AddFoundryResponses overloads that register an agent factory with a DI lifetime, and +/// AddFoundryAgentSessionStore, which registers session stores separately from agents. +/// +public sealed class FoundryAgentLifetimeTests +{ + private const string AgentName = "lifetime-agent"; + + [Fact] + public async Task CreateAsync_ScopedKeyedAgent_CreatesOnePerResponseAndDisposesItAsync() + { + // Arrange + var log = new AgentLifetimeLog(); + var services = new ServiceCollection(); + services.AddSingleton(new TestAgentSessionStore()); + services.AddKeyedScoped(AgentName, (_, _) => new TrackingAgent(log)); + AgentFrameworkResponseHandler handler = CreateHandler(services); + + // Act + await DrainAsync(handler, CreateRequest(AgentName), CreateContext()); + await DrainAsync(handler, CreateRequest(AgentName), CreateContext()); + + // Assert + Assert.Equal(2, log.Created.Count); + Assert.NotSame(log.Created.First(), log.Created.Last()); + Assert.All(log.Created, agent => Assert.True(agent.IsDisposed)); + Assert.All(log.Created, agent => Assert.False(agent.WasDisposedDuringRun)); + } + + [Fact] + public async Task CreateAsync_TransientKeyedAgent_IsDisposedWhenTheResponseEndsAsync() + { + // Arrange + var log = new AgentLifetimeLog(); + var services = new ServiceCollection(); + services.AddSingleton(new TestAgentSessionStore()); + services.AddKeyedTransient(AgentName, (_, _) => new TrackingAgent(log)); + AgentFrameworkResponseHandler handler = CreateHandler(services); + + // Act + await DrainAsync(handler, CreateRequest(AgentName), CreateContext()); + + // Assert: the handler only needs the keyed agent, and the root container does not keep it. + Assert.NotEmpty(log.Created); + Assert.All(log.Created, agent => Assert.True(agent.IsDisposed)); + } + + [Fact] + public async Task CreateAsync_SingletonKeyedAgent_IsReusedAndNotDisposedAsync() + { + // Arrange + var log = new AgentLifetimeLog(); + var agent = new TrackingAgent(log); + var services = new ServiceCollection(); + services.AddSingleton(new TestAgentSessionStore()); + services.AddKeyedSingleton(AgentName, agent); + AgentFrameworkResponseHandler handler = CreateHandler(services); + + // Act + await DrainAsync(handler, CreateRequest(AgentName), CreateContext()); + await DrainAsync(handler, CreateRequest(AgentName), CreateContext()); + + // Assert + Assert.Same(agent, Assert.Single(log.Created)); + Assert.Equal(2, agent.RunCount); + Assert.False(agent.IsDisposed); + } + + [Fact] + public async Task CreateAsync_ScopedAgent_StaysAliveUntilTheEventStreamEndsAsync() + { + // Arrange: the agent waits mid-run, as a long background run would. + var log = new AgentLifetimeLog(); + var gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var services = new ServiceCollection(); + services.AddSingleton(new TestAgentSessionStore()); + services.AddKeyedScoped(AgentName, (_, _) => new TrackingAgent(log, gate)); + AgentFrameworkResponseHandler handler = CreateHandler(services); + + // Act + Task run = Task.Run(() => DrainAsync(handler, CreateRequest(AgentName), CreateContext())); + TrackingAgent agent = await log.RunStarted.Task.WaitAsync(TimeSpan.FromSeconds(10)); + bool disposedWhileRunning = agent.IsDisposed; + gate.SetResult(); + await run.WaitAsync(TimeSpan.FromSeconds(10)); + + // Assert + Assert.False(disposedWhileRunning); + Assert.False(agent.WasDisposedDuringRun); + Assert.True(agent.IsDisposed); + } + + [Fact] + public async Task CreateAsync_EnumerationStoppedEarly_DisposesScopedAgentAsync() + { + // Arrange: a client disconnect or cancellation stops reading the stream early. + var log = new AgentLifetimeLog(); + var services = new ServiceCollection(); + services.AddSingleton(new TestAgentSessionStore()); + services.AddKeyedScoped(AgentName, (_, _) => new TrackingAgent(log)); + AgentFrameworkResponseHandler handler = CreateHandler(services); + IAsyncEnumerator enumerator = handler + .CreateAsync(CreateRequest(AgentName), CreateContext(), CancellationToken.None) + .GetAsyncEnumerator(); + + // Act + Assert.True(await enumerator.MoveNextAsync()); + TrackingAgent agent = Assert.Single(log.Created); + bool disposedBeforeStop = agent.IsDisposed; + await enumerator.DisposeAsync(); + + // Assert + Assert.False(disposedBeforeStop); + Assert.True(agent.IsDisposed); + } + + [Fact] + public async Task CreateAsync_ScopedKeyedSessionStore_IsCreatedPerResponseAndDisposedAsync() + { + // Arrange + var log = new AgentLifetimeLog(); + var stores = new ConcurrentQueue(); + var services = new ServiceCollection(); + services.AddKeyedSingleton(AgentName, new TrackingAgent(log)); + services.AddKeyedScoped(AgentName, (_, _) => + { + var store = new RecordingSessionStore(); + stores.Enqueue(store); + return store; + }); + AgentFrameworkResponseHandler handler = CreateHandler(services); + + // Act + await DrainAsync(handler, CreateRequest(AgentName), CreateContext()); + await DrainAsync(handler, CreateRequest(AgentName), CreateContext()); + + // Assert + Assert.Equal(2, stores.Count); + Assert.All(stores, store => Assert.True(store.IsDisposed)); + Assert.All(stores, store => Assert.NotEmpty(store.Identities)); + } + + [Theory] + [InlineData(ServiceLifetime.Singleton)] + [InlineData(ServiceLifetime.Scoped)] + [InlineData(ServiceLifetime.Transient)] + public async Task AddFoundryResponses_Factory_UsesTheInstanceOverloadStorageIdentityAsync(ServiceLifetime lifetime) + { + // Arrange: the instance overload stores sessions under the agent name. + var instanceStore = new RecordingSessionStore(); + var instanceServices = new ServiceCollection(); + instanceServices.AddLogging(); + instanceServices.AddFoundryResponses(new TrackingAgent(new AgentLifetimeLog()), instanceStore); + AgentFrameworkResponseHandler instanceHandler = ResolveHandler(instanceServices); + await DrainAsync(instanceHandler, CreateRequest(AgentName), CreateContext()); + string instanceIdentity = Assert.Single(instanceStore.Identities.Distinct()); + + var factoryStore = new RecordingSessionStore(); + var factoryServices = new ServiceCollection(); + factoryServices.AddLogging(); + factoryServices.AddFoundryAgentSessionStore(factoryStore); + factoryServices.AddFoundryResponses( + AgentName, + _ => new TrackingAgent(new AgentLifetimeLog()), + lifetime); + AgentFrameworkResponseHandler factoryHandler = ResolveHandler(factoryServices); + + // Act: one request names the agent and one does not. + await DrainAsync(factoryHandler, CreateRequest(AgentName), CreateContext()); + await DrainAsync(factoryHandler, CreateRequest(agentName: null), CreateContext()); + + // Assert: every lifetime and both request shapes read and write the same sessions. + Assert.Equal($"name:{AgentName}", instanceIdentity); + Assert.All(factoryStore.Identities, identity => Assert.Equal(instanceIdentity, identity)); + } + + [Theory] + [InlineData(ServiceLifetime.Scoped)] + [InlineData(ServiceLifetime.Transient)] + public async Task AddFoundryResponses_Factory_CreatesOneAgentPerResponseAsync(ServiceLifetime lifetime) + { + // Arrange + var log = new AgentLifetimeLog(); + var services = new ServiceCollection(); + services.AddLogging(); + services.AddFoundryAgentSessionStore(new TestAgentSessionStore()); + services.AddFoundryResponses(AgentName, _ => new TrackingAgent(log), lifetime); + AgentFrameworkResponseHandler handler = ResolveHandler(services); + + // Act + await DrainAsync(handler, CreateRequest(AgentName), CreateContext()); + await DrainAsync(handler, CreateRequest(agentName: null), CreateContext()); + + // Assert: the named and the default registration share one agent within a response run. + Assert.Equal(2, log.Created.Count); + Assert.All(log.Created, agent => Assert.Equal(1, agent.RunCount)); + Assert.All(log.Created, agent => Assert.True(agent.IsDisposed)); + } + + [Fact] + public async Task AddFoundryResponses_SingletonFactory_CreatesOneAgentForTheProcessAsync() + { + // Arrange + var log = new AgentLifetimeLog(); + var services = new ServiceCollection(); + services.AddLogging(); + services.AddFoundryAgentSessionStore(new TestAgentSessionStore()); + services.AddFoundryResponses(AgentName, _ => new TrackingAgent(log), ServiceLifetime.Singleton); + AgentFrameworkResponseHandler handler = ResolveHandler(services); + + // Act + await DrainAsync(handler, CreateRequest(AgentName), CreateContext()); + await DrainAsync(handler, CreateRequest(agentName: null), CreateContext()); + + // Assert + TrackingAgent agent = Assert.Single(log.Created); + Assert.Equal(2, agent.RunCount); + Assert.False(agent.IsDisposed); + } + + [Theory] + [InlineData(ServiceLifetime.Scoped)] + [InlineData(ServiceLifetime.Transient)] + public async Task AddFoundryAgentSessionStore_Factory_CreatesOneStorePerResponseIndependentlyOfTheAgentAsync(ServiceLifetime lifetime) + { + // Arrange: a singleton agent with a store that has its own, shorter lifetime. + var log = new AgentLifetimeLog(); + var stores = new ConcurrentQueue(); + var services = new ServiceCollection(); + services.AddLogging(); + services.AddFoundryResponses(AgentName, _ => new TrackingAgent(log)); + services.AddFoundryAgentSessionStore( + _ => + { + var store = new RecordingSessionStore(); + stores.Enqueue(store); + return store; + }, + lifetime); + AgentFrameworkResponseHandler handler = ResolveHandler(services); + + // Act + await DrainAsync(handler, CreateRequest(AgentName), CreateContext()); + await DrainAsync(handler, CreateRequest(agentName: null), CreateContext()); + + // Assert: each run uses one store, which is disposed when the run ends, while the agent lives on. + Assert.Equal(2, stores.Count); + Assert.All(stores, store => Assert.Equal(1, store.DisposeCount)); + Assert.All(stores, store => Assert.NotEmpty(store.Identities)); + Assert.All(stores.SelectMany(store => store.Identities), identity => Assert.Equal($"name:{AgentName}", identity)); + TrackingAgent agent = Assert.Single(log.Created); + Assert.False(agent.IsDisposed); + } + + [Fact] + public async Task AddFoundryAgentSessionStore_Factory_DefaultsToOneStoreForTheProcessAsync() + { + // Arrange: a scoped agent with a store registered without a lifetime. + var stores = new ConcurrentQueue(); + var services = new ServiceCollection(); + services.AddLogging(); + services.AddFoundryResponses(AgentName, _ => new TrackingAgent(new AgentLifetimeLog()), ServiceLifetime.Scoped); + services.AddFoundryAgentSessionStore(_ => + { + var store = new RecordingSessionStore(); + stores.Enqueue(store); + return store; + }); + AgentFrameworkResponseHandler handler = ResolveHandler(services); + + // Act + await DrainAsync(handler, CreateRequest(AgentName), CreateContext()); + await DrainAsync(handler, CreateRequest(agentName: null), CreateContext()); + + // Assert + RecordingSessionStore store = Assert.Single(stores); + Assert.False(store.IsDisposed); + } + + [Fact] + public async Task AddFoundryAgentSessionStore_Instance_IsNeverDisposedByTheContainerAsync() + { + // Arrange: a shared store object used by a scoped agent. + var store = new RecordingSessionStore(); + var services = new ServiceCollection(); + services.AddLogging(); + services.AddFoundryResponses(AgentName, _ => new TrackingAgent(new AgentLifetimeLog()), ServiceLifetime.Scoped); + services.AddFoundryAgentSessionStore(store); + ServiceProvider provider = services.BuildServiceProvider(new ServiceProviderOptions { ValidateScopes = true }); + var handler = Assert.IsType(provider.GetRequiredService()); + + // Act + await DrainAsync(handler, CreateRequest(AgentName), CreateContext()); + await DrainAsync(handler, CreateRequest(agentName: null), CreateContext()); + await provider.DisposeAsync(); + + // Assert: the caller owns the store, so neither the response runs nor shutdown dispose it. + Assert.NotEmpty(store.Identities); + Assert.Equal(0, store.DisposeCount); + } + + [Fact] + public async Task AddFoundryAgentSessionStore_FactoryReturnsNull_ThrowsInvalidOperationExceptionAsync() + { + // Arrange + var services = new ServiceCollection(); + services.AddLogging(); + services.AddFoundryResponses(AgentName, _ => new TrackingAgent(new AgentLifetimeLog()), ServiceLifetime.Scoped); + services.AddFoundryAgentSessionStore(_ => null!, ServiceLifetime.Scoped, AgentName); + AgentFrameworkResponseHandler handler = ResolveHandler(services); + + // Act & Assert + InvalidOperationException exception = await Assert.ThrowsAsync( + () => DrainAsync(handler, CreateRequest(AgentName), CreateContext())); + Assert.Contains("session store factory", exception.Message, StringComparison.Ordinal); + Assert.Contains(AgentName, exception.Message, StringComparison.Ordinal); + } + + [Fact] + public async Task AddFoundryResponses_FactoryReturnsNull_ThrowsInvalidOperationExceptionAsync() + { + // Arrange + var services = new ServiceCollection(); + services.AddLogging(); + services.AddFoundryAgentSessionStore(new TestAgentSessionStore()); + services.AddFoundryResponses(AgentName, _ => null!, ServiceLifetime.Scoped); + AgentFrameworkResponseHandler handler = ResolveHandler(services); + + // Act & Assert + InvalidOperationException exception = await Assert.ThrowsAsync( + () => DrainAsync(handler, CreateRequest(AgentName), CreateContext())); + Assert.Contains(AgentName, exception.Message, StringComparison.Ordinal); + } + + [Fact] + public void AddFoundryResponses_Factory_InvalidArguments_Throw() + { + // Arrange + var services = new ServiceCollection(); + + // Act & Assert + Assert.Throws(() => services.AddFoundryResponses(" ", _ => new TrackingAgent(new AgentLifetimeLog()))); + Assert.Throws(() => services.AddFoundryResponses(AgentName, null!)); + Assert.Throws(() => services.AddFoundryResponses( + AgentName, + _ => new TrackingAgent(new AgentLifetimeLog()), + (ServiceLifetime)42)); + } + + [Fact] + public void AddFoundryAgentSessionStore_InvalidArguments_Throw() + { + // Arrange + var services = new ServiceCollection(); + var store = new TestAgentSessionStore(); + + // Act & Assert + Assert.Throws(() => ((IServiceCollection)null!).AddFoundryAgentSessionStore(store)); + Assert.Throws(() => services.AddFoundryAgentSessionStore((AgentSessionStore)null!)); + Assert.Throws(() => services.AddFoundryAgentSessionStore((Func)null!)); + Assert.Throws(() => services.AddFoundryAgentSessionStore(store, " ")); + Assert.Throws(() => services.AddFoundryAgentSessionStore(_ => store, ServiceLifetime.Singleton, string.Empty)); + Assert.Throws(() => services.AddFoundryAgentSessionStore(_ => store, (ServiceLifetime)42)); + Assert.Throws(() => ((IHostApplicationBuilder)null!).AddFoundryAgentSessionStore(store)); + Assert.Empty(services); + } + + [Fact] + public void AddFoundryResponses_Factory_RegistersOnlyTheKeyedAgentWithTheRequestedLifetime() + { + // Arrange + var services = new ServiceCollection(); + + // Act + services.AddFoundryResponses(AgentName, _ => new TrackingAgent(new AgentLifetimeLog()), ServiceLifetime.Scoped); + + // Assert: one registration per created agent, so DI owns its disposal once, and no store of its own. + ServiceDescriptor keyed = Assert.Single(services, d => d.ServiceType == typeof(AIAgent)); + Assert.True(keyed.IsKeyedService); + Assert.Equal(AgentName, keyed.ServiceKey); + Assert.Equal(ServiceLifetime.Scoped, keyed.Lifetime); + Assert.DoesNotContain(services, d => d.ServiceType == typeof(AgentSessionStore) && d.IsKeyedService); + ServiceDescriptor defaultStore = Assert.Single(services, d => d.ServiceType == typeof(AgentSessionStore)); + Assert.Equal(ServiceLifetime.Singleton, defaultStore.Lifetime); + FoundryDefaultAgentName defaultAgentName = Assert.IsType( + Assert.Single(services, d => d.ServiceType == typeof(FoundryDefaultAgentName)).ImplementationInstance); + Assert.Equal(AgentName, defaultAgentName.Name); + } + + [Fact] + public void AddFoundryResponses_SecondFactoryAgent_IsNotRecordedAsTheDefault() + { + // Arrange + var services = new ServiceCollection(); + + // Act + services.AddFoundryResponses("first", _ => new TrackingAgent(new AgentLifetimeLog(), name: "first")); + services.AddFoundryResponses("second", _ => new TrackingAgent(new AgentLifetimeLog(), name: "second")); + + // Assert + FoundryDefaultAgentName defaultAgentName = Assert.IsType( + Assert.Single(services, d => d.ServiceType == typeof(FoundryDefaultAgentName)).ImplementationInstance); + Assert.Equal("first", defaultAgentName.Name); + } + + [Theory] + [InlineData(true)] + [InlineData(false)] + public void AddFoundryAgentSessionStore_ForAllAgents_ReplacesTheDefaultStoreInAnyOrder(bool registerStoreFirst) + { + // Arrange + var store = new TestAgentSessionStore(); + var services = new ServiceCollection(); + services.AddLogging(); + + // Act + if (registerStoreFirst) + { + services.AddFoundryAgentSessionStore(store); + services.AddFoundryResponses(AgentName, _ => new TrackingAgent(new AgentLifetimeLog())); + } + else + { + services.AddFoundryResponses(AgentName, _ => new TrackingAgent(new AgentLifetimeLog())); + services.AddFoundryAgentSessionStore(store); + } + + using ServiceProvider provider = services.BuildServiceProvider(); + + // Assert + Assert.Single(services, d => d.ServiceType == typeof(AgentSessionStore)); + Assert.Same(store, provider.GetRequiredService()); + } + + [Fact] + public void AddFoundryAgentSessionStore_CalledTwice_TheLastCallWins() + { + // Arrange + var first = new TestAgentSessionStore(); + var second = new TestAgentSessionStore(); + var firstForAgent = new TestAgentSessionStore(); + var secondForAgent = new TestAgentSessionStore(); + var services = new ServiceCollection(); + + // Act + services.AddFoundryAgentSessionStore(first); + services.AddFoundryAgentSessionStore(firstForAgent, AgentName); + services.AddFoundryAgentSessionStore(_ => second, ServiceLifetime.Scoped); + services.AddFoundryAgentSessionStore(secondForAgent, AgentName); + using ServiceProvider provider = services.BuildServiceProvider(); + using IServiceScope scope = provider.CreateScope(); + + // Assert: replacing the store for all agents keeps the store registered for the agent name. + Assert.Equal(2, services.Count(d => d.ServiceType == typeof(AgentSessionStore))); + Assert.Same(second, scope.ServiceProvider.GetRequiredService()); + Assert.Same(secondForAgent, scope.ServiceProvider.GetRequiredKeyedService(AgentName)); + } + + [Fact] + public void AddFoundryAgentSessionStore_HostBuilder_RegistersTheStoreWithTheRequestedLifetime() + { + // Arrange + HostApplicationBuilder builder = Host.CreateApplicationBuilder(); + var store = new TestAgentSessionStore(); + + // Act + builder.AddFoundryAgentSessionStore(store); + builder.AddFoundryAgentSessionStore(_ => new TestAgentSessionStore(), ServiceLifetime.Transient, AgentName); + + // Assert + ServiceDescriptor forAllAgents = Assert.Single(builder.Services, d => d.ServiceType == typeof(AgentSessionStore) && !d.IsKeyedService); + Assert.Same(store, forAllAgents.ImplementationInstance); + ServiceDescriptor forAgent = Assert.Single(builder.Services, d => d.ServiceType == typeof(AgentSessionStore) && d.IsKeyedService); + Assert.Equal(AgentName, forAgent.ServiceKey); + Assert.Equal(ServiceLifetime.Transient, forAgent.Lifetime); + } + + [Fact] + public async Task AddFoundryAgentSessionStore_ForAgentName_TakesPrecedenceOverTheStoreForAllAgentsAsync() + { + // Arrange + var storeForAll = new RecordingSessionStore(); + var storeForSecond = new RecordingSessionStore(); + var services = new ServiceCollection(); + services.AddLogging(); + services.AddFoundryResponses("first", _ => new TrackingAgent(new AgentLifetimeLog(), name: "first"), ServiceLifetime.Scoped); + services.AddFoundryResponses("second", _ => new TrackingAgent(new AgentLifetimeLog(), name: "second"), ServiceLifetime.Scoped); + services.AddFoundryAgentSessionStore(storeForAll); + services.AddFoundryAgentSessionStore(storeForSecond, "second"); + AgentFrameworkResponseHandler handler = ResolveHandler(services); + + // Act + string firstIdentity = await DrainAndGetIdentityAsync(handler, storeForAll, "first"); + string secondIdentity = await DrainAndGetIdentityAsync(handler, storeForSecond, "second"); + + // Assert + Assert.Equal("name:first", firstIdentity); + Assert.Equal("key:second", secondIdentity); + Assert.DoesNotContain("key:second", storeForAll.Identities); + Assert.DoesNotContain("name:first", storeForSecond.Identities); + } + + [Fact] + public async Task AddFoundryAgentSessionStore_ForAgentName_ReplacesTheStorePassedToTheInstanceOverloadAsync() + { + // Arrange + var instanceStore = new RecordingSessionStore(); + var replacement = new RecordingSessionStore(); + var services = new ServiceCollection(); + services.AddLogging(); + services.AddFoundryResponses(new TrackingAgent(new AgentLifetimeLog()), instanceStore); + services.AddFoundryAgentSessionStore(replacement, AgentName); + AgentFrameworkResponseHandler handler = ResolveHandler(services); + + // Act + await DrainAsync(handler, CreateRequest(AgentName), CreateContext()); + + // Assert + Assert.Empty(instanceStore.Identities); + Assert.NotEmpty(replacement.Identities); + } + + [Theory] + [InlineData(true)] + [InlineData(false)] + public async Task AddFoundryResponses_InstanceWithoutStore_UsesTheStoreForAllAgentsAsync(bool registerStoreFirst) + { + // Arrange + var storeForAll = new RecordingSessionStore(); + var services = new ServiceCollection(); + services.AddLogging(); + if (registerStoreFirst) + { + services.AddFoundryAgentSessionStore(storeForAll); + } + + services.AddFoundryResponses(new TrackingAgent(new AgentLifetimeLog())); + if (!registerStoreFirst) + { + services.AddFoundryAgentSessionStore(storeForAll); + } + + AgentFrameworkResponseHandler handler = ResolveHandler(services); + + // Act + string namedIdentity = await DrainAndGetIdentityAsync(handler, storeForAll, AgentName); + string unnamedIdentity = await DrainAndGetIdentityAsync(handler, storeForAll, agentName: null); + + // Assert + Assert.Equal($"name:{AgentName}", namedIdentity); + Assert.Equal(namedIdentity, unnamedIdentity); + } + + [Fact] + public async Task AddFoundryResponses_InstanceWithStore_DoesNotShareItsStoreWithOtherAgentsAsync() + { + // Arrange: a store passed for one agent, and another agent registered by hand without a store. + var instanceStore = new RecordingSessionStore(); + var storeForAll = new RecordingSessionStore(); + var services = new ServiceCollection(); + services.AddLogging(); + services.AddFoundryResponses(new TrackingAgent(new AgentLifetimeLog(), name: "first"), instanceStore); + services.AddKeyedSingleton("other", new TrackingAgent(new AgentLifetimeLog(), name: "other")); + services.AddFoundryAgentSessionStore(storeForAll); + AgentFrameworkResponseHandler handler = ResolveHandler(services); + + // Act + await DrainAsync(handler, CreateRequest("other"), CreateContext()); + await DrainAsync(handler, CreateRequest("first"), CreateContext()); + + // Assert + Assert.Equal("key:other", Assert.Single(storeForAll.Identities.Distinct())); + Assert.Equal("name:first", Assert.Single(instanceStore.Identities.Distinct())); + } + + [Fact] + public void AddFoundryResponses_InstanceWithStore_KeepsTheDefaultStoreForOtherAgents() + { + // Arrange + var instanceStore = new TestAgentSessionStore(); + var services = new ServiceCollection(); + services.AddLogging(); + + // Act + services.AddFoundryResponses(new TrackingAgent(new AgentLifetimeLog()), instanceStore); + using ServiceProvider provider = services.BuildServiceProvider(); + + // Assert + Assert.Same(instanceStore, provider.GetRequiredKeyedService(AgentName)); + Assert.IsType(provider.GetRequiredService()); + } + + [Fact] + public async Task CreateAsync_UnknownAgentName_UsesTheDefaultAgentsOwnStoreAsync() + { + // Arrange: a request naming no registered agent is answered by the default agent. + var instanceStore = new RecordingSessionStore(); + var storeForAll = new RecordingSessionStore(); + var services = new ServiceCollection(); + services.AddLogging(); + services.AddFoundryResponses(new TrackingAgent(new AgentLifetimeLog()), instanceStore); + services.AddFoundryAgentSessionStore(storeForAll); + AgentFrameworkResponseHandler handler = ResolveHandler(services); + + // Act + await DrainAsync(handler, CreateRequest("unknown"), CreateContext()); + await DrainAsync(handler, CreateRequest(agentName: null), CreateContext()); + + // Assert + Assert.Empty(storeForAll.Identities); + Assert.Equal($"name:{AgentName}", Assert.Single(instanceStore.Identities.Distinct())); + } + [Theory] + [InlineData(ServiceLifetime.Singleton)] + [InlineData(ServiceLifetime.Scoped)] + [InlineData(ServiceLifetime.Transient)] + public async Task AddFoundryResponses_TwoFactoryAgents_UseTheInstanceOverloadStorageIdentitiesAsync(ServiceLifetime lifetime) + { + // Arrange: the instance overload makes the first agent the default and keys the second by name. + var instanceStore = new RecordingSessionStore(); + var instanceServices = new ServiceCollection(); + instanceServices.AddLogging(); + instanceServices.AddFoundryResponses(new TrackingAgent(new AgentLifetimeLog(), name: "first"), instanceStore); + instanceServices.AddFoundryResponses(new TrackingAgent(new AgentLifetimeLog(), name: "second"), instanceStore); + AgentFrameworkResponseHandler instanceHandler = ResolveHandler(instanceServices); + + var factoryStore = new RecordingSessionStore(); + var factoryServices = new ServiceCollection(); + factoryServices.AddLogging(); + factoryServices.AddFoundryAgentSessionStore(factoryStore); + factoryServices.AddFoundryResponses("first", _ => new TrackingAgent(new AgentLifetimeLog(), name: "first"), lifetime); + factoryServices.AddFoundryResponses("second", _ => new TrackingAgent(new AgentLifetimeLog(), name: "second"), lifetime); + AgentFrameworkResponseHandler factoryHandler = ResolveHandler(factoryServices); + + // Act + string[] instanceIdentities = + [ + await DrainAndGetIdentityAsync(instanceHandler, instanceStore, "first"), + await DrainAndGetIdentityAsync(instanceHandler, instanceStore, "second"), + await DrainAndGetIdentityAsync(instanceHandler, instanceStore, agentName: null), + ]; + string[] factoryIdentities = + [ + await DrainAndGetIdentityAsync(factoryHandler, factoryStore, "first"), + await DrainAndGetIdentityAsync(factoryHandler, factoryStore, "second"), + await DrainAndGetIdentityAsync(factoryHandler, factoryStore, agentName: null), + ]; + + // Assert + Assert.Equal(["name:first", "key:second", "name:first"], instanceIdentities); + Assert.Equal(instanceIdentities, factoryIdentities); + } + + [Fact] + public async Task AddFoundryResponses_FactoryDefaultThenInstanceAgent_KeepsTheFactoryAgentAsDefaultAsync() + { + // Arrange: the instance overload must not take over the default slot the factory agent holds. + var log = new AgentLifetimeLog(); + var store = new RecordingSessionStore(); + var services = new ServiceCollection(); + services.AddLogging(); + services.AddFoundryAgentSessionStore(store); + services.AddFoundryResponses("first", _ => new TrackingAgent(log, name: "first"), ServiceLifetime.Scoped); + services.AddFoundryResponses(new TrackingAgent(new AgentLifetimeLog(), name: "second"), store); + AgentFrameworkResponseHandler handler = ResolveHandler(services); + + // Act + string unnamedIdentity = await DrainAndGetIdentityAsync(handler, store, agentName: null); + string secondIdentity = await DrainAndGetIdentityAsync(handler, store, "second"); + + // Assert + Assert.Equal("name:first", unnamedIdentity); + Assert.Equal("key:second", secondIdentity); + Assert.DoesNotContain(services, d => d.ServiceType == typeof(AIAgent) && !d.IsKeyedService); + } + + [Fact] + public async Task AddFoundryResponses_NonKeyedAgentRegisteredLater_DoesNotShareTheFactoryAgentSessionsAsync() + { + // Arrange: a different agent with the same name is registered as the non-keyed default after the + // factory overload, so it receives requests without an agent name. + var store = new RecordingSessionStore(); + var services = new ServiceCollection(); + services.AddLogging(); + services.AddSingleton(store); + services.AddFoundryResponses(AgentName, _ => new TrackingAgent(new AgentLifetimeLog()), ServiceLifetime.Scoped); + services.AddSingleton(new TrackingAgent(new AgentLifetimeLog())); + AgentFrameworkResponseHandler handler = ResolveHandler(services); + + // Act + string factoryIdentity = await DrainAndGetIdentityAsync(handler, store, AgentName); + string nonKeyedIdentity = await DrainAndGetIdentityAsync(handler, store, agentName: null); + + // Assert: the two agents never read or write each other's sessions. + Assert.Equal($"key:{AgentName}", factoryIdentity); + Assert.Equal($"name:{AgentName}", nonKeyedIdentity); + } + + [Theory] + [InlineData(ServiceLifetime.Scoped)] + [InlineData(ServiceLifetime.Transient)] + public async Task AddFoundryResponses_Factory_DisposesEachAgentAndStoreExactlyOnceAsync(ServiceLifetime lifetime) + { + // Arrange + var log = new AgentLifetimeLog(); + var stores = new ConcurrentQueue(); + var services = new ServiceCollection(); + services.AddLogging(); + services.AddFoundryResponses(AgentName, _ => new TrackingAgent(log), lifetime); + services.AddFoundryAgentSessionStore( + _ => + { + var store = new RecordingSessionStore(); + stores.Enqueue(store); + return store; + }, + lifetime, + AgentName); + ServiceProvider provider = services.BuildServiceProvider(new ServiceProviderOptions { ValidateScopes = true }); + var handler = Assert.IsType(provider.GetRequiredService()); + + // Act: a named request, an unnamed request, and the readiness enumeration of all agents. + await DrainAsync(handler, CreateRequest(AgentName), CreateContext()); + await DrainAsync(handler, CreateRequest(agentName: null), CreateContext()); + await using (AsyncServiceScope scope = provider.CreateAsyncScope()) + { + Assert.Single(ServiceCollectionExtensions.ResolveRegisteredAgents(scope.ServiceProvider)); + } + + await provider.DisposeAsync(); + + // Assert + Assert.Equal(3, log.Created.Count); + Assert.All(log.Created, agent => Assert.Equal(1, agent.DisposeCount)); + Assert.Equal(2, stores.Count); + Assert.All(stores, store => Assert.Equal(1, store.DisposeCount)); + } + + [Fact] + public async Task AddFoundryResponses_SingletonFactory_DisposesTheAgentAndStoreOnceAtShutdownAsync() + { + // Arrange + var log = new AgentLifetimeLog(); + var store = new RecordingSessionStore(); + var services = new ServiceCollection(); + services.AddLogging(); + services.AddFoundryResponses(AgentName, _ => new TrackingAgent(log), ServiceLifetime.Singleton); + services.AddFoundryAgentSessionStore(_ => store); + ServiceProvider provider = services.BuildServiceProvider(new ServiceProviderOptions { ValidateScopes = true }); + var handler = Assert.IsType(provider.GetRequiredService()); + + // Act + await DrainAsync(handler, CreateRequest(AgentName), CreateContext()); + await DrainAsync(handler, CreateRequest(agentName: null), CreateContext()); + bool disposedBeforeShutdown = store.IsDisposed || log.Created.Any(agent => agent.IsDisposed); + await provider.DisposeAsync(); + + // Assert + Assert.False(disposedBeforeShutdown); + Assert.Equal(1, Assert.Single(log.Created).DisposeCount); + Assert.Equal(1, store.DisposeCount); + } + + [Fact] + public async Task AddFoundryResponses_ExistingScopedKeyedAgent_IsNotCapturedBySingletonDefaultAsync() + { + // Arrange: the caller already registered a scoped agent under the name; the factory call uses + // the default singleton lifetime. + var log = new AgentLifetimeLog(); + var services = new ServiceCollection(); + services.AddLogging(); + services.AddKeyedScoped(AgentName, (_, _) => new TrackingAgent(log)); + services.AddFoundryAgentSessionStore(new TestAgentSessionStore()); + services.AddFoundryResponses(AgentName, _ => new TrackingAgent(new AgentLifetimeLog())); + AgentFrameworkResponseHandler handler = ResolveHandler(services); + + // Act + await DrainAsync(handler, CreateRequest(AgentName), CreateContext()); + await DrainAsync(handler, CreateRequest(AgentName), CreateContext()); + + // Assert: the existing scoped registration is used per run and nothing forwards to it. + Assert.DoesNotContain(services, d => d.ServiceType == typeof(AIAgent) && !d.IsKeyedService); + Assert.Equal(2, log.Created.Count); + Assert.All(log.Created, agent => Assert.Equal(1, agent.DisposeCount)); + } + + [Fact] + public void AddFoundryResponses_Factory_KeepsAnExistingKeyedRegistration() + { + // Arrange + var existing = new TrackingAgent(new AgentLifetimeLog()); + var services = new ServiceCollection(); + services.AddKeyedSingleton(AgentName, existing); + + // Act + services.AddFoundryResponses(AgentName, _ => new TrackingAgent(new AgentLifetimeLog()), ServiceLifetime.Scoped); + + // Assert: the existing registration wins, like the instance overload's TryAdd behavior. + ServiceDescriptor keyed = Assert.Single(services, d => d.ServiceType == typeof(AIAgent) && d.IsKeyedService); + Assert.Same(existing, keyed.KeyedImplementationInstance); + Assert.DoesNotContain(services, d => d.ServiceType == typeof(FoundryDefaultAgentName)); + } + + [Fact] + public void AddFoundryResponses_HostBuilderFactory_RegistersAgentWithTheRequestedLifetime() + { + // Arrange + HostApplicationBuilder builder = Host.CreateApplicationBuilder(); + + // Act + builder.AddFoundryResponses(AgentName, _ => new TrackingAgent(new AgentLifetimeLog()), ServiceLifetime.Transient); + using ServiceProvider provider = builder.Services.BuildServiceProvider(); + + // Assert + ServiceDescriptor keyed = Assert.Single(builder.Services, d => d.ServiceType == typeof(AIAgent) && d.IsKeyedService); + Assert.Equal(ServiceLifetime.Transient, keyed.Lifetime); + Assert.IsType(provider.GetRequiredService()); + } + + [Fact] + public async Task BackgroundResponse_ScopedAgent_OutlivesTheHttpRequestAsync() + { + // Arrange + var log = new AgentLifetimeLog(); + var gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + builder.WebHost.UseTestServer(); + builder.Services.AddFoundryResponses( + AgentName, + _ => new TrackingAgent(log, gate), + ServiceLifetime.Scoped); + builder.Services.AddFoundryAgentSessionStore(_ => new TestAgentSessionStore(), ServiceLifetime.Scoped); + builder.Services.AddLogging(); + await using WebApplication app = builder.Build(); + app.MapFoundryResponses(); + await app.StartAsync(); + using HttpClient client = ((TestServer)app.Services.GetRequiredService()).CreateClient(); + + // Act: the POST returns while the agent is still running in the background. + string body = JsonSerializer.Serialize(new { model = "test", input = "Hello", store = true, background = true }); + using HttpResponseMessage response = await client.PostAsync( + new Uri("/responses", UriKind.Relative), + new StringContent(body, Encoding.UTF8, "application/json")); + response.EnsureSuccessStatusCode(); + TrackingAgent agent = await log.RunStarted.Task.WaitAsync(TimeSpan.FromSeconds(10)); + bool disposedAfterHttpReturned = agent.IsDisposed; + gate.SetResult(); + await agent.Disposed.Task.WaitAsync(TimeSpan.FromSeconds(10)); + + // Assert + Assert.False(disposedAfterHttpReturned); + Assert.False(agent.WasDisposedDuringRun); + Assert.True(agent.IsDisposed); + } + + private static AgentFrameworkResponseHandler CreateHandler(IServiceCollection services) => + new( + services.BuildServiceProvider(new ServiceProviderOptions { ValidateScopes = true }), + NullLogger.Instance); + + private static AgentFrameworkResponseHandler ResolveHandler(IServiceCollection services) => + Assert.IsType( + services.BuildServiceProvider(new ServiceProviderOptions { ValidateScopes = true }) + .GetRequiredService()); + + private static CreateResponse CreateRequest(string? agentName) + { + var request = agentName is null + ? new CreateResponse { Model = "test" } + : new CreateResponse { Model = "test", AgentReference = new AgentReference(agentName) }; + request.Input = BinaryData.FromObjectAsJson(new[] + { + new + { + type = "message", + id = "msg_1", + status = "completed", + role = "user", + content = new[] { new { type = "input_text", text = "Hello" } }, + }, + }); + return request; + } + + private static ResponseContext CreateContext() + { + var context = new Mock("resp_" + new string('0', 46)) { CallBase = true }; + 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 DrainAsync(AgentFrameworkResponseHandler handler, CreateResponse request, ResponseContext context) + { + await foreach (ResponseStreamEvent _ in handler.CreateAsync(request, context, CancellationToken.None)) + { + } + } + + /// Runs one request and returns the single storage identity the store saw during it. + private static async Task DrainAndGetIdentityAsync( + AgentFrameworkResponseHandler handler, + RecordingSessionStore store, + string? agentName) + { + int before = store.Identities.Count; + await DrainAsync(handler, CreateRequest(agentName), CreateContext()); + return Assert.Single(store.Identities.Skip(before).Distinct()); + } + + private sealed class AgentLifetimeLog + { + public ConcurrentQueue Created { get; } = new(); + + public TaskCompletionSource RunStarted { get; } = + new(TaskCreationOptions.RunContinuationsAsynchronously); + } + + /// An agent that records when it is created, run, and disposed. + private sealed class TrackingAgent : AIAgent, IAsyncDisposable + { + private readonly AgentLifetimeLog _log; + private readonly TaskCompletionSource? _gate; + private readonly string _name; + private int _runCount; + private int _disposeCount; + + public TrackingAgent(AgentLifetimeLog log, TaskCompletionSource? gate = null, string name = AgentName) + { + this._log = log; + this._gate = gate; + this._name = name; + log.Created.Enqueue(this); + } + + public override string? Name => this._name; + + public int RunCount => this._runCount; + + public int DisposeCount => this._disposeCount; + + public bool IsDisposed { get; private set; } + + public bool WasDisposedDuringRun { get; private set; } + + public TaskCompletionSource Disposed { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public ValueTask DisposeAsync() + { + Interlocked.Increment(ref this._disposeCount); + this.IsDisposed = true; + this.Disposed.TrySetResult(); + return default; + } + + protected override async IAsyncEnumerable RunCoreStreamingAsync( + IEnumerable messages, + AgentSession? session, + AgentRunOptions? options, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + Interlocked.Increment(ref this._runCount); + this._log.RunStarted.TrySetResult(this); + if (this._gate is not null) + { + await this._gate.Task.WaitAsync(cancellationToken); + } + + this.WasDisposedDuringRun = this.IsDisposed; + yield return new AgentResponseUpdate { MessageId = "resp_msg_1", Contents = [new MeaiTextContent("ok")] }; + } + + protected override Task RunCoreAsync( + IEnumerable messages, + AgentSession? session, + AgentRunOptions? options, + CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + + protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) => + new(new TrackingSession()); + + protected override ValueTask SerializeSessionCoreAsync( + AgentSession session, + JsonSerializerOptions? jsonSerializerOptions, + CancellationToken cancellationToken = default) => + new(JsonDocument.Parse("{}").RootElement.Clone()); + + protected override ValueTask DeserializeSessionCoreAsync( + JsonElement serializedState, + JsonSerializerOptions? jsonSerializerOptions, + CancellationToken cancellationToken = default) => + new(new TrackingSession()); + + private sealed class TrackingSession : AgentSession; + } + + /// A session store that records the storage identity of every agent it serves. + private sealed class RecordingSessionStore : AgentSessionStore, IAsyncDisposable + { + public ConcurrentQueue Identities { get; } = new(); + + public bool IsDisposed => this.DisposeCount > 0; + + public int DisposeCount { get; private set; } + + public ValueTask DisposeAsync() + { + this.DisposeCount++; + return default; + } + + public override ValueTask GetSessionAsync( + AIAgent agent, + AgentSessionStoreKey key, + CancellationToken cancellationToken = default) + { + this.Identities.Enqueue(FoundryHostingAgent.GetSessionStorageIdentity(agent, allowInstanceId: true)); + return new((AgentSession?)null); + } + + public override ValueTask SaveSessionAsync( + AIAgent agent, + AgentSessionStoreKey key, + AgentSession session, + CancellationToken cancellationToken = default) + { + this.Identities.Enqueue(FoundryHostingAgent.GetSessionStorageIdentity(agent, allowInstanceId: true)); + return default; + } + + public override ValueTask DeleteSessionAsync( + AIAgent agent, + AgentSessionStoreKey key, + CancellationToken cancellationToken = default) => + default; + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedStoredOutputHealthCheckTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedStoredOutputHealthCheckTests.cs index af2164c5f63..b948c524019 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedStoredOutputHealthCheckTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedStoredOutputHealthCheckTests.cs @@ -195,6 +195,33 @@ public async Task CheckHealthAsync_FoundryAgent_IsProbeableAsync() Assert.Contains("foundry-agent", (List)result.Data["storingAgents"]); } + [Fact] + public async Task CheckHealthAsync_ScopedAgentRegistration_IsResolvedFromAScopeAsync() + { + // Arrange: a scoped agent cannot be resolved from the root container when scopes are validated. + var services = new ServiceCollection(); + services.AddKeyedScoped("scoped-agent", (_, _) => new ChatClientAgent( + NewSilentChatClient(), + new ChatClientAgentOptions + { + Name = "scoped-agent", + ChatOptions = new ChatOptions + { + RawRepresentationFactory = _ => new CreateResponseOptions { StoredOutputEnabled = true }, + }, + })); + var check = new HostedStoredOutputHealthCheck( + services.BuildServiceProvider(new ServiceProviderOptions { ValidateScopes = true }), + new FoundryResponsesServerOptions()); + + // Act + var result = await check.CheckHealthAsync(NewContext(), CancellationToken.None); + + // Assert: the agent was resolved and probed rather than failing the readiness probe. + Assert.Equal(HealthStatus.Unhealthy, result.Status); + Assert.Contains("scoped-agent", (List)result.Data["storingAgents"]); + } + private static HostedStoredOutputHealthCheck BuildCheckFor(AIAgent agent, FoundryResponsesServerOptions? hostingOptions = null) { var services = new ServiceCollection(); diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedWorkflowCheckpointingHealthCheckTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedWorkflowCheckpointingHealthCheckTests.cs index 30d928271d9..e1af9d66e8e 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedWorkflowCheckpointingHealthCheckTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/HostedWorkflowCheckpointingHealthCheckTests.cs @@ -120,6 +120,27 @@ public void ApplyWorkflowCheckpointing_NotInAFoundryContainer_UsesTheSdkLocalFal Assert.True(result.GetService()?.UsesOwnCheckpointStorage); } + [Fact] + public async Task CheckHealthAsync_ScopedWorkflowAgent_IsResolvedFromAScopeAsync() + { + // Arrange: a scoped agent cannot be resolved from the root container when scopes are validated. + var environment = InProcessExecution.OffThread.WithCheckpointing(CheckpointManager.CreateInMemory()); + var services = new ServiceCollection(); + services.AddKeyedScoped("workflow", (_, _) => BuildWorkflowAgent(environment)); + var check = new HostedWorkflowCheckpointingHealthCheck( + services.BuildServiceProvider(new ServiceProviderOptions { ValidateScopes = true })) + { + IsHosted = true, + }; + + // Act + var result = await check.CheckHealthAsync(NewContext(), CancellationToken.None); + + // Assert: the agent was resolved and inspected rather than failing the readiness probe. + Assert.Equal(HealthStatus.Unhealthy, result.Status); + Assert.Equal(["WorkflowAgent"], Assert.IsType>(result.Data["incompatibleAgents"])); + } + private static HostedWorkflowCheckpointingHealthCheck BuildCheckFor(AIAgent agent) { var services = new ServiceCollection(); diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/ResilientTwoLifetimeIntegrationTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/ResilientTwoLifetimeIntegrationTests.cs index 1027e24dd79..838175aa322 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/ResilientTwoLifetimeIntegrationTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/ResilientTwoLifetimeIntegrationTests.cs @@ -125,8 +125,11 @@ await coordinator.PhasePersisted.Task.WaitAsync( } } - [Fact] - public async Task StoppedHost_RecoversWorkflowWithCompleteOrderedOutputAsync() + [Theory] + [InlineData(null)] + [InlineData(ServiceLifetime.Scoped)] + [InlineData(ServiceLifetime.Transient)] + public async Task StoppedHost_RecoversWorkflowWithCompleteOrderedOutputAsync(ServiceLifetime? factoryLifetime) { // Arrange string stateRoot = Path.Combine( @@ -152,7 +155,9 @@ public async Task StoppedHost_RecoversWorkflowWithCompleteOrderedOutputAsync() Directory.CreateDirectory(checkpointRoot))) { WebApplication firstHost = await StartServerAsync( - BuildCountdownWorkflowAgent(coordinator, checkpointStore), + "countdown-workflow", + () => BuildCountdownWorkflowAgent(coordinator, checkpointStore), + factoryLifetime, new FoundryAgentSessionStore(storeName: sessionStoreName), coordinator); try @@ -193,7 +198,9 @@ await WaitForPersistedOutputCountAsync( using var recoveryCheckpointStore = new FileSystemJsonCheckpointStore( Directory.CreateDirectory(checkpointRoot)); await using WebApplication secondHost = await StartServerAsync( - BuildCountdownWorkflowAgent(coordinator, recoveryCheckpointStore), + "countdown-workflow", + () => BuildCountdownWorkflowAgent(coordinator, recoveryCheckpointStore), + factoryLifetime, new FoundryAgentSessionStore(storeName: sessionStoreName)); using HttpClient secondClient = GetClient(secondHost); JsonElement completed = await WaitForTerminalAsync( @@ -229,17 +236,37 @@ await WaitForPersistedOutputCountAsync( } } - private static async Task StartServerAsync( + private static Task StartServerAsync( AIAgent agent, AgentSessionStore sessionStore, + CountdownRecoveryCoordinator? recoveryCoordinator = null) => + StartServerAsync(agent.Name!, () => agent, factoryLifetime: null, sessionStore, recoveryCoordinator); + + /// + /// Starts a host that registers the agent instance with AddFoundryResponses(agent) when + /// is , or with the factory overload and + /// that lifetime otherwise. + /// + private static async Task StartServerAsync( + string agentName, + Func createAgent, + ServiceLifetime? factoryLifetime, + AgentSessionStore sessionStore, CountdownRecoveryCoordinator? recoveryCoordinator = null) { WebApplicationBuilder builder = WebApplication.CreateBuilder(); builder.WebHost.UseTestServer(); - builder.Services.AddFoundryResponses( - agent, - sessionStore, - new FoundryResponsesServerOptions { ResilientBackground = true }); + var options = new FoundryResponsesServerOptions { ResilientBackground = true }; + if (factoryLifetime is { } lifetime) + { + builder.Services.AddFoundryResponses(agentName, _ => createAgent(), lifetime, options); + builder.Services.AddFoundryAgentSessionStore(sessionStore); + } + else + { + builder.Services.AddFoundryResponses(createAgent(), sessionStore, options); + } + builder.Services.AddLogging(); if (recoveryCoordinator is not null) {